232 lines
7.5 KiB
C++
232 lines
7.5 KiB
C++
/*
|
|
Copyright (c) 2015 Microsoft Corporation. All rights reserved.
|
|
Released under Apache 2.0 license as described in the file LICENSE.
|
|
|
|
Author: Leonardo de Moura
|
|
*/
|
|
#pragma once
|
|
#include "util/rb_map.h"
|
|
#include "util/optional.h"
|
|
|
|
namespace lean {
|
|
/** \brief (template for) Union-find datastructure that "explains" implied equalities.
|
|
We use functional datastructures to be able to have a O(1) copy operation.
|
|
|
|
Each join/union is decorated with a justification.
|
|
|
|
\c cmp implements a total order on \c node. That is, it provides the method:
|
|
int operator()(node const & n1, node const & n2) const
|
|
s.t. the result is negative when n1 < n2, 0 if n1 == n2, and positive if n1 > n2.
|
|
|
|
The implementation also provides a method to traverse the elements of an equivalence
|
|
class. The implementation is based on a datastructure used in the Simplify theorem prover.
|
|
|
|
Since it provides extra functionality, it does not implement the O(n*alpha(n)) amortized time
|
|
per operation algorithm.
|
|
*/
|
|
template<typename node, typename jst, typename cmp>
|
|
class union_find : private cmp {
|
|
rb_map<node, node, cmp> m_root;
|
|
rb_map<node, node, cmp> m_next;
|
|
rb_map<node, unsigned, cmp> m_rank;
|
|
rb_map<node, pair<node, jst>, cmp> m_jst;
|
|
|
|
bool is_equal(node const & n1, node const & n2) const {
|
|
return cmp::operator()(n1, n2) == 0;
|
|
}
|
|
|
|
unsigned rank(node const & n) const {
|
|
if (auto r = m_rank.find(n))
|
|
return *r;
|
|
else
|
|
return 0;
|
|
}
|
|
void set_rank(node const & n, unsigned r) { m_rank.insert(n, r); }
|
|
|
|
node const & root(node const & n) const {
|
|
if (auto r = m_root.find(n))
|
|
return *r;
|
|
else
|
|
return n;
|
|
}
|
|
void set_root(node const & n, node const & r) { m_root.insert(n, r); }
|
|
|
|
node const & next(node const & n) const {
|
|
if (auto r = m_next.find(n))
|
|
return *r;
|
|
else
|
|
return n;
|
|
}
|
|
void set_next(node const & n, node const & nx) { m_next.insert(n, nx); }
|
|
void set_justification(node const & n, node const & t, jst const & j) { m_jst.insert(n, mk_pair(t, j)); }
|
|
|
|
// for debugging purposes only
|
|
bool check_inv(node const & n) const {
|
|
node r = root(n);
|
|
unsigned sz = size(r);
|
|
node it = n;
|
|
do {
|
|
lean_assert_eq(root(it), r);
|
|
lean_assert(reaches(it, r));
|
|
lean_assert(size(it), sz);
|
|
it = next(it);
|
|
} while (!is_equal(it, n));
|
|
return true;
|
|
}
|
|
|
|
void join_core(node const & n1, node r1, node const & n2, node r2, jst const & j) {
|
|
// r1 will be the root of the resulting equivalence class.
|
|
DEBUG_CODE(unsigned sz1 = size(n1); unsigned sz2 = size(n2););
|
|
// Step 1) update m_jst
|
|
//
|
|
// Given justification paths
|
|
// n1 -> ... -> r1
|
|
// n2 -> ... -> r2
|
|
// we generate the path
|
|
// r2 -> ... -> n2 -> n1 -> ... -> r1
|
|
buffer<pair<node, jst>> trace;
|
|
node it2 = n2;
|
|
while (pair<node, jst> const * p = m_jst.find(it2)) {
|
|
trace.push_back(*p);
|
|
it2 = p->first;
|
|
}
|
|
lean_assert(is_equal(it2, r2));
|
|
unsigned i = trace.size();
|
|
while (i > 1) {
|
|
--i;
|
|
set_justification(trace[i].first, trace[i-1].first, trace[i].second);
|
|
}
|
|
if (i > 0) {
|
|
set_justification(trace[0].first, n2, trace[0].second);
|
|
}
|
|
set_justification(n2, n1, j);
|
|
|
|
// Step 2) update m_root of nodes in n2 equivalence class to r1
|
|
it2 = n2;
|
|
do {
|
|
set_root(it2, r1);
|
|
it2 = next(it2);
|
|
} while (!is_equal(it2, n2));
|
|
|
|
// Step 3) update m_next of r1 and r2
|
|
node next1 = next(r1);
|
|
node next2 = next(r2);
|
|
set_next(r1, next2);
|
|
set_next(r2, next1);
|
|
|
|
lean_assert(check_inv(r1));
|
|
lean_assert_eq(size(n1), sz1 + sz2);
|
|
}
|
|
|
|
/** \brief Return true if \c s reaches \c r by following m_jst edges */
|
|
bool reaches(node const & s, node const & r) const {
|
|
node it = s;
|
|
while (true) {
|
|
if (is_equal(it, r))
|
|
return true;
|
|
pair<node, jst> const * p = m_jst.find(it);
|
|
if (p) {
|
|
it = p->first;
|
|
} else {
|
|
return false;
|
|
}
|
|
}
|
|
}
|
|
|
|
void explain_core(node const & n1, node const & n2, node const & r, buffer<jst> & js) const {
|
|
lean_assert(is_equal(root(n1), r));
|
|
lean_assert(is_equal(root(n2), r));
|
|
node it1 = n1;
|
|
while (true) {
|
|
if (reaches(n2, it1)) {
|
|
// it is the common in the paths n1 -> r and n2 -> r
|
|
node it2 = n2;
|
|
unsigned sz1 = js.size();
|
|
while (true) {
|
|
if (is_equal(it2, it1)) {
|
|
std::reverse(js.begin() + sz1, js.end());
|
|
return;
|
|
}
|
|
pair<node, jst> const * p = m_jst.find(it2);
|
|
lean_assert(p);
|
|
js.push_back(p->second);
|
|
it2 = p->first;
|
|
}
|
|
} else {
|
|
pair<node, jst> const * p = m_jst.find(it1);
|
|
lean_assert(p);
|
|
js.push_back(p->second);
|
|
it1 = p->first;
|
|
}
|
|
}
|
|
}
|
|
|
|
public:
|
|
union_find(cmp const & c = cmp()):cmp(c) {}
|
|
|
|
/** \brief Merge the equivalence class of \c n1 with \c n2 using justification \c j. */
|
|
void join(node const & n1, node const & n2, jst const & j) {
|
|
node const & r1 = root(n1);
|
|
node const & r2 = root(n2);
|
|
if (is_equal(r1, r2))
|
|
return;
|
|
unsigned k1 = rank(n1);
|
|
unsigned k2 = rank(n2);
|
|
if (k1 > k2) {
|
|
join_core(n1, r1, n2, r2, j);
|
|
} else if (k1 == k2) {
|
|
join_core(n1, r1, n2, r2, j);
|
|
set_rank(n1, k1+1);
|
|
} else {
|
|
join_core(n2, r2, n1, r1, j);
|
|
}
|
|
}
|
|
|
|
/** \brief Return the size of the equivalence class containing \c n */
|
|
unsigned size(node const & n) const {
|
|
unsigned r = 0;
|
|
node it = n;
|
|
do {
|
|
lean_assert(is_eq(it, n));
|
|
r++;
|
|
it = next(it);
|
|
} while (!is_equal(it, n));
|
|
return r;
|
|
}
|
|
|
|
/** \brief Return the representative for the equivalence class containing \c n. */
|
|
node rep(node const & n) const { return root(n); }
|
|
|
|
/** \brief Return true iff \c n1 and \c n2 are in the same equivalence class. */
|
|
bool is_eq(node const & n1, node const & n2) const { return is_equal(rep(n1), rep(n2)); }
|
|
|
|
/** \brief For each node \c m in the equivalence class of \c n, execute <tt>f(m)</tt> */
|
|
template<typename F>
|
|
void for_each(node const & n, F f) const {
|
|
node it = n;
|
|
do {
|
|
lean_assert(is_eq(it, n));
|
|
f(it);
|
|
it = next(it);
|
|
} while (!is_equal(it, n));
|
|
}
|
|
|
|
/** \brief If is_eq(n1, n2), then return true and store the justifications that can be used to produce
|
|
a transitivity+symmetry proof for n1 = n2 */
|
|
bool explain(node const & n1, node const & n2, buffer<jst> & js) const {
|
|
node r1 = root(n1);
|
|
node r2 = root(n2);
|
|
if (is_equal(r1, r2)) {
|
|
if (rank(r1) >= rank(r2)) {
|
|
explain_core(n1, n2, r1, js);
|
|
} else {
|
|
explain_core(n2, n1, r1, js);
|
|
std::reverse(js.begin(), js.end());
|
|
}
|
|
return true;
|
|
} else {
|
|
return false;
|
|
}
|
|
}
|
|
};
|
|
}
|