diff --git a/include/data/tools/rb_map.hpp b/include/data/tools/rb_map.hpp index 39030f1a..4dc29d06 100644 --- a/include/data/tools/rb_map.hpp +++ b/include/data/tools/rb_map.hpp @@ -6,12 +6,15 @@ #define DATA_MAP_RB #include +#include #include #include #include namespace data::tool { + template struct rb_map_iterator; + template struct rb_map { using entry = data::entry; @@ -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 begin() const; + + rb_map_iterator end() const; + + }; + + template + struct rb_map_iterator { + using map = milewski::okasaki::RBMap; + using node = milewski::okasaki::Node; + + ptr Map; + ptr Next; + linked_stack> 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 &operator*() const; + + bool operator==(const rb_map_iterator i) const; + int operator-(const rb_map_iterator& i) const; + + void go_left(); + + }; +} + +namespace std { + template + struct iterator_traits> { + using value_type = remove_const_t>; + using difference_type = int; + using pointer = const remove_reference_t>*; + using reference = const data::entry&; + using iterator_concept = input_iterator_tag; }; +} + +namespace data::tool { template std::ostream inline &operator<<(std::ostream& o, const rb_map& x) { @@ -176,6 +227,69 @@ namespace data::tool { return Size; } + template + rb_map_iterator rb_map::begin() const { + return rb_map_iterator{&Map}; + } + + template + rb_map_iterator rb_map::end() const { + return rb_map_iterator{&Map, static_cast(Size)}; + } + + template + rb_map_iterator rb_map_iterator::operator++(int) { + auto x = *this; + ++(*this); + return x; + } + + template + rb_map_iterator &rb_map_iterator::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 + const data::entry &rb_map_iterator::operator*() const { + return Next->_entry; + } + + template + bool rb_map_iterator::operator==(const rb_map_iterator i) const { + return Map == i.Map && Next == i.Next; + } + + template + int rb_map_iterator::operator-(const rb_map_iterator& i) const { + if (Map == i.Map) return Index - i.Index; + return 0; + } + + template + void rb_map_iterator::go_left() { + if (Next == nullptr) return; + while(Next->_lft != nullptr) { + Last = Last << Next; + Next = Next->_lft; + } + } + } #endif diff --git a/include/milewski/RBMap/RBMap.h b/include/milewski/RBMap/RBMap.h index aef55191..fcb0623e 100644 --- a/include/milewski/RBMap/RBMap.h +++ b/include/milewski/RBMap/RBMap.h @@ -4,6 +4,8 @@ #include #include +#include + namespace milewski::okasaki { enum Color { R, B }; @@ -11,24 +13,25 @@ namespace milewski::okasaki { // 1. No red node has a red child. // 2. Every path from rootKey to empty node contains the same // number of black nodes. + template + struct Node + { + Node(Color c, + std::shared_ptr const & lft, + const K key, V val, + std::shared_ptr const & rgt) + : _c(c), _lft(lft), _entry(key, val), _rgt(rgt) + {} + Color _c; + std::shared_ptr _lft; + data::entry _entry; + std::shared_ptr _rgt; + }; template - class RBMap + struct RBMap { - struct Node - { - Node(Color c, - std::shared_ptr const & lft, - const K key, V val, - std::shared_ptr const & rgt) - : _c(c), _lft(lft), _key(key), _val(val), _rgt(rgt) - {} - Color _c; - std::shared_ptr _lft; - const K _key; - const V _val; - std::shared_ptr _rgt; - }; + using Node = okasaki::Node; explicit RBMap(std::shared_ptr const & node) : _root(node) {} Color rootColor() const @@ -36,7 +39,7 @@ namespace milewski::okasaki { 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(_root->_val); + return const_cast(_root->_entry.Value); + } + + const data::entry& 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 _root; }; diff --git a/test/testMap.cpp b/test/testMap.cpp index 39c3e3c1..2b8fd19e 100644 --- a/test/testMap.cpp +++ b/test/testMap.cpp @@ -74,5 +74,27 @@ namespace data { EXPECT_EQ(m1r2, m2); } + + TEST(MapTest, TestIterate) { + map empty_map{}; + + EXPECT_EQ(empty_map.begin(), empty_map.end()); + EXPECT_EQ(empty_map.end() - empty_map.begin(), 0); + + entry first_entry{7, 7}; + map 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); + + } }