feat: implemented CandidateStorage

dot/parachain/prospective-parachains/fragment-chain/fragment_chain.go

@@ -2,6 +2,7 @@ package fragmentchain
 
 import (
 	"fmt"
+	"iter"
 
 	parachaintypes "github.com/ChainSafe/gossamer/dot/parachain/types"
 	inclusionemulator "github.com/ChainSafe/gossamer/dot/parachain/util/inclusion-emulator"
@@ -80,7 +81,7 @@ func NewCandidateEntry(
 type CandidateStorage struct {
 	byParentHead    map[common.Hash]map[parachaintypes.CandidateHash]any
 	byOutputHead    map[common.Hash]map[parachaintypes.CandidateHash]any
-	byCandidateHash map[parachaintypes.CandidateHash]CandidateEntry
+	byCandidateHash map[parachaintypes.CandidateHash]*CandidateEntry
 }
 
 func (c *CandidateStorage) AddPendingAvailabilityCandidate(
@@ -103,6 +104,8 @@ func (c *CandidateStorage) addCandidateEntry(candidate *CandidateEntry) error {
 	}
 
 	// updates the reference parent hash -> candidate
+	// we don't check the `ok` value since the key can
+	// exists in the map but pointing to a nil hashset
 	setOfCandidates := c.byParentHead[candidate.parentHeadDataHash]
 	if setOfCandidates == nil {
 		setOfCandidates = make(map[parachaintypes.CandidateHash]any)
@@ -118,4 +121,96 @@ func (c *CandidateStorage) addCandidateEntry(candidate *CandidateEntry) error {
 	setOfCandidates[candidate.candidateHash] = struct{}{}
 	c.byOutputHead[candidate.outputHeadDataHash] = setOfCandidates
 
+	return nil
+}
+
+func (c *CandidateStorage) removeCandidate(candidateHash parachaintypes.CandidateHash) {
+	entry, ok := c.byCandidateHash[candidateHash]
+	if !ok {
+		return
+	}
+
+	delete(c.byCandidateHash, candidateHash)
+
+	if setOfCandidates, ok := c.byParentHead[entry.parentHeadDataHash]; ok {
+		delete(setOfCandidates, candidateHash)
+		if len(setOfCandidates) == 0 {
+			delete(c.byParentHead, entry.parentHeadDataHash)
+		}
+	}
+
+	if setOfCandidates, ok := c.byOutputHead[entry.outputHeadDataHash]; ok {
+		delete(setOfCandidates, candidateHash)
+		if len(setOfCandidates) == 0 {
+			delete(c.byOutputHead, entry.outputHeadDataHash)
+		}
+	}
+}
+
+func (c *CandidateStorage) MarkBacked(candidateHash parachaintypes.CandidateHash) {
+	entry, ok := c.byCandidateHash[candidateHash]
+	if !ok {
+		fmt.Println("candidate not found while marking as backed")
+	}
+
+	entry.state = Backed
+	fmt.Println("candidate marked as backed")
+}
+
+func (c *CandidateStorage) Contains(candidateHash parachaintypes.CandidateHash) bool {
+	_, ok := c.byCandidateHash[candidateHash]
+	return ok
+}
+
+// Candidates returns an iterator over references to the stored candidates, in arbitrary order.
+func (c *CandidateStorage) Candidates() iter.Seq[*CandidateEntry] {
+	return func(yield func(*CandidateEntry) bool) {
+		for _, entry := range c.byCandidateHash {
+			if !yield(entry) {
+				return
+			}
+		}
+	}
+}
+
+func (c *CandidateStorage) HeadDataByHash(hash common.Hash) *parachaintypes.HeadData {
+	// first, search for candidates outputting this head data and extract the head data
+	// from their commitments if they exist.
+	// otherwise, search for candidates building upon this head data and extract the
+	// head data from their persisted validation data if they exist.
+
+	if setOfCandidateHashes, ok := c.byOutputHead[hash]; ok {
+		for candidateHash := range setOfCandidateHashes {
+			if candidate, ok := c.byCandidateHash[candidateHash]; ok {
+				return &candidate.candidate.Commitments.HeadData
+			}
+		}
+	}
+
+	if setOfCandidateHashes, ok := c.byParentHead[hash]; ok {
+		for candidateHash := range setOfCandidateHashes {
+			if candidate, ok := c.byCandidateHash[candidateHash]; ok {
+				return &candidate.candidate.PersistedValidationData.ParentHead
+			}
+		}
+	}
+
+	return nil
+}
+
+func (c *CandidateStorage) PossibleBackedParaChildren(parentHeadHash common.Hash) iter.Seq[*CandidateEntry] {
+	return func(yield func(*CandidateEntry) bool) {
+		seqOfCandidateHashes, ok := c.byParentHead[parentHeadHash]
+		if !ok {
+			return
+		}
+
+		for candidateHash := range seqOfCandidateHashes {
+			if entry, ok := c.byCandidateHash[candidateHash]; ok && entry.state == Backed {
+				if !yield(entry) {
+					return
+				}
+			}
+		}
+	}
 }

dot/parachain/prospective-parachains/fragment-chain/fragment_chain_test.go

@@ -0,0 +1,251 @@
+package fragmentchain
+
+import (
+	"testing"
+
+	parachaintypes "github.com/ChainSafe/gossamer/dot/parachain/types"
+	inclusionemulator "github.com/ChainSafe/gossamer/dot/parachain/util/inclusion-emulator"
+	"github.com/ChainSafe/gossamer/lib/common"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCandidateStorage_RemoveCandidate(t *testing.T) {
+	storage := &CandidateStorage{
+		byParentHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+		byOutputHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+		byCandidateHash: make(map[parachaintypes.CandidateHash]*CandidateEntry),
+	}
+
+	candidateHash := parachaintypes.CandidateHash{Value: common.Hash{1, 2, 3}}
+	parentHeadHash := common.Hash{4, 5, 6}
+	outputHeadHash := common.Hash{7, 8, 9}
+
+	entry := &CandidateEntry{
+		candidateHash:      candidateHash,
+		parentHeadDataHash: parentHeadHash,
+		outputHeadDataHash: outputHeadHash,
+		state:              Backed,
+	}
+
+	storage.byCandidateHash[candidateHash] = entry
+	storage.byParentHead[parentHeadHash] = map[parachaintypes.CandidateHash]any{candidateHash: struct{}{}}
+	storage.byOutputHead[outputHeadHash] = map[parachaintypes.CandidateHash]any{candidateHash: struct{}{}}
+
+	storage.removeCandidate(candidateHash)
+
+	_, exists := storage.byCandidateHash[candidateHash]
+	assert.False(t, exists, "candidate should be removed from byCandidateHash")
+
+	_, exists = storage.byParentHead[parentHeadHash]
+	assert.False(t, exists, "candidate should be removed from byParentHead")
+
+	_, exists = storage.byOutputHead[outputHeadHash]
+	assert.False(t, exists, "candidate should be removed from byOutputHead")
+}
+
+func TestCandidateStorage_MarkBacked(t *testing.T) {
+	storage := &CandidateStorage{
+		byParentHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+		byOutputHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+		byCandidateHash: make(map[parachaintypes.CandidateHash]*CandidateEntry),
+	}
+
+	candidateHash := parachaintypes.CandidateHash{Value: common.Hash{1, 2, 3}}
+	parentHeadHash := common.Hash{4, 5, 6}
+	outputHeadHash := common.Hash{7, 8, 9}
+
+	entry := &CandidateEntry{
+		candidateHash:      candidateHash,
+		parentHeadDataHash: parentHeadHash,
+		outputHeadDataHash: outputHeadHash,
+		state:              Seconded,
+	}
+
+	storage.byCandidateHash[candidateHash] = entry
+	storage.byParentHead[parentHeadHash] = map[parachaintypes.CandidateHash]any{candidateHash: struct{}{}}
+	storage.byOutputHead[outputHeadHash] = map[parachaintypes.CandidateHash]any{candidateHash: struct{}{}}
+
+	storage.MarkBacked(candidateHash)
+
+	assert.Equal(t, Backed, entry.state, "candidate state should be marked as backed")
+}
+
+func TestCandidateStorage_HeadDataByHash(t *testing.T) {
+	tests := map[string]struct {
+		setup    func() *CandidateStorage
+		hash     common.Hash
+		expected *parachaintypes.HeadData
+	}{
+		"find_head_data_of_first_candidate_using_output_head_data_hash": {
+			setup: func() *CandidateStorage {
+				storage := &CandidateStorage{
+					byParentHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+					byOutputHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+					byCandidateHash: make(map[parachaintypes.CandidateHash]*CandidateEntry),
+				}
+
+				candidateHash := parachaintypes.CandidateHash{Value: common.Hash{1, 2, 3}}
+				parentHeadHash := common.Hash{4, 5, 6}
+				outputHeadHash := common.Hash{7, 8, 9}
+				headData := parachaintypes.HeadData{Data: []byte{10, 11, 12}}
+
+				entry := &CandidateEntry{
+					candidateHash:      candidateHash,
+					parentHeadDataHash: parentHeadHash,
+					outputHeadDataHash: outputHeadHash,
+					candidate: inclusionemulator.ProspectiveCandidate{
+						Commitments: parachaintypes.CandidateCommitments{
+							HeadData: headData,
+						},
+					},
+				}
+
+				storage.byCandidateHash[candidateHash] = entry
+				storage.byParentHead[parentHeadHash] = map[parachaintypes.CandidateHash]any{candidateHash: struct{}{}}
+				storage.byOutputHead[outputHeadHash] = map[parachaintypes.CandidateHash]any{candidateHash: struct{}{}}
+
+				return storage
+			},
+			hash:     common.Hash{7, 8, 9},
+			expected: &parachaintypes.HeadData{Data: []byte{10, 11, 12}},
+		},
+		"find_head_data_using_parent_head_data_hash_from_second_candidate": {
+			setup: func() *CandidateStorage {
+				storage := &CandidateStorage{
+					byParentHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+					byOutputHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+					byCandidateHash: make(map[parachaintypes.CandidateHash]*CandidateEntry),
+				}
+
+				candidateHash := parachaintypes.CandidateHash{Value: common.Hash{13, 14, 15}}
+				parentHeadHash := common.Hash{16, 17, 18}
+				outputHeadHash := common.Hash{19, 20, 21}
+				headData := parachaintypes.HeadData{Data: []byte{22, 23, 24}}
+
+				entry := &CandidateEntry{
+					candidateHash:      candidateHash,
+					parentHeadDataHash: parentHeadHash,
+					outputHeadDataHash: outputHeadHash,
+					candidate: inclusionemulator.ProspectiveCandidate{
+						PersistedValidationData: parachaintypes.PersistedValidationData{
+							ParentHead: headData,
+						},
+					},
+				}
+
+				storage.byCandidateHash[candidateHash] = entry
+				storage.byParentHead[parentHeadHash] = map[parachaintypes.CandidateHash]any{candidateHash: struct{}{}}
+				storage.byOutputHead[outputHeadHash] = map[parachaintypes.CandidateHash]any{candidateHash: struct{}{}}
+
+				return storage
+			},
+			hash:     common.Hash{16, 17, 18},
+			expected: &parachaintypes.HeadData{Data: []byte{22, 23, 24}},
+		},
+		"use_nonexistent_hash_and_should_get_nil": {
+			setup: func() *CandidateStorage {
+				storage := &CandidateStorage{
+					byParentHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+					byOutputHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+					byCandidateHash: make(map[parachaintypes.CandidateHash]*CandidateEntry),
+				}
+				return storage
+			},
+			hash:     common.Hash{99, 99, 99},
+			expected: nil,
+		},
+		"insert_0_candidates_and_try_to_find_but_should_get_nil": {
+			setup: func() *CandidateStorage {
+				return &CandidateStorage{
+					byParentHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+					byOutputHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+					byCandidateHash: make(map[parachaintypes.CandidateHash]*CandidateEntry),
+				}
+			},
+			hash:     common.Hash{7, 8, 9},
+			expected: nil,
+		},
+	}
+
+	for name, tt := range tests {
+		tt := tt
+		t.Run(name, func(t *testing.T) {
+			storage := tt.setup()
+			result := storage.HeadDataByHash(tt.hash)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+func TestCandidateStorage_PossibleBackedParaChildren(t *testing.T) {
+	tests := map[string]struct {
+		setup    func() *CandidateStorage
+		hash     common.Hash
+		expected []*CandidateEntry
+	}{
+		"insert_2_candidates_for_same_parent_one_seconded_one_backed": {
+			setup: func() *CandidateStorage {
+				storage := &CandidateStorage{
+					byParentHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+					byOutputHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+					byCandidateHash: make(map[parachaintypes.CandidateHash]*CandidateEntry),
+				}
+
+				candidateHash1 := parachaintypes.CandidateHash{Value: common.Hash{1, 2, 3}}
+				parentHeadHash := common.Hash{4, 5, 6}
+				outputHeadHash1 := common.Hash{7, 8, 9}
+
+				candidateHash2 := parachaintypes.CandidateHash{Value: common.Hash{10, 11, 12}}
+				outputHeadHash2 := common.Hash{13, 14, 15}
+
+				entry1 := &CandidateEntry{
+					candidateHash:      candidateHash1,
+					parentHeadDataHash: parentHeadHash,
+					outputHeadDataHash: outputHeadHash1,
+					state:              Seconded,
+				}
+
+				entry2 := &CandidateEntry{
+					candidateHash:      candidateHash2,
+					parentHeadDataHash: parentHeadHash,
+					outputHeadDataHash: outputHeadHash2,
+					state:              Backed,
+				}
+
+				storage.byCandidateHash[candidateHash1] = entry1
+				storage.byCandidateHash[candidateHash2] = entry2
+				storage.byParentHead[parentHeadHash] = map[parachaintypes.CandidateHash]any{
+					candidateHash1: struct{}{},
+					candidateHash2: struct{}{},
+				}
+
+				return storage
+			},
+			hash:     common.Hash{4, 5, 6},
+			expected: []*CandidateEntry{{candidateHash: parachaintypes.CandidateHash{Value: common.Hash{10, 11, 12}}, parentHeadDataHash: common.Hash{4, 5, 6}, outputHeadDataHash: common.Hash{13, 14, 15}, state: Backed}},
+		},
+		"insert_nothing_and_call_function_should_return_nothing": {
+			setup: func() *CandidateStorage {
+				return &CandidateStorage{
+					byParentHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+					byOutputHead:    make(map[common.Hash]map[parachaintypes.CandidateHash]any),
+					byCandidateHash: make(map[parachaintypes.CandidateHash]*CandidateEntry),
+				}
+			},
+			hash:     common.Hash{4, 5, 6},
+			expected: nil,
+		},
+	}
+
+	for name, tt := range tests {
+		tt := tt
+		t.Run(name, func(t *testing.T) {
+			storage := tt.setup()
+			var result []*CandidateEntry
+			for entry := range storage.PossibleBackedParaChildren(tt.hash) {
+				result = append(result, entry)
+			}
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}

go.mod

@@ -16,7 +16,6 @@ require (
 	github.com/fatih/color v1.17.0
 	github.com/gammazero/deque v0.2.1
 	github.com/go-playground/validator/v10 v10.21.0
-	github.com/golang/mock v1.6.0
 	github.com/google/go-cmp v0.6.0
 	github.com/google/uuid v1.6.0
 	github.com/gorilla/mux v1.8.1
@@ -200,7 +199,7 @@ require (
 	lukechampine.com/blake3 v1.2.1 // indirect
 )
 
-go 1.21
+go 1.23.1
 
 replace github.com/tetratelabs/wazero => github.com/ChainSafe/wazero v0.0.0-20240319130522-78b21a59bd5f