Initial TreeSortedMap and SortedMap in C++ (#980)

* Prepare for TreeSortedMap
  * Factor out SortedMapBase
  * Move ArraySortedMap to impl
  * Factor out SortedMap testing utilities

* Add a minimal TreeSortedMap

* Add the public SortedMap type

14 files changed, 907 insertions, 178 deletions

diff --git a/Firestore/core/src/firebase/firestore/immutable/CMakeLists.txt b/Firestore/core/src/firebase/firestore/immutable/CMakeLists.txt
index 4079307..543f912 100644
--- a/Firestore/core/src/firebase/firestore/immutable/CMakeLists.txt
+++ b/Firestore/core/src/firebase/firestore/immutable/CMakeLists.txt
@@ -15,9 +15,13 @@
 cc_library(
   firebase_firestore_immutable
   SOURCES
-    array_sorted_map.cc
     array_sorted_map.h
+    llrb_node.h
     map_entry.h
+    sorted_map.h
+    sorted_map_base.h
+    sorted_map_base.cc
+    tree_sorted_map.h
   DEPENDS
     firebase_firestore_util
 )
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 4b8acde..56da9e9 100644
--- a/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h
+++ b/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h
@@ -25,45 +25,15 @@
 #include <utility>
 
 #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/firebase_assert.h"
 
 namespace firebase {
 namespace firestore {
 namespace immutable {
-
 namespace impl {
 
 /**
- * A base class for implementing ArraySortedMap, containing types and constants
- * that don't depend upon the template parameters to the main class.
- *
- * Note that this exists as a base class rather than as just a namespace in
- * order to make it possible for users of ArraySortedMap to avoid needing to
- * declare storage for each instantiation of the template.
- */
-class ArraySortedMapBase {
- public:
-  /**
-   * The type of size() methods on immutable collections. Note that this is not
-   * size_t specifically to save space in the TreeSortedMap implementation.
-   */
-  using size_type = uint32_t;
-
-  /**
-   * The maximum size of an ArraySortedMap.
-   *
-   * This is the size threshold where we use a tree backed sorted map instead of
-   * an array backed sorted map. This is a more or less arbitrary chosen value,
-   * that was chosen to be large enough to fit most of object kind of Firebase
-   * data, but small enough to not notice degradation in performance for
-   * inserting and lookups. Feel free to empirically determine this constant,
-   * but don't expect much gain in real world performance.
-   */
-  // TODO(wilhuff): actually use this for switching implementations.
-  static constexpr size_type kFixedSize = 25;
-};
-
-/**
  * A bounded-size array that allocates its contents directly in itself. This
  * saves a heap allocation when compared with std::vector (though std::vector
  * can resize itself while FixedArray cannot).
@@ -78,10 +48,10 @@ class ArraySortedMapBase {
  * @tparam T The type of an element in the array.
  * @tparam fixed_size the fixed size to use in creating the FixedArray.
  */
-template <typename T, ArraySortedMapBase::size_type fixed_size>
+template <typename T, SortedMapBase::size_type fixed_size>
 class FixedArray {
  public:
-  using size_type = ArraySortedMapBase::size_type;
+  using size_type = SortedMapBase::size_type;
   using array_type = std::array<T, fixed_size>;
   using iterator = typename array_type::iterator;
   using const_iterator = typename array_type::const_iterator;
@@ -144,14 +114,12 @@ class FixedArray {
   size_type size_ = 0;
 };
 
-}  // namespace impl
-
 /**
  * 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>>
-class ArraySortedMap : public impl::ArraySortedMapBase {
+class ArraySortedMap : public SortedMapBase {
  public:
   using key_comparator_type = KeyComparator<K, V, C>;
 
@@ -163,7 +131,7 @@ class ArraySortedMap : public impl::ArraySortedMapBase {
   /**
    * The type of the fixed-size array containing entries of value_type.
    */
-  using array_type = impl::FixedArray<value_type, kFixedSize>;
+  using array_type = FixedArray<value_type, kFixedSize>;
   using const_iterator = typename array_type::const_iterator;
 
   using array_pointer = std::shared_ptr<const array_type>;
@@ -172,7 +140,7 @@ class ArraySortedMap : public impl::ArraySortedMapBase {
    * Creates an empty ArraySortedMap.
    */
   explicit ArraySortedMap(const C& comparator = C())
-      : array_(EmptyArray()), key_comparator_(comparator) {
+      : array_{EmptyArray()}, key_comparator_{comparator} {
   }
 
   /**
@@ -180,8 +148,8 @@ class ArraySortedMap : public impl::ArraySortedMapBase {
    */
   ArraySortedMap(std::initializer_list<value_type> entries,
                  const C& comparator = C())
-      : array_(std::make_shared<array_type>(entries.begin(), entries.end())),
-        key_comparator_(comparator) {
+      : array_{std::make_shared<array_type>(entries.begin(), entries.end())},
+        key_comparator_{comparator} {
   }
 
   /**
@@ -211,7 +179,7 @@ class ArraySortedMap : public impl::ArraySortedMapBase {
     auto copy = std::make_shared<array_type>(begin(), pos);
 
     // Copy the value to be inserted.
-    copy->append(value_type(key, value));
+    copy->append({key, value});
 
     if (replacing_entry) {
       // Skip the thing at pos because it compares the same as the pair above.
@@ -297,11 +265,11 @@ class ArraySortedMap : public impl::ArraySortedMapBase {
 
   ArraySortedMap(const array_pointer& array,
                  const key_comparator_type& key_comparator) noexcept
-      : array_(array), key_comparator_(key_comparator) {
+      : array_{array}, key_comparator_{key_comparator} {
   }
 
   ArraySortedMap wrap(const array_pointer& array) const noexcept {
-    return ArraySortedMap(array, key_comparator_);
+    return ArraySortedMap{array, key_comparator_};
   }
 
   const_iterator LowerBound(const K& key) const {
@@ -312,6 +280,7 @@ class ArraySortedMap : public impl::ArraySortedMapBase {
   key_comparator_type key_comparator_;
 };
 
+}  // namespace impl
 }  // namespace immutable
 }  // namespace firestore
 }  // namespace firebase
diff --git a/Firestore/core/src/firebase/firestore/immutable/llrb_node.h b/Firestore/core/src/firebase/firestore/immutable/llrb_node.h
new file mode 100644
index 0000000..89fd8cf
--- /dev/null
+++ b/Firestore/core/src/firebase/firestore/immutable/llrb_node.h
@@ -0,0 +1,140 @@
+/*
+ * Copyright 2018 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_LLRB_NODE_H_
+#define FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_LLRB_NODE_H_
+
+#include <memory>
+#include <utility>
+
+#include "Firestore/core/src/firebase/firestore/immutable/sorted_map_base.h"
+
+namespace firebase {
+namespace firestore {
+namespace immutable {
+namespace impl {
+
+/**
+ * A Color of a tree node in a red-black tree.
+ */
+enum Color : unsigned int {
+  Black,
+  Red,
+};
+
+/**
+ * LlrbNode is a node in a TreeSortedMap.
+ */
+template <typename K, typename V>
+class LlrbNode : public SortedMapBase {
+ public:
+  using first_type = K;
+  using second_type = V;
+
+  /**
+   * The type of the entries stored in the map.
+   */
+  using value_type = std::pair<K, V>;
+
+  /**
+   * 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;
+  }
+
+  /** Returns the number of elements at this node or beneath it in the tree. */
+  size_type size() const {
+    return data_->size_;
+  }
+
+  /** Returns true if this is an empty node--a leaf node in the tree. */
+  bool empty() const {
+    return size() == 0;
+  }
+
+  /** Returns true if this node is red (as opposed to black). */
+  bool red() const {
+    return static_cast<bool>(data_->red_);
+  }
+
+  const value_type& entry() const {
+    return data_->contents_;
+  }
+  const K& key() const {
+    return entry().first;
+  }
+  const V& value() const {
+    return entry().second;
+  }
+  Color color() const {
+    return data_->red_ ? Color::Red : Color::Black;
+  }
+  const LlrbNode& left() const {
+    return data_->left_;
+  }
+  const LlrbNode& right() const {
+    return data_->right_;
+  }
+
+ private:
+  struct NodeData {
+    value_type contents_;
+
+    // Store the color in the high bit of the size to save memory.
+    size_type red_ : 1;
+    size_type size_ : 31;
+
+    LlrbNode left_;
+    LlrbNode right_;
+  };
+
+  explicit LlrbNode(NodeData&& data)
+      : data_{std::make_shared<NodeData>(std::move(data))} {
+  }
+
+  /**
+   * 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.
+   */
+  explicit LlrbNode(std::nullptr_t) : data_{nullptr} {
+  }
+
+  std::shared_ptr<NodeData> data_;
+};
+
+}  // namespace impl
+}  // namespace immutable
+}  // namespace firestore
+}  // namespace firebase
+
+#endif  // FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_LLRB_NODE_H_
diff --git a/Firestore/core/src/firebase/firestore/immutable/sorted_map.h b/Firestore/core/src/firebase/firestore/immutable/sorted_map.h
new file mode 100644
index 0000000..5ed16b3
--- /dev/null
+++ b/Firestore/core/src/firebase/firestore/immutable/sorted_map.h
@@ -0,0 +1,209 @@
+/*
+ * Copyright 2018 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_SORTED_MAP_H_
+#define FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_SORTED_MAP_H_
+
+#include <utility>
+
+#include "Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h"
+#include "Firestore/core/src/firebase/firestore/immutable/sorted_map_base.h"
+#include "Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h"
+#include "Firestore/core/src/firebase/firestore/util/comparison.h"
+
+namespace firebase {
+namespace firestore {
+namespace immutable {
+
+/**
+ * SortedMap 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 = util::Comparator<K>>
+class SortedMap : public impl::SortedMapBase {
+ public:
+  /** The type of the entries stored in the map. */
+  using value_type = std::pair<K, V>;
+  using array_type = impl::ArraySortedMap<K, V, C>;
+  using tree_type = impl::TreeSortedMap<K, V, C>;
+
+  /**
+   * Creates an empty SortedMap.
+   */
+  explicit SortedMap(const C& comparator = {})
+      : SortedMap{array_type{comparator}} {
+  }
+
+  /**
+   * Creates an SortedMap containing the given entries.
+   */
+  SortedMap(std::initializer_list<value_type> entries,
+            const C& comparator = {}) {
+    if (entries.size() <= kFixedSize) {
+      tag_ = Tag::Array;
+      new (&array_) array_type{entries, comparator};
+    } else {
+      // TODO(wilhuff): implement tree initialization
+      abort();
+    }
+  }
+
+  SortedMap(const SortedMap& other) : tag_{other.tag_} {
+    switch (tag_) {
+      case Tag::Array:
+        new (&array_) array_type{other.array_};
+        break;
+      case Tag::Tree:
+        new (&tree_) tree_type{other.tree_};
+        break;
+    }
+  }
+
+  SortedMap(SortedMap&& other) : tag_{other.tag_} {
+    switch (tag_) {
+      case Tag::Array:
+        new (&array_) array_type{std::move(other.array_)};
+        break;
+      case Tag::Tree:
+        new (&tree_) tree_type{std::move(other.tree_)};
+        break;
+    }
+  }
+
+  ~SortedMap() {
+    switch (tag_) {
+      case Tag::Array:
+        array_.~ArraySortedMap();
+        break;
+      case Tag::Tree:
+        tree_.~TreeSortedMap();
+        break;
+    }
+  }
+
+  SortedMap& operator=(const SortedMap& other) {
+    if (tag_ == other.tag_) {
+      switch (tag_) {
+        case Tag::Array:
+          array_ = other.array_;
+          break;
+        case Tag::Tree:
+          tree_ = other.tree_;
+          break;
+      }
+    } else {
+      this->~SortedMap();
+      new (this) SortedMap{other};
+    }
+    return *this;
+  }
+
+  SortedMap& operator=(SortedMap&& other) {
+    if (tag_ == other.tag_) {
+      switch (tag_) {
+        case Tag::Array:
+          array_ = std::move(other.array_);
+          break;
+        case Tag::Tree:
+          tree_ = std::move(other.tree_);
+          break;
+      }
+    } else {
+      this->~SortedMap();
+      new (this) SortedMap{std::move(other)};
+    }
+    return *this;
+  }
+
+  /**
+   * Creates a new map identical to this one, but with a key-value pair added or
+   * updated.
+   *
+   * @param key The key to insert/update.
+   * @param value The value to associate with the key.
+   * @return A new dictionary with the added/updated value.
+   */
+  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)};
+      case Tag::Tree:
+        return SortedMap{tree_.insert(key, value)};
+    }
+  }
+
+  /**
+   * Creates a new map identical to this one, but with a key removed from it.
+   *
+   * @param key The key to remove.
+   * @return A new map without that value.
+   */
+  SortedMap erase(const K& key) const {
+    switch (tag_) {
+      case Tag::Array:
+        return SortedMap{array_.erase(key)};
+      case Tag::Tree:
+        return SortedMap{tree_.erase(key)};
+    }
+  }
+
+  /** Returns true if the map contains no elements. */
+  bool empty() const {
+    switch (tag_) {
+      case Tag::Array:
+        return array_.empty();
+      case Tag::Tree:
+        return tree_.empty();
+    }
+  }
+
+  /** Returns the number of items in this map. */
+  size_type size() const {
+    switch (tag_) {
+      case Tag::Array:
+        return array_.size();
+      case Tag::Tree:
+        return tree_.size();
+    }
+  }
+
+ private:
+  explicit SortedMap(array_type&& array)
+      : tag_{Tag::Array}, array_{std::move(array)} {
+  }
+
+  explicit SortedMap(tree_type&& tree)
+      : tag_{Tag::Tree}, tree_{std::move(tree)} {
+  }
+
+  enum class Tag {
+    Array,
+    Tree,
+  };
+
+  Tag tag_;
+  union {
+    array_type array_;
+    tree_type tree_;
+  };
+};
+
+}  // namespace immutable
+}  // namespace firestore
+}  // namespace firebase
+
+#endif  // FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_SORTED_MAP_H_
diff --git a/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.cc b/Firestore/core/src/firebase/firestore/immutable/sorted_map_base.cc
index 48e7553..f2faa33 100644
--- a/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.cc
+++ b/Firestore/core/src/firebase/firestore/immutable/sorted_map_base.cc
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include "Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h"
+#include "Firestore/core/src/firebase/firestore/immutable/sorted_map_base.h"
 
 namespace firebase {
 namespace firestore {
@@ -22,7 +22,7 @@ namespace immutable {
 namespace impl {
 
 // Define external storage for constants:
-constexpr ArraySortedMapBase::size_type ArraySortedMapBase::kFixedSize;
+constexpr SortedMapBase::size_type SortedMapBase::kFixedSize;
 
 }  // namespace impl
 }  // namespace immutable
diff --git a/Firestore/core/src/firebase/firestore/immutable/sorted_map_base.h b/Firestore/core/src/firebase/firestore/immutable/sorted_map_base.h
new file mode 100644
index 0000000..accb5ef
--- /dev/null
+++ b/Firestore/core/src/firebase/firestore/immutable/sorted_map_base.h
@@ -0,0 +1,64 @@
+/*
+ * Copyright 2018 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_SORTED_MAP_BASE_H_
+#define FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_SORTED_MAP_BASE_H_
+
+#include <stdint.h>
+
+namespace firebase {
+namespace firestore {
+namespace immutable {
+namespace impl {
+
+/**
+ * A base class for implementing sorted maps, containing types and constants
+ * that don't depend upon the template parameters to the main class.
+ *
+ * Note that this exists as a base class rather than as just a namespace in
+ * order to make it possible for users of the SortedMap classes to avoid needing
+ * to declare storage for each instantiation of the template.
+ */
+class SortedMapBase {
+ public:
+  /**
+   * The type of size() methods on immutable collections. Note:
+   *   * This is not size_t specifically to save space in the TreeSortedMap
+   *     implementation.
+   *   * This remains unsigned for straightforward casting to size_t.
+   */
+  using size_type = uint32_t;
+
+  /**
+   * The maximum size of an ArraySortedMap.
+   *
+   * This is the size threshold where we use a tree backed sorted map instead of
+   * an array backed sorted map. This is a more or less arbitrary chosen value,
+   * that was chosen to be large enough to fit most of object kind of Firebase
+   * data, but small enough to not notice degradation in performance for
+   * inserting and lookups. Feel free to empirically determine this constant,
+   * but don't expect much gain in real world performance.
+   */
+  // TODO(wilhuff): actually use this for switching implementations.
+  static constexpr size_type kFixedSize = 25;
+};
+
+}  // namespace impl
+}  // namespace immutable
+}  // namespace firestore
+}  // namespace firebase
+
+#endif  // FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_SORTED_MAP_BASE_H_
diff --git a/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h b/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h
new file mode 100644
index 0000000..e3102e7
--- /dev/null
+++ b/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h
@@ -0,0 +1,121 @@
+/*
+ * Copyright 2018 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_TREE_SORTED_MAP_H_
+#define FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_TREE_SORTED_MAP_H_
+
+#include <assert.h>
+
+#include <algorithm>
+#include <functional>
+#include <memory>
+#include <utility>
+
+#include "Firestore/core/src/firebase/firestore/immutable/llrb_node.h"
+#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/comparator_holder.h"
+#include "Firestore/core/src/firebase/firestore/util/comparison.h"
+#include "Firestore/core/src/firebase/firestore/util/iterator_adaptors.h"
+
+namespace firebase {
+namespace firestore {
+namespace immutable {
+namespace impl {
+
+/**
+ * TreeSortedMap 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 = util::Comparator<K>>
+class TreeSortedMap : public SortedMapBase, private util::ComparatorHolder<C> {
+ public:
+  /**
+   * The type of the entries stored in the map.
+   */
+  using value_type = std::pair<K, V>;
+
+  /**
+   * The type of the node containing entries of value_type.
+   */
+  using node_type = LlrbNode<K, V>;
+
+  /**
+   * Creates an empty TreeSortedMap.
+   */
+  explicit TreeSortedMap(const C& comparator = {})
+      : util::ComparatorHolder<C>{comparator}, root_{node_type::Empty()} {
+  }
+
+  /**
+   * Creates a new map identical to this one, but with a key-value pair added or
+   * updated.
+   *
+   * @param key The key to insert/update.
+   * @param value The value to associate with the key.
+   * @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;
+  }
+
+  /**
+   * Creates a new map identical to this one, but with a key removed from it.
+   *
+   * @param key The key to remove.
+   * @return A new map without that value.
+   */
+  TreeSortedMap erase(const K& key) const {
+    // TODO(wilhuff): Actually implement erase
+    (void)key;
+    return *this;
+  }
+
+  /** Returns true if the map contains no elements. */
+  bool empty() const {
+    return root_.empty();
+  }
+
+  /** Returns the number of items in this map. */
+  size_type size() const {
+    return root_.size();
+  }
+
+  const node_type& root() const {
+    return root_;
+  }
+
+ private:
+  TreeSortedMap(node_type&& root, const C& comparator) noexcept
+      : util::ComparatorHolder<C>{comparator}, root_{std::move(root)} {
+  }
+
+  TreeSortedMap Wrap(node_type&& root) noexcept {
+    return TreeSortedMap{std::move(root), this->comparator()};
+  }
+
+  node_type root_;
+};
+
+}  // namespace impl
+}  // namespace immutable
+}  // namespace firestore
+}  // namespace firebase
+
+#endif  // FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_TREE_SORTED_MAP_H_
diff --git a/Firestore/core/src/firebase/firestore/util/CMakeLists.txt b/Firestore/core/src/firebase/firestore/util/CMakeLists.txt
index 7a080d4..eac7053 100644
--- a/Firestore/core/src/firebase/firestore/util/CMakeLists.txt
+++ b/Firestore/core/src/firebase/firestore/util/CMakeLists.txt
@@ -109,6 +109,7 @@ cc_library(
     autoid.h
     bits.cc
     bits.h
+    comparator_holder.h
     comparison.cc
     comparison.h
     config.h
diff --git a/Firestore/core/src/firebase/firestore/util/comparator_holder.h b/Firestore/core/src/firebase/firestore/util/comparator_holder.h
new file mode 100644
index 0000000..8641b0f
--- /dev/null
+++ b/Firestore/core/src/firebase/firestore/util/comparator_holder.h
@@ -0,0 +1,62 @@
+/*
+ * Copyright 2018 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_UTIL_COMPARATOR_HOLDER_H_
+#define FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_UTIL_COMPARATOR_HOLDER_H_
+
+#include <type_traits>
+
+namespace firebase {
+namespace firestore {
+namespace util {
+
+/**
+ * A holder of some comparator (e.g. std::less or util::Comparator) that
+ * implements the empty base optimization for the common case where the
+ * comparator occupies no storage.
+ */
+template <typename C, bool = std::is_empty<C>::value>
+class ComparatorHolder {
+ protected:
+  explicit ComparatorHolder(const C& comparator) noexcept
+      : comparator_(comparator) {
+  }
+
+  const C& comparator() const noexcept {
+    return comparator_;
+  }
+
+ private:
+  C comparator_;
+};
+
+// Implementation to use when C is empty.
+template <typename C>
+class ComparatorHolder<C, true> : private C {
+ protected:
+  explicit ComparatorHolder(const C&) noexcept {
+  }
+
+  const C& comparator() const noexcept {
+    return *this;
+  }
+};
+
+}  // namespace util
+}  // namespace firestore
+}  // namespace firebase
+
+#endif  // FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_UTIL_COMPARATOR_HOLDER_H_
diff --git a/Firestore/core/test/firebase/firestore/immutable/CMakeLists.txt b/Firestore/core/test/firebase/firestore/immutable/CMakeLists.txt
index e7b0b6e..753e2d0 100644
--- a/Firestore/core/test/firebase/firestore/immutable/CMakeLists.txt
+++ b/Firestore/core/test/firebase/firestore/immutable/CMakeLists.txt
@@ -16,6 +16,9 @@ cc_test(
   firebase_firestore_immutable_test
   SOURCES
     array_sorted_map_test.cc
+    testing.h
+    sorted_map_test.cc
+    tree_sorted_map_test.cc
   DEPENDS
     firebase_firestore_immutable
     firebase_firestore_util
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 260506b..fceab7d 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
@@ -20,127 +20,18 @@
 #include <random>
 
 #include "Firestore/core/src/firebase/firestore/util/secure_random.h"
+
+#include "Firestore/core/test/firebase/firestore/immutable/testing.h"
 #include "gtest/gtest.h"
 
 namespace firebase {
 namespace firestore {
 namespace immutable {
+namespace impl {
 
 typedef ArraySortedMap<int, int> IntMap;
 constexpr IntMap::size_type kFixedSize = IntMap::kFixedSize;
 
-template <typename K, typename V>
-testing::AssertionResult NotFound(const ArraySortedMap<K, V>& set,
-                                  const K& key) {
-  auto found = set.find(key);
-  if (found == set.end()) {
-    return testing::AssertionSuccess();
-  } else {
-    return testing::AssertionFailure()
-           << "Should not have found (" << found->first << ", " << found->second
-           << ") @ " << found;
-  }
-}
-
-template <typename K, typename V>
-testing::AssertionResult Found(const ArraySortedMap<K, V>& map,
-                               const K& key,
-                               const V& expected) {
-  auto found = map.find(key);
-  if (found == map.end()) {
-    return testing::AssertionFailure() << "Did not find key " << key;
-  }
-  if (found->second == expected) {
-    return testing::AssertionSuccess();
-  } else {
-    return testing::AssertionFailure() << "Found entry was (" << found->first
-                                       << ", " << found->second << ")";
-  }
-}
-
-/**
- * Creates a vector containing a sequence of integers from the given starting
- * element up to, but not including, the given end element, with values
- * incremented by the given step.
- *
- * If step is negative the sequence is in descending order (but still starting
- * at start and ending before end).
- */
-std::vector<int> Sequence(int start, int end, int step = 1) {
-  std::vector<int> result;
-  if (step > 0) {
-    for (int i = start; i < end; i += step) {
-      result.push_back(i);
-    }
-  } else {
-    for (int i = start; i > end; i += step) {
-      result.push_back(i);
-    }
-  }
-  return result;
-}
-
-/**
- * Creates a vector containing a sequence of integers with the given number of
- * elements, from zero up to, but not including the given value.
- */
-std::vector<int> Sequence(int num_elements) {
-  return Sequence(0, num_elements);
-}
-
-/**
- * Creates a copy of the given vector with contents shuffled randomly.
- */
-std::vector<int> Shuffled(const std::vector<int>& values) {
-  std::vector<int> result(values);
-  util::SecureRandom rng;
-  std::shuffle(result.begin(), result.end(), rng);
-  return result;
-}
-
-/**
- * Creates a copy of the given vector with contents sorted.
- */
-std::vector<int> Sorted(const std::vector<int>& values) {
-  std::vector<int> result(values);
-  std::sort(result.begin(), result.end());
-  return result;
-}
-
-/**
- * Creates a vector of pairs where each pair has the same first and second
- * corresponding to an element in the given vector.
- */
-std::vector<std::pair<int, int>> Pairs(const std::vector<int>& values) {
-  std::vector<std::pair<int, int>> result;
-  for (auto&& value : values) {
-    result.emplace_back(value, value);
-  }
-  return result;
-}
-
-/**
- * Creates an ArraySortedMap by inserting a pair for each value in the vector.
- * Each pair will have the same key and value.
- */
-IntMap ToMap(const std::vector<int>& values) {
-  IntMap result;
-  for (auto&& value : values) {
-    result = result.insert(value, value);
-  }
-  return result;
-}
-
-/**
- * Appends the contents of the given container to a new vector.
- */
-template <typename Container>
-std::vector<typename Container::value_type> Append(const Container& container) {
-  std::vector<typename Container::value_type> result;
-  result.insert(result.begin(), container.begin(), container.end());
-  return result;
-}
-
 // TODO(wilhuff): ReverseTraversal
 
 #define ASSERT_SEQ_EQ(x, y) ASSERT_EQ((x), Append(y));
@@ -167,7 +58,7 @@ TEST(ArraySortedMap, RemoveKeyValuePair) {
 }
 
 TEST(ArraySortedMap, MoreRemovals) {
-  IntMap map = IntMap()
+  IntMap map = IntMap{}
                    .insert(1, 1)
                    .insert(50, 50)
                    .insert(3, 3)
@@ -230,7 +121,7 @@ TEST(ArraySortedMap, Increasing) {
 }
 
 TEST(ArraySortedMap, Override) {
-  IntMap map = IntMap().insert(10, 10).insert(10, 8);
+  IntMap map = IntMap{}.insert(10, 10).insert(10, 8);
 
   ASSERT_TRUE(Found(map, 10, 8));
   ASSERT_FALSE(Found(map, 10, 10));
@@ -238,7 +129,7 @@ TEST(ArraySortedMap, Override) {
 
 TEST(ArraySortedMap, ChecksSize) {
   std::vector<int> to_insert = Sequence(kFixedSize);
-  IntMap map = ToMap(to_insert);
+  IntMap map = ToMap<IntMap>(to_insert);
 
   // Replacing an existing entry should not hit increase size
   map = map.insert(5, 10);
@@ -247,25 +138,11 @@ TEST(ArraySortedMap, ChecksSize) {
   ASSERT_ANY_THROW(map.insert(next, next));
 }
 
-TEST(ArraySortedMap, Empty) {
-  IntMap map = IntMap().insert(10, 10).erase(10);
-  EXPECT_TRUE(map.empty());
-  EXPECT_EQ(0u, map.size());
-  EXPECT_TRUE(NotFound(map, 1));
-  EXPECT_TRUE(NotFound(map, 10));
-}
-
 TEST(ArraySortedMap, EmptyGet) {
   IntMap map;
   EXPECT_TRUE(NotFound(map, 10));
 }
 
-TEST(ArraySortedMap, EmptySize) {
-  IntMap map;
-  EXPECT_TRUE(map.empty());
-  EXPECT_EQ(0u, map.size());
-}
-
 TEST(ArraySortedMap, EmptyRemoval) {
   IntMap map;
   IntMap new_map = map.erase(1);
@@ -281,7 +158,7 @@ TEST(ArraySortedMap, InsertionAndRemovalOfMaxItems) {
   std::vector<int> to_remove = Shuffled(to_insert);
 
   // Add them to the map
-  IntMap map = ToMap(to_insert);
+  IntMap map = ToMap<IntMap>(to_insert);
   ASSERT_EQ(expected_size, map.size())
       << "Check if all N objects are in the map";
 
@@ -297,7 +174,7 @@ TEST(ArraySortedMap, InsertionAndRemovalOfMaxItems) {
 TEST(ArraySortedMap, BalanceProblem) {
   std::vector<int> to_insert{1, 7, 8, 5, 2, 6, 4, 0, 3};
 
-  IntMap map = ToMap(to_insert);
+  IntMap map = ToMap<IntMap>(to_insert);
   ASSERT_SEQ_EQ(Pairs(Sorted(to_insert)), map);
 }
 
@@ -306,7 +183,7 @@ TEST(ArraySortedMap, BalanceProblem) {
 // TODO(wilhuff): IndexOf
 
 TEST(ArraySortedMap, AvoidsCopying) {
-  IntMap map = IntMap().insert(10, 20);
+  IntMap map = IntMap{}.insert(10, 20);
   auto found = map.find(10);
   ASSERT_NE(found, map.end());
   EXPECT_EQ(20, found->second);
@@ -321,6 +198,7 @@ TEST(ArraySortedMap, AvoidsCopying) {
   EXPECT_EQ(found, duped_found);
 }
 
+}  // namespace impl
 }  // namespace immutable
 }  // namespace firestore
 }  // namespace firebase
diff --git a/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc b/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc
new file mode 100644
index 0000000..44dca50
--- /dev/null
+++ b/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc
@@ -0,0 +1,76 @@
+/*
+ * Copyright 2018 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "Firestore/core/src/firebase/firestore/immutable/sorted_map.h"
+
+#include <numeric>
+#include <random>
+
+#include "Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h"
+#include "Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h"
+#include "Firestore/core/src/firebase/firestore/util/secure_random.h"
+
+#include "Firestore/core/test/firebase/firestore/immutable/testing.h"
+#include "gtest/gtest.h"
+
+using firebase::firestore::immutable::impl::SortedMapBase;
+
+namespace firebase {
+namespace firestore {
+namespace immutable {
+
+using SizeType = SortedMapBase::size_type;
+
+template <typename MapType>
+struct TestPolicy {
+  // TODO(mcg): increase beyond what ArraySortedMap supports
+  static const SizeType kFixedSize = 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);
+  }
+};
+
+typedef ::testing::Types<SortedMap<int, int>,
+                         impl::ArraySortedMap<int, int>,
+                         impl::TreeSortedMap<int, int>>
+    TestedTypes;
+TYPED_TEST_CASE(SortedMapTest, TestedTypes);
+
+TYPED_TEST(SortedMapTest, EmptySize) {
+  TypeParam map;
+  EXPECT_TRUE(map.empty());
+  EXPECT_EQ(0u, map.size());
+}
+
+TYPED_TEST(SortedMapTest, Empty) {
+  TypeParam map = TypeParam{}.insert(10, 10).erase(10);
+  EXPECT_TRUE(map.empty());
+  EXPECT_EQ(0u, map.size());
+
+  // TODO(wilhuff): re-add find
+  // EXPECT_TRUE(NotFound(map, 1));
+  // EXPECT_TRUE(NotFound(map, 10));
+}
+
+}  // namespace immutable
+}  // namespace firestore
+}  // namespace firebase
diff --git a/Firestore/core/test/firebase/firestore/immutable/testing.h b/Firestore/core/test/firebase/firestore/immutable/testing.h
new file mode 100644
index 0000000..337140f
--- /dev/null
+++ b/Firestore/core/test/firebase/firestore/immutable/testing.h
@@ -0,0 +1,162 @@
+/*
+ * Copyright 2018 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef FIRESTORE_CORE_TEST_FIREBASE_FIRESTORE_IMMUTABLE_TESTING_H_
+#define FIRESTORE_CORE_TEST_FIREBASE_FIRESTORE_IMMUTABLE_TESTING_H_
+
+#include <algorithm>
+#include <utility>
+#include <vector>
+
+#include "Firestore/core/src/firebase/firestore/util/secure_random.h"
+#include "gtest/gtest.h"
+
+namespace firebase {
+namespace firestore {
+namespace immutable {
+
+template <typename Container, typename K>
+testing::AssertionResult NotFound(const Container& map, const K& key) {
+  auto found = map.find(key);
+  if (found == map.end()) {
+    return testing::AssertionSuccess();
+  } else {
+    return testing::AssertionFailure()
+           << "Should not have found (" << found->first << ", " << found->second
+           << ")";
+  }
+}
+
+template <typename Container, typename K, typename V>
+testing::AssertionResult Found(const Container& map,
+                               const K& key,
+                               const V& expected) {
+  auto found = map.find(key);
+  if (found == map.end()) {
+    return testing::AssertionFailure() << "Did not find key " << key;
+  }
+  if (found->second == expected) {
+    return testing::AssertionSuccess();
+  } else {
+    return testing::AssertionFailure() << "Found entry was (" << found->first
+                                       << ", " << found->second << ")";
+  }
+}
+
+/** Creates an empty vector (for readability). */
+inline std::vector<int> Empty() {
+  return {};
+}
+
+/**
+ * Creates a vector containing a sequence of integers from the given starting
+ * element up to, but not including, the given end element, with values
+ * incremented by the given step.
+ *
+ * If step is negative the sequence is in descending order (but still starting
+ * at start and ending before end).
+ */
+inline std::vector<int> Sequence(int start, int end, int step = 1) {
+  std::vector<int> result;
+  if (step > 0) {
+    for (int i = start; i < end; i += step) {
+      result.push_back(i);
+    }
+  } else {
+    for (int i = start; i > end; i += step) {
+      result.push_back(i);
+    }
+  }
+  return result;
+}
+
+/**
+ * Creates a vector containing a sequence of integers with the given number of
+ * elements, from zero up to, but not including the given value.
+ */
+inline std::vector<int> Sequence(int num_elements) {
+  return Sequence(0, num_elements);
+}
+
+/**
+ * Creates a copy of the given vector with contents shuffled randomly.
+ */
+inline std::vector<int> Shuffled(const std::vector<int>& values) {
+  std::vector<int> result{values};
+  util::SecureRandom rng;
+  std::shuffle(result.begin(), result.end(), rng);
+  return result;
+}
+
+/**
+ * Creates a copy of the given vector with contents sorted.
+ */
+inline std::vector<int> Sorted(const std::vector<int>& values) {
+  std::vector<int> result{values};
+  std::sort(result.begin(), result.end());
+  return result;
+}
+
+/**
+ * Creates a copy of the given vector with contents reversed.
+ */
+inline std::vector<int> Reversed(const std::vector<int>& values) {
+  std::vector<int> result{values};
+  std::reverse(result.begin(), result.end());
+  return result;
+}
+
+/**
+ * Creates a vector of pairs where each pair has the same first and second
+ * corresponding to an element in the given vector.
+ */
+inline std::vector<std::pair<int, int>> Pairs(const std::vector<int>& values) {
+  std::vector<std::pair<int, int>> result;
+  for (auto&& value : values) {
+    result.emplace_back(value, value);
+  }
+  return result;
+}
+
+/**
+ * Creates a SortedMap by inserting a pair for each value in the vector.
+ * Each pair will have the same key and value.
+ */
+template <typename Container>
+Container ToMap(const std::vector<int>& values) {
+  Container result;
+  for (auto&& value : values) {
+    result = result.insert(value, value);
+  }
+  return result;
+}
+
+/**
+ * Appends the contents of the given container to a new vector.
+ */
+template <typename Container>
+std::vector<typename Container::value_type> Append(const Container& container) {
+  return {container.begin(), container.end()};
+}
+
+#define ASSERT_SEQ_EQ(x, y) ASSERT_EQ((x), Append(y));
+#define EXPECT_SEQ_EQ(x, y) EXPECT_EQ((x), Append(y));
+
+}  // namespace immutable
+}  // namespace firestore
+}  // namespace firebase
+
+#endif  // FIRESTORE_CORE_TEST_FIREBASE_FIRESTORE_IMMUTABLE_TESTING_H_
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
new file mode 100644
index 0000000..7a96b67
--- /dev/null
+++ b/Firestore/core/test/firebase/firestore/immutable/tree_sorted_map_test.cc
@@ -0,0 +1,40 @@
+/*
+ * Copyright 2018 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h"
+
+#include "Firestore/core/src/firebase/firestore/util/secure_random.h"
+#include "Firestore/core/test/firebase/firestore/immutable/testing.h"
+#include "gtest/gtest.h"
+
+namespace firebase {
+namespace firestore {
+namespace immutable {
+namespace impl {
+
+typedef TreeSortedMap<int, int> IntMap;
+
+TEST(TreeSortedMap, EmptySize) {
+  IntMap map;
+  EXPECT_TRUE(map.empty());
+  EXPECT_EQ(0u, map.size());
+  EXPECT_EQ(Color::Black, map.root().color());
+}
+
+}  // namespace impl
+}  // namespace immutable
+}  // namespace firestore
+}  // namespace firebase