From ebc4a94785202ed16545960fd23ab5967238faaf Mon Sep 17 00:00:00 2001
From: Gui Iribarren <gui@altermundi.net>
Date: Tue, 29 Oct 2024 18:43:22 +0100
Subject: [PATCH] arbo: add CalculateProofNodes

---
 tree/arbo/circomproofs.go | 18 ++++++++++++++++++
 tree/arbo/proof.go        | 22 ++++++++++++++++++----
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/tree/arbo/circomproofs.go b/tree/arbo/circomproofs.go
index 94a4aacbb..9166df4f5 100644
--- a/tree/arbo/circomproofs.go
+++ b/tree/arbo/circomproofs.go
@@ -1,7 +1,9 @@
 package arbo
 
 import (
+	"bytes"
 	"encoding/json"
+	"slices"
 )
 
 // CircomVerifierProof contains the needed data to check a Circom Verifier Proof
@@ -89,3 +91,19 @@ func (t *Tree) GenerateCircomVerifierProof(k []byte) (*CircomVerifierProof, erro
 
 	return &cp, nil
 }
+
+// CalculateProofNodes calculates the chain of hashes in the path of the proof.
+// In the returned list, first item is the root, and last item is the hash of the leaf.
+func (cvp CircomVerifierProof) CalculateProofNodes(hashFunc HashFunction) ([][]byte, error) {
+	paddedSiblings := slices.Clone(cvp.Siblings)
+	for k, v := range paddedSiblings {
+		if bytes.Equal(v, []byte{0}) {
+			paddedSiblings[k] = make([]byte, hashFunc.Len())
+		}
+	}
+	packedSiblings, err := PackSiblings(hashFunc, paddedSiblings)
+	if err != nil {
+		return nil, err
+	}
+	return CalculateProofNodes(hashFunc, cvp.Key, cvp.Value, packedSiblings)
+}
diff --git a/tree/arbo/proof.go b/tree/arbo/proof.go
index ad19f671d..7aede4cf0 100644
--- a/tree/arbo/proof.go
+++ b/tree/arbo/proof.go
@@ -160,19 +160,31 @@ func bytesToBitmap(b []byte) []bool {
 // CheckProof verifies the given proof. The proof verification depends on the
 // HashFunction passed as parameter.
 func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
-	siblings, err := UnpackSiblings(hashFunc, packedSiblings)
+	hashes, err := CalculateProofNodes(hashFunc, k, v, packedSiblings)
 	if err != nil {
 		return false, err
 	}
+	return bytes.Equal(hashes[0], root), nil
+}
+
+// CalculateProofNodes calculates the chain of hashes in the path of the given proof.
+// In the returned list, first item is the root, and last item is the hash of the leaf.
+func CalculateProofNodes(hashFunc HashFunction, k, v, packedSiblings []byte) ([][]byte, error) {
+	siblings, err := UnpackSiblings(hashFunc, packedSiblings)
+	if err != nil {
+		return nil, err
+	}
 
 	keyPath := make([]byte, int(math.Ceil(float64(len(siblings))/float64(8))))
 	copy(keyPath, k)
 
 	key, _, err := newLeafValue(hashFunc, k, v)
 	if err != nil {
-		return false, err
+		return nil, err
 	}
 
+	hashes := [][]byte{key}
+
 	path := getPath(len(siblings), keyPath)
 	for i, sibling := range slices.Backward(siblings) {
 		if path[i] {
@@ -181,8 +193,10 @@ func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool,
 			key, _, err = newIntermediate(hashFunc, key, sibling)
 		}
 		if err != nil {
-			return false, err
+			return nil, err
 		}
+		hashes = append(hashes, key)
 	}
-	return bytes.Equal(key, root), nil
+	slices.Reverse(hashes)
+	return hashes, nil
 }