github.com/evdatsion/aphelion-dpos-bft@v0.32.1/libs/db/prefix_db.go

package db

import (
	"bytes"
	"fmt"
	"sync"
)

// IteratePrefix is a convenience function for iterating over a key domain
// restricted by prefix.
func IteratePrefix(db DB, prefix []byte) Iterator {
	var start, end []byte
	if len(prefix) == 0 {
		start = nil
		end = nil
	} else {
		start = cp(prefix)
		end = cpIncr(prefix)
	}
	return db.Iterator(start, end)
}

/*
TODO: Make test, maybe rename.
// Like IteratePrefix but the iterator strips the prefix from the keys.
func IteratePrefixStripped(db DB, prefix []byte) Iterator {
	start, end := ...
	return newPrefixIterator(prefix, start, end, IteratePrefix(db, prefix))
}
*/

//----------------------------------------
// prefixDB

type prefixDB struct {
	mtx    sync.Mutex
	prefix []byte
	db     DB
}

// NewPrefixDB lets you namespace multiple DBs within a single DB.
func NewPrefixDB(db DB, prefix []byte) *prefixDB {
	return &prefixDB{
		prefix: prefix,
		db:     db,
	}
}

// Implements atomicSetDeleter.
func (pdb *prefixDB) Mutex() *sync.Mutex {
	return &(pdb.mtx)
}

// Implements DB.
func (pdb *prefixDB) Get(key []byte) []byte {
	pdb.mtx.Lock()
	defer pdb.mtx.Unlock()

	pkey := pdb.prefixed(key)
	value := pdb.db.Get(pkey)
	return value
}

// Implements DB.
func (pdb *prefixDB) Has(key []byte) bool {
	pdb.mtx.Lock()
	defer pdb.mtx.Unlock()

	return pdb.db.Has(pdb.prefixed(key))
}

// Implements DB.
func (pdb *prefixDB) Set(key []byte, value []byte) {
	pdb.mtx.Lock()
	defer pdb.mtx.Unlock()

	pkey := pdb.prefixed(key)
	pdb.db.Set(pkey, value)
}

// Implements DB.
func (pdb *prefixDB) SetSync(key []byte, value []byte) {
	pdb.mtx.Lock()
	defer pdb.mtx.Unlock()

	pdb.db.SetSync(pdb.prefixed(key), value)
}

// Implements DB.
func (pdb *prefixDB) Delete(key []byte) {
	pdb.mtx.Lock()
	defer pdb.mtx.Unlock()

	pdb.db.Delete(pdb.prefixed(key))
}

// Implements DB.
func (pdb *prefixDB) DeleteSync(key []byte) {
	pdb.mtx.Lock()
	defer pdb.mtx.Unlock()

	pdb.db.DeleteSync(pdb.prefixed(key))
}

// Implements DB.
func (pdb *prefixDB) Iterator(start, end []byte) Iterator {
	pdb.mtx.Lock()
	defer pdb.mtx.Unlock()

	var pstart, pend []byte
	pstart = append(cp(pdb.prefix), start...)
	if end == nil {
		pend = cpIncr(pdb.prefix)
	} else {
		pend = append(cp(pdb.prefix), end...)
	}
	return newPrefixIterator(
		pdb.prefix,
		start,
		end,
		pdb.db.Iterator(
			pstart,
			pend,
		),
	)
}

// Implements DB.
func (pdb *prefixDB) ReverseIterator(start, end []byte) Iterator {
	pdb.mtx.Lock()
	defer pdb.mtx.Unlock()

	var pstart, pend []byte
	pstart = append(cp(pdb.prefix), start...)
	if end == nil {
		pend = cpIncr(pdb.prefix)
	} else {
		pend = append(cp(pdb.prefix), end...)
	}
	ritr := pdb.db.ReverseIterator(pstart, pend)
	return newPrefixIterator(
		pdb.prefix,
		start,
		end,
		ritr,
	)
}

// Implements DB.
// Panics if the underlying DB is not an
// atomicSetDeleter.
func (pdb *prefixDB) NewBatch() Batch {
	pdb.mtx.Lock()
	defer pdb.mtx.Unlock()

	return newPrefixBatch(pdb.prefix, pdb.db.NewBatch())
}

