Merge pull request #2828 from yahoo/toregge/minor-refactor-btreenode

Toregge/minor refactor btreenode

3 files changed, 37 insertions, 51 deletions

diff --git a/searchlib/src/tests/btree/btreeaggregation_test.cpp b/searchlib/src/tests/btree/btreeaggregation_test.cpp
index e07de5cd9d5..50ef7f21c4c 100644
--- a/searchlib/src/tests/btree/btreeaggregation_test.cpp
+++ b/searchlib/src/tests/btree/btreeaggregation_test.cpp
@@ -588,11 +588,11 @@ Test::requireThatNodeStealWorks()
                                "[60:160,70:170,80:180[min=160,max=180]]"
                                "[min=110,max=180]]", t));
         t.remove(60);
-        EXPECT_TRUE(assertTree("[40:"
-                               "[10:110,20:120,30:130,40:140"
-                               "[min=110,max=140]]"
+        EXPECT_TRUE(assertTree("[30:"
+                               "[10:110,20:120,30:130"
+                               "[min=110,max=130]]"
                                ",80:"
-                               "[50:150,70:170,80:180[min=150,max=180]]"
+                               "[40:140,50:150,70:170,80:180[min=140,max=180]]"
                                "[min=110,max=180]]", t));
     }
     { // steal some from right
@@ -615,12 +615,12 @@ Test::requireThatNodeStealWorks()
                                "[min=150,max=190]]"
                                "[min=110,max=190]]", t));
         t.remove(20);
-        EXPECT_TRUE(assertTree("[50:"
-                               "[10:110,30:130,50:150"
-                               "[min=110,max=150]]"
+        EXPECT_TRUE(assertTree("[60:"
+                               "[10:110,30:130,50:150,60:160"
+                               "[min=110,max=160]]"
                                ",90:"
-                               "[60:160,70:170,80:180,90:190"
-                               "[min=160,max=190]]"
+                               "[70:170,80:180,90:190"
+                               "[min=170,max=190]]"
                                "[min=110,max=190]]", t));
     }
 }
diff --git a/searchlib/src/vespa/searchlib/btree/btreenode.h b/searchlib/src/vespa/searchlib/btree/btreenode.h
index 13fccf72623..94b0016e1b5 100644
--- a/searchlib/src/vespa/searchlib/btree/btreenode.h
+++ b/searchlib/src/vespa/searchlib/btree/btreenode.h
@@ -398,6 +398,9 @@ private:
         return *this;
     }
 
+    template <typename NodeAllocatorType>
+    uint32_t countValidLeaves(uint32_t start, uint32_t end, NodeAllocatorType &allocator);
+
 public:
     BTreeNode::Ref getChild(uint32_t idx) const { return getData(idx); }
     void setChild(uint32_t idx, BTreeNode::Ref child) { setData(idx, child); }
diff --git a/searchlib/src/vespa/searchlib/btree/btreenode.hpp b/searchlib/src/vespa/searchlib/btree/btreenode.hpp
index f307623a8e7..3523641705b 100644
--- a/searchlib/src/vespa/searchlib/btree/btreenode.hpp
+++ b/searchlib/src/vespa/searchlib/btree/btreenode.hpp
@@ -170,7 +170,7 @@ stealSomeFromLeftNode(NodeType *victim)
     assert(validSlots() + victim->validSlots() >= NodeType::minSlots());
     assert(!getFrozen());
     assert(!victim->getFrozen());
