M treeinterface/arbotree/go.mod → treeinterface/arbotree/go.mod

@@ -5,7 +5,7 @@
 // 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
 )

M treeinterface/arbotree/go.sum → treeinterface/arbotree/go.sum

@@ -105,6 +105,8 @@ github.com/arnaucube/arbo v0.0.0-20210508132613-cb373f881c7c h1:H1EVaXwUfqOlGvznEjDxridyUqfSJ87Qq16Yc1oDSqw=
 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=

M treeinterface/arbotree/wrapper.go → treeinterface/arbotree/wrapper.go

@@ -28,6 +28,7 @@ storage, err := pebble.NewPebbleStorage(dbDir, false) // TODO TMP
 	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 @@ return t.Tree.ImportDump(data)
 }
 
 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 {

M treeinterface/arbotree/wrapper_test.go → treeinterface/arbotree/wrapper_test.go

@@ -20,3 +20,36 @@ if !ok {
 		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")
+		}
+	}
+}