From 0b9be1069fdde1810e17bec7ac9e264a82182440 Mon Sep 17 00:00:00 2001
From: Ethan Frey <ethanfrey@users.noreply.github.com>
Date: Thu, 3 Aug 2017 16:58:37 +0200
Subject: [PATCH] Cleanup some todos

---
 app/store.go     |  4 ++--
 state/bonsai.go  | 13 +++++++++++--
 state/kvcache.go |  7 ++++++-
 state/kvstore.go |  7 ++++++-
 state/merkle.go  |  4 ++++
 5 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/app/store.go b/app/store.go
index 4c3bc8ada..adf512806 100644
--- a/app/store.go
+++ b/app/store.go
@@ -112,7 +112,7 @@ func (s *Store) Info() abci.ResponseInfo {
 		"height", s.height,
 		"hash", fmt.Sprintf("%X", s.hash))
 	return abci.ResponseInfo{
-		Data:             cmn.Fmt("size:%v", s.State.Committed().Size()),
+		Data:             cmn.Fmt("size:%v", s.State.Size()),
 		LastBlockHeight:  s.height - 1,
 		LastBlockAppHash: s.hash,
 	}
@@ -144,7 +144,7 @@ func (s *Store) Commit() abci.Result {
 		return abci.NewError(abci.CodeType_InternalError, "AppHash is incorrect")
 	}
 
-	if s.State.Committed().Size() == 0 {
+	if s.State.Size() == 0 {
 		return abci.NewResultOK(nil, "Empty hash for empty tree")
 	}
 	return abci.NewResultOK(s.hash, "")
diff --git a/state/bonsai.go b/state/bonsai.go
index 2a9873298..9800ba54f 100644
--- a/state/bonsai.go
+++ b/state/bonsai.go
@@ -11,8 +11,8 @@ type nonce int64
 
 // Bonsai is a deformed tree forced to fit in a small pot
 type Bonsai struct {
-	id nonce
-	merkle.Tree
+	id   nonce
+	Tree merkle.Tree
 }
 
 var _ SimpleDB = &Bonsai{}
@@ -31,6 +31,11 @@ func (b *Bonsai) Get(key []byte) []byte {
 	return value
 }
 
+// Get matches the signature of KVStore
+func (b *Bonsai) Has(key []byte) bool {
+	return b.Tree.Has(key)
+}
+
 // Set matches the signature of KVStore
 func (b *Bonsai) Set(key, value []byte) {
 	b.Tree.Set(key, value)
@@ -41,6 +46,10 @@ func (b *Bonsai) Remove(key []byte) (value []byte) {
 	return
 }
 
+func (b *Bonsai) Proof(key []byte) (value []byte, proof []byte, exists bool) {
+	return b.Tree.Proof(key)
+}
+
 func (b *Bonsai) List(start, end []byte, limit int) []Model {
 	res := []Model{}
 	stopAtCount := func(key []byte, value []byte) (stop bool) {
diff --git a/state/kvcache.go b/state/kvcache.go
index 1141d20e3..6fcb176f3 100644
--- a/state/kvcache.go
+++ b/state/kvcache.go
@@ -117,7 +117,12 @@ func (c *MemKVCache) Commit(sub SimpleDB) error {
 	if !ok {
 		return ErrNotASubTransaction()
 	}
-	// TODO: see if it points to us
+
+	// see if it points to us
+	ref, ok := cache.store.(*MemKVCache)
+	if !ok || ref != c {
+		return ErrNotASubTransaction()
+	}
 
 	// apply the cached data to us
 	cache.applyCache()
diff --git a/state/kvstore.go b/state/kvstore.go
index 08dd85f21..9af9b5753 100644
--- a/state/kvstore.go
+++ b/state/kvstore.go
@@ -149,7 +149,12 @@ func (m *MemKVStore) Commit(sub SimpleDB) error {
 	if !ok {
 		return ErrNotASubTransaction()
 	}
-	// TODO: see if it points to us
+
+	// see if it points to us
+	ref, ok := cache.store.(*MemKVStore)
+	if !ok || ref != m {
+		return ErrNotASubTransaction()
+	}
 
 	// apply the cached data to us
 	cache.applyCache()
diff --git a/state/merkle.go b/state/merkle.go
index f383fd0d5..666cac0bc 100644
--- a/state/merkle.go
+++ b/state/merkle.go
@@ -24,6 +24,10 @@ func NewState(tree merkle.Tree, persistent bool) State {
 	}
 }
 
+func (s State) Size() int {
+	return s.committed.Tree.Size()
+}
+
 func (s State) Committed() *Bonsai {
 	return s.committed
 }