Implement iterators for our immutable maps (#1132)

* Add a minimal LlrbNodeIterator

* Remove fixed_size type parameter from FixedArray

The parameter wasn't that useful and caused problems in trying to define
dependent iterator types.

* Add begin()/end() to SortedMap.

7 files changed, 434 insertions, 8 deletions

diff --git a/Firestore/core/src/firebase/firestore/immutable/CMakeLists.txt b/Firestore/core/src/firebase/firestore/immutable/CMakeLists.txt
index 543f912..accf5c5 100644
--- a/Firestore/core/src/firebase/firestore/immutable/CMakeLists.txt
+++ b/Firestore/core/src/firebase/firestore/immutable/CMakeLists.txt
@@ -17,10 +17,12 @@ cc_library(
   SOURCES
     array_sorted_map.h
     llrb_node.h
+    llrb_node_iterator.h
     map_entry.h
     sorted_map.h
     sorted_map_base.h
     sorted_map_base.cc
+    sorted_map_iterator.h
     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 adba962..6808297 100644
--- a/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h
+++ b/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h
@@ -47,13 +47,12 @@ namespace impl {
  * to a FixedArray.
  *
  * @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, SortedMapBase::size_type fixed_size>
+template <typename T>
 class FixedArray {
  public:
   using size_type = SortedMapBase::size_type;
-  using array_type = std::array<T, fixed_size>;
+  using array_type = std::array<T, SortedMapBase::kFixedSize>;
   using iterator = typename array_type::iterator;
   using const_iterator = typename array_type::const_iterator;
 
@@ -73,7 +72,7 @@ class FixedArray {
   void append(SourceIterator src_begin, SourceIterator src_end) {
     auto appending = static_cast<size_type>(src_end - src_begin);
     auto new_size = size_ + appending;
-    FIREBASE_ASSERT(new_size <= fixed_size);
+    FIREBASE_ASSERT(new_size <= SortedMapBase::kFixedSize);
 
     std::copy(src_begin, src_end, end());
     size_ = new_size;
@@ -84,7 +83,7 @@ class FixedArray {
    */
   void append(T&& value) {
     size_type new_size = size_ + 1;
-    FIREBASE_ASSERT(new_size <= fixed_size);
+    FIREBASE_ASSERT(new_size <= SortedMapBase::kFixedSize);
 
     *end() = std::move(value);
     size_ = new_size;
@@ -132,7 +131,7 @@ class ArraySortedMap : public SortedMapBase {
   /**
    * The type of the fixed-size array containing entries of value_type.
    */
-  using array_type = FixedArray<value_type, kFixedSize>;
+  using array_type = FixedArray<value_type>;
   using const_iterator = typename array_type::const_iterator;
 
   using array_pointer = std::shared_ptr<const array_type>;
diff --git a/Firestore/core/src/firebase/firestore/immutable/llrb_node.h b/Firestore/core/src/firebase/firestore/immutable/llrb_node.h
index bfd0f8a..0294efe 100644
--- a/Firestore/core/src/firebase/firestore/immutable/llrb_node.h
+++ b/Firestore/core/src/firebase/firestore/immutable/llrb_node.h
@@ -20,6 +20,7 @@
 #include <memory>
 #include <utility>
 
+#include "Firestore/core/src/firebase/firestore/immutable/llrb_node_iterator.h"
 #include "Firestore/core/src/firebase/firestore/immutable/sorted_map_base.h"
 
 namespace firebase {
@@ -48,6 +49,7 @@ class LlrbNode : public SortedMapBase {
    * The type of the entries stored in the map.
    */
   using value_type = std::pair<K, V>;
+  using const_iterator = LlrbNodeIterator<LlrbNode<K, V>>;
 
   /**
    * Constructs an empty node.
diff --git a/Firestore/core/src/firebase/firestore/immutable/llrb_node_iterator.h b/Firestore/core/src/firebase/firestore/immutable/llrb_node_iterator.h
new file mode 100644
index 0000000..d113d34
--- /dev/null
+++ b/Firestore/core/src/firebase/firestore/immutable/llrb_node_iterator.h
@@ -0,0 +1,181 @@
+/*
+ * 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_ITERATOR_H_
+#define FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_LLRB_NODE_ITERATOR_H_
+
+#include <iterator>
+#include <stack>
+#include <utility>
+
+#include "Firestore/core/src/firebase/firestore/immutable/llrb_node.h"
+#include "Firestore/core/src/firebase/firestore/util/comparison.h"
+#include "Firestore/core/src/firebase/firestore/util/firebase_assert.h"
+
+namespace firebase {
+namespace firestore {
+namespace immutable {
+namespace impl {
+
+/**
+ * A forward iterator for traversing LlrbNodes. LlrbNodes represent the nodes
+ * in a tree implementing a sorted map so iterating with LlrbNodeIterator is
+ * an in-order traversal of the map.
+ *
+ * ## Complexity
+ *
+ * LlrbNode is an immutable tree, where mutations create new trees without
+ * invalidating any of the old instances. This means the tree cannot contain
+ * parent pointers and thus this iterator implementation must keep an explicit
+ * stack.
+ *
+ * For an underlying tree of size `n`:
+ *
+ *   * LlrbNodeIterator uses `O(lg(n))` memory for its stack, and
+ *   * incrementing an iterator is an `O(lg(n))` operation.
+ *
+ * ## Invalidation and Comparison
+ *
+ * LlrbNodeIterators compare based on the identity of the nodes themselves,
+ * not based on the values of the keys in the nodes. When adding and removing
+ * the same key and iterator obtained before and after will not compare equal.
+ *
+ * LlrbNodeIterators are not invalidated in any conventional sense because
+ * mutations of the underlying tree create new trees. Together this means that
+ * any given version of the tree can be iterated over from the same iterator
+ * repeatedly, but a "mutable view" of the tree kept by replacing the pointer
+ * to the root is effectively invalidated on each mutation.
+ *
+ * Note: LlrbNodeIterator does not extend the lifetime of its underlying tree.
+ */
+template <typename N>
+class LlrbNodeIterator {
+ public:
+  using node_type = N;
+  using key_type = typename node_type::first_type;
+
+  using stack_type = std::stack<const node_type*>;
+
+  using iterator_category = std::forward_iterator_tag;
+  using value_type = typename node_type::value_type;
+
+  using pointer = typename node_type::value_type const*;
+  using reference = typename node_type::value_type const&;
+  using difference_type = std::ptrdiff_t;
+
+  /**
+   * Constructs an iterator starting at the first node in the iteration
+   * sequence of the tree represented by the given root node (i.e. it points at
+   * the left-most node).
+   */
+  static LlrbNodeIterator Begin(const node_type* root) {
+    stack_type stack;
+    AccumulateLeft(root, &stack);
+    return LlrbNodeIterator{std::move(stack)};
+  }
+
+  /**
+   * Constructs an iterator pointing at the end of the iteration sequence of the
+   * tree pointed to by the given node (i.e. one past the right-most node)
+   */
+  static LlrbNodeIterator End() {
+    return LlrbNodeIterator{stack_type{}};
+  }
+
+  // Default constructor to conform to the requirements of ForwardIterator
+  LlrbNodeIterator() {
+  }
+
+  /**
+   * Returns true if this iterator points at the end of the iteration sequence.
+   */
+  bool is_end() const {
+    return stack_.empty();
+  }
+
+  /**
+   * Returns the address of the entry in the node that this iterator points to.
+   * This can only be called if `end()` is false.
+   */
+  pointer get() const {
+    FIREBASE_ASSERT(!is_end());
+    return &(stack_.top()->entry());
+  }
+
+  reference operator*() const {
+    return *get();
+  }
+
+  pointer operator->() const {
+    return get();
+  }
+
+  LlrbNodeIterator& operator++() {
+    FIREBASE_ASSERT(!is_end());
+
+    // Pop the stack, moving the currently pointed to node to the parent.
+    const node_type* node = stack_.top();
+    stack_.pop();
+
+    // If the popped node has a right subtree that has to precede the parent in
+    // the iteration order so push those on.
+    node = &node->right();
+    AccumulateLeft(node, &stack_);
+
+    return *this;
+  }
+
+  LlrbNodeIterator operator++(int /*unused*/) {
+    LlrbNodeIterator result = *this;
+    ++*this;
+    return result;
+  }
+
+  friend bool operator==(const LlrbNodeIterator& a, const LlrbNodeIterator& b) {
+    if (a.is_end()) {
+      return b.is_end();
+    } else if (b.is_end()) {
+      return false;
+    } else {
+      const key_type& left_key = a.get()->first;
+      const key_type& right_key = b.get()->first;
+      return left_key == right_key;
+    }
+  }
+
+  bool operator!=(const LlrbNodeIterator& b) const {
+    return !(*this == b);
+  }
+
+ private:
+  explicit LlrbNodeIterator(stack_type&& stack) : stack_(std::move(stack)) {
+  }
+
+  static void AccumulateLeft(const node_type* node, stack_type* stack) {
+    for (; !node->empty(); node = &node->left()) {
+      stack->push(node);
+    }
+  }
+
+  stack_type stack_;
+};
+
+}  // namespace impl
+}  // namespace immutable
+}  // namespace firestore
+}  // namespace firebase
+
+#endif  // FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_LLRB_NODE_ITERATOR_H_
diff --git a/Firestore/core/src/firebase/firestore/immutable/sorted_map.h b/Firestore/core/src/firebase/firestore/immutable/sorted_map.h
index b7729b3..7c8c832 100644
--- a/Firestore/core/src/firebase/firestore/immutable/sorted_map.h
+++ b/Firestore/core/src/firebase/firestore/immutable/sorted_map.h
@@ -21,6 +21,7 @@
 
 #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/sorted_map_iterator.h"
 #include "Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h"
 #include "Firestore/core/src/firebase/firestore/util/comparison.h"
 
@@ -40,6 +41,11 @@ class SortedMap : public impl::SortedMapBase {
   using array_type = impl::ArraySortedMap<K, V, C>;
   using tree_type = impl::TreeSortedMap<K, V, C>;
 
+  using const_iterator = impl::SortedMapIterator<
+      value_type,
+      typename impl::FixedArray<value_type>::const_iterator,
+      typename impl::LlrbNode<K, V>::const_iterator>;
+
   /**
    * Creates an empty SortedMap.
    */
@@ -71,7 +77,7 @@ class SortedMap : public impl::SortedMapBase {
     }
   }
 
-  SortedMap(SortedMap&& other) : tag_{other.tag_} {
+  SortedMap(SortedMap&& other) noexcept : tag_{other.tag_} {
     switch (tag_) {
       case Tag::Array:
         new (&array_) array_type{std::move(other.array_)};
@@ -110,7 +116,7 @@ class SortedMap : public impl::SortedMapBase {
     return *this;
   }
 
-  SortedMap& operator=(SortedMap&& other) {
+  SortedMap& operator=(SortedMap&& other) noexcept {
     if (tag_ == other.tag_) {
       switch (tag_) {
         case Tag::Array:
@@ -195,6 +201,33 @@ class SortedMap : public impl::SortedMapBase {
     FIREBASE_UNREACHABLE();
   }
 
+  /**
+   * Returns an iterator pointing to the first entry in the map. If there are
+   * no entries in the map, begin() == end().
+   */
+  const_iterator begin() const {
+    switch (tag_) {
+      case Tag::Array:
+        return const_iterator{array_.begin()};
+      case Tag::Tree:
+        return const_iterator{tree_.begin()};
+    }
+    FIREBASE_UNREACHABLE();
+  }
+
+  /**
+   * Returns an iterator pointing past the last entry in the map.
+   */
+  const_iterator end() const {
+    switch (tag_) {
+      case Tag::Array:
+        return const_iterator{array_.end()};
+      case Tag::Tree:
+        return const_iterator{tree_.end()};
+    }
+    FIREBASE_UNREACHABLE();
+  }
+
  private:
   explicit SortedMap(array_type&& array)
       : tag_{Tag::Array}, array_{std::move(array)} {
diff --git a/Firestore/core/src/firebase/firestore/immutable/sorted_map_iterator.h b/Firestore/core/src/firebase/firestore/immutable/sorted_map_iterator.h
new file mode 100644
index 0000000..36f3967
--- /dev/null
+++ b/Firestore/core/src/firebase/firestore/immutable/sorted_map_iterator.h
@@ -0,0 +1,191 @@
+/*
+ * 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_ITERATOR_H_
+#define FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_SORTED_MAP_ITERATOR_H_
+
+#include <utility>
+
+#include "Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h"
+#include "Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h"
+
+namespace firebase {
+namespace firestore {
+namespace immutable {
+namespace impl {
+
+template <typename V, typename ArrayIter, typename TreeIter>
+class SortedMapIterator {
+ public:
+  using iterator_category = std::forward_iterator_tag;
+  using value_type = V;
+  using pointer = const value_type*;
+  using reference = const value_type&;
+  using difference_type = std::ptrdiff_t;
+
+ public:
+  // Default constructor to conform to the requirements of ForwardIterator
+  SortedMapIterator() : tag_{Tag::Array}, array_iter_{} {
+  }
+
+  explicit SortedMapIterator(ArrayIter&& delegate)
+      : tag_{Tag::Array}, array_iter_{std::move(delegate)} {
+  }
+
+  explicit SortedMapIterator(TreeIter&& delegate)
+      : tag_{Tag::Tree}, tree_iter_{std::move(delegate)} {
+  }
+
+  SortedMapIterator(const SortedMapIterator& other) : tag_(other.tag_) {
+    switch (tag_) {
+      case Tag::Array:
+        new (&array_iter_) ArrayIter{other.array_iter_};
+        break;
+      case Tag::Tree:
+        new (&tree_iter_) TreeIter{other.tree_iter_};
+        break;
+    }
+  }
+
+  SortedMapIterator(SortedMapIterator&& other) : tag_(other.tag_) {
+    switch (tag_) {
+      case Tag::Array:
+        new (&array_iter_) ArrayIter{std::move(other.array_iter_)};
+        break;
+      case Tag::Tree:
+        new (&tree_iter_) TreeIter{std::move(other.tree_iter_)};
+        break;
+    }
+  }
+
+  ~SortedMapIterator() {
+    switch (tag_) {
+      case Tag::Array:
+        array_iter_.~ArrayIter();
+        break;
+      case Tag::Tree:
+        tree_iter_.~TreeIter();
+        break;
+    }
+  }
+
+  SortedMapIterator& operator=(const SortedMapIterator& other) {
+    if (tag_ == other.tag_) {
+      switch (tag_) {
+        case Tag::Array:
+          array_iter_ = other.array_iter_;
+          break;
+        case Tag::Tree:
+          tree_iter_ = other.tree_iter_;
+          break;
+      }
+    } else {
+      this->~SortedMapIterator();
+      new (this) SortedMapIterator(other);
+    }
+    return *this;
+  }
+
+  SortedMapIterator& operator=(SortedMapIterator&& other) {
+    if (tag_ == other.tag_) {
+      switch (tag_) {
+        case Tag::Array:
+          array_iter_ = std::move(other.array_iter_);
+          break;
+        case Tag::Tree:
+          tree_iter_ = std::move(other.tree_iter_);
+          break;
+      }
+    } else {
+      this->~SortedMapIterator();
+      new (this) SortedMapIterator(std::move(other));
+    }
+    return *this;
+  }
+
+  pointer get() const {
+    switch (tag_) {
+      case Tag::Array:
+        return array_iter_;
+      case Tag::Tree:
+        return tree_iter_.get();
+    }
+  }
+
+  reference operator*() const {
+    return *get();
+  }
+
+  pointer operator->() const {
+    return get();
+  }
+
+  SortedMapIterator& operator++() {
+    switch (tag_) {
+      case Tag::Array:
+        ++array_iter_;
+        break;
+      case Tag::Tree:
+        ++tree_iter_;
+        break;
+    }
+    return *this;
+  }
+
+  SortedMapIterator operator++(int /*unused*/) {
+    SortedMapIterator result = *this;
+    ++*this;
+    return result;
+  }
+
+  friend bool operator==(const SortedMapIterator& a,
+                         const SortedMapIterator& b) {
+    if (a.tag_ != b.tag_) {
+      return false;
+    }
+
+    switch (a.tag_) {
+      case Tag::Array:
+        return a.array_iter_ == b.array_iter_;
+      case Tag::Tree:
+        return a.tree_iter_ == b.tree_iter_;
+    }
+    FIREBASE_UNREACHABLE();
+  }
+
+  bool operator!=(const SortedMapIterator& b) const {
+    return !(*this == b);
+  }
+
+ private:
+  enum class Tag {
+    Array,
+    Tree,
+  };
+
+  Tag tag_;
+  union {
+    ArrayIter array_iter_;
+    TreeIter tree_iter_;
+  };
+};
+
+}  // namespace impl
+}  // namespace immutable
+}  // namespace firestore
+}  // namespace firebase
+
+#endif  // FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_IMMUTABLE_SORTED_MAP_ITERATOR_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
index ff8c9f9..c5eddc2 100644
--- a/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h
+++ b/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h
@@ -51,6 +51,7 @@ class TreeSortedMap : public SortedMapBase, private util::ComparatorHolder<C> {
    * The type of the node containing entries of value_type.
    */
   using node_type = LlrbNode<K, V>;
+  using const_iterator = typename node_type::const_iterator;
 
   /**
    * Creates an empty TreeSortedMap.
@@ -110,6 +111,23 @@ class TreeSortedMap : public SortedMapBase, private util::ComparatorHolder<C> {
     return root_;
   }
 
+  /**
+   * Returns a forward iterator pointing to the first entry in the map. If there
+   * are no entries in the map, begin() == end().
+   *
+   * See LlrbNodeIterator for details
+   */
+  const_iterator begin() const {
+    return const_iterator::Begin(&root_);
+  }
+
+  /**
+   * Returns an iterator pointing past the last entry in the map.
+   */
+  const_iterator end() const {
+    return const_iterator::End();
+  }
+
  private:
   TreeSortedMap(node_type&& root, const C& comparator) noexcept
       : util::ComparatorHolder<C>{comparator}, root_{std::move(root)} {