source: src/Common/PersistentDisjointSet.h@ eff03a94

//
// Cforall Version 1.0.0 Copyright (C) 2015 University of Waterloo
//
// The contents of this file are covered under the licence agreement in the
// file "LICENCE" distributed with Cforall.
//
// PersistentDisjointSet.h --
//
// Author           : Aaron B. Moss
// Created On       : Wed Jun 13 16:31:00 2018
// Last Modified By : Aaron B. Moss
// Last Modified On : Wed Jun 13 16:31:00 2018
// Update Count     : 1
//

#pragma once

#include <cassert>
#include <functional>
#include <unordered_map>
#include <utility>

#include "GC.h"

/// Persistent disjoint-set data structure based on the persistent array in 
/// Conchon & Filliatre "A Persistent Union-Find Data Structure". Path 
/// compression is not performed (to lower cost of persistent rollback). 
/// Auxilliary lists are kept for efficient retrieval of class members.
/// Find root should still operate in O(log k), for k the size of an 
/// equivalence class.

template<typename Elm, typename Hash = std::hash<Elm>, typename Eq = std::equal_to<Elm>>
class PersistentDisjointSet : public GC_Object {
public:
        /// Type of this class
        using Self = PersistentDisjointSet<Elm, Hash, Eq>;

        /// Types of version nodes
        enum Mode { 
                BASE,    ///< Root node of version tree
                ADD,     ///< Add key to set
                REM,     ///< Reverse add operation
                ADDTO,   ///< Merge one class root under another
                REMFROM  ///< Reverse addTo operation
        };

private:
        /// Type of node height
        using Height = unsigned char;

        /// Disjoint-set node
        struct Node {
                Elm parent;     ///< Parent node in equivalence class
                Elm next;       ///< Next node in equivalence class
                Height height;  ///< Tree height of the node

                template<typename E>
                Node(E&& e) : parent(e), next(std::forward<E>(e)), height(0) {}

                template<typename E, typename F>
                Node(E&& p, F&& n, Height h) 
                        : parent(std::forward<E>(p)), next(std::forward<F>(n)), height(h) {}
        };

        /// Type of class map
        using Base = std::unordered_map<Elm, Node, Hash, Eq>;

        /// Node inserted into underlying map as new equivalence class
        struct Add {
                Self* base;  ///< Modified map
                Elm root;    ///< Element added

                template<typename E>
                Add(Self* b, E&& r) : base(b), root(std::forward<E>(r)) {}
        };

        /// Two classes merged
        struct AddTo {
                Self* base;       ///< Modified map
                Elm root;         ///< Root node
                Elm child;        ///< Child node, formerly root of own class
                bool new_height;  ///< Did the root node's height change?

                template<typename R, typename C>
                AddTo(Self* b, R&& r, C&& c, bool h) 
                        : base(b), root(std::forward<R>(r)), child(std::forward<C>(c)), new_height(h) {}
        };

        /// Underlying storage
        union Data {
                char none;
                Base base;
                Add add;
                AddTo add_to;

                Data() : none('\0') {}
                ~Data() {}
        } data;

        /// Type of version node
        mutable Mode mode;

        /// get mutable reference as T
        template<typename T>
        T& as() { return reinterpret_cast<T&>(data); }

        /// get const reference as T
        template<typename T>
        const T& as() const { return reinterpret_cast<const T&>(data); }

        /// get rvalue reference as T
        template<typename T>
        T&& take_as() { return std::move(as<T>()); }

        /// initialize as T
        template<typename T, typename... Args>
        void init( Args&&... args ) {
                new( &as<T>() ) T { std::forward<Args>(args)... };
        }

        /// reset according to current mode
        void reset() {
                switch( mode ) {
                        case BASE:                as<Base>().~Base();   break;
                        case ADD: case REM:       as<Add>().~Add();     break;
                        case ADDTO: case REMFROM: as<AddTo>().~AddTo(); break;
                        default: assertf(false, "invalid mode");
                }
        }

        /// reset as base
        void reset_as_base() {
                assertf( mode == BASE, "can only reset_as_base() on BASE" );
                as<Base>().~Base();
        }

