diff --git a/src/tests/util/CMakeLists.txt b/src/tests/util/CMakeLists.txt index 7ec77af46..c47a4b1b1 100644 --- a/src/tests/util/CMakeLists.txt +++ b/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) diff --git a/src/tests/util/rb_map.cpp b/src/tests/util/rb_map.cpp new file mode 100644 index 000000000..ebb1756b2 --- /dev/null +++ b/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 +#include +#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 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; +} diff --git a/src/util/CMakeLists.txt b/src/util/CMakeLists.txt index b941c05e5..b5e828eef 100644 --- a/src/util/CMakeLists.txt +++ b/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}) diff --git a/src/util/rb_map.cpp b/src/util/rb_map.cpp new file mode 100644 index 000000000..7450bb36b --- /dev/null +++ b/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 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 }, + {"contains", safe_function}, + {"size", safe_function}, + {"empty", safe_function}, + {"insert", safe_function}, + {"erase", safe_function}, + {"find", safe_function}, + {"copy", safe_function}, + {"for_each", safe_function}, + {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"); +} +} diff --git a/src/util/rb_map.h b/src/util/rb_map.h new file mode 100644 index 000000000..11283f23d --- /dev/null +++ b/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 +#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 +class rb_map : public CMP { +public: + typedef std::pair 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 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 + 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 +rb_map insert(rb_map const & m, K const & k, T const & v) { + auto r = m; + r.insert(k, v); + return r; +} +template +rb_map erase(rb_map const & m, K const & k) { + auto r = m; + r.erase(k); + return r; +} +template +void for_each(rb_map const & m, F f) { + return m.for_each(f); +} +void open_rb_map(lua_State * L); +} diff --git a/src/util/rb_tree.h b/src/util/rb_tree.h index 1b7f6248f..2d8670b69 100644 --- a/src/util/rb_tree.h +++ b/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());