From aff9bcc18604d447c52d005ff066cad2b19c2cbc Mon Sep 17 00:00:00 2001
From: arnaucube <root@arnaucube.com>
Date: Tue, 18 May 2021 11:00:17 +0200
Subject: [PATCH] arbotree: Add tree.Size, update arbo version

---
 treeinterface/arbotree/go.mod          |  2 +-
 treeinterface/arbotree/go.sum          |  2 ++
 treeinterface/arbotree/wrapper.go      | 12 ++++++++--
 treeinterface/arbotree/wrapper_test.go | 33 ++++++++++++++++++++++++++
 4 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/treeinterface/arbotree/go.mod b/treeinterface/arbotree/go.mod
index 9ffbdef..e80732e 100644
--- a/treeinterface/arbotree/go.mod
+++ b/treeinterface/arbotree/go.mod
@@ -5,7 +5,7 @@ go 1.16
 // replace github.com/arnaucube/arbo => ../../../arbo
 
 require (
-	github.com/arnaucube/arbo v0.0.0-20210508150807-6dcbbdf4a51a // indirect
+	github.com/arnaucube/arbo v0.0.0-20210518085706-03bb9f744792 // indirect
 	github.com/iden3/go-merkletree v0.0.0-20210308143313-8b63ca866189
 	go.vocdoni.io/dvote v0.6.1-0.20210508104027-72380d23da7f
 )
diff --git a/treeinterface/arbotree/go.sum b/treeinterface/arbotree/go.sum
index e6bc769..b95806b 100644
--- a/treeinterface/arbotree/go.sum
+++ b/treeinterface/arbotree/go.sum
@@ -105,6 +105,8 @@ github.com/arnaucube/arbo v0.0.0-20210508132613-cb373f881c7c h1:H1EVaXwUfqOlGvzn
 github.com/arnaucube/arbo v0.0.0-20210508132613-cb373f881c7c/go.mod h1:ZlejgAnqIb9ysj5z9Husdqant+3arJbac8CBStess/s=
 github.com/arnaucube/arbo v0.0.0-20210508150807-6dcbbdf4a51a h1:gBPYofXmm/ehWn071mLbjZKv6N6Zhl4huWHfZE1gZhA=
 github.com/arnaucube/arbo v0.0.0-20210508150807-6dcbbdf4a51a/go.mod h1:ZlejgAnqIb9ysj5z9Husdqant+3arJbac8CBStess/s=
+github.com/arnaucube/arbo v0.0.0-20210518085706-03bb9f744792 h1:0lwF3FXiztk0Dh6wpdjYGczOtG5HAB8reIgeuHBuuqo=
+github.com/arnaucube/arbo v0.0.0-20210518085706-03bb9f744792/go.mod h1:ZlejgAnqIb9ysj5z9Husdqant+3arJbac8CBStess/s=
 github.com/arnaucube/go-blindsecp256k1 v0.0.0-20210203222605-876755a714c3/go.mod h1:10oM1DQMpukFfSzNsjSQG2rE6UjoGsXT4iIxWIiVbu0=
 github.com/aryann/difflib v0.0.0-20170710044230-e206f873d14a/go.mod h1:DAHtR1m6lCRdSC2Tm3DSWRPvIPr6xNKyeHdqDQSQT+A=
 github.com/aws/aws-lambda-go v1.13.3/go.mod h1:4UKl9IzQMoD+QF79YdCuzCwp8VbmG4VAQwij/eHl5CU=
diff --git a/treeinterface/arbotree/wrapper.go b/treeinterface/arbotree/wrapper.go
index eb8cf30..b7fda7a 100644
--- a/treeinterface/arbotree/wrapper.go
+++ b/treeinterface/arbotree/wrapper.go
@@ -28,6 +28,7 @@ func (t *Tree) Init(name, storageDir string) error {
 	if err != nil {
 		return err
 	}
+	// storage := memory.NewMemoryStorage()
 
 	mt, err := arbo.NewTree(storage, 140, arbo.HashFunctionBlake2b) // TODO here the hash function would depend on the usage
 	if err != nil {
@@ -125,14 +126,21 @@ func (t *Tree) ImportDump(data []byte) error {
 }
 
 func (t *Tree) Size(root []byte) (int64, error) {
-	// TODO
-	return 0, nil
+	count := 0
+	err := t.Tree.Iterate(func(k, v []byte) {
+		if v[0] != arbo.PrefixValueLeaf {
+			return
+		}
+		count++
+	})
+	return int64(count), err
 }
 
 func (t *Tree) Snapshot(root []byte) (censustree.Tree, error) {
 	// TODO
 	return t, nil
 }
+
 func (t *Tree) HashExists(hash []byte) (bool, error) {
 	_, _, err := t.Tree.Get(hash)
 	if err != nil {
diff --git a/treeinterface/arbotree/wrapper_test.go b/treeinterface/arbotree/wrapper_test.go
index 8f36b81..9820640 100644
--- a/treeinterface/arbotree/wrapper_test.go
+++ b/treeinterface/arbotree/wrapper_test.go
@@ -20,3 +20,36 @@ func TestInterface(t *testing.T) {
 		t.Fatal("censustree interface not matched by arbotree wrapper")
 	}
 }
+
+func TestGenProof(t *testing.T) {
+	storage := t.TempDir()
+	tr1 := &Tree{}
+	err := tr1.Init("test1", storage)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var keys, values [][]byte
+	for i := 0; i < 10; i++ {
+		keys = append(keys, []byte{byte(i)})
+		values = append(values, []byte{byte(i)})
+		err = tr1.Add([]byte{byte(i)}, []byte{byte(i)})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	for i := 0; i < 10; i++ {
+		p, err := tr1.GenProof(keys[i], values[i])
+		if err != nil {
+			t.Fatal(err)
+		}
+		v, err := tr1.CheckProof(keys[i], values[i], tr1.Root(), p)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !v {
+			t.Fatal("CheckProof failed")
+		}
+	}
+}