rb map is iterable.

include/data/tools/rb_map.hpp

@@ -6,12 +6,15 @@
 #define DATA_MAP_RB
 
 #include <data/tools/ordered_list.hpp>
+#include <data/tools/linked_stack.hpp>
 #include <data/functional/map.hpp>
 #include <data/fold.hpp>
 #include <milewski/RBMap/RBMap.h>
 
 namespace data::tool {
 
+    template <typename K, typename V> struct rb_map_iterator;
+
     template <typename K, typename V>
     struct rb_map {
         using entry = data::entry<K, V>;
@@ -67,11 +70,59 @@ namespace data::tool {
 
         bool operator==(const rb_map& map) const;
 
-        bool operator!=(const rb_map& map) const {
-            return !(*this == map);
+        rb_map_iterator<K, V> begin() const;
+
+        rb_map_iterator<K, V> end() const;
+
+    };
+
+    template <typename K, typename V> 
+    struct rb_map_iterator {
+        using map = milewski::okasaki::RBMap<K, V>;
+        using node = milewski::okasaki::Node<K, V>;
+
+        ptr<const node> Map;
+        ptr<const node> Next;
+        linked_stack<ptr<const node>> Last;
+        int Index;
+
+        // we need this constructor in order to satisfy some 
+        // std concepts but it's not really good for anything. 
+        rb_map_iterator() : Map{nullptr}, Next{}, Last{}, Index{0} {}
+
+        // constructor for the end of a map. 
+        rb_map_iterator(const map *m, int size) : Map{m->_root}, Next{}, Last{}, Index{size} {}
+
+        // constructor for the beginning of a map. 
+        rb_map_iterator(const map *m) : Map{m->_root}, Next{m->_root}, Last{}, Index{0} {
+            go_left();
         }
 
+        rb_map_iterator operator++(int);
+        rb_map_iterator &operator++();
+
+        const data::entry<K, V> &operator*() const;
+
+        bool operator==(const rb_map_iterator i) const;
+        int operator-(const rb_map_iterator& i) const;
+
+        void go_left();
+
+    };
+}
+
+namespace std {
+    template <typename K, typename V> 
+    struct iterator_traits<data::tool::rb_map_iterator<K, V>> {
+        using value_type = remove_const_t<data::entry<K, V>>;
+        using difference_type = int;
+        using pointer = const remove_reference_t<data::entry<K, V>>*;
+        using reference = const data::entry<K, V>&;
+        using iterator_concept = input_iterator_tag;
     };
+}
+
+namespace data::tool {
 
     template <typename K, typename V>
     std::ostream inline &operator<<(std::ostream& o, const rb_map<K, V>& x) {
@@ -176,6 +227,69 @@ namespace data::tool {
         return Size;
     }
 
+    template <typename K, typename V>
+    rb_map_iterator<K, V> rb_map<K, V>::begin() const {
+        return rb_map_iterator<K, V>{&Map};
+    }
+
+    template <typename K, typename V>
+    rb_map_iterator<K, V> rb_map<K, V>::end() const {
+        return rb_map_iterator<K, V>{&Map, static_cast<int>(Size)};
+    }
+
+    template <typename K, typename V>
+    rb_map_iterator<K, V> rb_map_iterator<K, V>::operator++(int) {
+        auto x = *this;
+        ++(*this);
+        return x;
+    }
+
+    template <typename K, typename V>
+    rb_map_iterator<K, V> &rb_map_iterator<K, V>::operator++() {
+        if (Next == nullptr) return *this;
+        Index++;
+
+        if (Next->_rgt != nullptr) {
+            Next = Next->_rgt;
+            go_left();
+            return *this;
+        } 
+
+        if (!data::empty(Last)) {
+            Next = Last.first();
+            Last = Last.rest();
+            return *this;
+        }
+
+        Next = nullptr;
+        return *this;
+    }
+
+    template <typename K, typename V>
+    const data::entry<K, V> &rb_map_iterator<K, V>::operator*() const {
+        return Next->_entry;
+    }
+
+    template <typename K, typename V>
+    bool rb_map_iterator<K, V>::operator==(const rb_map_iterator i) const {
+        return Map == i.Map && Next == i.Next;
+    }
+
+    template <typename K, typename V>
+    int rb_map_iterator<K, V>::operator-(const rb_map_iterator& i) const {
+        if (Map == i.Map) return Index - i.Index;
+        return 0;
+    }
+
+    template <typename K, typename V>
+    void rb_map_iterator<K, V>::go_left() {
+        if (Next == nullptr) return;
+        while(Next->_lft != nullptr) {
+            Last = Last << Next;
+            Next = Next->_lft;
+        }
+    }
+
 }
 
 #endif

include/milewski/RBMap/RBMap.h

@@ -4,39 +4,42 @@
 #include <cassert>
 #include <memory>
 
+#include <data/functional/map.hpp>
+
 namespace milewski::okasaki {
 
     enum Color { R, B };
 
     // 1. No red node has a red child.
     // 2. Every path from rootKey to empty node contains the same
     // number of black nodes.
+    template<class K, class V>
+    struct Node
+    {
+        Node(Color c,
+            std::shared_ptr<const Node> const & lft,
+            const K key, V val,
+            std::shared_ptr<const Node> const & rgt)
+            : _c(c), _lft(lft), _entry(key, val), _rgt(rgt)
+        {}
+        Color _c;
+        std::shared_ptr<const Node> _lft;
+        data::entry<K, V> _entry;
+        std::shared_ptr<const Node> _rgt;
+    };
 
     template<class K, class V>
-    class RBMap
+    struct RBMap
     {
-        struct Node
-        {
-            Node(Color c,
-                std::shared_ptr<const Node> const & lft,
-                const K key, V val,
-                std::shared_ptr<const Node> const & rgt)
-                : _c(c), _lft(lft), _key(key), _val(val), _rgt(rgt)
-            {}
-            Color _c;
-            std::shared_ptr<const Node> _lft;
-            const K _key;
-            const V _val;
-            std::shared_ptr<const Node> _rgt;
-        };
+        using Node = okasaki::Node<K, V>;
 
         explicit RBMap(std::shared_ptr<const Node> const & node) : _root(node) {}
         Color rootColor() const
         {
             assert(!isEmpty());
             return _root->_c;
         }
-    public:
+        
         RBMap() : _root{nullptr} {}
 
         RBMap(Color c, RBMap const & lft, const K& key, const V& val, RBMap const & rgt)
@@ -51,13 +54,19 @@ namespace milewski::okasaki {
         const K& rootKey() const
         {
             assert(!isEmpty());
-            return _root->_key;
+            return _root->_entry.Key;
         }
 
         V& rootValue() const
         {
             assert(!isEmpty());
-            return const_cast<V&>(_root->_val);
+            return const_cast<V&>(_root->_entry.Value);
+        }
+
+        const data::entry<K, V>& root() const
+        {
+            assert(!isEmpty());
+            return _root->_entry;
         }
 
         RBMap left() const
@@ -142,7 +151,6 @@ namespace milewski::okasaki {
             return (rootColor() == B) ? 1 + lft : lft;
         }
 
-    private:
         RBMap ins(const K& x, const V& v) const
         {
             assert1();
@@ -252,7 +260,7 @@ namespace milewski::okasaki {
             assert(!isEmpty());
             return RBMap(c, left(), rootKey(), rootValue(), right());
         }
-    private:
+        
         std::shared_ptr<const Node> _root;
     };
 

test/testMap.cpp

@@ -74,5 +74,27 @@ namespace data {
 
         EXPECT_EQ(m1r2, m2);
     }
+
+    TEST(MapTest, TestIterate) {
+        map<int, int> empty_map{};
+
+        EXPECT_EQ(empty_map.begin(), empty_map.end());
+        EXPECT_EQ(empty_map.end() - empty_map.begin(), 0);
+
+        entry<int, int> first_entry{7, 7};
+        map<int, int> small_map{first_entry};
+
+        auto small_begin = small_map.begin();
+        auto small_end = small_map.end();
+
+        EXPECT_NE(small_begin, small_end);
+        EXPECT_EQ(small_end - small_begin, 1);
+
+        EXPECT_EQ(*small_begin, first_entry);
+
+        ++small_begin;
+        EXPECT_EQ(small_begin, small_end);
+
+    }
 }