feat/btree: add BTree with scale #3536
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,187 @@ | ||
| // Copyright 2023 ChainSafe Systems (ON) | ||
| // SPDX-License-Identifier: LGPL-3.0-only | ||
|
|
||
| package btree | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "io" | ||
| "reflect" | ||
|
|
||
| "github.com/ChainSafe/gossamer/pkg/scale" | ||
|
|
||
| "golang.org/x/exp/constraints" | ||
|
|
||
| "github.com/tidwall/btree" | ||
| ) | ||
|
|
||
| type Codec interface { | ||
| MarshalSCALE() ([]byte, error) | ||
| UnmarshalSCALE(reader io.Reader) error | ||
| } | ||
|
|
||
| // Tree is a wrapper around tidwall/btree.BTree that also stores the comparator function and the type of the items | ||
| // stored in the BTree. This is needed during decoding because the Tree item is a generic type, and we need to know it | ||
| // at the time of decoding. | ||
| type Tree struct { | ||
| *btree.BTree | ||
| Comparator func(a, b interface{}) bool | ||
| ItemType reflect.Type | ||
| } | ||
|
|
||
| // MarshalSCALE encodes the Tree using SCALE. | ||
| func (bt Tree) MarshalSCALE() ([]byte, error) { | ||
| encodedLen, err := scale.Marshal(uint(bt.Len())) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to encode BTree length: %w", err) | ||
| } | ||
|
|
||
| var encodedItems []byte | ||
| bt.Ascend(nil, func(item interface{}) bool { | ||
| var encodedItem []byte | ||
| encodedItem, err = scale.Marshal(item) | ||
| if err != nil { | ||
| return false | ||
| } | ||
|
|
||
| encodedItems = append(encodedItems, encodedItem...) | ||
| return true | ||
| }) | ||
|
|
||
| return append(encodedLen, encodedItems...), err | ||
| } | ||
|
|
||
| // UnmarshalSCALE decodes the Tree using SCALE. | ||
| func (bt Tree) UnmarshalSCALE(reader io.Reader) error { | ||
| if bt.Comparator == nil { | ||
| return fmt.Errorf("comparator not found") | ||
| } | ||
|
|
||
| sliceType := reflect.SliceOf(bt.ItemType) | ||
| slicePtr := reflect.New(sliceType) | ||
| encodedItems, err := io.ReadAll(reader) | ||
| if err != nil { | ||
| return fmt.Errorf("read BTree items: %w", err) | ||
| } | ||
| err = scale.Unmarshal(encodedItems, slicePtr.Interface()) | ||
| if err != nil { | ||
| return fmt.Errorf("decode BTree items: %w", err) | ||
| } | ||
|
|
||
| for i := 0; i < slicePtr.Elem().Len(); i++ { | ||
| item := slicePtr.Elem().Index(i).Interface() | ||
| bt.Set(item) | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| // Copy returns a copy of the Tree. | ||
| func (bt Tree) Copy() *Tree { | ||
| return &Tree{ | ||
| BTree: bt.BTree.Copy(), | ||
| Comparator: bt.Comparator, | ||
| ItemType: bt.ItemType, | ||
| } | ||
| } | ||
|
|
||
| // NewTree creates a new Tree with the given comparator function. | ||
| func NewTree[T any](comparator func(a, b any) bool) Tree { | ||
| elementType := reflect.TypeOf((*T)(nil)).Elem() | ||
| return Tree{ | ||
| BTree: btree.New(comparator), | ||
| Comparator: comparator, | ||
| ItemType: elementType, | ||
| } | ||
| } | ||
|
|
||
| var _ Codec = (*Tree)(nil) | ||
|
|
||
| // Map is a wrapper around tidwall/btree.Map | ||
| type Map[K constraints.Ordered, V any] struct { | ||
| *btree.Map[K, V] | ||
| Degree int | ||
| } | ||
|
Comment on lines
+100
to
+103
There was a problem hiding this comment. Why do we need to hold There was a problem hiding this comment. It is so that the user can set the degree beforehand and not worry about filtering the data after unmarshal. There was a problem hiding this comment. Let me rephrase. It gives the option for the user to set the Degree. One usecase that I see it supporting is, if there are 100 entries in the db and I care about only top 10, I can set that during the creation. The usecase leads to lesser memory usage. There was a problem hiding this comment. But yeah, I guess I don't need to store it. I'm not familiar with your implementation style, feels like a lot of repetitive code. |
||
|
|
||
| type mapItem[K constraints.Ordered, V any] struct { | ||
| Key K | ||
| Value V | ||
| } | ||
|
|
||
| // MarshalSCALE encodes the Map using SCALE. | ||
| func (btm Map[K, V]) MarshalSCALE() ([]byte, error) { | ||
| encodedLen, err := scale.Marshal(uint(btm.Len())) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to encode Map length: %w", err) | ||
| } | ||
|
|
||
| var ( | ||
| pivot K | ||
| encodedItems []byte | ||
| ) | ||
| btm.Ascend(pivot, func(key K, value V) bool { | ||
| var ( | ||
| encodedKey []byte | ||
| encodedValue []byte | ||
| ) | ||
| encodedKey, err = scale.Marshal(key) | ||
| if err != nil { | ||
| return false | ||
| } | ||
|
|
||
| encodedValue, err = scale.Marshal(value) | ||
| if err != nil { | ||
| return false | ||
| } | ||
|
|
||
| encodedItems = append(encodedItems, encodedKey...) | ||
| encodedItems = append(encodedItems, encodedValue...) | ||
| return true | ||
| }) | ||
|
|
||
| return append(encodedLen, encodedItems...), err | ||
| } | ||
|
|
||
| // UnmarshalSCALE decodes the Map using SCALE. | ||
| func (btm Map[K, V]) UnmarshalSCALE(reader io.Reader) error { | ||
| if btm.Degree == 0 { | ||
| return fmt.Errorf("nothing to decode into") | ||
| } | ||
|
|
||
| if btm.Map == nil { | ||
| btm.Map = btree.NewMap[K, V](btm.Degree) | ||
| } | ||
|
|
||
| sliceType := reflect.SliceOf(reflect.TypeOf((*mapItem[K, V])(nil)).Elem()) | ||
| slicePtr := reflect.New(sliceType) | ||
| encodedItems, err := io.ReadAll(reader) | ||
| if err != nil { | ||
| return fmt.Errorf("read Map items: %w", err) | ||
| } | ||
| err = scale.Unmarshal(encodedItems, slicePtr.Interface()) | ||
| if err != nil { | ||
| return fmt.Errorf("decode Map items: %w", err) | ||
| } | ||
|
|
||
| for i := 0; i < slicePtr.Elem().Len(); i++ { | ||
| item := slicePtr.Elem().Index(i).Interface().(mapItem[K, V]) | ||
| btm.Map.Set(item.Key, item.Value) | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| // Copy returns a copy of the Map. | ||
| func (btm Map[K, V]) Copy() Map[K, V] { | ||
| return Map[K, V]{ | ||
| Map: btm.Map.Copy(), | ||
| } | ||
| } | ||
|
|
||
| // NewMap creates a new Map with the given degree. | ||
| func NewMap[K constraints.Ordered, V any](degree int) Map[K, V] { | ||
| return Map[K, V]{ | ||
| Map: btree.NewMap[K, V](degree), | ||
| Degree: degree, | ||
| } | ||
| } | ||
|
|
||
| var _ Codec = (*Map[int, string])(nil) | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,85 @@ | ||
| // Copyright 2023 ChainSafe Systems (ON) | ||
| // SPDX-License-Identifier: LGPL-3.0-only | ||
|
|
||
| package btree | ||
|
|
||
| import ( | ||
| "testing" | ||
|
|
||
| "github.com/ChainSafe/gossamer/pkg/scale" | ||
|
|
||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| type dummy struct { | ||
| Field1 uint32 | ||
| Field2 [32]byte | ||
| } | ||
|
|
||
| func TestBTree_Codec(t *testing.T) { | ||
| comparator := func(a, b interface{}) bool { | ||
| v1 := a.(dummy) | ||
| v2 := b.(dummy) | ||
| return v1.Field1 < v2.Field1 | ||
| } | ||
|
|
||
| // Create a Tree with 3 dummy items | ||
| tree := NewTree[dummy](comparator) | ||
| tree.Set(dummy{Field1: 1}) | ||
| tree.Set(dummy{Field1: 2}) | ||
| tree.Set(dummy{Field1: 3}) | ||
| encoded, err := scale.Marshal(tree) | ||
| require.NoError(t, err) | ||
|
|
||
| //let mut btree = Map::<u32, Hash>::new(); | ||
| //btree.insert(1, Hash::zero()); | ||
| //btree.insert(2, Hash::zero()); | ||
| //btree.insert(3, Hash::zero()); | ||
| //let encoded = btree.encode(); | ||
| //println!("encoded: {:?}", encoded); | ||
| expectedEncoded := []byte{12, | ||
| 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, | ||
| 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, | ||
| 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, | ||
| } | ||
| require.Equal(t, expectedEncoded, encoded) | ||
|
|
||
| expected := NewTree[dummy](comparator) | ||
| err = scale.Unmarshal(expectedEncoded, &expected) | ||
| require.NoError(t, err) | ||
|
|
||
| // Check that the expected Tree has the same items as the original | ||
| require.Equal(t, tree.Len(), expected.Len()) | ||
| require.Equal(t, tree.ItemType, expected.ItemType) | ||
| require.Equal(t, tree.Min(), expected.Min()) | ||
| require.Equal(t, tree.Max(), expected.Max()) | ||
| require.Equal(t, tree.Get(dummy{Field1: 1}), expected.Get(dummy{Field1: 1})) | ||
| require.Equal(t, tree.Get(dummy{Field1: 2}), expected.Get(dummy{Field1: 2})) | ||
| require.Equal(t, tree.Get(dummy{Field1: 3}), expected.Get(dummy{Field1: 3})) | ||
| } | ||
|
|
||
| func TestBTreeMap_Codec(t *testing.T) { | ||
| btreeMap := NewMap[uint32, dummy](32) | ||
| btreeMap.Set(uint32(1), dummy{Field1: 1}) | ||
| btreeMap.Set(uint32(2), dummy{Field1: 2}) | ||
| btreeMap.Set(uint32(3), dummy{Field1: 3}) | ||
| encoded, err := scale.Marshal(btreeMap) | ||
| require.NoError(t, err) | ||
|
|
||
| //let mut btree = Map::<u32, (u32, Hash)>::new(); | ||
| //btree.insert(1, (1, Hash::zero())); | ||
| //btree.insert(2, (2, Hash::zero())); | ||
| //btree.insert(3, (3, Hash::zero())); | ||
| //let encoded = btree.encode(); | ||
| //println!("encoded: {:?}", encoded); | ||
| expectedEncoded := []byte{12, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, | ||
| 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, | ||
| 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, | ||
| 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, | ||
| } | ||
| require.Equal(t, expectedEncoded, encoded) | ||
| expected := NewMap[uint32, dummy](32) | ||
| err = scale.Unmarshal(expectedEncoded, &expected) | ||
| require.NoError(t, err) | ||
| require.Equal(t, btreeMap.Len(), expected.Len()) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not sure why you need this attribute given it's not used in any of the
MarshalSCALEorUnmarshalSCALEmethods.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please find the follow up questions below to get a better understanding.
Do you mean the comparator? Is the comment not descriptive enough?