feat(util): add maps based on red-black trees

Signed-off-by: Leonardo de Moura <leonardo@microsoft.com>

src/tests/util/CMakeLists.txt

@@ -31,6 +31,9 @@ add_test(memory ${CMAKE_CURRENT_BINARY_DIR}/memory)
 add_executable(rb_tree rb_tree.cpp)
 target_link_libraries(rb_tree ${EXTRA_LIBS})
 add_test(rb_tree ${CMAKE_CURRENT_BINARY_DIR}/rb_tree)
+add_executable(rb_map rb_map.cpp)
+target_link_libraries(rb_map ${EXTRA_LIBS})
+add_test(rb_map ${CMAKE_CURRENT_BINARY_DIR}/rb_map)
 add_executable(splay_tree splay_tree.cpp)
 target_link_libraries(splay_tree ${EXTRA_LIBS})
 add_test(splay_tree ${CMAKE_CURRENT_BINARY_DIR}/splay_tree)

src/tests/util/rb_map.cpp

@@ -0,0 +1,49 @@
+/*
+Copyright (c) 2013 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Author: Leonardo de Moura
+*/
+#include <iostream>
+#include <sstream>
+#include "util/test.h"
+#include "util/rb_map.h"
+#include "util/name.h"
+using namespace lean;
+
+struct int_cmp { int operator()(int i1, int i2) const { return i1 < i2 ? -1 : (i1 > i2 ? 1 : 0); } };
+
+typedef rb_map<int, name, int_cmp> int2name;
+
+static void tst0() {
+    int2name m1;
+    m1[10] = name("t1");
+    m1[20] = name("t2");
+    int2name m2(m1);
+    m2[10] = name("t3");
+    lean_assert(m1[10] == name("t1"));
+    lean_assert(m1[20] == name("t2"));
+    lean_assert(m2[10] == name("t3"));
+    lean_assert(m2[20] == name("t2"));
+    lean_assert(m2.size() == 2);
+    lean_assert(m2[100] == name());
+    lean_assert(m2.size() == 3);
+    lean_assert(m2[100] == name());
+    lean_assert(m2.size() == 3);
+}
+
+static void tst1() {
+    int2name m1, m2;
+    m1[10] = name("t1");
+    lean_assert(m1.size() == 1);
+    lean_assert(m2.size() == 0);
+    swap(m1, m2);
+    lean_assert(m2.size() == 1);
+    lean_assert(m1.size() == 0);
+}
+
+int main() {
+    tst0();
+    tst1();
+    return has_violations() ? 1 : 0;
+}

src/util/CMakeLists.txt

@@ -7,7 +7,7 @@ endif()
 add_library(util trace.cpp debug.cpp name.cpp name_set.cpp name_generator.cpp
   exception.cpp interrupt.cpp hash.cpp escaped.cpp bit_tricks.cpp
   safe_arith.cpp ascii.cpp memory.cpp shared_mutex.cpp realpath.cpp
-  script_state.cpp script_exception.cpp splay_map.cpp lua.cpp
+  script_state.cpp script_exception.cpp splay_map.cpp rb_map.cpp lua.cpp
   luaref.cpp stackinfo.cpp lean_path.cpp serializer.cpp ${THREAD_CPP})
 
 target_link_libraries(util ${LEAN_LIBS})

src/util/rb_map.cpp

@@ -0,0 +1,100 @@
+/*
+Copyright (c) 2014 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Author: Leonardo de Moura
+*/
+#include "util/lua.h"
+#include "util/luaref.h"
+#include "util/rb_map.h"
+
+namespace lean {
+typedef rb_map<luaref, luaref, luaref_lt_proc> lua_rb_map;
+DECL_UDATA(lua_rb_map)
+
+static int mk_lua_rb_map(lua_State * L) {
+    lua_rb_map r;
+    return push_lua_rb_map(L, r);
+}
+
+static int lua_rb_map_size(lua_State * L) {
+    lua_pushinteger(L, to_lua_rb_map(L, 1).size());
+    return 1;
+}
+
+static int lua_rb_map_contains(lua_State * L) {
+    lua_pushboolean(L, to_lua_rb_map(L, 1).contains(luaref(L, 2)));
+    return 1;
+}
+
+static int lua_rb_map_empty(lua_State * L) {
+    lua_pushboolean(L, to_lua_rb_map(L, 1).empty());
+    return 1;
+}
+
+static int lua_rb_map_insert(lua_State * L) {
+    to_lua_rb_map(L, 1).insert(luaref(L, 2), luaref(L, 3));
+    return 0;
+}
+
+static int lua_rb_map_erase(lua_State * L) {
+    to_lua_rb_map(L, 1).erase(luaref(L, 2));
+    return 0;
+}
+
+static int lua_rb_map_find(lua_State * L) {
+    lua_rb_map & m = to_lua_rb_map(L, 1);
+    luaref const * val = m.find(luaref(L, 2));
+    if (val) {
+        lean_assert(val->get_state() == L);
+        val->push();
+    } else {
+        lua_pushnil(L);
+    }
+    return 1;
+}
+
+static int lua_rb_map_copy(lua_State * L) {
+    return push_lua_rb_map(L, to_lua_rb_map(L, 1));
+}
+
+static int lua_rb_map_for_each(lua_State * L) {
+    // Remark: we take a copy of the map to make sure
+    // for_each will not crash if the map is updated while being
+    // traversed.
+    // The copy operation is very cheap O(1).
+    lua_rb_map m(to_lua_rb_map(L, 1)); // map
+    luaL_checktype(L, 2, LUA_TFUNCTION); // user-fun
+    m.for_each([&](luaref const & k, luaref const & v) {
+        lua_pushvalue(L, 2); // push user-fun
+        k.push();
+        v.push();
+        pcall(L, 2, 0, 0);
+        });
+    return 0;
+}
+
+static const struct luaL_Reg lua_rb_map_m[] = {
+    {"__gc",            lua_rb_map_gc}, // never throws
+    {"__len",           safe_function<lua_rb_map_size> },
+    {"contains",        safe_function<lua_rb_map_contains>},
+    {"size",            safe_function<lua_rb_map_size>},
+    {"empty",           safe_function<lua_rb_map_empty>},
+    {"insert",          safe_function<lua_rb_map_insert>},
+    {"erase",           safe_function<lua_rb_map_erase>},
+    {"find",            safe_function<lua_rb_map_find>},
+    {"copy",            safe_function<lua_rb_map_copy>},
+    {"for_each",        safe_function<lua_rb_map_for_each>},
+    {0, 0}
+};
+
+void open_rb_map(lua_State * L) {
+    luaL_newmetatable(L, lua_rb_map_mt);
+    lua_pushvalue(L, -1);
+    lua_setfield(L, -2, "__index");
+    setfuncs(L, lua_rb_map_m, 0);
+
+    SET_GLOBAL_FUN(mk_lua_rb_map,          "rb_map");
+    SET_GLOBAL_FUN(lua_rb_map_pred,        "is_rb_map");
+}
+}

