vm/storage: add prefix helper functions

Signed-off-by: Gyuho Lee <[email protected]>

Author: gyuho

vm/storage/prefix.go

@@ -0,0 +1,50 @@
+package storage
+
+import (
+	"bytes"
+	"errors"
+)
+
+var (
+	delimiter   = byte(0x2f) // '/'
+	noPrefixEnd = []byte{0}
+)
+
+var (
+	ErrInvalidKeyDelimiter = errors.New("key has unexpected delimiters; only flat key or sub-key is supported")
+)
+
+// getPrefix returns the prefixed key and range query end key for list calls.
+// put "foo" becomes "foo/" for its own namespace, and range ends with "foo0"
+// put "fop" becomes "fop/" for its own namespace, and range ends with "fop0"
+// put "foo1" becomes "foo1/" for its own namespace, and range ends with "foo10"
+func getPrefix(key []byte) (pfx []byte, sub []byte, end []byte, err error) {
+	if bytes.Count(key, []byte{delimiter}) > 1 {
+		return nil, nil, nil, ErrInvalidKeyDelimiter
+	}
+
+	idx := bytes.IndexRune(key, rune(delimiter))
+	switch {
+	case idx == -1: // "foo"
+		pfx = append(key, delimiter)
+	case idx == len(key)-1: // "foo/"
+		pfx, sub = key, nil
+	default: // "a/b"
+		ss := bytes.Split(key, []byte{delimiter})
+		pfx, sub = append(ss[0], delimiter), ss[1]
+	}
+
+	// next lexicographical key (range end) for prefix queries
+	end = make([]byte, len(pfx))
+	copy(end, pfx)
+	for i := len(end) - 1; i >= 0; i-- {
+		if end[i] < 0xff {
+			end[i] = end[i] + 1
+			end = end[:i+1]
+			return pfx, sub, end, nil
+		}
+	}
+	// next prefix does not exist (e.g., 0xffff);
+	// default to special end key
+	return pfx, sub, noPrefixEnd, nil
+}

vm/storage/prefix_test.go

@@ -0,0 +1,60 @@
+package storage
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestPrefix(t *testing.T) {
+	tt := []struct {
+		key []byte
+		pfx []byte
+		sub []byte
+		end []byte
+		err error
+	}{
+		{
+			key: []byte("foo"),
+			pfx: []byte("foo/"),
+			sub: nil,
+			end: []byte("foo0"),
+			err: nil,
+		},
+		{
+			key: []byte("foo/"),
+			pfx: []byte("foo/"),
+			sub: nil,
+			end: []byte("foo0"),
+			err: nil,
+		},
+		{
+			key: []byte("a/b"),
+			pfx: []byte("a/"),
+			sub: []byte("b"),
+			end: []byte("a0"),
+			err: nil,
+		},
+		{
+			key: []byte("a/b/"),
+			pfx: nil,
+			sub: nil,
+			end: nil,
+			err: ErrInvalidKeyDelimiter,
+		},
+	}
+	for i, tv := range tt {
+		pfx, sub, end, err := getPrefix(tv.key)
+		if !bytes.Equal(pfx, tv.pfx) {
+			t.Fatalf("#%d: pfx expected %q, got %q", i, string(tv.pfx), string(pfx))
+		}
+		if !bytes.Equal(sub, tv.sub) {
+			t.Fatalf("#%d: sub expected %q, got %q", i, string(tv.sub), string(sub))
+		}
+		if !bytes.Equal(end, tv.end) {
+			t.Fatalf("#%d: end expected %q, got %q", i, string(tv.end), string(end))
+		}
+		if err != tv.err {
+			t.Fatalf("#%d: err expected %v, got %v", i, tv.err, err)
+		}
+	}
+}

vm/storage/storage.go

@@ -4,28 +4,101 @@
 package storage
 
 import (
+	"errors"
+
 	"github.com/ava-labs/avalanchego/database"
 	"github.com/ava-labs/avalanchego/database/prefixdb"
 	"github.com/ava-labs/avalanchego/database/versiondb"
+	"github.com/ava-labs/quarkvm/crypto/ed25519"
+)
+
+var (
+	blockStatePrefix = []byte("block")
+
+	// (singleton prefix info)
+	infoPrefix = []byte("info")
+	// (tx hashes)
+	txPrefix = []byte("tx")
+	// (prefix keys)
+	keyPrefix = []byte("key")
 )
 
 type Storage interface {
-	Block() *prefixdb.Database
+	Block() database.Database
+	Info() database.Database
+	Tx() database.Database
+	Key() database.Database
+
+	// Put writes key-value pair with prefix checks.
+	Put(k []byte, v []byte, opts ...OpOption) error
 }
 
 type storage struct {
 	baseDB  *versiondb.Database
 	blockDB *prefixdb.Database
+	infoDB  *prefixdb.Database
+	txDB    *prefixdb.Database
+	keyDB   *prefixdb.Database
 }
 
-var (
-	blockStatePrefix = []byte("block")
-)
-
 func New(db database.Database) Storage {
 	baseDB := versiondb.New(db)
 	return &storage{
 		baseDB:  baseDB,
 		blockDB: prefixdb.New(blockStatePrefix, baseDB),
+		infoDB:  prefixdb.New(infoPrefix, baseDB),
+		txDB:    prefixdb.New(txPrefix, baseDB),
+		keyDB:   prefixdb.New(keyPrefix, baseDB),
+	}
+}
+
+func (s *storage) Block() database.Database {
+	return s.blockDB
+}
+
+func (s *storage) Info() database.Database {
+	return s.infoDB
+}
+
+func (s *storage) Tx() database.Database {
+	return s.txDB
+}
+
+func (s *storage) Key() database.Database {
+	return s.keyDB
+}
+
+func (s *storage) Put(k []byte, v []byte, opts ...OpOption) error {
+	ret := &Op{}
+	ret.applyOpts(opts)
+	if ret.overwrite && ret.owner == nil {
+		return errors.New("caller must provide the public key to make modification or overwrite")
 	}
+
+	// TODO: look up key
+	// if exists, only overwrite when conditions are met
+	// implement other business logic
+
+	return s.keyDB.Put(k, v)
+}
+
+type Op struct {
+	overwrite bool
+	owner     ed25519.PublicKey
+}
+
+type OpOption func(*Op)
+
+func (op *Op) applyOpts(opts []OpOption) {
+	for _, opt := range opts {
+		opt(op)
+	}
+}
+
+func WithOverwrite(b bool) OpOption {
+	return func(op *Op) { op.overwrite = b }
+}
+
+func WithOwner(owner ed25519.PublicKey) OpOption {
+	return func(op *Op) { op.owner = owner }
 }