Per-row counts of valid values in ellpack multimap are stored in separate vector

src/TNL/Experimental/Multimaps/EllpackIndexMultimap.h

@@ -61,6 +61,7 @@ class EllpackIndexMultimap
 
    protected:
       Containers::Vector< IndexType, DeviceType, IndexType > values;
+      Containers::Vector< LocalIndexType, DeviceType, IndexType > valuesCounts;
 
       IndexType keysRange;
       LocalIndexType maxValuesCount;

src/TNL/Experimental/Multimaps/EllpackIndexMultimapValues.h

@@ -10,6 +10,7 @@
 
 #pragma once
 
+#include <type_traits>
 #include <ostream>
 
 namespace TNL {
@@ -59,18 +60,25 @@ class EllpackIndexMultimapValues
 
       bool operator==( const ThisType& other ) const;
 
+      bool operator!=( const ThisType& other ) const;
+
       void print( std::ostream& str ) const;
 
    protected:
+      using ValuesCountType = typename std::conditional< std::is_const< IndexType >::value,
+                                                         typename std::add_const< LocalIndexType >::type,
+                                                         LocalIndexType >::type;
+
       EllpackIndexMultimapValues( IndexType* values,
+                                  ValuesCountType* valuesCounts,
                                   const IndexType& input,
                                   const LocalIndexType& allocatedSize );
 
       IndexType* values;
 
-      // TODO: step is unused
-//      LocalIndexType step;
+      ValuesCountType* valuesCount;
 
+      // TODO: this is useless for a const-accessor (without setSize etc.)
       LocalIndexType allocatedSize;
 
       friend EllpackIndexMultimap< IndexType, DeviceType, LocalIndexType >;

src/TNL/Experimental/Multimaps/EllpackIndexMultimapValues_impl.h

@@ -21,7 +21,7 @@ template< typename Index,
           typename LocalIndex >
 EllpackIndexMultimapValues< Index, Device, LocalIndex >::
 EllpackIndexMultimapValues()
-: values( nullptr ), allocatedSize( 0 )
+: values( nullptr ), valuesCount( nullptr ), allocatedSize( 0 )
 {
 }
 
@@ -30,9 +30,10 @@ template< typename Index,
           typename LocalIndex >
 EllpackIndexMultimapValues< Index, Device, LocalIndex >::
 EllpackIndexMultimapValues( ThisType&& other )
-: values( other.values ), allocatedSize( other.allocatedSize )
+: values( other.values ), valuesCount( other.valuesCount ), allocatedSize( other.allocatedSize )
 {
    other.values = nullptr;
+   other.valuesCount = nullptr;
    other.allocatedSize = 0;
 }
 
@@ -59,8 +60,10 @@ EllpackIndexMultimapValues< Index, Device, LocalIndex >::
 operator=( ThisType&& other )
 {
    this->values = other.values;
+   this->valuesCount = other.valuesCount;
    this->allocatedSize = other.allocatedSize;
    other.values = nullptr;
+   other.valuesCount = nullptr;
    other.allocatedSize = 0;
    return *this;
 }
@@ -73,6 +76,7 @@ EllpackIndexMultimapValues< Index, Device, LocalIndex >::
 bind( const ThisType& other )
 {
    this->values = other.values;
+   this->valuesCount = other.valuesCount;
    this->allocatedSize = other.allocatedSize;
 }
 
@@ -81,11 +85,14 @@ template< typename Index,
           typename LocalIndex >
 EllpackIndexMultimapValues< Index, Device, LocalIndex >::
 EllpackIndexMultimapValues( IndexType* values,
+                            ValuesCountType* valuesCounts,
                             const IndexType& input,
                             const LocalIndexType& allocatedSize )
 {
-   this->values = &values[ input * ( allocatedSize + 1 ) ];
+   this->values = &values[ input * allocatedSize ];
+   this->valuesCount = &valuesCounts[ input ];
    this->allocatedSize = allocatedSize;
+   Assert( *(this->valuesCount) <= allocatedSize, );
 }
 
 template< typename Index,
@@ -95,9 +102,9 @@ bool
 EllpackIndexMultimapValues< Index, Device, LocalIndex >::
 setSize( const LocalIndexType& size )
 {
-   if( ! this->values || size > this->allocatedSize )
+   if( ! this->valuesCount || size > this->allocatedSize )
       return false;
-   this->values[ this->allocatedSize ] = size;
+   *valuesCount = size;
    return true;
 }
 
@@ -108,9 +115,9 @@ LocalIndex
 EllpackIndexMultimapValues< Index, Device, LocalIndex >::
 getSize() const
 {
-   if( ! this->values )
+   if( ! valuesCount )
       return 0;
-   return this->values[ this->allocatedSize ];
+   return *valuesCount;
 }
 
 template< typename Index,
@@ -195,6 +202,16 @@ operator==( const ThisType& other ) const
    return true;
 }
 
