Old changes

proofs.go

@@ -327,7 +327,7 @@ func verifyProofWithUpdates(
 	root, key, value []byte,
 	spec *TrieSpec,
 ) (bool, [][][]byte, error) {
-	// Retrieve the trie path for the give key
+	// Retrieve the trie path for the key being proven
 	path := spec.ph.Path(key)
 
 	if err := proof.validateBasic(spec); err != nil {
@@ -338,31 +338,31 @@ func verifyProofWithUpdates(
 
 	// Determine what the leaf hash should be.
 	var currentHash, currentData []byte
-	if bytes.Equal(value, defaultEmptyValue) { // Non-membership proof.
-		if proof.NonMembershipLeafData == nil { // Leaf is a placeholder value.
+	if bytes.Equal(value, defaultEmptyValue) {
+		// Non-membership proof if `value` is empty.
+		if proof.NonMembershipLeafData == nil {
+			// Leaf is a placeholder value.
 			currentHash = spec.placeholder()
-		} else { // Leaf is an unrelated leaf.
+		} else {
+			// Leaf is an unrelated leaf.
 			var actualPath, valueHash []byte
 			actualPath, valueHash = spec.parseLeafNode(proof.NonMembershipLeafData)
 			if bytes.Equal(actualPath, path) {
 				// This is not an unrelated leaf; non-membership proof failed.
 				return false, nil, errors.Join(ErrBadProof, errors.New("non-membership proof on related leaf"))
 			}
 			currentHash, currentData = spec.digestLeaf(actualPath, valueHash)
-
-			update := make([][]byte, 2)
-			update[0], update[1] = currentHash, currentData
-			updates = append(updates, update)
 		}
 	} else {
-		// Membership proof.
+		// Membership proof if `value` is non-empty.
 		valueHash := spec.valueHash(value)
 		currentHash, currentData = spec.digestLeaf(path, valueHash)
-		update := make([][]byte, 2)
-		update[0], update[1] = currentHash, currentData
-		updates = append(updates, update)
 	}
 
+	update := make([][]byte, 2)
+	update[0], update[1] = currentHash, currentData
+	updates = append(updates, update)
+
 	// Recompute root.
 	for i := 0; i < len(proof.SideNodes); i++ {
 		node := make([]byte, spec.hashSize())

smst_example_test.go

@@ -32,14 +32,17 @@ func TestExampleSMST(t *testing.T) {
 	err = smst.Update([]byte("bin"), []byte("nib"), 3)
 	require.NoError(t, err)
 	dataMap["bin"] = "nib"
+	err = smst.Update([]byte("binn"), []byte("nnib"), 3)
+	require.NoError(t, err)
+	dataMap["binn"] = "nnib"
 
 	// Commit the changes to the nodeStore
 	err = smst.Commit()
 	require.NoError(t, err)
 
 	// Calculate the total sum of the trie
 	sum := smst.Sum()
-	require.Equal(t, uint64(20), sum)
+	require.Equal(t, uint64(23), sum)
 
 	// Generate a Merkle proof for "foo"
 	proof1, err := smst.Prove([]byte("foo"))
@@ -114,6 +117,7 @@ func exportToCSV(
 		}
 	*/
 
+	fmt.Println("Root sum", smst.Sum())
 	helper(t, smst, nodeStore, smst.Root())
 }
 
@@ -131,12 +135,19 @@ func helper(t *testing.T, smst SparseMerkleSumTrie, nodeStore kvstore.MapStore,
 		helper(t, smst, nodeStore, childData)
 		return
 	} else if isLeafNode(node) {
-		path, value := smst.Spec().parseLeafNode(node)
-		fmt.Println("leaf node sum", 0)
-		fmt.Println(path, value)
+		path, valueHash, weight := smst.Spec().parseSumLeafNode(node)
+		fmt.Println("leaf node weight", weight)
+		fmt.Println(path, valueHash)
 	} else if isInnerNode(node) {
 		leftData, rightData, sum := smst.Spec().th.parseSumInnerNode(node)
 		fmt.Println("inner node sum", sum)
+
+		// leftChild, err := nodeStore.Get(leftData)
+		// require.NoError(t, err)
+
+		// rightChild, err := nodeStore.Get(rightData)
+		// require.NoError(t, err)
+
 		helper(t, smst, nodeStore, leftData)
 		helper(t, smst, nodeStore, rightData)
 	}

trie_spec.go

@@ -237,6 +237,22 @@ func (spec *TrieSpec) parseExtNode(data []byte) (pathBounds, path, childData []b
 	return
 }
 
+// parseSumLeafNode parses a leafNode and returns its weight as well
+func (spec *TrieSpec) parseSumLeafNode(data []byte) (path, value []byte, weight uint64) {
+	// panics if not a leaf node
+	checkPrefix(data, leafNodePrefix)
+
+	path = data[prefixLen : prefixLen+spec.ph.PathSize()]
+	value = data[prefixLen+spec.ph.PathSize():]
+
+	// Extract the sum from the encoded node data
+	var weightBz [sumSizeBits]byte
+	copy(weightBz[:], value[len(value)-sumSizeBits:])
+	binary.BigEndian.PutUint64(weightBz[:], weight)
+
+	return
+}
+
 // parseSumExtNode parses the pathBounds, path, child data and sum from the encoded extension node data
 func (spec *TrieSpec) parseSumExtNode(data []byte) (pathBounds, path, childData []byte, sum uint64) {
 	// panics if not an extension node