/* NOTE: Uncomment to use memBatch instead of prefixBatch
// Implements atomicSetDeleter.
func (pdb *prefixDB) SetNoLock(key []byte, value []byte) {
	pdb.db.(atomicSetDeleter).SetNoLock(pdb.prefixed(key), value)
}

// Implements atomicSetDeleter.
func (pdb *prefixDB) SetNoLockSync(key []byte, value []byte) {
	pdb.db.(atomicSetDeleter).SetNoLockSync(pdb.prefixed(key), value)
}

// Implements atomicSetDeleter.
func (pdb *prefixDB) DeleteNoLock(key []byte) {
	pdb.db.(atomicSetDeleter).DeleteNoLock(pdb.prefixed(key))
}

// Implements atomicSetDeleter.
func (pdb *prefixDB) DeleteNoLockSync(key []byte) {
	pdb.db.(atomicSetDeleter).DeleteNoLockSync(pdb.prefixed(key))
}
*/

// Implements DB.
func (pdb *prefixDB) Close() {
	pdb.mtx.Lock()
	defer pdb.mtx.Unlock()

	pdb.db.Close()
}

// Implements DB.
func (pdb *prefixDB) Print() {
	fmt.Printf("prefix: %X\n", pdb.prefix)

	itr := pdb.Iterator(nil, nil)
	defer itr.Close()
	for ; itr.Valid(); itr.Next() {
		key := itr.Key()
		value := itr.Value()
		fmt.Printf("[%X]:\t[%X]\n", key, value)
	}
}

// Implements DB.
func (pdb *prefixDB) Stats() map[string]string {
	stats := make(map[string]string)
	stats["prefixdb.prefix.string"] = string(pdb.prefix)
	stats["prefixdb.prefix.hex"] = fmt.Sprintf("%X", pdb.prefix)
	source := pdb.db.Stats()
	for key, value := range source {
		stats["prefixdb.source."+key] = value
	}
	return stats
}

func (pdb *prefixDB) prefixed(key []byte) []byte {
	return append(cp(pdb.prefix), key...)
}

//----------------------------------------
// prefixBatch

type prefixBatch struct {
	prefix []byte
	source Batch
}

func newPrefixBatch(prefix []byte, source Batch) prefixBatch {
	return prefixBatch{
		prefix: prefix,
		source: source,
	}
}

func (pb prefixBatch) Set(key, value []byte) {
	pkey := append(cp(pb.prefix), key...)
	pb.source.Set(pkey, value)
}

func (pb prefixBatch) Delete(key []byte) {
	pkey := append(cp(pb.prefix), key...)
	pb.source.Delete(pkey)
}

func (pb prefixBatch) Write() {
	pb.source.Write()
}

func (pb prefixBatch) WriteSync() {
	pb.source.WriteSync()
}

func (pb prefixBatch) Close() {
	pb.source.Close()
}

//----------------------------------------
// prefixIterator

var _ Iterator = (*prefixIterator)(nil)

// Strips prefix while iterating from Iterator.
type prefixIterator struct {
	prefix []byte
	start  []byte
	end    []byte
	source Iterator
	valid  bool
}

func newPrefixIterator(prefix, start, end []byte, source Iterator) *prefixIterator {
	if !source.Valid() || !bytes.HasPrefix(source.Key(), prefix) {
		return &prefixIterator{
			prefix: prefix,
			start:  start,
			end:    end,
			source: source,
			valid:  false,
		}
	} else {
		return &prefixIterator{
			prefix: prefix,
			start:  start,
			end:    end,
			source: source,
			valid:  true,
		}
	}
}

func (itr *prefixIterator) Domain() (start []byte, end []byte) {
	return itr.start, itr.end
}

func (itr *prefixIterator) Valid() bool {
	return itr.valid && itr.source.Valid()
}

func (itr *prefixIterator) Next() {
	if !itr.valid {
		panic("prefixIterator invalid, cannot call Next()")
	}
	itr.source.Next()
	if !itr.source.Valid() || !bytes.HasPrefix(itr.source.Key(), itr.prefix) {
		itr.valid = false
		return
	}
}

func (itr *prefixIterator) Key() (key []byte) {
	if !itr.valid {
		panic("prefixIterator invalid, cannot call Key()")
	}
	return stripPrefix(itr.source.Key(), itr.prefix)
}

func (itr *prefixIterator) Value() (value []byte) {
	if !itr.valid {
		panic("prefixIterator invalid, cannot call Value()")
	}
	return itr.source.Value()
}

func (itr *prefixIterator) Close() {
	itr.source.Close()
}

//----------------------------------------

func stripPrefix(key []byte, prefix []byte) (stripped []byte) {
	if len(key) < len(prefix) {
		panic("should not happen")
	}
	if !bytes.Equal(key[:len(prefix)], prefix) {
		panic("should not happne")
	}
	return key[len(prefix):]
}