implement cow on leaves

Signed-off-by: Ignacio Hagopian <[email protected]>

Author: jsign

proof_test.go

@@ -34,6 +34,8 @@ import (
 )
 
 func TestProofVerifyTwoLeaves(t *testing.T) {
+	cfg := GetConfig()
+
 	root := New()
 	root.Insert(zeroKeyTest, zeroKeyTest, nil)
 	root.Insert(oneKeyTest, zeroKeyTest, nil)
@@ -42,7 +44,6 @@ func TestProofVerifyTwoLeaves(t *testing.T) {
 
 	proof, cis, zis, yis, _ := MakeVerkleMultiProof(root, [][]byte{ffx32KeyTest}, map[string][]byte{string(ffx32KeyTest): zeroKeyTest})
 
-	cfg := GetConfig()
 	if !VerifyVerkleProof(proof, cis, zis, yis, cfg) {
 		t.Fatalf("could not verify verkle proof: %s", ToDot(root))
 	}
@@ -176,6 +177,7 @@ func TestProofOfAbsenceLeafVerify(t *testing.T) {
 		t.Fatal("could not verify verkle proof")
 	}
 }
+
 func TestProofOfAbsenceLeafVerifyOtherSuffix(t *testing.T) {
 	root := New()
 	root.Insert(zeroKeyTest, zeroKeyTest, nil)
@@ -212,6 +214,7 @@ func TestProofOfAbsenceStemVerify(t *testing.T) {
 }
 
 func BenchmarkProofCalculation(b *testing.B) {
+	_ = GetConfig()
 	keys := make([][]byte, 100000)
 	root := New()
 	for i := 0; i < 100000; i++ {
@@ -220,6 +223,7 @@ func BenchmarkProofCalculation(b *testing.B) {
 		keys[i] = key
 		root.Insert(key, zeroKeyTest, nil)
 	}
+	root.Commit()
 
 	b.ResetTimer()
 	b.ReportAllocs()
@@ -370,7 +374,6 @@ func TestProofDeserialize(t *testing.T) {
 }
 
 func TestProofDeserializeErrors(t *testing.T) {
-
 	deserialized, err := DeserializeProof([]byte{0}, nil)
 	if err == nil {
 		t.Fatal("deserializing invalid proof didn't cause an error")

stateless_test.go

@@ -48,7 +48,7 @@ func TestStatelessChildren(t *testing.T) {
 		t.Fatal("invalid list length")
 	}
 
-	var emptycount = 0
+	emptycount := 0
 	for _, v := range list {
 		if _, ok := v.(Empty); ok {
 			emptycount++
@@ -377,6 +377,7 @@ func TestStatelessDeserializeDepth2(t *testing.T) {
 }
 
 func TestStatelessGetProofItems(t *testing.T) {
+	_ = GetConfig()
 	insertedKeys := [][]byte{zeroKeyTest, oneKeyTest, ffx32KeyTest}
 	provenKeys := [][]byte{zeroKeyTest, fourtyKeyTest}
 

tree.go

@@ -29,6 +29,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"sync"
 
 	"github.com/crate-crypto/go-ipa/banderwagon"
 )
@@ -187,6 +188,7 @@ type (
 
 		commitment *Point
 		c1, c2     *Point
+		cow        map[byte][]byte
 
 		depth byte
 	}
@@ -219,23 +221,6 @@ func NewLeafNode(stem []byte, values [][]byte) *LeafNode {
 		c2:     Generator(),
 	}
 
-	// Initialize the commitment with the extension tree
-	// marker and the stem.
-	cfg := GetConfig()
-	count := 0
-	var poly, c1poly, c2poly [256]Fr
-	poly[0].SetUint64(1)
-	StemFromBytes(&poly[1], leaf.stem)
-
-	count = fillSuffixTreePoly(c1poly[:], values[:128])
-	leaf.c1 = cfg.CommitToPoly(c1poly[:], 256-count)
-	toFr(&poly[2], leaf.c1)
-	count = fillSuffixTreePoly(c2poly[:], values[128:])
-	leaf.c2 = cfg.CommitToPoly(c2poly[:], 256-count)
-	toFr(&poly[3], leaf.c2)
-
-	leaf.commitment = cfg.CommitToPoly(poly[:], 252)
-
 	return leaf
 }
 
@@ -844,7 +829,7 @@ func (n *LeafNode) updateC(index byte, c *Point, oldc *Fr) {
 	n.commitment.Add(n.commitment, &diff)
 }
 
-func (n *LeafNode) updateCn(index byte, value []byte, c *Point) {
+func (n *LeafNode) updateCn(index byte, oldValue []byte, c *Point) {
 	var (
 		old, newH [2]Fr
 		diff      Point
@@ -857,8 +842,8 @@ func (n *LeafNode) updateCn(index byte, value []byte, c *Point) {
 	// do not include it. The result should be the same,
 	// but the computation time should be faster as one doesn't need to
 	// compute 1 - 1 mod N.
-	leafToComms(old[:], n.values[index])
-	leafToComms(newH[:], value)
+	leafToComms(old[:], oldValue)
+	leafToComms(newH[:], n.values[index])
 
 	newH[0].Sub(&newH[0], &old[0])
 	poly[2*(index%128)] = newH[0]
@@ -873,42 +858,36 @@ func (n *LeafNode) updateCn(index byte, value []byte, c *Point) {
 }
 
 func (n *LeafNode) updateLeaf(index byte, value []byte) {
-	c, oldc := n.getOldCn(index)
-
-	n.updateCn(index, value, c)
+	// If we haven't calculated a commitment for this node, we don't need to create the cow map since all the
+	// previous values are empty. If we already have a calculated commitment, then we track new values in
+	// cow so we can do diff-updating in the next Commit().
+	if n.commitment != nil {
+		// If cow was never setup, then initialize the map.
+		if n.cow == nil {
+			n.cow = make(map[byte][]byte)
+		}
 
-	n.updateC(index, c, oldc)
+		// If we are touching an value in an index for the first time,
+		// we save the original value for future use to update commitments.
+		if _, ok := n.cow[index]; !ok {
+			if n.values[index] == nil {
+				n.cow[index] = nil
+			} else {
+				n.cow[index] = make([]byte, 32)
+				copy(n.cow[index], n.values[index])
+			}
+		}
+	}
 
 	n.values[index] = value
 }
 
 func (n *LeafNode) updateMultipleLeaves(values [][]byte) {
-	var c1, c2 *Point
-	var old1, old2 *Fr
-	for i, v := range values {
-		if len(v) != 0 && !bytes.Equal(v, n.values[i]) {
-			if i < 128 {
-				if c1 == nil {
-					c1, old1 = n.getOldCn(byte(i))
-				}
-				n.updateCn(byte(i), v, c1)
-			} else {
-				if c2 == nil {
-					c2, old2 = n.getOldCn(byte(i))
-				}
-				n.updateCn(byte(i), v, c2)
-			}
-
-			n.values[i] = v
+	for i := range values {
+		if values[i] != nil {
+			n.updateLeaf(byte(i), values[i])
 		}
 	}
-
-	if c1 != nil {
-		n.updateC(0, c1, old1)
-	}
-	if c2 != nil {
-		n.updateC(128, c2, old2)
-	}
 }
 
 func (n *LeafNode) InsertOrdered(key []byte, value []byte, _ NodeFlushFn) error {
@@ -957,8 +936,77 @@ func (n *LeafNode) Commitment() *Point {
 	return n.commitment
 }
 
-func (n *LeafNode) Commit() *Point {
-	return n.commitment
+var frPool = sync.Pool{
+	New: func() any {
+		ret := make([]Fr, NodeWidth)
+		return &ret
+	},
+}
+
+func (leaf *LeafNode) Commit() *Point {
+	// If we've never calculated a commitment for this leaf node, we calculate the commitment
+	// in a single shot considering all the values.
+	if leaf.commitment == nil {
+		// Initialize the commitment with the extension tree
+		// marker and the stem.
+		count := 0
+		c1polyp := frPool.Get().(*[]Fr)
+		c1poly := *c1polyp
+		defer func() {
+			for i := 0; i < 256; i++ {
+				c1poly[i] = Fr{}
+			}
+			frPool.Put(c1polyp)
+		}()
+
+		count = fillSuffixTreePoly(c1poly, leaf.values[:128])
+		leaf.c1 = cfg.CommitToPoly(c1poly, 256-count)
+
+		for i := 0; i < 256; i++ {
+			c1poly[i] = Fr{}
+		}
+		count = fillSuffixTreePoly(c1poly, leaf.values[128:])
+		leaf.c2 = cfg.CommitToPoly(c1poly, 256-count)
+
+		for i := 0; i < 256; i++ {
+			c1poly[i] = Fr{}
+		}
+		c1poly[0].SetUint64(1)
+		StemFromBytes(&c1poly[1], leaf.stem)
+
+		toFrMultiple([]*Fr{&c1poly[2], &c1poly[3]}, []*Point{leaf.c1, leaf.c2})
+		leaf.commitment = cfg.CommitToPoly(c1poly, 252)
+
+	} else if len(leaf.cow) != 0 {
+		// If we've already have a calculated commitment, and there're touched leaf values, we do a diff update.
+		var c1, c2 *Point
+		var old1, old2 *Fr
+		for i, oldValue := range leaf.cow {
+			if !bytes.Equal(oldValue, leaf.values[i]) {
+				if i < 128 {
+					if c1 == nil {
+						c1, old1 = leaf.getOldCn(i)
+					}
+					leaf.updateCn(i, oldValue, c1)
+				} else {
+					if c2 == nil {
+						c2, old2 = leaf.getOldCn(i)
+					}
+					leaf.updateCn(i, oldValue, c2)
+				}
+			}
+		}
+
+		if c1 != nil {
+			leaf.updateC(0, c1, old1)
+		}
+		if c2 != nil {
+			leaf.updateC(128, c2, old2)
+		}
+		leaf.cow = nil
+	}
+
+	return leaf.commitment
 }
 
 // fillSuffixTreePoly takes one of the two suffix tree and

tree_test.go

@@ -33,11 +33,17 @@ import (
 	"errors"
 	"fmt"
 	mRand "math/rand"
+	"os"
 	"sort"
 	"testing"
 	"time"
 )
 
+func TestMain(m *testing.M) {
+	_ = GetConfig()
+	os.Exit(m.Run())
+}
+
 // a 32 byte value, as expected in the tree structure
 var testValue = []byte("0123456789abcdef0123456789abcdef")
 
@@ -106,7 +112,6 @@ func TestInsertTwoLeavesLastLevel(t *testing.T) {
 	if !bytes.Equal(leaf.values[0], testValue) {
 		t.Fatalf("did not find correct value in trie %x != %x", testValue, leaf.values[0])
 	}
-
 }
 
 func TestGetTwoLeaves(t *testing.T) {
@@ -413,6 +418,7 @@ func TestDeleteUnequalPath(t *testing.T) {
 		t.Fatalf("didn't catch the deletion of non-existing key, err =%v", err)
 	}
 }
+
 func TestDeleteResolve(t *testing.T) {
 	key1, _ := hex.DecodeString("0105000000000000000000000000000000000000000000000000000000000000")
 	key2, _ := hex.DecodeString("0107000000000000000000000000000000000000000000000000000000000000")
@@ -719,7 +725,7 @@ func isLeafEqual(a, b *LeafNode) bool {
 
 func TestGetResolveFromHash(t *testing.T) {
 	var count uint
-	var dummyError = errors.New("dummy")
+	dummyError := errors.New("dummy")
 	var serialized []byte
 	getter := func([]byte) ([]byte, error) {
 		count++
@@ -813,7 +819,7 @@ func TestInsertIntoHashedNode(t *testing.T) {
 		t.Fatalf("error detecting a decoding error after resolution: %v", err)
 	}
 
-	var randomResolverError = errors.New("'clef' was mispronounced")
+	randomResolverError := errors.New("'clef' was mispronounced")
 	// Check that the proper error is raised if the resolver returns an error
 	erroringResolver := func(h []byte) ([]byte, error) {
 		return nil, randomResolverError
@@ -879,9 +885,9 @@ func TestLeafToCommsLessThan16(*testing.T) {
 }
 
 func TestGetProofItemsNoPoaIfStemPresent(t *testing.T) {
-
 	root := New()
 	root.Insert(ffx32KeyTest, zeroKeyTest, nil)
+	root.Commit()
 
 	// insert two keys that differ from the inserted stem
 	// by one byte.