        /// Non-initializing constructor; should call init() before use
        PersistentDisjointSet( Mode m ) : data(), mode(m) {}

        PersistentDisjointSet( Mode m, Base&& b ) : data(), mode(m) {
                assertf(m == BASE, "invalid mode");
                init<Base>(std::move(b));
        }

        template<typename R>
        PersistentDisjointSet( Mode m, const Self* b, R&& r ) : data(), mode(m) {
                assertf(m == ADD || m == REM, "invalid mode");
                init<Add>(b, std::forward<R>(r));
        }

        template<typename R, typename C>
        PersistentDisjointSet( Mode m, const Self* b, R&& r, C&& c, bool h ) : data(), mode(m) {
                assertf(m == ADDTO || m == REMFROM, "invalid mode");
                init<AddTo>(b, std::forward<R>(r), std::forward<C>(c), h);
        }

        /// Adds (also removes) graph edges.
        /// * `from.parent` updated to `new_root`, 
        /// * `from.next` and `to.next` swapped (splices or un-splices class lists)
        /// * `to.height` adjusted by change
        template<typename R>
        static void addEdge( Node& from, Node& to, R&& new_root, Height change ) {
                from.parent = std::forward<R>(new_root);
                std::swap(from.next, to.next);
                to.height += change;
        }

protected:
        void trace( const GC& gc ) const {
                switch( mode ) {
                        case BASE: {
                                for (const auto& entry : as<Base>()) {
                                        gc.maybe_trace( entry.first );
                                }
                                return;
                        }
                        case ADD: case REM: {
                                const Add& self = as<Add>();
                                gc << self.base; 
                                gc.maybe_trace( self.root );
                                return;
                        }
                        case ADDTO: case REMFROM: {
                                const AddTo& self = as<AddTo>();
                                gc << self.base;
                                gc.maybe_trace( self.root, self.child );
                                return;
                        }
                        default: assertf(false, "invalid mode");
                }
        }

public:
        using size_type = std::size_t;

        using iterator = typename Base::const_iterator;

        PersistentDisjointSet() : data(), mode(BASE) { init<Base>(); }

        PersistentDisjointSet( const Self& o ) = delete;

        Self& operator= ( const Self& o ) = delete;

        ~PersistentDisjointSet() { reset(); }

        /// reroot persistent data structure at current node
        void reroot() const {
                if ( mode == BASE ) return;

                // reroot base
                Self* mut_this = const_cast<Self*>(this);
                Self* base = ( mode == ADD || mode == REM ) ? 
                        mut_this->as<Add>().base : 
                        mut_this->as<AddTo>().base;
                base->reroot();
                assertf(base->mode == BASE, "reroot results in base");

                // take map out of base
                Base base_map = base->take_as<Base>();
                base->reset_as_base();

                // switch base to inverse of self and mutate base map
                switch ( mode ) {
                        case ADD: {
                                Add& self = mut_this->as<Add>();

                                base->init<Add>( mut_this, self.root );
                                base->mode = REM;

                                base_map.emplace( self.root, Node{ std::move(self.root) } );
                        } break;
                        case REM: {
                                Add& self = mut_this->as<Add>();

                                base->init<Add>( mut_this, self.root );
                                base->mode = ADD;

                                base_map.erase( self.root );
                        } break;
                        case ADDTO: {
                                AddTo& self = mut_this->as<AddTo>();

                                base->init<AddTo>( mut_this, self.root, self.child, self.new_height );
                                base->mode = REMFROM;

                                auto child_it = base_map.find( self.child );
                                auto root_it = base_map.find( self.root );
                                assertf(child_it != base_map.end() && root_it != base_map.end(), 
                                        "nodes must exist in base");
                                Node& child = child_it->second;
                                Node& root = root_it->second;

                                addEdge( child, root, std::move(self.root), Height(self.new_height) );
                        } break;
                        case REMFROM: {
                                AddTo& self = mut_this->as<AddTo>();

                                base->init<AddTo>( mut_this, self.root, self.child, self.new_height );
                                base->mode = ADDTO;

                                auto child_it = base_map.find( self.child );
                                auto root_it = base_map.find( self.root );
                                assertf(child_it != base_map.end() && root_it != base_map.end(), 
                                        "nodes must exist in base");
                                Node& child = child_it->second;
                                Node& root = root_it->second;

                                addEdge( child, root, std::move(self.child), Height(-1 * self.new_height) );
                        } break;
                        default: assertf(false, "invalid mode");
                }

                // set base map into self
                mut_this->reset();
                mut_this->init<Base>( std::move(base_map) );
                mode = BASE;
        }

private:
        /// Gets the base after rerooting at the current node 
        const Base& rerooted() const {
                reroot();
                assertf(mode == BASE, "reroot results in base");
                return as<Base>();
        }

public:
        /// true if the set of sets is empty
        bool empty() const { return rerooted().empty(); }

