From 582917c1f51ab8d54f9d30f8b5ffa684a5d8f16c Mon Sep 17 00:00:00 2001
From: Sebastian Messmer <messmer@cryfs.org>
Date: Fri, 15 Jul 2016 16:23:00 +0200
Subject: [PATCH] Improve code decreasing tree depth

---
 .../datatreestore/impl/LeafTraverser.cpp      | 25 +++++++++++++++----
 .../datatreestore/impl/LeafTraverser.h        |  2 ++
 .../DataTreeTest_Performance.cpp              |  2 +-
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/src/blobstore/implementations/onblocks/datatreestore/impl/LeafTraverser.cpp b/src/blobstore/implementations/onblocks/datatreestore/impl/LeafTraverser.cpp
index 48b10b43..0ec9ecfc 100644
--- a/src/blobstore/implementations/onblocks/datatreestore/impl/LeafTraverser.cpp
+++ b/src/blobstore/implementations/onblocks/datatreestore/impl/LeafTraverser.cpp
@@ -10,6 +10,7 @@ using std::vector;
 using boost::none;
 using cpputils::Data;
 using cpputils::unique_ref;
+using cpputils::dynamic_pointer_move;
 using blobstore::onblocks::datanodestore::DataNodeStore;
 using blobstore::onblocks::datanodestore::DataNode;
 using blobstore::onblocks::datanodestore::DataInnerNode;
@@ -196,16 +197,30 @@ namespace blobstore {
             unique_ref<DataNode> LeafTraverser::_whileRootHasOnlyOneChildReplaceRootWithItsChild(unique_ref<DataNode> root) {
                 DataInnerNode *inner = dynamic_cast<DataInnerNode*>(root.get());
                 if (inner != nullptr && inner->numChildren() == 1) {
-                    auto child = _nodeStore->load(inner->getChild(0)->key());
-                    ASSERT(child != none, "Couldn't load first child of root node");
-                    auto newRoot = _nodeStore->overwriteNodeWith(std::move(root), **child);
-                    _nodeStore->remove(std::move(*child));
-                    return _whileRootHasOnlyOneChildReplaceRootWithItsChild(std::move(newRoot));
+                    auto newRoot = _whileRootHasOnlyOneChildRemoveRootReturnChild(inner->getChild(0)->key());
+                    auto result = _nodeStore->overwriteNodeWith(std::move(root), *newRoot);
+                    _nodeStore->remove(std::move(newRoot));
+                    return result;
                 } else {
                     return root;
                 }
             }
 
+            unique_ref<DataNode> LeafTraverser::_whileRootHasOnlyOneChildRemoveRootReturnChild(const blockstore::Key &key) {
+                auto current = _nodeStore->load(key);
+                ASSERT(current != none, "Node not found");
+                auto inner = dynamic_pointer_move<DataInnerNode>(*current);
+                if (inner == none) {
+                    return std::move(*current);
+                } else if ((*inner)->numChildren() == 1) {
+                    auto result = _whileRootHasOnlyOneChildRemoveRootReturnChild((*inner)->getChild(0)->key());
+                    _nodeStore->remove(std::move(*inner));
+                    return result;
+                } else {
+                    return std::move(*inner);
+                }
+            }
+
         }
     }
 }
diff --git a/src/blobstore/implementations/onblocks/datatreestore/impl/LeafTraverser.h b/src/blobstore/implementations/onblocks/datatreestore/impl/LeafTraverser.h
index db10330c..56e0e8fe 100644
--- a/src/blobstore/implementations/onblocks/datatreestore/impl/LeafTraverser.h
+++ b/src/blobstore/implementations/onblocks/datatreestore/impl/LeafTraverser.h
@@ -5,6 +5,7 @@
 #include <cpp-utils/macros.h>
 #include <cpp-utils/pointer/unique_ref.h>
 #include <cpp-utils/data/Data.h>
+#include <blockstore/utils/Key.h>
 
 namespace blobstore {
     namespace onblocks {
@@ -50,6 +51,7 @@ namespace blobstore {
                 uint32_t _maxLeavesForTreeDepth(uint8_t depth) const;
                 std::function<cpputils::Data (uint32_t index)> _createMaxSizeLeaf() const;
                 cpputils::unique_ref<datanodestore::DataNode> _whileRootHasOnlyOneChildReplaceRootWithItsChild(cpputils::unique_ref<datanodestore::DataNode> root);
+                cpputils::unique_ref<datanodestore::DataNode> _whileRootHasOnlyOneChildRemoveRootReturnChild(const blockstore::Key &key);
 
                 DISALLOW_COPY_AND_ASSIGN(LeafTraverser);
             };
diff --git a/test/blobstore/implementations/onblocks/datatreestore/DataTreeTest_Performance.cpp b/test/blobstore/implementations/onblocks/datatreestore/DataTreeTest_Performance.cpp
index b4ce81f0..3a5518e4 100644
--- a/test/blobstore/implementations/onblocks/datatreestore/DataTreeTest_Performance.cpp
+++ b/test/blobstore/implementations/onblocks/datatreestore/DataTreeTest_Performance.cpp
@@ -464,6 +464,6 @@ TEST_F(DataTreeTest_Performance, ResizeNumBytes_DecreaseTreeDepth_2to0) {
     EXPECT_EQ(5u, blockStore->loadedBlocks().size()); // load new last leaf (+inner node), load second inner node to remove its subtree, then 2x load first child of root to replace root with its child.
     EXPECT_EQ(0u, blockStore->createdBlocks());
     EXPECT_EQ(3u + maxChildrenPerInnerNode, blockStore->removedBlocks().size());
-    EXPECT_EQ(2u, blockStore->distinctWrittenBlocks().size()); // 2x rewrite root node to be a leaf
+    EXPECT_EQ(2u, blockStore->distinctWrittenBlocks().size()); // remove children from inner node and rewrite root node to be a leaf
     EXPECT_EQ(0u, blockStore->resizedBlocks().size());
 }