From 21a2ca0b72ed9f392a515c23f57d6579e5be6da3 Mon Sep 17 00:00:00 2001
From: Christian Zimmermann <christian.zimmermann@ur.de>
Date: Wed, 13 Mar 2024 01:38:13 +0100
Subject: [PATCH] WIP: rrange

---
 src/opt/mpi/include/rrange.cc.h | 154 ++++++++++++++++++++++++++++++--
 src/opt/mpi/include/rrange.h    |  15 ++--
 2 files changed, 151 insertions(+), 18 deletions(-)

diff --git a/src/opt/mpi/include/rrange.cc.h b/src/opt/mpi/include/rrange.cc.h
index 50a815c..7ace931 100644
--- a/src/opt/mpi/include/rrange.cc.h
+++ b/src/opt/mpi/include/rrange.cc.h
@@ -46,19 +46,12 @@ namespace CNORXZ
 	template <class IndexI, class IndexK>
 	RIndex<IndexI,IndexK>::RIndex(const RangePtr& global, SizeT lexpos = 0) :
 	    mRange(rangeCast<RangeType>(global)),
-	    mI(std::make_shared<Index>(mRange->local())),
-	    mK(std::make_shared<YIndex>(mRange->geom()))
+	    mI(std::make_shared<IndexI>(mRange->local())),
+	    mK(std::make_shared<IndexK>(mRange->geom()))
 	{
 	    *this = lexpos;
 	}
 
-	template <class IndexI, class IndexK>
-	RIndex<IndexI,IndexK>::RIndex(const Sptr<Index>& local) :
-	{
-	    CXZ_ERROR("not implemented");
-	    //!!!
-	}
-	    
 	template <class IndexI, class IndexK>
 	RIndex<IndexI,IndexK>& RIndex<IndexI,IndexK>::operator=(SizeT pos)
 	{
@@ -299,6 +292,11 @@ namespace CNORXZ
 	    mRI(ri),
 	    mRK(rk)
 	{
+	    int s = 0;
+	    MPI_Comm_size(MPI_COMM_WORLD, &s);
+	    CXZ_ASSERT(rk->size() == static_cast<SizeT>(s),
+		       "geometry rank size ( = " << rk->size()
+		       << ") does not match number of ranks ( = " << s <<  ")");
 	    if constexpr(has_static_sub<typename RangeI::IndexType>::value and
 			 has_static_sub<typename RangeK::IndexType>::value) {
 		static_assert(typename RangeI::NR == typename RangeK::NR,
@@ -323,6 +321,144 @@ namespace CNORXZ
 	    }
 	}
 
+	/*==============+
+	 |    RRange    |
+	 +==============*/
+
+	template <class RangeI, class RangeK>
+	RangePtr RRange<RangeI,RangeK>::sub(SizeT num) const
+	{
+	    if(num == 0){
+		return mGeom;
+	    }
+	    else if(num == 1){
+		return mLocal;
+	    }
+	    else {
+		CXZ_ERROR("out of range: " << num);
+		return nullptr;
+	    }
+	}
+	
+	template <class RangeI, class RangeK>
+	MArray<RangePtr> RRange<RangeI,RangeK>::sub() const
+	{
+	    RangePtr sr = SRangeFactory<SizeT,2>( Arr<SizeT,2> { 0, 1 } ).create();
+	    return MArray<RangePtr>( sr, Vector<RangePtr>( { mGeom, mLocal } ) );
+	}
+	
+	template <class RangeI, class RangeK>
+	SizeT RRange<RangeI,RangeK>::size() const
+	{
+	    return mGeom->size() * mLocal->size();
+	}
+	
+	template <class RangeI, class RangeK>
+	SizeT RRange<RangeI,RangeK>::dim() const
+	{
+	    return mGeom->dim() + mLocal->dim();
+	}
+	
+	template <class RangeI, class RangeK>
+	String RRange<RangeI,RangeK>::stringMeta(SizeT pos) const
+	{
+	    return (this->begin()+pos).stringMeta();
+	}
+	
+	template <class RangeI, class RangeK>
+	const TypeInfo& RRange<RangeI,RangeK>::type() const
+	{
+	    return typeid(RRange<RangeI,RangeK>);
+	}
+	
+	template <class RangeI, class RangeK>
+	const TypeInfo& RRange<RangeI,RangeK>::metaType() const
+	{
+	    return typeid(MetaType);
+	}
+	
+	template <class RangeI, class RangeK>
+	RangePtr RRange<RangeI,RangeK>::extend(const RangePtr& r) const
+	{
+	    CXZ_ERROR("not implemented");
+	    return nullptr;
+	}
+
+	template <class RangeI, class RangeK>
+	Sptr<RangeI> RRange<RangeI,RangeK>::local() const
+	{
+	    return mLocal;
+	}
+
+	template <class RangeI, class RangeK>
+	Sptr<RangeK> RRange<RangeI,RangeK>::geom() const
+	{
+	    return mGeom;
+	}
+	    
+	template <class RangeI, class RangeK>
+	const MetaType RRange<RangeI,RangeK>::get(SizeT pos) const
+	{
+	    return (this->begin()+pos)->meta();
+	}
+
+	template <class RangeI, class RangeK>
+	SizeT RRange<RangeI,RangeK>::getMeta(const MetaType& metaPos) const
+	{
+	    auto i = this->begin();
+	    return i.at(metaPos).pos();
+	}
+
+	template <class RangeI, class RangeK>
+	SizeT RRange<RangeI,RangeK>::getRank(SizeT pos) const
+	{
+	    return (this->begin()+pos)->rank();
+	}
+	
+	/*=========================+
+	 |    RRange (protected)   |
+	 +=========================*/
+
+	template <class RangeI, class RangeK>
+	RRange<RangeI,RangeK>::RRange(const Sptr<RangeI>& loc, const Sptr<RangeK>& geom) :
+	    mLocal(loc),
+	    mGeom(geom)
+	{}
+
+	/*===========================+
+	 |    non-member functions   |
+	 +===========================*/
+
+	template <class RangeJ, class RangeK>
+	RangePtr rrange(const Sptr<RangeJ>& global, const Sptr<RangeK>& geom)
+	{
+	    if constexpr(has_static_sub<typename RangeJ::IndexType>::value and
+			 has_static_sub<typename RangeK::IndexType>::value) {
+		static_assert(typename RangeI::NR == typename RangeK::NR,
+			      "ranges have to be of same dimension");
+
+		constexpr SizeT N = typename RangeI::NR;
+		auto mr = ifor<0,N>( [&](auto mu) {
+		    return split( global->space()[CSizeT<mu>{}], geom->space()[CSizeT<mu>{}] );
+		}, [](const auto&... r) { return xplMrange(r,...); } );
+		typedef std::remove_reference<decltype(mr)>::type RangeI;
+		return RRangeFactory<RangeI,RangeK>(mr, geom).create(); // explicit range type!!!
+	    }
+	    // other cases!!!
+	    else {
+		CXZ_ASSERT(global->dim() == geom->dim(),
+			   "ranges have to be of same dimension, got "
+			   << global->dim() << " and " << geom->dim());
+		Vector<RangePtr> o(global->dim());
+		for(SizeT mu = 0; mu != global->dim(); ++mu){
+		    o[mu] = split( global->space()[mu], geom->space()[mu] );
+		}
+		const Sptr<YRange> yr = xplYrange(o);
+		return RRangeFactory<YRange,RangeK>(yr, geom).create();
+	    }
+	}
+	
+	
     } // namespace mpi
 } // namespace CNORXZ
 
