Mana vector fixes

packages/core/epoch/types.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"github.com/cockroachdb/errors"
 	"github.com/iotaledger/hive.go/core/generics/set"
+	"github.com/iotaledger/hive.go/core/generics/shrinkingmap"
 	"github.com/iotaledger/hive.go/core/identity"
 	"strings"
 	"time"
@@ -227,7 +228,9 @@ func ComputeECR(tangleRoot, stateMutationRoot, stateRoot, manaRoot MerkleRoot) E
 
 // region NodesActivityLog //////////////////////////////////////////////////////////////////////////////////////////////////
 
-type NodesActivityLog map[Index]*ActivityLog
+type NodesActivityLog struct {
+	shrinkingmap.ShrinkingMap[Index, *ActivityLog]
+}
 
 func (al *NodesActivityLog) FromBytes(data []byte) (err error) {
 	_, err = serix.DefaultAPI.Decode(context.Background(), data, al, serix.WithValidation())
@@ -246,38 +249,42 @@ func (al *NodesActivityLog) Bytes() []byte {
 	return objBytes
 }
 
-// endregion ///////////////////////////////////////////////////////////////////////////////////////////////////////////
+func NewNodesActivityLog() *NodesActivityLog {
+	return &NodesActivityLog{*shrinkingmap.New[Index, *ActivityLog]()}
+}
+
+// endregi`azon ///////////////////////////////////////////////////////////////////////////////////////////////////////////
 
 // region ActivityLog //////////////////////////////////////////////////////////////////////////////////////////////////
 
 // ActivityLog is a time-based log of node activity. It stores information when a node is active and provides
 // functionality to query for certain timeframes.
 type ActivityLog struct {
-	SetEpochs *set.AdvancedSet[identity.ID] `serix:"0,lengthPrefixType=uint32"`
+	setEpochs *set.AdvancedSet[identity.ID] `serix:"0,lengthPrefixType=uint32"`
 }
 
 // NewActivityLog is the constructor for ActivityLog.
 func NewActivityLog() *ActivityLog {
 	a := &ActivityLog{
-		SetEpochs: set.NewAdvancedSet[identity.ID](),
+		setEpochs: set.NewAdvancedSet[identity.ID](),
 	}
 
 	return a
 }
 
 // Add adds a node to the activity log.
 func (a *ActivityLog) Add(nodeID identity.ID) (added bool) {
-	return a.SetEpochs.Add(nodeID)
+	return a.setEpochs.Add(nodeID)
 }
 
 // Remove removes a node from the activity log.
 func (a *ActivityLog) Remove(nodeID identity.ID) (removed bool) {
-	return a.SetEpochs.Delete(nodeID)
+	return a.setEpochs.Delete(nodeID)
 }
 
 // Active returns true if the provided node was active.
 func (a *ActivityLog) Active(nodeID identity.ID) (active bool) {
-	if a.SetEpochs.Has(nodeID) {
+	if a.setEpochs.Has(nodeID) {
 		return true
 	}
 
@@ -287,8 +294,8 @@ func (a *ActivityLog) Active(nodeID identity.ID) (active bool) {
 // String returns a human-readable version of ActivityLog.
 func (a *ActivityLog) String() string {
 	var builder strings.Builder
-	builder.WriteString(fmt.Sprintf("ActivityLog(len=%d, elements=", a.SetEpochs.Size()))
-	a.SetEpochs.ForEach(func(nodeID identity.ID) (err error) {
+	builder.WriteString(fmt.Sprintf("ActivityLog(len=%d, elements=", a.setEpochs.Size()))
+	a.setEpochs.ForEach(func(nodeID identity.ID) (err error) {
 		builder.WriteString(fmt.Sprintf("%s, ", nodeID.String()))
 		return
 	})
@@ -299,10 +306,26 @@ func (a *ActivityLog) String() string {
 // Clone clones the ActivityLog.
 func (a *ActivityLog) Clone() *ActivityLog {
 	clone := NewActivityLog()
-	clone.SetEpochs = a.SetEpochs.Clone()
+	clone.setEpochs = a.setEpochs.Clone()
 	return clone
 }
 
+// ForEach iterates through the activity set and calls the callback for every element.
+func (a *ActivityLog) ForEach(callback func(nodeID identity.ID) (err error)) (err error) {
+	a.setEpochs.ForEach(func(nodeID identity.ID) (err error) {
+		if err = callback(nodeID); err != nil {
+			return err
+		}
+		return nil
+	})
+
+	return err
+}
+
+func (a *ActivityLog) Size() int {
+	return a.Size()
+}
+
 // endregion ///////////////////////////////////////////////////////////////////////////////////////////////////////////
 
 // SnapshotEpochActivity is the data structure to store node activity for the snapshot.

packages/core/tangleold/cmanaweightprovider.go

@@ -4,14 +4,14 @@ import (
 	"fmt"
 	"github.com/iotaledger/goshimmer/packages/core/epoch"
 	"github.com/iotaledger/hive.go/core/generics/shrinkingmap"
+	"github.com/iotaledger/hive.go/core/serix"
 	"github.com/iotaledger/hive.go/core/types"
 	"sync"
 	"time"
 
 	"github.com/cockroachdb/errors"
 	"github.com/iotaledger/hive.go/core/identity"
 	"github.com/iotaledger/hive.go/core/kvstore"
-	"github.com/iotaledger/hive.go/core/serix"
 )
 
 func init() {
@@ -38,7 +38,7 @@ type ActivityUpdatesCount map[identity.ID]uint64
 type CManaWeightProvider struct {
 	store                       kvstore.KVStore
 	mutex                       sync.RWMutex
-	activityLog                 epoch.NodesActivityLog
+	activityLog                 *epoch.NodesActivityLog
 	updatedActivityCount        *shrinkingmap.ShrinkingMap[epoch.Index, ActivityUpdatesCount]
 	manaRetrieverFunc           ManaRetrieverFunc
 	timeRetrieverFunc           TimeRetrieverFunc
@@ -48,7 +48,7 @@ type CManaWeightProvider struct {
 // NewCManaWeightProvider is the constructor for CManaWeightProvider.
 func NewCManaWeightProvider(manaRetrieverFunc ManaRetrieverFunc, timeRetrieverFunc TimeRetrieverFunc, confirmedEpochRetrieverFunc ConfirmedEpochRetrieverFunc, store ...kvstore.KVStore) (cManaWeightProvider *CManaWeightProvider) {
 	cManaWeightProvider = &CManaWeightProvider{
-		activityLog:                 make(epoch.NodesActivityLog),
+		activityLog:                 epoch.NewNodesActivityLog(),
 		updatedActivityCount:        shrinkingmap.New[epoch.Index, ActivityUpdatesCount](shrinkingmap.WithShrinkingThresholdCount(100)),
 		manaRetrieverFunc:           manaRetrieverFunc,
 		timeRetrieverFunc:           timeRetrieverFunc,
@@ -83,10 +83,10 @@ func (c *CManaWeightProvider) Update(ei epoch.Index, nodeID identity.ID) {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
 
-	a, exists := c.activityLog[ei]
+	a, exists := c.activityLog.Get(ei)
 	if !exists {
 		a = epoch.NewActivityLog()
-		c.activityLog[ei] = a
+		c.activityLog.Set(ei, a)
 	}
 
 	a.Add(nodeID)
@@ -108,7 +108,7 @@ func (c *CManaWeightProvider) Remove(ei epoch.Index, nodeID identity.ID, updated
 	}
 	// if that was the last activity for this node in the ei epoch, then remove it from activity list
 	if epochUpdatesCount[nodeID] == 0 {
-		if a, exists := c.activityLog[ei]; exists {
+		if a, exists := c.activityLog.Get(ei); exists {
 			a.Remove(nodeID)
 			return true
 		}
@@ -136,11 +136,11 @@ func (c *CManaWeightProvider) WeightsOfRelevantVoters() (weights map[identity.ID
 	// nodes mana is counted only once for total weight calculation
 	totalWeightOnce := make(map[identity.ID]types.Empty)
 	for ei := lowerBoundEpoch; ei <= upperBoundEpoch; ei++ {
-		al, exists := c.activityLog[ei]
+		al, exists := c.activityLog.Get(ei)
 		if !exists {
 			continue
 		}
-		al.SetEpochs.ForEach(func(nodeID identity.ID) error {
+		al.ForEach(func(nodeID identity.ID) error {
 			nodeMana := mana[nodeID]
 			// Do this check after determining whether a node was active because otherwise we would never clean up
 			// the ActivityLog of nodes lower than the threshold.
@@ -170,8 +170,8 @@ func (c *CManaWeightProvider) SnapshotEpochActivity(epochDiffIndex epoch.Index)
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
 
-	for ei, al := range c.activityLog {
-		al.SetEpochs.ForEach(func(nodeID identity.ID) error {
+	c.activityLog.ForEach(func(ei epoch.Index, activity *epoch.ActivityLog) bool {
+		activity.ForEach(func(nodeID identity.ID) error {
 			// we save only activity log up to epochDiffIndex as it is the last snapshotted epoch
 			if ei > epochDiffIndex {
 				return nil
@@ -186,7 +186,9 @@ func (c *CManaWeightProvider) SnapshotEpochActivity(epochDiffIndex epoch.Index)
 			}
 			return nil
 		})
-	}
+		return true
+	})
+
 	return
 }
 
@@ -204,11 +206,14 @@ func (c *CManaWeightProvider) LoadActiveNodes(loadedActiveNodes epoch.SnapshotEp
 	defer c.mutex.Unlock()
 
 	for ei, epochActivity := range loadedActiveNodes {
-		if _, ok := c.activityLog[ei]; !ok {
-			c.activityLog[ei] = epoch.NewActivityLog()
+		var activityLog *epoch.ActivityLog
+		var ok bool
+		if activityLog, ok = c.activityLog.Get(ei); !ok {
+			activityLog = epoch.NewActivityLog()
+			c.activityLog.Set(ei, activityLog)
 		}
 		for nodeID, activityCount := range epochActivity.NodesLog() {
-			c.activityLog[ei].Add(nodeID)
+			activityLog.Add(nodeID)
 			c.updateActivityCount(ei, nodeID, activityCount)
 		}
 	}
@@ -224,15 +229,16 @@ type TimeRetrieverFunc func() time.Time
 type ConfirmedEpochRetrieverFunc func() epoch.Index
 
 // activeNodes returns the map of the active nodes.
-func (c *CManaWeightProvider) activeNodes() (activeNodes epoch.NodesActivityLog) {
-	activeNodes = make(epoch.NodesActivityLog)
+func (c *CManaWeightProvider) activeNodes() (activeNodes *epoch.NodesActivityLog) {
+	activeNodes = epoch.NewNodesActivityLog()
 
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
 
-	for nodeID, al := range c.activityLog {
-		activeNodes[nodeID] = al.Clone()
-	}
+	c.activityLog.ForEach(func(ei epoch.Index, activity *epoch.ActivityLog) bool {
+		activeNodes.Set(ei, activity.Clone())
+		return true
+	})
 
 	return activeNodes
 }
@@ -249,11 +255,12 @@ func (c *CManaWeightProvider) activityBoundaries() (lowerBoundEpoch, upperBoundE
 
 // clean removes all activity logs for epochs lower than provided bound.
 func (c *CManaWeightProvider) clean(cutoffEpoch epoch.Index) {
-	for ei := range c.activityLog {
+	c.activityLog.ForEachKey(func(ei epoch.Index) bool {
 		if ei < cutoffEpoch {
-			delete(c.activityLog, ei)
+			c.activityLog.Delete(ei)
 		}
-	}
+		return true
+	})
 	// clean also the updates counting map
 	c.updatedActivityCount.ForEach(func(ei epoch.Index, count ActivityUpdatesCount) bool {
 		if ei < cutoffEpoch {

packages/core/tangleold/cmanaweightprovider_test.go

@@ -1,6 +1,7 @@
 package tangleold
 
 import (
+	"math/rand"
 	"testing"
 	"time"
 
@@ -19,22 +20,39 @@ func TestActiveNodesMarshalling(t *testing.T) {
 		"node3": identity.GenerateIdentity().ID(),
 	}
 
-	activeNodes := make(epoch.NodesActivityLog)
+	activeNodes := epoch.NewNodesActivityLog()
+
+	for i := 0; i < 100; i++ {
+		ei := epoch.Index(i)
+		al := epoch.NewActivityLog()
+		activeNodes.Set(ei, al)
+		weight := 1.0
+		for _, nodeID := range nodes {
+			if rand.Float64() < 0.1*weight {
+				al.Add(nodeID)
+			}
+			weight += 1
+		}
+	}
 
 	for _, nodeID := range nodes {
 		for i := 0; i < crypto.Randomness.Intn(100); i++ {
-			activeNodes[epoch.Index(i)] = epoch.NewActivityLog()
-			activeNodes[epoch.Index(i)].Add(nodeID)
+			al := epoch.NewActivityLog()
+			al.Add(nodeID)
+			activeNodes.Set(epoch.Index(i), al)
 		}
 	}
 	activeNodesBytes := activeNodes.Bytes()
-	activeNodes2 := make(epoch.NodesActivityLog)
+	activeNodes2 := epoch.NewNodesActivityLog()
 	err := activeNodes2.FromBytes(activeNodesBytes)
 	require.NoError(t, err)
 
-	for nodeID, a := range activeNodes {
-		assert.EqualValues(t, a.SetEpochs.Size(), activeNodes2[nodeID].SetEpochs.Size())
-	}
+	activeNodes.ForEach(func(ei epoch.Index, activity *epoch.ActivityLog) bool {
+		activity2, exists := activeNodes2.Get(ei)
+		require.True(t, exists)
+		assert.EqualValues(t, activity.Size(), activity2.Size())
+		return true
+	})
 }
 
 func TestCManaWeightProvider(t *testing.T) {