From 219d4e8427b2ddb0e9146ca7319bce2e1f22a6ce Mon Sep 17 00:00:00 2001
From: Alexis Sellier <alexis@cloudhead.io>
Date: Thu, 28 Sep 2017 17:35:01 +0200
Subject: [PATCH] Add IteratorPrefix method to Iterator

---
 db/db.go          | 1 +
 db/go_level_db.go | 5 +++++
 db/mem_db.go      | 9 ++++++++-
 3 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/db/db.go b/db/db.go
index 51283ac2a..8156c1e92 100644
--- a/db/db.go
+++ b/db/db.go
@@ -11,6 +11,7 @@ type DB interface {
 	Close()
 	NewBatch() Batch
 	Iterator() Iterator
+	IteratorPrefix([]byte) Iterator
 
 	// For debugging
 	Print()
diff --git a/db/go_level_db.go b/db/go_level_db.go
index 54ae1149f..31c85ce86 100644
--- a/db/go_level_db.go
+++ b/db/go_level_db.go
@@ -7,6 +7,7 @@ import (
 	"github.com/syndtr/goleveldb/leveldb"
 	"github.com/syndtr/goleveldb/leveldb/errors"
 	"github.com/syndtr/goleveldb/leveldb/opt"
+	"github.com/syndtr/goleveldb/leveldb/util"
 
 	. "github.com/tendermint/tmlibs/common"
 )
@@ -119,6 +120,10 @@ func (db *GoLevelDB) Iterator() Iterator {
 	return db.db.NewIterator(nil, nil)
 }
 
+func (db *GoLevelDB) IteratorPrefix(prefix []byte) Iterator {
+	return db.db.NewIterator(util.BytesPrefix(prefix), nil)
+}
+
 func (db *GoLevelDB) NewBatch() Batch {
 	batch := new(leveldb.Batch)
 	return &goLevelDBBatch{db, batch}
diff --git a/db/mem_db.go b/db/mem_db.go
index 04d48cac0..561633291 100644
--- a/db/mem_db.go
+++ b/db/mem_db.go
@@ -2,6 +2,7 @@ package db
 
 import (
 	"fmt"
+	"strings"
 	"sync"
 )
 
@@ -108,6 +109,10 @@ func (it *memDBIterator) Error() error {
 }
 
 func (db *MemDB) Iterator() Iterator {
+	return db.IteratorPrefix([]byte{})
+}
+
+func (db *MemDB) IteratorPrefix(prefix []byte) Iterator {
 	it := newMemDBIterator()
 	it.db = db
 	it.last = -1
@@ -117,7 +122,9 @@ func (db *MemDB) Iterator() Iterator {
 
 	// unfortunately we need a copy of all of the keys
 	for key, _ := range db.db {
-		it.keys = append(it.keys, key)
+		if strings.HasPrefix(key, string(prefix)) {
+			it.keys = append(it.keys, key)
+		}
 	}
 	return it
 }