diff --git a/src/opt/mpi/include/rrange.h b/src/opt/mpi/include/rrange.h
index 03ebae3..cab6b30 100644
--- a/src/opt/mpi/include/rrange.h
+++ b/src/opt/mpi/include/rrange.h
@@ -51,14 +51,11 @@ namespace CNORXZ
 	    RIndex& operator=(const RIndex& in);
 
 	    /** Construct from global range and format.
-		@param global Pointer to global range; this can be a RRange or the bare range (will be converted to a RRange).
+		@param global Pointer to global range (RRange).
 		@param lexpos Start position.
 	     */
 	    RIndex(const RangePtr& global, SizeT lexpos = 0);
 
-	    /** Construct from local index. */
-	    RIndex(const Sptr<IndexI,IndexK>& local);
-	    
 	    /** @copydoc IndexInterface::operator=(SizeT) */
 	    RIndex& operator=(SizeT pos);
 
@@ -191,7 +188,7 @@ namespace CNORXZ
 	    typedef RIndex<typename RangeI::IndexType,typename RangeK::IndexType> IndexType;
 	    typedef typename RangeI::MetaType MetaType;
 
-	    friend RRangeFactory<Range>;
+	    friend RRangeFactory<RangeI,RangeK>;
 
 	    virtual RangePtr sub(SizeT num) const override final;
 	    virtual MArray<RangePtr> sub() const override final;
@@ -221,7 +218,7 @@ namespace CNORXZ
 	    /** Get rank from lexicographic meta data position.
 		@param pos Lexicographic meta data position.
 	    */
-	    int getRank(SizeT pos) const;
+	    SizeT getRank(SizeT pos) const;
 
 	protected:
 
@@ -235,7 +232,7 @@ namespace CNORXZ
 		@param loc Local range.
 		@param geom Rank geometry range.
 	    */
-	    RRange(const Sptr<Range>& loc, const Sptr<YRange>& geom);
+	    RRange(const Sptr<RangeI>& loc, const Sptr<RangeK>& geom);
 	
 	    Sptr<RangeI> mLocal; /**< Local range of THIS rank. */
 	    Sptr<RangeK> mGeom; /**< Rank geometry range. */
@@ -246,8 +243,8 @@ namespace CNORXZ
 	    @param global Global range.
 	    @param geom Rank geometry.
 	*/
-	template <class GRange>
-	RangePtr rrange(const Sptr<GRange>& global, const Sptr<YRange>& geom);
+	template <class RangeJ, class RangeK>
+	RangePtr rrange(const Sptr<RangeJ>& global, const Sptr<RangeK>& geom);
 	
     } // namespace mpi
 } // namespace CNORXZ