Add item specific callbacks to OnChangeMap (#442)

Author: muXxer

core/generics/onchangemap/onchangemap.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/iotaledger/hive.go/core/generics/constraints"
 	"github.com/iotaledger/hive.go/core/generics/lo"
+	"github.com/iotaledger/hive.go/core/generics/options"
 )
 
 // Item represents an item in the OnChangeMap.
@@ -14,53 +15,105 @@ type Item[K comparable, C constraints.ComparableStringer[K]] interface {
 	Clone() Item[K, C]
 }
 
-// OnChangeMap is a map that executes a callback if the map or an item is modified,
+// OnChangeMap is a map that executes callbacks if the map or an item is modified,
 // in case callbackEnabled is true.
 type OnChangeMap[K comparable, C constraints.ComparableStringer[K], I Item[K, C]] struct {
 	mutex sync.RWMutex
 
-	m               map[K]I
-	callback        func([]I) error
-	callbackEnabled bool
+	m                map[K]I
+	callbacksEnabled bool
+
+	changedCallback      func([]I) error
+	itemAddedCallback    func(I) error
+	itemModifiedCallback func(I) error
+	itemDeletedCallback  func(I) error
 }
 
-// NewOnChangeMap creates a new OnChangeMap.
-func NewOnChangeMap[K comparable, C constraints.ComparableStringer[K], I Item[K, C]](callback func([]I) error) *OnChangeMap[K, C, I] {
-	return &OnChangeMap[K, C, I]{
-		m:               make(map[K]I),
-		callback:        callback,
-		callbackEnabled: false,
+// WithChangedCallback is triggered when something in the OnChangeMap is changed (added/modified/deleted).
+func WithChangedCallback[K comparable, C constraints.ComparableStringer[K], I Item[K, C]](changedCallback func([]I) error) options.Option[OnChangeMap[K, C, I]] {
+	return func(r *OnChangeMap[K, C, I]) {
+		r.changedCallback = changedCallback
+	}
+}
+
+// WithItemAddedCallback is triggered when a new item is added.
+func WithItemAddedCallback[K comparable, C constraints.ComparableStringer[K], I Item[K, C]](itemAddedCallback func(I) error) options.Option[OnChangeMap[K, C, I]] {
+	return func(r *OnChangeMap[K, C, I]) {
+		r.itemAddedCallback = itemAddedCallback
+	}
+}
+
+// WithItemModifiedCallback is triggered when an item is modified.
+func WithItemModifiedCallback[K comparable, C constraints.ComparableStringer[K], I Item[K, C]](itemModifiedCallback func(I) error) options.Option[OnChangeMap[K, C, I]] {
+	return func(r *OnChangeMap[K, C, I]) {
+		r.itemModifiedCallback = itemModifiedCallback
+	}
+}
+
+// WithItemDeletedCallback is triggered when an item is deleted.
+func WithItemDeletedCallback[K comparable, C constraints.ComparableStringer[K], I Item[K, C]](itemDeletedCallback func(I) error) options.Option[OnChangeMap[K, C, I]] {
+	return func(r *OnChangeMap[K, C, I]) {
+		r.itemDeletedCallback = itemDeletedCallback
 	}
 }
 
-// CallbackEnabled sets whether executing the callback on change is active or not.
-func (r *OnChangeMap[K, C, I]) CallbackEnabled(enabled bool) {
-	r.callbackEnabled = enabled
+// NewOnChangeMap creates a new OnChangeMap.
+func NewOnChangeMap[K comparable, C constraints.ComparableStringer[K], I Item[K, C]](opts ...options.Option[OnChangeMap[K, C, I]]) *OnChangeMap[K, C, I] {
+	return options.Apply(&OnChangeMap[K, C, I]{
+		m:                    make(map[K]I),
+		callbacksEnabled:     false,
+		changedCallback:      nil,
+		itemAddedCallback:    nil,
+		itemModifiedCallback: nil,
+		itemDeletedCallback:  nil,
+	}, opts)
+}
+
+// CallbacksEnabled sets whether executing the callbacks on change is active or not.
+func (r *OnChangeMap[K, C, I]) CallbacksEnabled(enabled bool) {
+	r.callbacksEnabled = enabled
 }
 
-// executeCallbackWithoutLocking calls the callback if callbackEnabled is true.
-func (r *OnChangeMap[K, C, I]) executeCallbackWithoutLocking() error {
-	if !r.callbackEnabled {
+// executeChangedCallback calls the changedCallback if callbackEnabled is true.
+func (r *OnChangeMap[K, C, I]) executeChangedCallback() error {
+	if !r.callbacksEnabled {
 		return nil
 	}
 
-	if r.callback == nil {
+	if r.changedCallback != nil {
+		if err := r.changedCallback(lo.Values(r.m)); err != nil {
+			return fmt.Errorf("failed to execute callback in OnChangeMap: %w", err)
+		}
+	}
+
+	return nil
+}
+
+// executeItemCallback calls the given callback if callbackEnabled is true.
+func (r *OnChangeMap[K, C, I]) executeItemCallback(callback func(I) error, item I) error {
+	if !r.callbacksEnabled {
 		return nil
 	}
 
-	if err := r.callback(lo.Values(r.m)); err != nil {
-		return fmt.Errorf("failed to execute callback in OnChangeMap: %w", err)
+	if err := r.executeChangedCallback(); err != nil {
+		return err
+	}
+
+	if callback != nil {
+		if err := callback(item); err != nil {
+			return fmt.Errorf("failed to execute item callback in OnChangeMap: %w", err)
+		}
 	}
 
 	return nil
 }
 
-// ExecuteCallback calls the callback if callbackEnabled is true.
-func (r *OnChangeMap[K, C, I]) ExecuteCallback() error {
+// ExecuteChangedCallback calls the changedCallback if callbackEnabled is true.
+func (r *OnChangeMap[K, C, I]) ExecuteChangedCallback() error {
 	r.mutex.RLock()
 	defer r.mutex.RUnlock()
 
-	return r.executeCallbackWithoutLocking()
+	return r.executeChangedCallback()
 }
 
 // All returns a copy of all items.
@@ -100,7 +153,7 @@ func (r *OnChangeMap[K, C, I]) Add(item I) error {
 
 	r.m[item.ID().Key()] = item
 
-	return r.executeCallbackWithoutLocking()
+	return r.executeItemCallback(r.itemAddedCallback, item)
 }
 
 // Modify modifies an item in the map and returns a copy.
@@ -118,19 +171,20 @@ func (r *OnChangeMap[K, C, I]) Modify(id C, callback func(item I) bool) (I, erro
 		return item.Clone().(I), nil
 	}
 
-	return item.Clone().(I), r.executeCallbackWithoutLocking()
+	return item.Clone().(I), r.executeItemCallback(r.itemModifiedCallback, item)
 }
 
 // Delete removes an item from the map.
 func (r *OnChangeMap[K, C, I]) Delete(id C) error {
 	r.mutex.Lock()
 	defer r.mutex.Unlock()
 
-	if _, exists := r.m[id.Key()]; !exists {
+	item, exists := r.m[id.Key()]
+	if !exists {
 		return fmt.Errorf("unable to remove item: \"%s\" does not exist in map", id)
 	}
 
 	delete(r.m, id.Key())
 
-	return r.executeCallbackWithoutLocking()
+	return r.executeItemCallback(r.itemDeletedCallback, item)
 }

core/generics/onchangemap/onchangemap_test.go

@@ -42,18 +42,39 @@ func (i *testItem) Clone() Item[int, comparableInteger] {
 
 func TestOnChangeMap(t *testing.T) {
 	var storedItems []*testItem
-
-	onChangeMap := NewOnChangeMap[int, comparableInteger](func(items []*testItem) error {
-		storedItems = items
-		return nil
-	})
+	var itemAdded *testItem
+	var itemModified *testItem
+	var itemDeleted *testItem
+
+	onChangeMap := NewOnChangeMap(
+		WithChangedCallback[int, comparableInteger](func(items []*testItem) error {
+			storedItems = items
+			return nil
+		}),
+		WithItemAddedCallback[int, comparableInteger](func(item *testItem) error {
+			itemAdded = item
+			return nil
+		}),
+		WithItemModifiedCallback[int, comparableInteger](func(item *testItem) error {
+			itemModified = item
+			return nil
+		}),
+		WithItemDeletedCallback[int, comparableInteger](func(item *testItem) error {
+			itemDeleted = item
+			return nil
+		}),
+	)
 
 	require.NotNil(t, onChangeMap)
 
 	item1 := newTestItem(1, "one")
 	item2 := newTestItem(2, "two")
 	item3 := newTestItem(3, "three")
 
+	require.Nil(t, itemAdded)
+	require.Nil(t, itemModified)
+	require.Nil(t, itemDeleted)
+
 	// add and get an item
 	err := onChangeMap.Add(item1)
 	require.NoError(t, err)
@@ -74,28 +95,39 @@ func TestOnChangeMap(t *testing.T) {
 
 	// check store on change
 	require.Equal(t, len(storedItems), 0)
-	err = onChangeMap.ExecuteCallback()
+	err = onChangeMap.ExecuteChangedCallback()
 	require.NoError(t, err)
 	require.Equal(t, len(storedItems), 0)
 
 	// enable store on changed
-	onChangeMap.CallbackEnabled(true)
-	err = onChangeMap.ExecuteCallback()
+	onChangeMap.CallbacksEnabled(true)
+	err = onChangeMap.ExecuteChangedCallback()
 	require.NoError(t, err)
 	require.Equal(t, len(storedItems), 1)
 
 	// add duplicate
 	err = onChangeMap.Add(item1)
 	require.Error(t, err)
+	require.Nil(t, itemAdded)
+	require.Nil(t, itemModified)
+	require.Nil(t, itemDeleted)
+	itemAdded = nil
 
 	// add second item
 	err = onChangeMap.Add(item2)
 	require.NoError(t, err)
 	require.Equal(t, len(storedItems), 2)
+	require.NotNil(t, itemAdded)
+	require.Nil(t, itemModified)
+	require.Nil(t, itemDeleted)
+	itemAdded = nil
 
 	// modify non existing item
 	_, err = onChangeMap.Modify(3, nil)
 	require.Error(t, err)
+	require.Nil(t, itemAdded)
+	require.Nil(t, itemModified)
+	require.Nil(t, itemDeleted)
 
 	// modify existing item
 	item2Copy, err := onChangeMap.Modify(2, func(item *testItem) bool {
@@ -105,6 +137,10 @@ func TestOnChangeMap(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, len(storedItems), 2)
 	require.Equal(t, item2Copy.value, item3.value)
+	require.Nil(t, itemAdded)
+	require.NotNil(t, itemModified)
+	require.Nil(t, itemDeleted)
+	itemModified = nil
 
 	// get modified item
 	item2Copy, err = onChangeMap.Get(2)
@@ -115,8 +151,15 @@ func TestOnChangeMap(t *testing.T) {
 	err = onChangeMap.Delete(2)
 	require.NoError(t, err)
 	require.Equal(t, len(storedItems), 1)
+	require.Nil(t, itemAdded)
+	require.Nil(t, itemModified)
+	require.NotNil(t, itemDeleted)
+	itemDeleted = nil
 
 	// delete non-existing item
 	err = onChangeMap.Delete(2)
 	require.Error(t, err)
+	require.Nil(t, itemAdded)
+	require.Nil(t, itemModified)
+	require.Nil(t, itemDeleted)
 }