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