When shrinking a tree, LeafTraverser takes care of decreasing tree depth, not DataTree anymore. This way, LeafTraverser is now responsible for increasing and decreasing depth as necessary.

src/blobstore/implementations/onblocks/datatreestore/DataTree.cpp

@@ -46,17 +46,6 @@ DataTree::DataTree(DataNodeStore *nodeStore, unique_ref<DataNode> rootNode)
 DataTree::~DataTree() {
 }
 
-void DataTree::whileRootHasOnlyOneChildReplaceRootWithItsChild() {
-  DataInnerNode *rootNode = dynamic_cast<DataInnerNode*>(_rootNode.get());
-  if (rootNode != nullptr && rootNode->numChildren() == 1) {
-    auto child = _nodeStore->load(rootNode->getChild(0)->key());
-    ASSERT(child != none, "Couldn't load first child of root node");
-    _rootNode = _nodeStore->overwriteNodeWith(std::move(_rootNode), **child);
-    _nodeStore->remove(std::move(*child));
-    whileRootHasOnlyOneChildReplaceRootWithItsChild();
-  }
-}
-
 const Key &DataTree::key() const {
   return _rootNode->key();
 }
@@ -184,7 +173,6 @@ void DataTree::resizeNumBytes(uint64_t newNumBytes) {
   };
 
   _traverseLeaves(newNumLeaves - 1, newNumLeaves, onExistingLeaf, onCreateLeaf, onBacktrackFromSubtree);
-  whileRootHasOnlyOneChildReplaceRootWithItsChild();
   _numLeavesCache = newNumLeaves;
   ASSERT(newNumBytes == _numStoredBytes(), "We resized to the wrong number of bytes ("+std::to_string(numStoredBytes())+" instead of "+std::to_string(newNumBytes)+")");
 }

src/blobstore/implementations/onblocks/datatreestore/DataTree.h

@@ -50,8 +50,6 @@ private:
   cpputils::unique_ref<datanodestore::DataNode> releaseRootNode();
   friend class DataTreeStore;
 
-  void whileRootHasOnlyOneChildReplaceRootWithItsChild();
-
   //TODO Use underscore for private methods
   void _traverseLeaves(uint32_t beginIndex, uint32_t endIndex,
                        std::function<void (uint32_t index, datanodestore::DataLeafNode* leaf)> onExistingLeaf,

src/blobstore/implementations/onblocks/datatreestore/impl/LeafTraverser.cpp

@@ -45,7 +45,9 @@ namespace blobstore {
                     auto newRoot = _increaseTreeDepth(std::move(root));
                     return traverseAndReturnRoot(std::move(newRoot), std::max(beginIndex, maxLeavesForDepth), endIndex, onExistingLeaf, onCreateLeaf, onBacktrackFromSubtree);
                 } else {
-                    return std::move(root);
+                    // Once we're done growing the tree and done with the traversal, we might have to decrease tree depth,
+                    // because the callbacks could have deleted nodes (this happens for example when shrinking the tree using a traversal).
+                    return whileRootHasOnlyOneChildReplaceRootWithItsChild(std::move(root));
                 }
             }
 
@@ -192,6 +194,19 @@ 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));
+                } else {
+                    return root;
+                }
+            }
+
         }
     }
 }

src/blobstore/implementations/onblocks/datatreestore/impl/LeafTraverser.h

@@ -44,6 +44,7 @@ namespace blobstore {
                                                                                 std::function<void (datanodestore::DataInnerNode *node)> onBacktrackFromSubtree);
                 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);
 
                 DISALLOW_COPY_AND_ASSIGN(LeafTraverser);
             };