WIP. Implementation of Postgres Storage for MerkleTree. Changes in how storage works in general.

db/memory/memory.go

@@ -2,64 +2,72 @@ package memory
 
 import (
 	"bytes"
+	"github.com/iden3/go-merkletree"
 	"sort"
-
-	"github.com/iden3/go-merkletree/db"
 )
 
 // Storage implements the db.Storage interface
 type Storage struct {
-	prefix []byte
-	kv     db.KvMap
+	prefix      []byte
+	kv          merkletree.KvMap
+	currentRoot *merkletree.Hash
 }
 
 // StorageTx implements the db.Tx interface
 type StorageTx struct {
-	s  *Storage
-	kv db.KvMap
+	s           *Storage
+	kv          merkletree.KvMap
+	currentRoot *merkletree.Hash
 }
 
 // NewMemoryStorage returns a new Storage
 func NewMemoryStorage() *Storage {
-	kvmap := make(db.KvMap)
-	return &Storage{[]byte{}, kvmap}
+	kvmap := make(merkletree.KvMap)
+	return &Storage{[]byte{}, kvmap, nil}
 }
 
 // WithPrefix implements the method WithPrefix of the interface db.Storage
-func (m *Storage) WithPrefix(prefix []byte) db.Storage {
-	return &Storage{db.Concat(m.prefix, prefix), m.kv}
+func (m *Storage) WithPrefix(prefix []byte) merkletree.Storage {
+	return &Storage{merkletree.Concat(m.prefix, prefix), m.kv, nil}
 }
 
 // NewTx implements the method NewTx of the interface db.Storage
-func (m *Storage) NewTx() (db.Tx, error) {
-	return &StorageTx{m, make(db.KvMap)}, nil
+func (m *Storage) NewTx() (merkletree.Tx, error) {
+	return &StorageTx{m, make(merkletree.KvMap), nil}, nil
 }
 
-// Get retreives a value from a key in the db.Storage
-func (m *Storage) Get(key []byte) ([]byte, error) {
-	if v, ok := m.kv.Get(db.Concat(m.prefix, key[:])); ok {
-		return v, nil
+// Get retrieves a value from a key in the db.Storage
+func (m *Storage) Get(key []byte) (*merkletree.Node, error) {
+	if v, ok := m.kv.Get(merkletree.Concat(m.prefix, key[:])); ok {
+		return &v, nil
 	}
-	return nil, db.ErrNotFound
+	return nil, merkletree.ErrNotFound
+}
+
+func (m *Storage) GetRoot() (*merkletree.Hash, error) {
+	if m.currentRoot != nil {
+		return m.currentRoot, nil
+	}
+	return nil, merkletree.ErrNotFound
 }
 
 // Iterate implements the method Iterate of the interface db.Storage
-func (m *Storage) Iterate(f func([]byte, []byte) (bool, error)) error {
-	kvs := make([]db.KV, 0)
+func (m *Storage) Iterate(f func([]byte, *merkletree.Node) (bool, error)) error {
+	kvs := make([]merkletree.KV, 0)
 	for _, v := range m.kv {
 		if len(v.K) < len(m.prefix) ||
 			!bytes.Equal(v.K[:len(m.prefix)], m.prefix) {
 			continue
 		}
 		localkey := v.K[len(m.prefix):]
-		kvs = append(kvs, db.KV{K: localkey, V: v.V})
+		kvs = append(kvs, merkletree.KV{K: localkey, V: v.V})
 	}
 	sort.SliceStable(kvs, func(i, j int) bool {
 		return bytes.Compare(kvs[i].K, kvs[j].K) < 0
 	})
 
 	for _, kv := range kvs {
-		if cont, err := f(kv.K, kv.V); err != nil {
+		if cont, err := f(kv.K, &kv.V); err != nil {
 			return err
 		} else if !cont {
 			break
@@ -69,20 +77,37 @@ func (m *Storage) Iterate(f func([]byte, []byte) (bool, error)) error {
 }
 
 // Get implements the method Get of the interface db.Tx
-func (tx *StorageTx) Get(key []byte) ([]byte, error) {
-	if v, ok := tx.kv.Get(db.Concat(tx.s.prefix, key)); ok {
-		return v, nil
+func (tx *StorageTx) Get(key []byte) (*merkletree.Node, error) {
+	if v, ok := tx.kv.Get(merkletree.Concat(tx.s.prefix, key)); ok {
+		return &v, nil
 	}
-	if v, ok := tx.s.kv.Get(db.Concat(tx.s.prefix, key)); ok {
-		return v, nil
+	if v, ok := tx.s.kv.Get(merkletree.Concat(tx.s.prefix, key)); ok {
+		return &v, nil
 	}
 
-	return nil, db.ErrNotFound
+	return nil, merkletree.ErrNotFound
 }
 
 // Put implements the method Put of the interface db.Tx
-func (tx *StorageTx) Put(k, v []byte) error {
-	tx.kv.Put(db.Concat(tx.s.prefix, k), v)
+func (tx *StorageTx) Put(k []byte, v *merkletree.Node) error {
+	tx.kv.Put(merkletree.Concat(tx.s.prefix, k), *v)
+	return nil
+}
+
+func (tx *StorageTx) GetRoot() (*merkletree.Hash, error) {
+	if tx.currentRoot != nil {
+		hash := merkletree.Hash{}
+		copy(tx.currentRoot[:], hash[:])
+		return &hash, nil
+	}
+	return nil, merkletree.ErrNotFound
+}
+
+// SetRoot sets a hash of merkle tree root in the interface db.Tx
+func (tx *StorageTx) SetRoot(hash *merkletree.Hash) error {
+	root := &merkletree.Hash{}
+	copy(root[:], hash[:])
+	tx.currentRoot = root
 	return nil
 }
 
@@ -96,7 +121,7 @@ func (tx *StorageTx) Commit() error {
 }
 
 // Add implements the method Add of the interface db.Tx
-func (tx *StorageTx) Add(atx db.Tx) error {
+func (tx *StorageTx) Add(atx merkletree.Tx) error {
 	mstx := atx.(*StorageTx)
 	for _, v := range mstx.kv {
 		tx.kv.Put(v.K, v.V)
@@ -114,10 +139,10 @@ func (m *Storage) Close() {
 }
 
 // List implements the method List of the interface db.Storage
-func (m *Storage) List(limit int) ([]db.KV, error) {
-	ret := []db.KV{}
-	err := m.Iterate(func(key []byte, value []byte) (bool, error) {
-		ret = append(ret, db.KV{K: db.Clone(key), V: db.Clone(value)})
+func (m *Storage) List(limit int) ([]merkletree.KV, error) {
+	ret := []merkletree.KV{}
+	err := m.Iterate(func(key []byte, value *merkletree.Node) (bool, error) {
+		ret = append(ret, merkletree.KV{K: merkletree.Clone(key), V: *value})
 		if len(ret) == limit {
 			return false, nil
 		}

db/memory/memory_test.go

@@ -1,15 +1,15 @@
 package memory
 
 import (
+	"github.com/iden3/go-merkletree"
 	"testing"
 
-	"github.com/iden3/go-merkletree/db"
 	"github.com/iden3/go-merkletree/db/test"
 	"github.com/stretchr/testify/require"
 )
 
 func TestMemoryStorageInterface(t *testing.T) {
-	var db db.Storage //nolint:gosimple
+	var db merkletree.Storage //nolint:gosimple
 
 	db = NewMemoryStorage()
 	require.NotNil(t, db)