Snap/merge range and proof

core/trie/key.go

@@ -153,3 +153,60 @@ func (k *Key) RemoveLastBit() {
 		inUseBytes[0] = (inUseBytes[0] << unusedBitsCount) >> unusedBitsCount
 	}
 }
+
+// CmpAligned is Cmp as if the value is bigendian bytes of key of the same length
+func (k Key) CmpAligned(other *Key) int {
+	// No its not aligned, so need to convert to bigint then left shift it so that the MSB is of the same index
+	height := k.len
+	if other.len > height {
+		height = other.len
+	}
+
+	b1i := k.alignedBitInt(height)
+	b2i := other.alignedBitInt(height)
+	return b1i.Cmp(b2i)
+}
+
+func (k Key) alignedBitInt(height uint8) *big.Int {
+	bigint := &big.Int{}
+	bigint = bigint.SetBytes(k.bitset[:])
+	if k.len < height {
+		bigint = bigint.Lsh(bigint, uint(height-k.len))
+	}
+
+	return bigint
+}
+
+func (k *Key) AppendBitMut(flag bool) {
+	bit := k.len
+	byteIdx := bit / 8
+	byteAtIdx := k.bitset[len(k.bitset)-int(byteIdx)-1]
+	bitIdx := bit % 8
+
+	mask := uint8(1 << bitIdx)
+	if flag {
+		byteAtIdx |= mask // set bit
+	} else {
+		byteAtIdx &= ^mask // clear bit
+	}
+
+	k.len++
+	k.bitset[len(k.bitset)-int(byteIdx)-1] = byteAtIdx
+}
+
+func (k Key) Append(otherKey *Key) Key {
+	result := NewKey(otherKey.len, otherKey.bitset[:])
+
+	// I'm sure someone will make this faster
+	for i := uint8(0); i < k.len; i++ {
+		result.AppendBitMut(k.Test(i))
+	}
+
+	return result
+}
+
+func (k Key) AppendBit(flag bool) Key {
+	result := NewKey(0, []byte{})
+	result.AppendBitMut(flag)
+	return k.Append(&result)
+}

core/trie/key_test.go

@@ -2,6 +2,7 @@ package trie_test
 
 import (
 	"bytes"
+	"fmt"
 	"testing"
 
 	"github.com/NethermindEth/juno/core/felt"
@@ -153,3 +154,120 @@ func TestTruncate(t *testing.T) {
 		})
 	}
 }
+
+func Test_cmp(t *testing.T) {
+	tests := []struct {
+		n1       int
+		n2       int
+		isHigher bool
+	}{
+		{
+			n1:       10,
+			n2:       0,
+			isHigher: true,
+		},
+		{
+			n1:       5,
+			n2:       0,
+			isHigher: true,
+		},
+		{
+			n1:       5,
+			n2:       4,
+			isHigher: true,
+		},
+		{
+			n1:       5,
+			n2:       5,
+			isHigher: false,
+		},
+		{
+			n1:       4,
+			n2:       5,
+			isHigher: false,
+		},
+		{
+			n1:       0,
+			n2:       5,
+			isHigher: false,
+		},
+		{
+			n1:       300,
+			n2:       1,
+			isHigher: true,
+		},
+		{
+			n1:       1,
+			n2:       300,
+			isHigher: false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("%d %d %v", test.n1, test.n2, test.isHigher), func(t *testing.T) {
+			k1 := numToKey(test.n1)
+			k2 := numToKey(test.n2)
+
+			assert.Equal(t,
+				k1.CmpAligned(&k2) > 0,
+				test.isHigher)
+		})
+	}
+}
+
+func numToKey(num int) trie.Key {
+	return trie.NewKey(8, []byte{byte(num)})
+}
+
+func TestKeyAppend(t *testing.T) {
+	tests := map[string]struct {
+		Key1        trie.Key
+		Key2        trie.Key
+		ExpectedKey trie.Key
+	}{
+		"no append": {
+			Key1:        trie.NewKey(1, []byte{0x01}),
+			Key2:        trie.NewKey(0, []byte{0x00}),
+			ExpectedKey: trie.NewKey(1, []byte{0x01}),
+		},
+		"from zero append": {
+			Key1:        trie.NewKey(0, []byte{0x00}),
+			Key2:        trie.NewKey(1, []byte{0x01}),
+			ExpectedKey: trie.NewKey(1, []byte{0x01}),
+		},
+		"append shift": {
+			Key1:        trie.NewKey(1, []byte{0x01}),
+			Key2:        trie.NewKey(7, []byte{0x00}),
+			ExpectedKey: trie.NewKey(8, []byte{0x80}),
+		},
+		"append to a new byte": {
+			Key1:        trie.NewKey(8, []byte{0xff}),
+			Key2:        trie.NewKey(1, []byte{0x01}),
+			ExpectedKey: trie.NewKey(9, []byte{0x01, 0xff}),
+		},
+		"append multi byte": {
+			Key1:        trie.NewKey(11, []byte{0x00, 0xff}),       //  000 1111 1111
+			Key2:        trie.NewKey(12, []byte{0x00, 0xff}),       // 0000 1111 1111
+			ExpectedKey: trie.NewKey(23, []byte{0x0f, 0xf0, 0xff}), //  000 1111 1111 0000 1111 1111
+		},
+	}
+
+	for desc, test := range tests {
+		t.Run(desc, func(t *testing.T) {
+			appended := test.Key1.Append(&test.Key2)
+			assert.Equal(t, test.ExpectedKey, appended)
+		})
+	}
+}
+
+func TestKeyAppendBit(t *testing.T) {
+	k1 := trie.NewKey(1, []byte{0x01})
+	k2 := k1.AppendBit(true)
+	expected := trie.NewKey(2, []byte{0x03})
+	assert.Equal(t, k2, expected)
+
+	k1 = trie.NewKey(1, []byte{0x00})
+	k2 = k1.AppendBit(true)
+	expected = trie.NewKey(2, []byte{0x01})
+	assert.Equal(t, k2, expected)
+}

