fix api

[bytom/bytom.git] / database / leveldb / mem_db.go

diff --git a/database/leveldb/mem_db.go b/database/leveldb/mem_db.go
index 62f40fc..e70aa9d 100644 (file)
--- a/database/leveldb/mem_db.go
+++ b/database/leveldb/mem_db.go
@@ -1,6 +1,7 @@
 package db
 
 import (
+       "bytes"
        "fmt"
        "sort"
        "strings"
@@ -78,13 +79,29 @@ func (db *MemDB) Stats() map[string]string {
 type memDBIterator struct {
        last int
        keys []string
-       db   *MemDB
+       db   DB
+
+       start []byte
 }
 
 func newMemDBIterator() *memDBIterator {
        return &memDBIterator{}
 }
 
+// Keys is expected to be in reverse order for reverse iterators.
+func newMemDBIteratorWithArgs(db DB, keys []string, start []byte) *memDBIterator {
+       itr := &memDBIterator{
+               db:    db,
+               keys:  keys,
+               start: start,
+               last:  -1,
+       }
+       if start != nil {
+               itr.Seek(start)
+       }
+       return itr
+}
+
 func (it *memDBIterator) Next() bool {
        if it.last >= len(it.keys)-1 {
                return false
@@ -143,10 +160,38 @@ func (db *MemDB) IteratorPrefix(prefix []byte) Iterator {
        return it
 }
 
+func (db *MemDB) IteratorPrefixWithStart(Prefix, start []byte, isReverse bool) Iterator {
+       db.mtx.Lock()
+       defer db.mtx.Unlock()
+
+       keys := db.getSortedKeys(start, isReverse)
+       return newMemDBIteratorWithArgs(db, keys, start)
+}
+
 func (db *MemDB) NewBatch() Batch {
        return &memDBBatch{db, nil}
 }
 
+func (db *MemDB) getSortedKeys(start []byte, reverse bool) []string {
+       keys := []string{}
+       for key := range db.db {
+               if bytes.Compare([]byte(key), start) < 0 {
+                       continue
+               }
+               keys = append(keys, key)
+       }
+       sort.Strings(keys)
+       if reverse {
+               nkeys := len(keys)
+               for i := 0; i < nkeys/2; i++ {
+                       temp := keys[i]
+                       keys[i] = keys[nkeys-i-1]
+                       keys[nkeys-i-1] = temp
+               }
+       }
+       return keys
+}
+
 //--------------------------------------------------------------------------------
 
 type memDBBatch struct {