From 19326669db623e0bd8be90e33da6b8bd77d9ade5 Mon Sep 17 00:00:00 2001
From: p4u <p4u@dabax.net>
Date: Thu, 13 Dec 2018 18:43:45 +0100
Subject: [PATCH] First implementation of tree package

Signed-off-by: p4u <p4u@dabax.net>
---
 tree/README.md | 22 ++++++++++++++++++
 tree/tree.go   | 62 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+)
 create mode 100644 tree/README.md
 create mode 100644 tree/tree.go

diff --git a/tree/README.md b/tree/README.md
new file mode 100644
index 0000000..3e9af9d
--- /dev/null
+++ b/tree/README.md
@@ -0,0 +1,22 @@
+#### Dvote Tree
+
+Implementation of dvote tree structure. Currently based on iden3 merkle tree.
+
+Example of usage:
+
+```
+  T := tree.Tree {namespace: "vocdoni"}
+  if T.init() != nil { fmt.Println("Cannot create tree database") }
+  err := T.addClaim([]byte("Hello you!"))
+  if err != nil {
+    fmt.Println("Claim already exist")
+  }
+  mpHex, err := T.genProof([]byte("Hello you!"))
+  fmt.Println(mpHex)
+  fmt.Println(T.checkProof([]byte("Hello you!"), mpHex))
+  T.close()
+```
+
+### To-Do
+
++ Add export/import methods
diff --git a/tree/tree.go b/tree/tree.go
new file mode 100644
index 0000000..d1578d2
--- /dev/null
+++ b/tree/tree.go
@@ -0,0 +1,62 @@
+package tree
+
+import (
+  "github.com/iden3/go-iden3/db"
+  "github.com/iden3/go-iden3/merkletree"
+  mkcore "github.com/iden3/go-iden3/core"
+  common3 "github.com/iden3/go-iden3/common"
+  "os/user"
+)
+
+type Tree struct {
+    namespace string
+    storage string
+    tree *merkletree.MerkleTree
+}
+
+func (t *Tree) init() error {
+  if len(t.storage) < 1 {
+    usr, err := user.Current()
+    if err == nil {
+      t.storage = usr.HomeDir + "/.dvote/Tree"
+    } else { t.storage = "./dvoteTree" }
+  }
+  mtdb, err := db.NewLevelDbStorage(t.storage, false)
+	if err != nil {
+    return err
+	}
+	mt, err := merkletree.New(mtdb, 140)
+	if err != nil {
+	  return err
+	}
+  t.tree = mt
+  return nil
+}
+
+func (t *Tree) close() {
+	defer t.tree.Storage().Close()
+}
+
+func (t *Tree) addClaim(data []byte) error {
+  claim := mkcore.NewGenericClaim(t.namespace, "default", data, nil)
+  return t.tree.Add(claim)
+}
+
+func (t *Tree) genProof(data []byte) (string, error) {
+  claim := mkcore.NewGenericClaim(t.namespace, "default", data, nil)
+  mp, err := t.tree.GenerateProof(claim.Hi())
+  if err!=nil {
+    return "", err
+  }
+  mpHex := common3.BytesToHex(mp)
+  return mpHex, nil
+}
+
+func (t *Tree) checkProof(data []byte, mpHex string) (bool, error) {
+  mp, err := common3.HexToBytes(mpHex)
+  if err != nil {
+    return false, err
+  }
+  claim := mkcore.NewGenericClaim(t.namespace, "default", data, nil)
+  return merkletree.CheckProof(t.tree.Root(), mp, claim.Hi(), claim.Ht(), t.tree.NumLevels()), nil
+}