src/util/rb_map.h

@@ -0,0 +1,95 @@
+/*
+Copyright (c) 2014 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Author: Leonardo de Moura
+*/
+#pragma once
+#include <utility>
+#include "util/lua.h"
+#include "util/pair.h"
+#include "util/rb_tree.h"
+
+namespace lean {
+/**
+   \brief Wrapper for implementing maps using red black trees.
+*/
+template<typename K, typename T, typename CMP>
+class rb_map : public CMP {
+public:
+    typedef std::pair<K, T> entry;
+private:
+    struct entry_cmp : public CMP {
+        entry_cmp(CMP const & c):CMP(c) {}
+        int operator()(entry const & e1, entry const & e2) const { return CMP::operator()(e1.first, e2.first); }
+    };
+    rb_tree<entry, entry_cmp> m_map;
+public:
+    rb_map(CMP const & cmp = CMP()):m_map(entry_cmp(cmp)) {}
+    friend void swap(rb_map & a, rb_map & b) { swap(a.m_map, b.m_map); }
+    bool empty() const { return m_map.empty(); }
+    void clear() { m_map.clear(); }
+    bool is_eqp(rb_map const & m) const { return m_map.is_eqp(m); }
+    unsigned size() const { return m_map.size(); }
+    void insert(K const & k, T const & v) { m_map.insert(mk_pair(k, v)); }
+    T const * find(K const & k) const { auto e = m_map.find(mk_pair(k, T())); return e ? &(e->second) : nullptr; }
+    bool contains(K const & k) const { return m_map.contains(mk_pair(k, T())); }
+    void erase(K const & k) { m_map.erase(mk_pair(k, T())); }
+
+    class ref {
+        rb_map & m_map;
+        K const &   m_key;
+    public:
+        ref(rb_map & m, K const & k):m_map(m), m_key(k) {}
+        ref & operator=(T const & v) { m_map.insert(m_key, v); return *this; }
+        operator T const &() const {
+            T const * e = m_map.find(m_key);
+            if (e) {
+                return *e;
+            } else {
+                m_map.insert(m_key, T());
+                return *(m_map.find(m_key));
+            }
+        }
+    };
+
+    /**
+       \brief Returns a reference to the value that is mapped to a key equivalent to key,
+       performing an insertion if such key does not already exist.
+    */
+    ref operator[](K const & k) { return ref(*this, k); }
+
+    template<typename F>
+    void for_each(F f) const {
+        auto f_prime = [&](entry const & e) { f(e.first, e.second); };
+        return m_map.for_each(f_prime);
+    }
+
+    /** \brief (For debugging) Display the content of this splay map. */
+    friend std::ostream & operator<<(std::ostream & out, rb_map const & m) {
+        out << "{";
+        m.for_each([&out](K const & k, T const & v) {
+                out << k << " |-> " << v << "; ";
+            });
+        out << "}";
+        return out;
+    }
+};
+template<typename K, typename T, typename CMP>
+rb_map<K, T, CMP> insert(rb_map<K, T, CMP> const & m, K const & k, T const & v) {
+    auto r = m;
+    r.insert(k, v);
+    return r;
+}
+template<typename K, typename T, typename CMP>
+rb_map<K, T, CMP> erase(rb_map<K, T, CMP> const & m, K const & k) {
+    auto r = m;
+    r.erase(k);
+    return r;
+}
+template<typename K, typename T, typename CMP, typename F>
+void for_each(rb_map<K, T, CMP> const & m, F f) {
+    return m.for_each(f);
+}
+void open_rb_map(lua_State * L);
+}

src/util/rb_tree.h

@@ -278,6 +278,13 @@ class rb_tree : public CMP {
     node m_root;
 
 public:
+    rb_tree(CMP const & cmp = CMP()):CMP(cmp) {}
+    rb_tree(rb_tree const & s):CMP(s), m_root(s.m_root) {}
+    rb_tree(rb_tree && s):CMP(s), m_root(s.m_root) {}
+
+    rb_tree & operator=(rb_tree const & s) { m_root = s.m_root; return *this; }
+    rb_tree & operator=(rb_tree && s) { m_root = s.m_root; return *this; }
+
     void insert(T const & v) {
         m_root = set_black(insert(m_root.steal(), v));
         lean_assert(check_invariant());