1package memory
  2
  3import (
  4	"bytes"
  5	"sort"
  6
  7	"github.com/iden3/go-merkletree/db"
  8)
  9
 10// Storage implements the db.Storage interface
 11type Storage struct {
 12	prefix []byte
 13	kv     db.KvMap
 14}
 15
 16// StorageTx implements the db.Tx interface
 17type StorageTx struct {
 18	s  *Storage
 19	kv db.KvMap
 20}
 21
 22// NewMemoryStorage returns a new Storage
 23func NewMemoryStorage() *Storage {
 24	kvmap := make(db.KvMap)
 25	return &Storage{[]byte{}, kvmap}
 26}
 27
 28// WithPrefix implements the method WithPrefix of the interface db.Storage
 29func (m *Storage) WithPrefix(prefix []byte) db.Storage {
 30	return &Storage{db.Concat(m.prefix, prefix), m.kv}
 31}
 32
 33// NewTx implements the method NewTx of the interface db.Storage
 34func (m *Storage) NewTx() (db.Tx, error) {
 35	return &StorageTx{m, make(db.KvMap)}, nil
 36}
 37
 38// Get retreives a value from a key in the db.Storage
 39func (m *Storage) Get(key []byte) ([]byte, error) {
 40	if v, ok := m.kv.Get(db.Concat(m.prefix, key[:])); ok {
 41		return v, nil
 42	}
 43	return nil, db.ErrNotFound
 44}
 45
 46// Iterate implements the method Iterate of the interface db.Storage
 47func (m *Storage) Iterate(f func([]byte, []byte) (bool, error)) error {
 48	kvs := make([]db.KV, 0)
 49	for _, v := range m.kv {
 50		if len(v.K) < len(m.prefix) ||
 51			!bytes.Equal(v.K[:len(m.prefix)], m.prefix) {
 52			continue
 53		}
 54		localkey := v.K[len(m.prefix):]
 55		kvs = append(kvs, db.KV{K: localkey, V: v.V})
 56	}
 57	sort.SliceStable(kvs, func(i, j int) bool {
 58		return bytes.Compare(kvs[i].K, kvs[j].K) < 0
 59	})
 60
 61	for _, kv := range kvs {
 62		if cont, err := f(kv.K, kv.V); err != nil {
 63			return err
 64		} else if !cont {
 65			break
 66		}
 67	}
 68	return nil
 69}
 70
 71// Get implements the method Get of the interface db.Tx
 72func (tx *StorageTx) Get(key []byte) ([]byte, error) {
 73	if v, ok := tx.kv.Get(db.Concat(tx.s.prefix, key)); ok {
 74		return v, nil
 75	}
 76	if v, ok := tx.s.kv.Get(db.Concat(tx.s.prefix, key)); ok {
 77		return v, nil
 78	}
 79
 80	return nil, db.ErrNotFound
 81}
 82
 83// Put implements the method Put of the interface db.Tx
 84func (tx *StorageTx) Put(k, v []byte) error {
 85	tx.kv.Put(db.Concat(tx.s.prefix, k), v)
 86	return nil
 87}
 88
 89// Commit implements the method Commit of the interface db.Tx
 90func (tx *StorageTx) Commit() error {
 91	for _, v := range tx.kv {
 92		tx.s.kv.Put(v.K, v.V)
 93	}
 94	tx.kv = nil
 95	return nil
 96}
 97
 98// Add implements the method Add of the interface db.Tx
 99func (tx *StorageTx) Add(atx db.Tx) error {
100	mstx := atx.(*StorageTx)
101	for _, v := range mstx.kv {
102		tx.kv.Put(v.K, v.V)
103	}
104	return nil
105}
106
107// Close implements the method Close of the interface db.Tx
108func (tx *StorageTx) Close() {
109	tx.kv = nil
110}
111
112// Close implements the method Close of the interface db.Storage
113func (m *Storage) Close() {
114}
115
116// List implements the method List of the interface db.Storage
117func (m *Storage) List(limit int) ([]db.KV, error) {
118	ret := []db.KV{}
119	err := m.Iterate(func(key []byte, value []byte) (bool, error) {
120		ret = append(ret, db.KV{K: db.Clone(key), V: db.Clone(value)})
121		if len(ret) == limit {
122			return false, nil
123		}
124		return true, nil
125	})
126	return ret, err
127}