memdb: add fixed-key comparison via nodeData

leveldb/memdb/memdb.go

@@ -13,6 +13,7 @@ import (
 	"sync"
 	"unsafe"
 
+	"bytes"
 	"github.com/syndtr/goleveldb/leveldb/comparer"
 	"github.com/syndtr/goleveldb/leveldb/errors"
 	"github.com/syndtr/goleveldb/leveldb/iterator"
@@ -176,8 +177,9 @@ func (i *dbIter) Release() {
 
 // DB is an in-memory key/value database.
 type DB struct {
-	cmp comparer.BasicComparer
-	rnd *rand.Rand
+	cmp      comparer.BasicComparer
+	rnd      *rand.Rand
+	quickCmp bool
 
 	mu        sync.RWMutex
 	kvData    []byte
@@ -198,14 +200,31 @@ func (p *DB) randHeight() (h int) {
 
 // Must hold RW-lock if prev == true, as it use shared prevNode slice.
 func (p *DB) findGE(key []byte, prev bool) (int, bool) {
-	node := 0
-	h := p.maxHeight - 1
+	var (
+		node = 0
+		h    = p.maxHeight - 1
+		qKey []byte
+	)
+	if p.quickCmp {
+		// If we're using quickCmp, the key we compare against needs to
+		// be padded to full 8 bytes.
+		qKey = padKey(key)
+	}
 	for {
 		next := p.nodeAt(node).nextAt(h)
 		cmp := 1
 		if next != 0 {
 			o := p.nodeAt(next)
-			cmp = p.cmp.Compare(p.kvData[o.kStart():o.kEnd()], key)
+			if p.quickCmp {
+				// If the default comparer is used, can use bytes.Compare
+				if cmp = o.quickCmp(qKey); cmp == 0 {
+					// Deep compare
+					cmp = bytes.Compare(p.kvData[o.kStart():o.kEnd()], key)
+				}
+			} else {
+				// Deep compare
+				cmp = p.cmp.Compare(p.kvData[o.kStart():o.kEnd()], key)
+			}
 		}
 		if cmp < 0 {
 			// Keep searching in this list
@@ -292,7 +311,7 @@ func (p *DB) Put(key []byte, value []byte) error {
 	p.kvData = append(p.kvData, value...)
 	// Node
 	node := len(p.nodeData)
-	newN := newNode(kvOffset, len(key), len(value), h)
+	newN := newNode(kvOffset, len(value), h, key)
 	for i, n := range p.prevNode[:h] {
 		prev := p.nodeAt(n)
 		newN.setNextAt(i, prev.nextAt(i))
@@ -439,7 +458,7 @@ func (p *DB) Reset() {
 
 	p.nodeData = p.nodeData[:0]
 	// Add empty first element
-	zero := newNode(0, 0, 0, tMaxHeight)
+	zero := newNode(0, 0, tMaxHeight, nil)
 	for n := 0; n < tMaxHeight; n++ {
 		zero.setNextAt(n, 0)
 		p.prevNode[n] = 0
@@ -459,12 +478,13 @@ func (p *DB) Reset() {
 func New(cmp comparer.BasicComparer, capacity int) *DB {
 	p := &DB{
 		cmp:       cmp,
+		quickCmp:  (cmp == comparer.DefaultComparer),
 		rnd:       rand.New(rand.NewSource(0xdeadbeef)),
 		maxHeight: 1,
 		kvData:    make([]byte, 0, capacity),
 	}
 	// Add empty first element
-	zero := newNode(0, 0, 0, tMaxHeight)
+	zero := newNode(0, 0, tMaxHeight, nil)
 	for n := 0; n < tMaxHeight; n++ {
 		zero.setNextAt(n, 0)
 	}

leveldb/memdb/node.go

@@ -7,6 +7,20 @@
 // Package memdb provides in-memory key/value database implementation.
 package memdb
 
+import "encoding/binary"
+
+// padKey pads, if necessary, the given key to be sufficiently long
+// to work with the quickCmp function.
+func padKey(key []byte) []byte {
+	if len(key) >= 8 {
+		return key
+	}
+	// pad
+	k := make([]byte, 8)
+	copy(k, key)
+	return k
+}
+
 // The backing representation for the nodeData slice
 type nodeInt int
 
@@ -17,22 +31,22 @@ type nodeInt int
 // is not known before reading the height).
 // Node data is laid out as follows:
 // [0]         : KV offset
-// [1]         : Key length
+// [1]         : Key length | height (8 bits)
 // [2]         : Value length
-// [3]         : Height
+// [3]         : Key first 8 bytes (padded if necessary)
 // [3..height] : Next nodes
 type node []nodeInt
 
 // newNode constructs a new node. Be careful -- this allocates a new slice,
 // along with space to store the next, according to the height given.
 // This node later needs to be written to the backing slice, making the original
 // instance moot.
-func newNode(kvOffset, kLen, vLen, height int) node {
+func newNode(kvOffset, vLen, height int, key []byte) node {
 	buf := make([]nodeInt, 4+height)
 	buf[0] = nodeInt(kvOffset)
-	buf[1] = nodeInt(kLen)
+	buf[1] = nodeInt(len(key))<<8 | nodeInt(height&0xff)
 	buf[2] = nodeInt(vLen)
-	buf[3] = nodeInt(height)
+	buf[3] = nodeInt(binary.BigEndian.Uint64(padKey(key)))
 	return node(buf)
 }
 
@@ -43,22 +57,22 @@ func (n node) kStart() int {
 
 // kEnd returns the start + length for the key.
 func (n node) kEnd() int {
-	return int(n[0] + n[1])
+	return int(n[0] + n[1]>>8)
 }
 
 // kLen return the key length.
 func (n node) kLen() int {
-	return int(n[1])
+	return int(n[1] >> 8)
 }
 
 // vStart return the offset for the value.
 func (n node) vStart() int {
-	return int(n[0] + n[1])
+	return int(n[0] + n[1]>>8)
 }
 
 // vEnd return the offset + length for value.
 func (n node) vEnd() int {
-	return int(n[0] + n[1] + n[2])
+	return int(n[0] + n[1]>>8 + n[2])
 }
 
 // vLen returns the value length.
@@ -80,7 +94,7 @@ func (n node) setVLen(size int) node {
 
 // height return the size of the next-tower.
 func (n node) height() int {
-	return int(n[3])
+	return int(n[1] & 0xff)
 }
 
 // nextAt return the item at the given height.
@@ -97,3 +111,16 @@ func (n node) setNextAt(height int, node int) {
 func (p *DB) nodeAt(idx int) node {
 	return node(p.nodeData[idx:])
 }
+
+// quickCmp compares the first N bytes of the key in this  node with the first N bytes
+// of the given key. If the comparison returns 0, that means a deeper comparison is
+// needed.
+func (n node) quickCmp(key []byte) int {
+	other := binary.BigEndian.Uint64(key)
+	if uint64(n[3]) < other {
+		return -1
+	} else if uint64(n[3]) > other {
+		return 1
+	}
+	return 0
+}