Implement TreeSortedMap::insert (#1081)

* Make LlrbNode Rep more explicit, share empty node

* SortedMap::insert converts implementations

* Implement LlrbNode::insert

* Remove TestPolicy<SortedMap>

8 files changed, 494 insertions, 69 deletions

diff --git a/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h b/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h
index 92fd823..adba962 100644
--- a/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h
+++ b/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h
@@ -26,6 +26,7 @@
 
 #include "Firestore/core/src/firebase/firestore/immutable/map_entry.h"
 #include "Firestore/core/src/firebase/firestore/immutable/sorted_map_base.h"
+#include "Firestore/core/src/firebase/firestore/util/comparison.h"
 #include "Firestore/core/src/firebase/firestore/util/firebase_assert.h"
 
 namespace firebase {
@@ -118,7 +119,7 @@ class FixedArray {
  * ArraySortedMap is a value type containing a map. It is immutable, but has
  * methods to efficiently create new maps that are mutations of it.
  */
-template <typename K, typename V, typename C = std::less<K>>
+template <typename K, typename V, typename C = util::Comparator<K>>
 class ArraySortedMap : public SortedMapBase {
  public:
   using key_comparator_type = KeyComparator<K, V, C>;
@@ -229,6 +230,10 @@ class ArraySortedMap : public SortedMapBase {
     }
   }
 
+  const key_comparator_type& comparator() const {
+    return key_comparator_;
+  }
+
   /** Returns true if the map contains no elements. */
   bool empty() const {
     return size() == 0;
diff --git a/Firestore/core/src/firebase/firestore/immutable/llrb_node.h b/Firestore/core/src/firebase/firestore/immutable/llrb_node.h
index 89fd8cf..bfd0f8a 100644
--- a/Firestore/core/src/firebase/firestore/immutable/llrb_node.h
+++ b/Firestore/core/src/firebase/firestore/immutable/llrb_node.h
@@ -52,26 +52,12 @@ class LlrbNode : public SortedMapBase {
   /**
    * Constructs an empty node.
    */
-  // TODO(wilhuff): move this into NodeData if that structure is to live on.
-  LlrbNode()
-      : LlrbNode{NodeData{{},
-                          Color::Black,
-                          /*size=*/0u,
-                          LlrbNode{nullptr},
-                          LlrbNode{nullptr}}} {
-  }
-
-  /**
-   * Returns a shared Empty node, to cut down on allocations in the base case.
-   */
-  static const LlrbNode& Empty() {
-    static const LlrbNode empty_node{};
-    return empty_node;
+  LlrbNode() : LlrbNode{EmptyRep()} {
   }
 
   /** Returns the number of elements at this node or beneath it in the tree. */
   size_type size() const {
-    return data_->size_;
+    return rep_->size_;
   }
 
   /** Returns true if this is an empty node--a leaf node in the tree. */
@@ -81,11 +67,11 @@ class LlrbNode : public SortedMapBase {
 
   /** Returns true if this node is red (as opposed to black). */
   bool red() const {
-    return static_cast<bool>(data_->red_);
+    return static_cast<bool>(rep_->color_);
   }
 
   const value_type& entry() const {
-    return data_->contents_;
+    return rep_->entry_;
   }
   const K& key() const {
     return entry().first;
@@ -94,44 +80,247 @@ class LlrbNode : public SortedMapBase {
     return entry().second;
   }
   Color color() const {
-    return data_->red_ ? Color::Red : Color::Black;
+    return static_cast<Color>(rep_->color_);
   }
   const LlrbNode& left() const {
-    return data_->left_;
+    return rep_->left_;
   }
   const LlrbNode& right() const {
-    return data_->right_;
+    return rep_->right_;
   }
 
+  /** Returns a tree node with the given key-value pair set/updated. */
+  template <typename Comparator>
+  LlrbNode insert(const K& key,
+                  const V& value,
+                  const Comparator& comparator) const;
+
  private:
-  struct NodeData {
-    value_type contents_;
+  struct Rep {
+    Rep(value_type&& entry,
+        size_type color,
+        size_type size,
+        LlrbNode left,
+        LlrbNode right)
+        : entry_{std::move(entry)},
+          color_{color},
+          size_{size},
+          left_{std::move(left)},
+          right_{std::move(right)} {
+    }
+
+    Rep(value_type&& entry, size_type color, LlrbNode left, LlrbNode right)
+        : entry_{std::move(entry)},
+          color_{color},
+          size_{left.size() + 1 + right.size()},
+          left_{std::move(left)},
+          right_{std::move(right)} {
+    }
+
+    value_type entry_;
 
     // Store the color in the high bit of the size to save memory.
-    size_type red_ : 1;
+    size_type color_ : 1;
     size_type size_ : 31;
 
     LlrbNode left_;
     LlrbNode right_;
   };
 
-  explicit LlrbNode(NodeData&& data)
-      : data_{std::make_shared<NodeData>(std::move(data))} {
+  explicit LlrbNode(Rep rep) : rep_{std::make_shared<Rep>(std::move(rep))} {
+  }
+
+  explicit LlrbNode(const std::shared_ptr<Rep>& rep) : rep_{rep} {
+  }
+
+  explicit LlrbNode(std::shared_ptr<Rep>&& rep) : rep_{std::move(rep)} {
   }
 
   /**
-   * Constructs a dummy node that's a child of the empty node. This exists so
-   * that every node can have non-optional left and right children, despite the
-   * fact that these don't actually get visited.
-   *
-   * This should only be called when constructing the empty node.
+   * Returns a shared Empty node, to cut down on allocations in the base case.
    */
-  explicit LlrbNode(std::nullptr_t) : data_{nullptr} {
+  static const std::shared_ptr<Rep>& EmptyRep() {
+    static const std::shared_ptr<Rep> empty_rep = [] {
+      auto empty = std::make_shared<Rep>(Rep{std::pair<K, V>{}, Color::Black,
+                                             /* size= */ 0u, LlrbNode{nullptr},
+                                             LlrbNode{nullptr}});
+
+      // Set up the empty Rep such that you can traverse infinitely down left
+      // and right links.
+      empty->left_.rep_ = empty;
+      empty->right_.rep_ = empty;
+      return empty;
+    }();
+    return empty_rep;
   }
 
-  std::shared_ptr<NodeData> data_;
+  /**
+   * Creates a new copy of this node, duplicating the Rep but without
+   * duplicating the left_ and right_ children.
+   */
+  LlrbNode Clone() const {
+    return LlrbNode{*rep_};
+  }
+
+  void set_size(size_type size) {
+    rep_->size_ = size;
+  }
+
+  void set_entry(const value_type& entry) {
+    rep_->entry_ = entry;
+  }
+  void set_entry(value_type&& entry) {
+    rep_->entry_ = std::move(entry);
+  }
+  void set_value(const V& value) {
+    rep_->entry_.second = value;
+  }
+  void set_color(size_type color) {
+    rep_->color_ = color;
+  }
+  void set_left(const LlrbNode& left) {
+    rep_->left_ = left;
+  }
+  void set_left(LlrbNode&& left) {
+    rep_->left_ = std::move(left);
+  }
+  void set_right(const LlrbNode& right) {
+    rep_->right_ = right;
+  }
+  void set_right(LlrbNode&& right) {
+    rep_->right_ = std::move(right);
+  }
+
+  template <typename Comparator>
+  LlrbNode InnerInsert(const K& key,
+                       const V& value,
+                       const Comparator& comparator) const;
+
+  void FixUp();
+
+  void RotateLeft();
+  void RotateRight();
+  void FlipColor();
+
+  size_type OppositeColor() const noexcept {
+    return rep_->color_ == Color::Red ? Color::Black : Color::Red;
+  }
+
+  std::shared_ptr<Rep> rep_;
 };
 
+template <typename K, typename V>
+template <typename Comparator>
+LlrbNode<K, V> LlrbNode<K, V>::insert(const K& key,
+                                      const V& value,
+                                      const Comparator& comparator) const {
+  LlrbNode root = InnerInsert(key, value, comparator);
+  // The root must always be black
+  if (root.red()) {
+    root.rep_->color_ = Color::Black;
+  }
+  return root;
+}
+
+template <typename K, typename V>
+template <typename Comparator>
+LlrbNode<K, V> LlrbNode<K, V>::InnerInsert(const K& key,
+                                           const V& value,
+                                           const Comparator& comparator) const {
+  if (empty()) {
+    return LlrbNode{Rep{{key, value}, Color::Red, LlrbNode{}, LlrbNode{}}};
+  }
+
+  // Inserting is going to result in a copy but we can save some allocations by
+  // creating the copy once and fixing that up, rather than copying and
+  // re-copying the result.
+  LlrbNode result = Clone();
+
+  const K& this_key = this->key();
+  bool descending = comparator(key, this_key);
+  if (descending) {
+    result.set_left(result.left().InnerInsert(key, value, comparator));
+    result.FixUp();
+
+  } else {
+    bool ascending = comparator(this_key, key);
+    if (ascending) {
+      result.set_right(result.right().InnerInsert(key, value, comparator));
+      result.FixUp();
+
+    } else {
+      // keys are equal so update the value.
+      result.set_value(value);
+    }
+  }
+  return result;
+}
+
+template <typename K, typename V>
+void LlrbNode<K, V>::FixUp() {
+  set_size(left().size() + 1 + right().size());
+
+  if (right().red() && !left().red()) {
+    RotateLeft();
+  }
+  if (left().red() && left().left().red()) {
+    RotateRight();
+  }
+  if (left().red() && right().red()) {
+    FlipColor();
+  }
+}
+
+// Rotates left:
+//
+//      X              R
+//    /   \          /   \
+//   L     R   =>   X    RR
+//        / \      / \
+//       RL RR     L RL
+template <typename K, typename V>
+void LlrbNode<K, V>::RotateLeft() {
+  LlrbNode new_left{
+      Rep{std::move(rep_->entry_), Color::Red, left(), right().left()}};
+
+  // size_ and color remain unchanged after a rotation.
+  set_entry(right().entry());
+  set_left(std::move(new_left));
+  set_right(right().right());
+}
+
+// Rotates right:
+//
+//      X              L
+//    /   \          /   \
+//   L     R   =>   LL    X
+//  / \                  / \
+// LL LR                LR R
+template <typename K, typename V>
+void LlrbNode<K, V>::RotateRight() {
+  LlrbNode new_right{
+      Rep{std::move(rep_->entry_), Color::Red, left().right(), right()}};
+
+  // size_ remains unchanged after a rotation. Preserve color too.
+  set_entry(left().entry());
+  set_left(left().left());
+  set_right(std::move(new_right));
+}
+
+template <typename K, typename V>
+void LlrbNode<K, V>::FlipColor() {
+  LlrbNode new_left = left().Clone();
+  new_left.set_color(left().OppositeColor());
+
+  LlrbNode new_right = right().Clone();
+  new_right.set_color(right().OppositeColor());
+
+  // Preserve contents_ and size_
+  set_color(OppositeColor());
+  set_left(std::move(new_left));
+  set_right(std::move(new_right));
+}
+
 }  // namespace impl
 }  // namespace immutable
 }  // namespace firestore
diff --git a/Firestore/core/src/firebase/firestore/immutable/map_entry.h b/Firestore/core/src/firebase/firestore/immutable/map_entry.h
index 1022b06..b2b3c97 100644
--- a/Firestore/core/src/firebase/firestore/immutable/map_entry.h
+++ b/Firestore/core/src/firebase/firestore/immutable/map_entry.h
@@ -51,6 +51,10 @@ struct KeyComparator {
     return key_comparator_(lhs.first, rhs.first);
   }
 
+  const C& comparator() const {
+    return key_comparator_;
+  }
+
  private:
   C key_comparator_;
 };
diff --git a/Firestore/core/src/firebase/firestore/immutable/sorted_map.h b/Firestore/core/src/firebase/firestore/immutable/sorted_map.h
index ef6f54e..b7729b3 100644
--- a/Firestore/core/src/firebase/firestore/immutable/sorted_map.h
+++ b/Firestore/core/src/firebase/firestore/immutable/sorted_map.h
@@ -56,8 +56,7 @@ class SortedMap : public impl::SortedMapBase {
       tag_ = Tag::Array;
       new (&array_) array_type{entries, comparator};
     } else {
-      // TODO(wilhuff): implement tree initialization
-      abort();
+      new (&tree_) tree_type{tree_type::Create(entries, comparator)};
     }
   }
 
@@ -139,8 +138,19 @@ class SortedMap : public impl::SortedMapBase {
   SortedMap insert(const K& key, const V& value) const {
     switch (tag_) {
       case Tag::Array:
-        // TODO(wilhuff): convert to TreeSortedMap
-        return SortedMap{array_.insert(key, value)};
+        if (array_.size() >= kFixedSize) {
+          // Strictly speaking this conversion is more eager than it needs to
+          // be since we could be replacing an existing key. However, the
+          // benefit of using the array for small maps doesn't really depend on
+          // exactly where this cut-off happens and just unconditionally
+          // converting if the next insertion could overflow keeps things
+          // simpler.
+          const C& comparator = array_.comparator().comparator();
+          tree_type tree = tree_type::Create(array_, comparator);
+          return SortedMap{tree.insert(key, value)};
+        } else {
+          return SortedMap{array_.insert(key, value)};
+        }
       case Tag::Tree:
         return SortedMap{tree_.insert(key, value)};
     }
diff --git a/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h b/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h
index dfe270d..ff8c9f9 100644
--- a/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h
+++ b/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h
@@ -56,7 +56,19 @@ class TreeSortedMap : public SortedMapBase, private util::ComparatorHolder<C> {
    * Creates an empty TreeSortedMap.
    */
   explicit TreeSortedMap(const C& comparator = {})
-      : util::ComparatorHolder<C>{comparator}, root_{node_type::Empty()} {
+      : util::ComparatorHolder<C>{comparator} {
+  }
+
+  /**
+   * Creates a TreeSortedMap from a range of pairs to insert.
+   */
+  template <typename Range>
+  static TreeSortedMap Create(const Range& range, const C& comparator) {
+    node_type node;
+    for (auto&& element : range) {
+      node = node.insert(element.first, element.second, comparator);
+    }
+    return TreeSortedMap{std::move(node), comparator};
   }
 
   /**
@@ -68,10 +80,8 @@ class TreeSortedMap : public SortedMapBase, private util::ComparatorHolder<C> {
    * @return A new dictionary with the added/updated value.
    */
   TreeSortedMap insert(const K& key, const V& value) const {
-    // TODO(wilhuff): Actually implement insert
-    (void)key;
-    (void)value;
-    return *this;
+    const C& comparator = this->comparator();
+    return TreeSortedMap{root_.insert(key, value, comparator), comparator};
   }
 
   /**
@@ -83,7 +93,7 @@ class TreeSortedMap : public SortedMapBase, private util::ComparatorHolder<C> {
   TreeSortedMap erase(const K& key) const {
     // TODO(wilhuff): Actually implement erase
     (void)key;
-    return *this;
+    return TreeSortedMap{this->comparator()};
   }
 
   /** Returns true if the map contains no elements. */
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