core/trie/node.go

@@ -9,9 +9,11 @@
 
 // A Node represents a node in the [Trie]
 type Node struct {
-	Value *felt.Felt
-	Left  *Key
-	Right *Key
+	Value     *felt.Felt
+	Left      *Key
+	Right     *Key
+	LeftHash  *felt.Felt
+	RightHash *felt.Felt
 }
 
 // Hash calculates the hash of a [Node]
@@ -30,6 +32,12 @@
 	return hash.Add(hash, &pathFelt)
 }
 
+// Hash calculates the hash of a [Node]
+func (n *Node) HashFromParent(parnetKey, nodeKey *Key, hashFunc hashFunc) *felt.Felt {
+	path := path(nodeKey, parnetKey)
+	return n.Hash(&path, hashFunc)
+}
+
 func (n *Node) WriteTo(buf *bytes.Buffer) (int64, error) {
 	if n.Value == nil {
 		return 0, errors.New("cannot marshal node with nil value")
@@ -45,18 +53,38 @@
 	}
 
 	if n.Left != nil {
-		wrote, err := n.Left.WriteTo(buf)
+		wrote, errInner := n.Left.WriteTo(buf)
 		totalBytes += wrote
 		if err != nil {
-			return totalBytes, err
+			return totalBytes, errInner
 		}
-		wrote, err = n.Right.WriteTo(buf) // n.Right is non-nil by design
+		wrote, errInner = n.Right.WriteTo(buf) // n.Right is non-nil by design
 		totalBytes += wrote
 		if err != nil {
-			return totalBytes, err
+			return totalBytes, errInner
 		}
 	}
 
+	if n.LeftHash == nil && n.RightHash == nil {
+		return totalBytes, nil
+	} else if (n.LeftHash != nil && n.RightHash == nil) || (n.LeftHash == nil && n.RightHash != nil) {
+		return totalBytes, errors.New("cannot store only one lefthash or righthash")
+	}
+
+	leftHashB := n.LeftHash.Bytes()
+	wrote, err = buf.Write(leftHashB[:])
+	totalBytes += int64(wrote)
+	if err != nil {
+		return totalBytes, err
+	}
+
+	rightHashB := n.RightHash.Bytes()
+	wrote, err = buf.Write(rightHashB[:])
+	totalBytes += int64(wrote)
+	if err != nil {
+		return totalBytes, err
+	}
+
 	return totalBytes, nil
 }
 
@@ -74,6 +102,8 @@
 	if len(data) == 0 {
 		n.Left = nil
 		n.Right = nil
+		n.LeftHash = nil
+		n.RightHash = nil
 		return nil
 	}
 
@@ -85,5 +115,26 @@
 	if err := n.Left.UnmarshalBinary(data); err != nil {
 		return err
 	}
-	return n.Right.UnmarshalBinary(data[n.Left.EncodedLen():])
+	data = data[n.Left.EncodedLen():]
+	if err := n.Right.UnmarshalBinary(data); err != nil {
+		return err
+	}
+	data = data[n.Right.EncodedLen():]
+
+	if n.LeftHash == nil {
+		n.LeftHash = new(felt.Felt)
+	}
+	if n.RightHash == nil {
+		n.RightHash = new(felt.Felt)
+	}
+	if len(data) == 0 {
+		return nil
+	}
+	if len(data) != 2*felt.Bytes {
+		return errors.New("the node does not contain both left and right hash")
+	}
+	n.LeftHash.SetBytes(data[:felt.Bytes])
+	data = data[felt.Bytes:]
+	n.RightHash.SetBytes(data[:felt.Bytes])
+	return nil
 }