+template< typename Index,
+          typename Device,
+          typename LocalIndex >
+bool
+EllpackIndexMultimapValues< Index, Device, LocalIndex >::
+operator!=( const ThisType& other ) const
+{
+   return ! ( *this == other );
+}
+
 template< typename Index,
           typename Device,
           typename LocalIndex >

src/TNL/Experimental/Multimaps/EllpackIndexMultimap_impl.h

@@ -79,11 +79,11 @@ allocate( const LocalIndexType& maxValuesCount )
 {
    Assert( maxValuesCount >= 0, );
    this->maxValuesCount = maxValuesCount;
-   if( ! this->values.setSize( this->keysRange * ( this->maxValuesCount + 1 ) ) )
+   if( ! this->values.setSize( this->keysRange * this->maxValuesCount ) )
       return false;
-   // TODO: maybe the local sizes should be stored differently?
-   for( IndexType i = 0; i < this->keysRange; i++ )
-      this->getValues( i ).setSize( maxValuesCount );
+   if( ! this->valuesCounts.setSize( this->keysRange ) )
+      return false;
+   this->valuesCounts.setValue( maxValuesCount );
    return true;
 }
 
@@ -101,11 +101,11 @@ allocate( const ValuesAllocationVectorType& valuesCounts )
  
    Assert( this->maxValuesCount >= 0,
               std::cerr << "this->maxValuesCount = " << this->maxValuesCount << std::endl; );
-   if( ! this->values.setSize( this->keysRange * ( this->maxValuesCount + 1 ) ) )
+   if( ! this->values.setSize( this->keysRange * this->maxValuesCount ) )
       return false;
-   // TODO: maybe the local sizes should be stored differently?
-   for( IndexType i = 0; i < this->keysRange; i++ )
-      this->getValues( i ).setSize( valuesCounts[ i ] );
+   if( ! this->valuesCounts.setSize( this->keysRange ) )
+      return false;
+   this->valuesCounts = valuesCounts;
    return true;
 }
 
@@ -116,7 +116,7 @@ typename EllpackIndexMultimap< Index, Device, LocalIndex >::ValuesAccessorType
 EllpackIndexMultimap< Index, Device, LocalIndex >::
 getValues( const IndexType& inputIndex )
 {
-   return ValuesAccessorType( this->values.getData(), inputIndex, this->maxValuesCount );
+   return ValuesAccessorType( this->values.getData(), this->valuesCounts.getData(), inputIndex, this->maxValuesCount );
 }
 
 template< typename Index,
@@ -126,7 +126,7 @@ typename EllpackIndexMultimap< Index, Device, LocalIndex >::ConstValuesAccessorT
 EllpackIndexMultimap< Index, Device, LocalIndex >::
 getValues( const IndexType& inputIndex ) const
 {
-   return ConstValuesAccessorType( this->values.getData(), inputIndex, this->maxValuesCount );
+   return ConstValuesAccessorType( this->values.getData(), this->valuesCounts.getData(), inputIndex, this->maxValuesCount );
 }
 
 template< typename Index,
@@ -136,9 +136,15 @@ bool
 EllpackIndexMultimap< Index, Device, LocalIndex >::
 operator==( const EllpackIndexMultimap< Index, Device, LocalIndex >& other ) const
 {
-   return ( this->keysRange == other.keysRange &&
-            this->maxValuesCount == other.maxValuesCount &&
-            this->values == other.values );
+   if( ! ( this->keysRange == other.keysRange &&
+           this->maxValuesCount == other.maxValuesCount &&
+           this->valuesCounts == other.valuesCounts ) )
+      return false;
+   // compare values for each key separately - the sizes may vary
+   for( IndexType i = 0; i < this->keysRange; i++ )
+      if( this->getValues( i ) != other.getValues( i ) )
+         return false;
+   return true;
 }
 
 template< typename Index,
@@ -156,6 +162,8 @@ save( File& file ) const
       return false;
    if( ! this->values.save( file ) )
       return false;
+   if( ! this->valuesCounts.save( file ) )
+      return false;
    return true;
 }
 
@@ -174,6 +182,8 @@ load( File& file )
       return false;
    if( ! this->values.load( file ) )
       return false;
+   if( ! this->valuesCounts.load( file ) )
+      return false;
    return true;
 }