From 39f0736eaf9f2b531fa870d027c84f79f4d4ae84 Mon Sep 17 00:00:00 2001
From: arnaucube <root@arnaucube.com>
Date: Wed, 1 Jul 2020 14:57:44 +0200
Subject: [PATCH] Add proof Siblings methods for circom compatibility

---
 merkletree.go      | 37 +++++++++++++++++++++++++++++++++++++
 merkletree_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/merkletree.go b/merkletree.go
index 26ed0d4..659bec4 100644
--- a/merkletree.go
+++ b/merkletree.go
@@ -2,6 +2,7 @@ package merkletree
 
 import (
 	"bytes"
+	"encoding/hex"
 	"errors"
 	"fmt"
 	"io"
@@ -57,6 +58,9 @@ type Hash [32]byte
 func (h Hash) String() string {
 	return new(big.Int).SetBytes(h[:]).String()
 }
+func (h Hash) Hex() string {
+	return hex.EncodeToString(h[:])
+}
 
 // BigInt returns the *big.Int representation of the *Hash
 func (h *Hash) BigInt() *big.Int {
@@ -392,6 +396,39 @@ func (p *Proof) Bytes() []byte {
 	return bs
 }
 
+// SiblingsFromProof returns all the siblings of the proof. This function is used to generate the siblings input for the circom circuits.
+func SiblingsFromProof(proof *Proof) []*Hash {
+	sibIdx := 0
+	var siblings []*Hash
+	for lvl := 0; lvl < int(proof.depth); lvl++ {
+		if common.TestBitBigEndian(proof.notempties[:], uint(lvl)) {
+			siblings = append(siblings, proof.Siblings[sibIdx])
+			sibIdx++
+		} else {
+			siblings = append(siblings, &HashZero)
+		}
+	}
+	return siblings
+}
+
+func (p *Proof) AllSiblings() []*Hash {
+	return SiblingsFromProof(p)
+}
+
+func (p *Proof) AllSiblingsCircom(levels int) []*big.Int {
+	siblings := p.AllSiblings()
+	// Add the rest of empty levels to the siblings
+	for i := len(siblings); i < levels; i++ {
+		siblings = append(siblings, &HashZero)
+	}
+	siblings = append(siblings, &HashZero) // add extra level for circom compatibility
+	siblingsBigInt := make([]*big.Int, len(siblings))
+	for i, sibling := range siblings {
+		siblingsBigInt[i] = sibling.BigInt()
+	}
+	return siblingsBigInt
+}
+
 // GenerateProof generates the proof of existence (or non-existence) of an
 // Entry's hash Index for a Merkle Tree given the root.
 // If the rootKey is nil, the current merkletree root is used
diff --git a/merkletree_test.go b/merkletree_test.go
index 2a23c08..27224c0 100644
--- a/merkletree_test.go
+++ b/merkletree_test.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/iden3/go-iden3-core/db"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestNewTree(t *testing.T) {
@@ -32,3 +33,31 @@ func TestNewTree(t *testing.T) {
 	assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(33), big.NewInt(44)))
 	assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(33), big.NewInt(45)))
 }
+
+func TestSiblingsFromProof(t *testing.T) {
+	mt, err := NewMerkleTree(db.NewMemoryStorage(), 140)
+	require.Nil(t, err)
+	defer mt.db.Close()
+
+	for i := 0; i < 64; i++ {
+		k := big.NewInt(int64(i))
+		v := big.NewInt(0)
+		if err := mt.Add(k, v); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	proof, err := mt.GenerateProof(big.NewInt(4), nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	siblings := SiblingsFromProof(proof)
+	assert.Equal(t, 6, len(siblings))
+	assert.Equal(t, "26bc69cfd3c982eba7b45cd2e6a2c75f218c546089f115777df47ab06f1fdb23", siblings[0].Hex())
+	assert.Equal(t, "ba7fbb5841bc8fa65193c124fbda4843aaa8d64d4132c79e7176cbed4de65621", siblings[1].Hex())
+	assert.Equal(t, "1339329b2d15e467ec3734ec06f6f4ae5a7c8c0b6ba98c26558b5a4db3e9a804", siblings[2].Hex())
+	assert.Equal(t, "b2261e6f47d9feb6fac2480928cc0fe03204d946cbc0a7b4de250d3e1384f40f", siblings[3].Hex())
+	assert.Equal(t, "b6c905f21c9928efa19a2c4e55d88d5fd0af493032f5b84620eb872e48ff5d01", siblings[4].Hex())
+	assert.Equal(t, "69873a951f49bbaff71039d672ffe52d73605aed6b40c1ce7ab068ad86a44d1e", siblings[5].Hex())
+}