From b67f6a0dc6fcc09b219cc2773f589d6171368820 Mon Sep 17 00:00:00 2001
From: Alexander Block <ablock84@gmail.com>
Date: Thu, 4 Apr 2019 10:03:56 +0200
Subject: [PATCH] Implement CDBTransactionIterator

This iterator allows merged iteration of the key/values from the parent and
the not-yet-committed key/values from the transaction. This also works for
nested transactions (as used in CEvoDB).

It's interface mimics CDBIterator.
---
 src/dbwrapper.h | 160 +++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 158 insertions(+), 2 deletions(-)

diff --git a/src/dbwrapper.h b/src/dbwrapper.h
index 94cfc9cfb7bab..94db323dc70c8 100644
--- a/src/dbwrapper.h
+++ b/src/dbwrapper.h
@@ -159,9 +159,8 @@ class CDBIterator
     void Next();
 
     template<typename K> bool GetKey(K& key) {
-        leveldb::Slice slKey = piter->key();
         try {
-            CDataStream ssKey(slKey.data(), slKey.data() + slKey.size(), SER_DISK, CLIENT_VERSION);
+            CDataStream ssKey = GetKey();
             ssKey >> key;
         } catch (const std::exception&) {
             return false;
@@ -169,6 +168,11 @@ class CDBIterator
         return true;
     }
 
+    CDataStream GetKey() {
+        leveldb::Slice slKey = piter->key();
+        return CDataStream(slKey.data(), slKey.data() + slKey.size(), SER_DISK, CLIENT_VERSION);
+    }
+
     unsigned int GetKeySize() {
         return piter->key().size();
     }
@@ -389,8 +393,153 @@ class CDBWrapper
 
 };
 
+template<typename CDBTransaction>
+class CDBTransactionIterator
+{
+private:
+    CDBTransaction& transaction;
+
+    typedef typename std::remove_pointer<decltype(transaction.parent.NewIterator())>::type ParentIterator;
+
+    // We maintain 2 iterators, one for the transaction and one for the parent
+    // At all times, only one of both provides the current value. The decision is made by comparing the current keys
+    // of both iterators, so that always the smaller key is the current one. On Next(), the previously chosen iterator
+    // is advanced.
+    typename CDBTransaction::WritesMap::iterator transactionIt;
+    std::unique_ptr<ParentIterator> parentIt;
+    CDataStream parentKey;
+    bool curIsParent{false};
+
+public:
+    CDBTransactionIterator(CDBTransaction& _transaction) :
+            transaction(_transaction),
+            parentKey(SER_DISK, CLIENT_VERSION)
+    {
+        transactionIt = transaction.writes.end();
+        parentIt = std::unique_ptr<ParentIterator>(transaction.parent.NewIterator());
+    }
+
+    void SeekToFirst() {
+        transactionIt = transaction.writes.begin();
+        parentIt->SeekToFirst();
+        SkipDeletedAndOverwritten();
+        DecideCur();
+    }
+
+    template<typename K>
+    void Seek(const K& key) {
+        Seek(CDBTransaction::KeyToDataStream(key));
+    }
+
+    void Seek(const CDataStream& ssKey) {
+        transactionIt = transaction.writes.lower_bound(ssKey);
+        parentIt->Seek(ssKey);
+        SkipDeletedAndOverwritten();
+        DecideCur();
+    }
+
+    bool Valid() {
+        return transactionIt != transaction.writes.end() || parentIt->Valid();
+    }
+
+    void Next() {
+        if (transactionIt == transaction.writes.end() && !parentIt->Valid()) {
+            return;
+        }
+        if (curIsParent) {
+            assert(parentIt->Valid());
+            parentIt->Next();
+            SkipDeletedAndOverwritten();
+        } else {
+            assert(transactionIt != transaction.writes.end());
+            ++transactionIt;
+        }
+        DecideCur();
+    }
+
+    template<typename K>
+    bool GetKey(K& key) {
+        if (!Valid()) {
+            return false;
+        }
+
+        if (curIsParent) {
+            return parentIt->GetKey(key);
+        } else {
+            try {
+                // TODO try to avoid this copy (we need a stream that allows reading from external buffers)
+                CDataStream ssKey = transactionIt->first;
+                ssKey >> key;
+            } catch (const std::exception&) {
+                return false;
+            }
+            return true;
+        }
+    }
+
+    CDataStream GetKey() {
+        if (!Valid()) {
+            return CDataStream(SER_DISK, CLIENT_VERSION);
+        }
+        if (curIsParent) {
+            return parentIt->GetKey();
+        } else {
+            return transactionIt->first;
+        }
+    }
+
+    unsigned int GetKeySize() {
+        if (!Valid()) {
+            return 0;
+        }
+        if (curIsParent) {
+            return parentIt->GetKeySize();
+        } else {
+            return transactionIt->first.vKey.size();
+        }
+    }
+
+    template<typename V>
+    bool GetValue(V& value) {
+        if (!Valid()) {
+            return false;
+        }
+        if (curIsParent) {
+            return transaction.Read(parentKey, value);
+        } else {
+            return transaction.Read(transactionIt->first, value);
+        }
+    };
+
+private:
+    void SkipDeletedAndOverwritten() {
+        while (parentIt->Valid()) {
+            parentKey = parentIt->GetKey();
+            if (!transaction.deletes.count(parentKey) && !transaction.writes.count(parentKey)) {
+                break;
+            }
+        }
+    }
+
+    void DecideCur() {
+        if (transactionIt != transaction.writes.end() && !parentIt->Valid()) {
+            curIsParent = false;
+        } else if (transactionIt == transaction.writes.end() && parentIt->Valid()) {
+            curIsParent = true;
+        } else if (transactionIt != transaction.writes.end() && parentIt->Valid()) {
+            if (CDBTransaction::DataStreamCmp::less(transactionIt->first, parentKey)) {
+                curIsParent = false;
+            } else {
+                curIsParent = true;
+            }
+        }
+    }
+};
+
 template<typename Parent, typename CommitTarget>
 class CDBTransaction {
+    friend class CDBTransactionIterator<CDBTransaction>;
+
 protected:
     Parent &parent;
     CommitTarget &commitTarget;
@@ -520,6 +669,13 @@ class CDBTransaction {
     bool IsClean() {
         return writes.empty() && deletes.empty();
     }
+
+    CDBTransactionIterator<CDBTransaction>* NewIterator() {
+        return new CDBTransactionIterator<CDBTransaction>(*this);
+    }
+    std::unique_ptr<CDBTransactionIterator<CDBTransaction>> NewIteratorUniquePtr() {
+        return std::make_unique<CDBTransactionIterator<CDBTransaction>>(*this);
+    }
 };
 
 template<typename Parent, typename CommitTarget>