1// Package arbotree provides the functions for creating and managing an arbo
  2// merkletree adapted to the CensusTree interface
  3package arbotree
  4
  5import (
  6	"bytes"
  7	"encoding/hex"
  8	"fmt"
  9	"path/filepath"
 10	"strings"
 11	"sync/atomic"
 12	"time"
 13
 14	"github.com/arnaucube/arbo"
 15	"github.com/iden3/go-merkletree/db/pebble"
 16	"go.vocdoni.io/dvote/censustree"
 17)
 18
 19type Tree struct {
 20	Tree           *arbo.Tree
 21	public         uint32
 22	lastAccessUnix int64 // a unix timestamp, used via sync/atomic
 23}
 24
 25func (t *Tree) Init(name, storageDir string) error {
 26	dbDir := filepath.Join(storageDir, "arbotree.db."+strings.TrimSpace(name))
 27	storage, err := pebble.NewPebbleStorage(dbDir, false) // TODO TMP
 28	if err != nil {
 29		return err
 30	}
 31	// storage := memory.NewMemoryStorage()
 32
 33	mt, err := arbo.NewTree(storage, 140, arbo.HashFunctionBlake2b) // TODO here the hash function would depend on the usage
 34	if err != nil {
 35		return err
 36	}
 37	t.Tree = mt
 38	return nil
 39}
 40
 41func (t *Tree) MaxKeySize() int {
 42	return t.Tree.HashFunction().Len()
 43}
 44
 45func (t *Tree) LastAccess() int64 {
 46	return atomic.LoadInt64(&t.lastAccessUnix)
 47}
 48
 49func (t *Tree) updateAccessTime() {
 50	atomic.StoreInt64(&t.lastAccessUnix, time.Now().Unix())
 51}
 52
 53// Publish makes a merkle tree available for queries.
 54// Application layer should check IsPublish() before considering the Tree available.
 55func (t *Tree) Publish() {
 56	atomic.StoreUint32(&t.public, 1)
 57}
 58
 59// UnPublish makes a merkle tree not available for queries
 60func (t *Tree) UnPublish() {
 61	atomic.StoreUint32(&t.public, 0)
 62}
 63
 64// IsPublic returns true if the tree is available
 65func (t *Tree) IsPublic() bool {
 66	return atomic.LoadUint32(&t.public) == 1
 67}
 68
 69func (t *Tree) Add(index, value []byte) error {
 70	t.updateAccessTime()
 71	return t.Tree.Add(index, value)
 72}
 73
 74func (t *Tree) AddBatch(indexes, values [][]byte) ([]int, error) {
 75	t.updateAccessTime()
 76	return t.Tree.AddBatch(indexes, values)
 77}
 78
 79func (t *Tree) GenProof(index, value []byte) ([]byte, error) {
 80	t.updateAccessTime()
 81	v, siblings, err := t.Tree.GenProof(index)
 82	if err != nil {
 83		return nil, err
 84	}
 85	if !bytes.Equal(v, value) {
 86		return nil, fmt.Errorf("value does not match %s!=%s", hex.EncodeToString(v), hex.EncodeToString(value))
 87	}
 88	return siblings, nil
 89}
 90
 91func (t *Tree) CheckProof(index, value, root, mproof []byte) (bool, error) {
 92	t.updateAccessTime()
 93	if root == nil {
 94		root = t.Root()
 95	}
 96	return arbo.CheckProof(t.Tree.HashFunction(), index, value, root, mproof)
 97}
 98
 99func (t *Tree) Root() []byte {
100	t.updateAccessTime()
101	return t.Tree.Root()
102}
103
104func (t *Tree) Dump(root []byte) ([]byte, error) {
105	return t.Tree.Dump() // TODO pass root once arbo is updated
106}
107
108func (t *Tree) DumpPlain(root []byte) ([][]byte, [][]byte, error) {
109	t.updateAccessTime()
110	var indexes, values [][]byte
111	// TODO pass root once arbo is updated
112	err := t.Tree.Iterate(func(k, v []byte) {
113		if v[0] != arbo.PrefixValueLeaf {
114			return
115		}
116		leafK, leafV := arbo.ReadLeafValue(v)
117		indexes = append(indexes, leafK)
118		values = append(values, leafV)
119	})
120	return indexes, values, err
121}
122
123func (t *Tree) ImportDump(data []byte) error {
124	t.updateAccessTime()
125	return t.Tree.ImportDump(data)
126}
127
128func (t *Tree) Size(root []byte) (int64, error) {
129	count := 0
130	err := t.Tree.Iterate(func(k, v []byte) {
131		if v[0] != arbo.PrefixValueLeaf {
132			return
133		}
134		count++
135	})
136	return int64(count), err
137}
138
139func (t *Tree) Snapshot(root []byte) (censustree.Tree, error) {
140	// TODO
141	return t, nil
142}
143
144func (t *Tree) HashExists(hash []byte) (bool, error) {
145	_, _, err := t.Tree.Get(hash)
146	if err != nil {
147		return false, err
148	}
149	return true, nil
150}