Implement TreeSortedMap::insert (#1081)

* Make LlrbNode Rep more explicit, share empty node

* SortedMap::insert converts implementations

* Implement LlrbNode::insert

* Remove TestPolicy<SortedMap>

3 files changed, 232 insertions, 25 deletions

diff --git a/Firestore/core/test/firebase/firestore/immutable/array_sorted_map_test.cc b/Firestore/core/test/firebase/firestore/immutable/array_sorted_map_test.cc
index 6758dd5..9f18f2d 100644
--- a/Firestore/core/test/firebase/firestore/immutable/array_sorted_map_test.cc
+++ b/Firestore/core/test/firebase/firestore/immutable/array_sorted_map_test.cc
@@ -32,11 +32,6 @@ namespace impl {
 using IntMap = ArraySortedMap<int, int>;
 constexpr IntMap::size_type kFixedSize = IntMap::kFixedSize;
 
-// TODO(wilhuff): ReverseTraversal
-
-#define ASSERT_SEQ_EQ(x, y) ASSERT_EQ((x), Append(y));
-#define EXPECT_SEQ_EQ(x, y) EXPECT_EQ((x), Append(y));
-
 TEST(ArraySortedMap, SearchForSpecificKey) {
   IntMap map{{1, 3}, {2, 4}};
 
@@ -105,21 +100,6 @@ TEST(ArraySortedMap, RemovesMiddle) {
   ASSERT_TRUE(Found(s1, 3, 3));
 }
 
-TEST(ArraySortedMap, Increasing) {
-  auto total = static_cast<int>(kFixedSize);
-  IntMap map;
-
-  for (int i = 0; i < total; i++) {
-    map = map.insert(i, i);
-  }
-  ASSERT_EQ(kFixedSize, map.size());
-
-  for (int i = 0; i < total; i++) {
-    map = map.erase(i);
-  }
-  ASSERT_EQ(0u, map.size());
-}
-
 TEST(ArraySortedMap, Override) {
   IntMap map = IntMap{}.insert(10, 10).insert(10, 8);
 
diff --git a/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc b/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc
index 747c66b..3253509 100644
--- a/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc
+++ b/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc
@@ -18,6 +18,7 @@
 
 #include <numeric>
 #include <random>
+#include <unordered_set>
 
 #include "Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h"
 #include "Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h"
@@ -36,16 +37,24 @@ using SizeType = SortedMapBase::size_type;
 
 template <typename MapType>
 struct TestPolicy {
-  // TODO(mcg): increase beyond what ArraySortedMap supports
-  static const SizeType kFixedSize = SortedMapBase::kFixedSize;
+  static const SizeType kLargeSize = 100;
+};
+
+template <>
+struct TestPolicy<impl::ArraySortedMap<int, int>> {
+  // ArraySortedMap cannot insert more than this number
+  static const SizeType kLargeSize = SortedMapBase::kFixedSize;
 };
 
 template <typename IntMap>
 class SortedMapTest : public ::testing::Test {
  public:
-  template <typename Integer = SizeType>
-  Integer fixed_size() {
-    return static_cast<Integer>(TestPolicy<IntMap>::kFixedSize);
+  SortedMapBase::size_type large_size() const {
+    return TestPolicy<IntMap>::kLargeSize;
+  }
+
+  int large_number() const {
+    return static_cast<int>(large_size());
   }
 };
 
@@ -72,6 +81,37 @@ TYPED_TEST(SortedMapTest, Empty) {
   // EXPECT_TRUE(NotFound(map, 10));
 }
 
+TYPED_TEST(SortedMapTest, Size) {
+  std::mt19937 rand;
+  std::uniform_int_distribution<int> dist(0, 999);
+
+  std::unordered_set<int> expected;
+
+  TypeParam map;
+  auto n = this->large_number();
+  for (int i = 0; i < n; ++i) {
+    int value = dist(rand);
+
+    // The random number sequence can generate duplicates, so the expected size
+    // won't necessarily depend upon `i`.
+    expected.insert(value);
+
+    map = map.insert(value, value);
+    EXPECT_EQ(expected.size(), map.size());
+  }
+}
+
+TYPED_TEST(SortedMapTest, Increasing) {
+  std::vector<int> to_insert = Sequence(this->large_number());
+  TypeParam map = ToMap<TypeParam>(to_insert);
+  ASSERT_EQ(this->large_size(), map.size());
+
+  for (int i : to_insert) {
+    map = map.erase(i);
+  }
+  ASSERT_EQ(0u, map.size());
+}
+
 }  // namespace immutable
 }  // namespace firestore
 }  // namespace firebase
diff --git a/Firestore/core/test/firebase/firestore/immutable/tree_sorted_map_test.cc b/Firestore/core/test/firebase/firestore/immutable/tree_sorted_map_test.cc
index c03dc6c..4f9ad3e 100644
--- a/Firestore/core/test/firebase/firestore/immutable/tree_sorted_map_test.cc
+++ b/Firestore/core/test/firebase/firestore/immutable/tree_sorted_map_test.cc
@@ -34,6 +34,193 @@ TEST(TreeSortedMap, EmptySize) {
   EXPECT_EQ(Color::Black, map.root().color());
 }
 
+TEST(TreeSortedMap, EmptyHasEmptyChildren) {
+  IntMap map;
+  IntMap::node_type left = map.root().left();
+  ASSERT_TRUE(left.empty());
+
+  IntMap::node_type right = map.root().right();
+  ASSERT_TRUE(right.empty());
+}
+
+TEST(TreeSortedMap, PropertiesForEmpty) {
+  IntMap empty;
+  EXPECT_TRUE(empty.empty());
+  EXPECT_EQ(0, empty.root().value());
+
+  EXPECT_EQ(Color::Black, empty.root().color());
+  EXPECT_FALSE(empty.root().red());
+}
+
+TEST(TreeSortedMap, PropertiesForNonEmpty) {
+  IntMap empty;
+
+  IntMap non_empty = empty.insert(1, 2);
+  EXPECT_FALSE(non_empty.empty());
+  EXPECT_EQ(1, non_empty.root().key());
+  EXPECT_EQ(2, non_empty.root().value());
+
+  // Root nodes are always black
+  EXPECT_EQ(Color::Black, non_empty.root().color());
+  EXPECT_FALSE(non_empty.root().red());
+  EXPECT_TRUE(non_empty.root().left().empty());
+  EXPECT_TRUE(non_empty.root().right().empty());
+}
+
+TEST(TreeSortedMap, RotatesLeft) {
+  IntMap map;
+  EXPECT_EQ(Color::Black, map.root().color());
+
+  // root node, with two empty children
+  map = map.insert(1, 1);
+  EXPECT_EQ(Color::Black, map.root().color());
+  EXPECT_EQ(Color::Black, map.root().left().color());
+  EXPECT_EQ(Color::Black, map.root().right().color());
+
+  // insert successor, leans left, rotation required
+  map = map.insert(2, 2);
+  EXPECT_EQ(Color::Black, map.root().color());
+  EXPECT_EQ(Color::Red, map.root().left().color());
+  EXPECT_EQ(Color::Black, map.root().right().color());
+
+  // insert successor, balanced, color flip required
+  map = map.insert(3, 3);
+  EXPECT_EQ(2, map.root().key());
+  EXPECT_EQ(Color::Black, map.root().color());
+  EXPECT_EQ(Color::Black, map.root().left().color());
+  EXPECT_EQ(Color::Black, map.root().right().color());
+}
+
+TEST(TreeSortedMap, RotatesLeftWithSubtree) {
+  // Build an initially balanced, all black tree
+  IntMap map;
+  map = map.insert(5, 5);
+  map = map.insert(3, 3);
+  map = map.insert(7, 7);
+  EXPECT_EQ(Color::Black, map.root().color());
+  EXPECT_EQ(Color::Black, map.root().left().color());
+  EXPECT_EQ(Color::Black, map.root().right().color());
+
+  // left child of right, no rotation yet
+  map = map.insert(6, 6);
+  EXPECT_EQ(5, map.root().key());
+  EXPECT_EQ(6, map.root().right().left().key());
+  EXPECT_EQ(Color::Red, map.root().right().left().color());
+
+  // right child of right, triggers a color flip in the right node and forces
+  // a left rotation of the root
+  map = map.insert(8, 8);
+  EXPECT_EQ(7, map.root().key());
+  EXPECT_EQ(Color::Black, map.root().color());
+
+  EXPECT_EQ(5, map.root().left().key());
+  EXPECT_EQ(Color::Red, map.root().left().color());
+
+  EXPECT_EQ(3, map.root().left().left().key());
+  EXPECT_EQ(Color::Black, map.root().left().left().color());
+
+  EXPECT_EQ(6, map.root().left().right().key());
+  EXPECT_EQ(Color::Black, map.root().left().right().color());
+
+  EXPECT_EQ(8, map.root().right().key());
+  EXPECT_EQ(Color::Black, map.root().right().color());
+}
+
+TEST(TreeSortedMap, RotatesRight) {
+  IntMap map;
+  EXPECT_EQ(Color::Black, map.root().color());
+
+  // root node, with two empty children
+  map = map.insert(3, 3);
+  EXPECT_EQ(Color::Black, map.root().color());
+  EXPECT_EQ(Color::Black, map.root().left().color());
+  EXPECT_EQ(Color::Black, map.root().right().color());
+
+  // insert predecessor, leans left, no rotation
+  map = map.insert(2, 2);
+  EXPECT_EQ(Color::Black, map.root().color());
+  EXPECT_EQ(Color::Red, map.root().left().color());
+  EXPECT_EQ(Color::Black, map.root().right().color());
+
+  // insert predecessor, rotation required
+  map = map.insert(1, 2);
+  EXPECT_EQ(2, map.root().key());
+  EXPECT_EQ(Color::Black, map.root().color());
+  EXPECT_EQ(Color::Black, map.root().left().color());
+  EXPECT_EQ(Color::Black, map.root().right().color());
+}
+
+TEST(TreeSortedMap, RotatesRightWithSubtree) {
+  // Build an initially balanced, all black tree
+  IntMap map;
+  map = map.insert(5, 5);
+  map = map.insert(3, 3);
+  map = map.insert(7, 7);
+  EXPECT_EQ(Color::Black, map.root().color());
+  EXPECT_EQ(Color::Black, map.root().left().color());
+  EXPECT_EQ(Color::Black, map.root().right().color());
+
+  // insert left.left, no rotation yet
+  map = map.insert(1, 1);
+  EXPECT_EQ(5, map.root().key());
+  EXPECT_EQ(1, map.root().left().left().key());
+  EXPECT_EQ(Color::Red, map.root().left().left().color());
+
+  // insert left.right, triggers a color flip in left but no rotation
+  map = map.insert(4, 4);
+  EXPECT_EQ(5, map.root().key());
+  EXPECT_EQ(Color::Red, map.root().left().color());
+  EXPECT_EQ(Color::Black, map.root().left().left().color());
+  EXPECT_EQ(Color::Black, map.root().left().right().color());
+
+  // insert left.left.left; still no rotation
+  map = map.insert(0, 0);
+  EXPECT_EQ(5, map.root().key());
+  EXPECT_EQ(Color::Black, map.root().color());
+  EXPECT_EQ(Color::Red, map.root().left().color());
+  EXPECT_EQ(Color::Black, map.root().left().left().color());
+  EXPECT_EQ(Color::Red, map.root().left().left().left().color());
+
+  EXPECT_EQ(Color::Black, map.root().right().color());
+
+  // insert left.left.right:
+  //   * triggers a color flip on left.left => Red
+  //   * triggers right rotation at the root because left and left.left are Red
+  //   * triggers a color flip on root => whole tree black
+  map = map.insert(2, 2);
+  EXPECT_EQ(3, map.root().key());
+  EXPECT_EQ(Color::Black, map.root().color());
+
+  EXPECT_EQ(1, map.root().left().key());
+  EXPECT_EQ(Color::Black, map.root().left().color());
+
+  EXPECT_EQ(0, map.root().left().left().key());
+  EXPECT_EQ(Color::Black, map.root().left().left().color());
+
+  EXPECT_EQ(2, map.root().left().right().key());
+  EXPECT_EQ(Color::Black, map.root().left().right().color());
+
+  EXPECT_EQ(5, map.root().right().key());
+  EXPECT_EQ(Color::Black, map.root().right().color());
+
+  EXPECT_EQ(4, map.root().right().left().key());
+  EXPECT_EQ(Color::Black, map.root().right().left().color());
+
+  EXPECT_EQ(7, map.root().right().right().key());
+  EXPECT_EQ(Color::Black, map.root().right().right().color());
+}
+
+TEST(TreeSortedMap, InsertIsImmutable) {
+  IntMap original = IntMap{}.insert(3, 3);
+
+  IntMap modified = original.insert(2, 2).insert(1, 1);
+  EXPECT_EQ(3, original.root().key());
+  EXPECT_EQ(3, original.root().value());
+  EXPECT_EQ(Color::Black, original.root().color());
+  EXPECT_TRUE(original.root().left().empty());
+  EXPECT_TRUE(original.root().right().empty());
+}
+
 }  // namespace impl
 }  // namespace immutable
 }  // namespace firestore