        /// Get number of entries in the map
        size_type size() const { return rerooted().size(); }

        /// Get begin iterator for map; may be invalidated by calls to non-iteration functions 
        /// or functions on other maps in the same chain
        iterator begin() const { return rerooted().begin(); }

        /// Get end iterator for map; may be invalidated by calls to non-iteration functions 
        /// or functions on other maps in the same chain
        iterator end() const { return rerooted().end(); }

        /// Check if value is present
        size_type count(Elm i) const { return rerooted().count( i ); }

        /// Finds root for element i, undefined behaviour if i is not present
        Elm find(Elm i) const {
                const Base& self = rerooted();
                
                auto it = self.find( i );
                while (true) {
                        assertf(it != self.end(), "find target not present");

                        if ( it->first == it->second.parent ) return it->first;

                        it = self.find( it->second.parent );
                }
        }

        /// Finds root for element i, or default if i is not present
        template<typename E>
        Elm find_or_default(Elm i, E&& d) const {
                const Base& self = rerooted();

                auto it = self.find( i );
                if ( it == self.end() ) return d;

                while ( it->first != it->second.parent ) {
                        it = self.find( it->second.parent );

                        assertf(it != self.end(), "find target not present");
                }
                return it->first;
        }

        /// Adds fresh class including only one item; returns updated map (or self if no change)
        template<typename E>
        Self* add(E&& i) {
                reroot();

                // add new element to node
                Base base_map = take_as<Base>();
                bool added = base_map.emplace( i, Node{ i } ).second;

                // fail early on node already present
                if ( ! added ) {
                        as<Base>() = std::move(base_map);
                        return this;
                }

                // make new return node and reset self as REM node
                Self* ret = new Self{ BASE, std::move(base_map) };
                reset_as_base();
                init<Add>( ret, i );
                mode = REM;

                return ret;
        }

