Paging (#234)

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

diff --git a/database/leveldb/mem_db.go b/database/leveldb/mem_db.go
index 62f40fc..63e2b76 100644 (file)
--- a/database/leveldb/mem_db.go
+++ b/database/leveldb/mem_db.go
@@ -1,6 +1,7 @@
-package db
+package leveldb
 
 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
@@ -94,6 +111,9 @@ func (it *memDBIterator) Next() bool {
 }
 
 func (it *memDBIterator) Key() []byte {
+       if it.last < 0 {
+               return []byte("")
+       }
        return []byte(it.keys[it.last])
 }
 
@@ -143,10 +163,30 @@ func (db *MemDB) IteratorPrefix(prefix []byte) Iterator {
        return it
 }
 
+func (db *MemDB) IteratorPrefixWithStart(Prefix, start []byte) Iterator {
+       db.mtx.Lock()
+       defer db.mtx.Unlock()
+
+       keys := db.getSortedKeys(start)
+       return newMemDBIteratorWithArgs(db, keys, start)
+}
+
 func (db *MemDB) NewBatch() Batch {
        return &memDBBatch{db, nil}
 }
 
+func (db *MemDB) getSortedKeys(start []byte) []string {
+       keys := []string{}
+       for key := range db.db {
+               if bytes.Compare([]byte(key), start) < 0 {
+                       continue
+               }
+               keys = append(keys, key)
+       }
+       sort.Strings(keys)
+       return keys
+}
+
 //--------------------------------------------------------------------------------
 
 type memDBBatch struct {