-    uint32_t median = (validSlots() + victim->validSlots()) / 2;
+    uint32_t median = (validSlots() + victim->validSlots() + 1) / 2;
     uint32_t steal = median - validSlots();
     _validSlots += steal;
     for (int32_t i = validSlots() - 1; i >= static_cast<int32_t>(steal); --i) {
@@ -193,7 +193,7 @@ stealSomeFromRightNode(NodeType *victim)
     assert(validSlots() + victim->validSlots() >= NodeType::minSlots());
     assert(!getFrozen());
     assert(!victim->getFrozen());
-    uint32_t median = (validSlots() + victim->validSlots()) / 2;
+    uint32_t median = (validSlots() + victim->validSlots() + 1) / 2;
     uint32_t steal = median - validSlots();
     for (uint32_t i = 0; i < steal; ++i) {
         _keys[validSlots() + i] = victim->_keys[i];
@@ -307,6 +307,19 @@ stealAllFromRightNode(const BTreeInternalNode *victim)
     _validLeaves += victim->_validLeaves;
 }
 
+template <typename KeyT, typename AggrT, uint32_t NumSlots>
+template <typename NodeAllocatorType>
+uint32_t
+BTreeInternalNode<KeyT, AggrT, NumSlots>::countValidLeaves(uint32_t start, uint32_t end, NodeAllocatorType &allocator)
+{
+    assert(start <= end);
+    assert(end <= validSlots());
+    uint32_t leaves = 0;
+    for (uint32_t i = start; i < end; ++i) {
+        leaves += allocator.validLeaves(getData(i));
+    }
+    return leaves;
+}
 
 template <typename KeyT, typename AggrT, uint32_t NumSlots>
 template <typename NodeAllocatorType>
@@ -314,26 +327,11 @@ void
 BTreeInternalNode<KeyT, AggrT, NumSlots>::
 stealSomeFromLeftNode(BTreeInternalNode *victim, NodeAllocatorType &allocator)
 {
-    assert(validSlots() + victim->validSlots() >= BTreeInternalNode::minSlots());
-    assert(!getFrozen());
-    assert(!victim->getFrozen());
-    uint32_t median = (validSlots() + victim->validSlots()) / 2;
-    uint32_t steal = median - validSlots();
-    _validSlots += steal;
-    for (int32_t i = validSlots() - 1; i >= static_cast<int32_t>(steal); --i) {
-        _keys[i] = _keys[i - steal];
-        setData(i, getData(i - steal));
-    }
-    uint32_t stolenLeaves = 0;
-    for (uint32_t i = 0; i < steal; ++i) {
-        _keys[i] = victim->_keys[victim->validSlots() - steal + i];
-        setData(i, victim->getData(victim->validSlots() - steal + i));
-        stolenLeaves += allocator.validLeaves(getData(i));
-    }
-    _validLeaves += stolenLeaves;
-    victim->_validLeaves -= stolenLeaves;
-    victim->cleanRange(victim->validSlots() - steal, victim->validSlots());
-    victim->_validSlots -= steal;
+    uint16_t oldValidSlots = validSlots();
+    ParentType::stealSomeFromLeftNode(victim);
+    uint32_t stolenLeaves = countValidLeaves(0, validSlots() - oldValidSlots, allocator);
+    incValidLeaves(stolenLeaves);
+    victim->decValidLeaves(stolenLeaves);
 }
 
 
@@ -343,26 +341,11 @@ void
 BTreeInternalNode<KeyT, AggrT, NumSlots>::
 stealSomeFromRightNode(BTreeInternalNode *victim, NodeAllocatorType &allocator)
 {
-    assert(validSlots() + victim->validSlots() >= BTreeInternalNode::minSlots());
-    assert(!getFrozen());
-    assert(!victim->getFrozen());
-    uint32_t median = (validSlots() + victim->validSlots()) / 2;
-    uint32_t steal = median - validSlots();
-    uint32_t stolenLeaves = 0;
-    for (uint32_t i = 0; i < steal; ++i) {
-        _keys[validSlots() + i] = victim->_keys[i];
-        setData(validSlots() + i, victim->getData(i));
-        stolenLeaves += allocator.validLeaves(victim->getData(i));
-    }
-    _validSlots += steal;
-    _validLeaves += stolenLeaves;
-    victim->_validLeaves -= stolenLeaves;
-    for (uint32_t i = steal; i < victim->validSlots(); ++i) {
-        victim->_keys[i - steal] = victim->_keys[i];
-        victim->setData(i - steal, victim->getData(i));
-    }
-    victim->cleanRange(victim->validSlots() - steal, victim->validSlots());
-    victim->_validSlots -= steal;
+    uint16_t oldValidSlots = validSlots();
+    ParentType::stealSomeFromRightNode(victim);
+    uint32_t stolenLeaves = countValidLeaves(oldValidSlots, validSlots(), allocator);
+    incValidLeaves(stolenLeaves);
+    victim->decValidLeaves(stolenLeaves);
 }