        /// Merges two classes given by their roots; returns updated map.
        /// If two classes have same height, `i` is new root.
        Self* merge(Elm i, Elm j) {
                reroot();

                // transfer map to new node
                Self* ret = new Self{ BASE, take_as<Base>() };
                reset_as_base();

                // find set nodes
                Base& base_map = ret->as<Base>();
                auto it = base_map.find( i );
                auto jt = base_map.find( j );
                assertf(it != base_map.end() && jt != base_map.end(), "nodes must exist in base");
                Node& in = it->second;
                Node& jn = jt->second;

                // update returned map and set self to appropriate REMFROM node
                if ( in.height < jn.height ) {
                        addEdge( in, jn, j, 0 );
                        init<AddTo>( ret, j, i, false );
                } else if ( jn.height < in.height ) {
                        addEdge( jn, in, i, 0 );
                        init<AddTo>( ret, i, j, false );
                } else /* if ( jn.height == in.height ) */ {
                        addEdge( jn, in, i, 1 );
                        init<AddTo>( ret, i, j, true );
                }
                mode = REMFROM;

                return ret;
        }

private:
        /// Removes the children of the tree rooted at `root` as `root_node`, returning a new 
        /// uninitialized node and editing `next_edit` to point toward this new node.
        static Self* remove_children(Elm root, Node& root_node, Base& base_map, Self* next_edit) {
                // Invariant: root.next is *always* a pointer to a leaf of the most-recently added
                // subtree.
                //
                // Proof: By induction on height:
                // * height == 0: root.next == root, trivially true
                // * added.height < root.height: true by ind. hyp.
                // * added.height == root.height: no node may have a child the same height, ergo 
                //   property holds for added, therefore root.
                //
                // Corollary: next of most-recently-added subtree root is previously added subtree

                // remove all subtrees
                while ( root != root_node.next ) {
                        // find added child
                        auto it = base_map.find( root_node.next );
                        while ( it->second.parent != root ) it = base_map.find( it->second.parent );
                        Elm added = it->first;
                        Node& added_node = it->second;

                        // unlink subtree and set up previous map as ADDTO new map
                        std::swap( root_node.next, added_node.next );
                        Self* new_edit = new Self{ BASE };
                        next_edit->init<AddTo>( new_edit, root, added, false );  // assume unchanged root height
                        next_edit->mode = ADDTO;

                        // remove subtree children from map
                        next_edit = remove_children( added, added_node, base_map, new_edit );

                        // remove subtree root from map and set up previous map as ADD to new map
                        base_map.erase( added );
                        new_edit = new Self{ BASE };
                        next_edit->init<Add>( new_edit, added );
                        next_edit->mode = ADD;

                        next_edit = new_edit;
                }

                return next_edit;
        }

public:
        /// Removes all elements of a class given by its root; returns updated map.
        Self* remove_class(Elm root) {
                reroot();

                // remove map from self
                Base base_map = take_as<Base>();
                reset_as_base();

                // find root node and remove children
                auto it = base_map.find( root );
                assertf(it != base_map.end(), "root node must exist in map");
                Self* next_edit = remove_children( root, it->second, base_map, this);

                // root is alone in class, remove from map
                base_map.erase( root );
                Self* ret = new Self{ BASE, std::move(base_map) };
                // reinitialize previous node as ADD to new map
                next_edit->init<Add>( ret, root );
                next_edit->mode = ADD;

                return ret;
        }

        /// Applies `f` to all members of a class containing `i` (undefined behaviour if not present).
        /// `f` should take members by const reference or copy
        template<typename F>
        void for_class(Elm i, F&& f) const {
                const Base& self = rerooted();

                // exit early if class not present
                auto it = self.find( i );
                if ( it == self.end() ) return;

                // apply f to each member of class
                f( i );
                for ( Elm crnt = it->second.next; crnt != i; crnt = it->second.next ) {
                        f( crnt );
                        it = self.find( crnt );
                        assertf( it != self.end(), "current node must exist in base" );
                }
        }

        /// Get version node type
        Mode get_mode() const { return mode; }

        /// Get next version up the revision tree (self if base node)
        const Self* get_base() const {
                switch ( mode ) {
                        case BASE:                return this;
                        case ADD: case REM:       return as<Add>().base;
                        case ADDTO: case REMFROM: return as<AddTo>().base;
                        default: assertf(false, "invalid mode");
                }
        }

        /// Get root of new class formed/removed/split (undefined if called on base)
        Elm get_root() const {
                switch ( mode ) {
                        case ADD: case REM:       return as<Add>().root;
                        case ADDTO: case REMFROM: return as<AddTo>().root;
                        default: assertf(false, "invalid mode for get_root()");
                }
        }

        /// Get child root of new class formed/split (undefined if called on base or add/remove node)
        Elm get_child() const {
                switch ( mode ) {
                        case ADDTO: case REMFROM: return as<AddTo>().child;
                        default: assertf(false, "invalid mode for get_child()");
                }
        }

        /// Gets acted-upon key for new class (root unless for add/remove node, child for add to/remove 
        /// from node, undefined otherwise)
        Elm get_key() const {
                switch ( mode ) {
                        case ADD: case REM:       return as<Add>().root;
                        case ADDTO: case REMFROM: return as<AddTo>().child;
                        default: assertf(false, "invalid mode for get_key()");
                }
        }
};

// Local Variables: //
// tab-width: 4 //
// mode: c++ //
// compile-command: "make install" //
// End: //