TriMesh.h

#pragma once
 
#include <wmtk/utils/VectorUtils.h>
#include <wmtk/AttributeCollection.hpp>
#include <wmtk/Types.hpp>
#include <wmtk/simplex/Simplex.hpp>
#include <wmtk/simplex/SimplexCollection.hpp>
#include <wmtk/threading/enumerable_thread_specific.hpp>
#include <wmtk/threading/vertex_mutex.hpp>
#include <wmtk/utils/Logger.hpp>
 
#include <algorithm>
#include <array>
#include <atomic>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <map>
#include <memory>
#include <optional>
#include <set>
#include <vector>
 
namespace wmtk {
 
class TriMesh
{
public:
    // Cell Tuple Navigator
    class Tuple
    {
    private:
        size_t m_vid = -1;
        size_t m_eid = -1;
        size_t m_fid = -1;
        size_t m_hash = -1;
 
        void update_hash(const TriMesh& m);
 
    public:
        void print_info();
 
        //         v2        /
        //       /    \      /
        //  e1  /      \  e0 /
        //     v0 - - - v1   /
        //         e2        /
        Tuple() {}
        Tuple(size_t vid, size_t eid, size_t fid, const TriMesh& m)
            : m_vid(vid)
            , m_eid(eid)
            , m_fid(fid)
        {
            update_hash(m);
        }
 
 
        inline size_t vid(const TriMesh&) const { return m_vid; }
 
        inline size_t fid(const TriMesh&) const { return m_fid; }
 
 
        size_t eid(const TriMesh& m) const;
 
        size_t local_eid(const TriMesh& m) const { return m_eid; };
        Tuple switch_vertex(const TriMesh& m) const;
        Tuple switch_edge(const TriMesh& m) const;
        std::optional<Tuple> switch_face(const TriMesh& m) const;
 
        std::vector<Tuple> switch_faces(const TriMesh& m) const;
 
        bool is_valid(const TriMesh& m) const;
 
        std::array<Tuple, 3> oriented_tri_vertices(const TriMesh& m) const;
        friend bool operator<(const Tuple& a, const Tuple& t)
        {
            return (
                std::tie(a.m_vid, a.m_eid, a.m_fid, a.m_hash) <
                std::tie(t.m_vid, t.m_eid, t.m_fid, t.m_hash));
        }
    };
 
    class SmartTuple
    {
        const TriMesh& m_mesh;
        Tuple m_tuple;
 
    public:
        SmartTuple(const TriMesh& mesh, const Tuple& t)
            : m_mesh(mesh)
            , m_tuple(t)
        {}
 
        const Tuple& tuple() const { return m_tuple; }
        const TriMesh& mesh() const { return m_mesh; }
 
        SmartTuple& operator=(const SmartTuple& t)
        {
            m_tuple = t.m_tuple;
            return *this;
        }
 
        bool is_valid() const { return m_tuple.is_valid(m_mesh); }
        size_t vid() const { return m_tuple.vid(m_mesh); }
        size_t eid() const { return m_tuple.eid(m_mesh); }
        size_t fid() const { return m_tuple.fid(m_mesh); }
        SmartTuple switch_vertex() const { return {m_mesh, m_tuple.switch_vertex(m_mesh)}; }
        SmartTuple switch_edge() const { return {m_mesh, m_tuple.switch_edge(m_mesh)}; }
        std::optional<SmartTuple> switch_face() const
        {
            const std::optional<Tuple> t = m_tuple.switch_face(m_mesh);
            if (t) {
                return std::optional<SmartTuple>({m_mesh, t.value()});
            }
            return {};
        }
    };
 
    class VertexConnectivity
    {
    public:
        std::vector<size_t> m_conn_tris;
        bool m_is_removed = false;
 
        inline size_t& operator[](const size_t index)
        {
            assert(index < m_conn_tris.size());
            return m_conn_tris[index];
        }
 
        inline size_t operator[](const size_t index) const
        {
            assert(index < m_conn_tris.size());
            return m_conn_tris[index];
        }
    };
 
    class TriangleConnectivity
    {
    public:
        std::array<size_t, 3> m_indices;
        bool m_is_removed = false;
        size_t hash = 0;
 
        inline size_t& operator[](size_t index)
        {
            assert(index < 3);
            return m_indices[index];
        }
 
        inline size_t operator[](size_t index) const
        {
            assert(index < 3);
            return m_indices[index];
        }
        inline int find(size_t v_id) const
        {
            for (int j = 0; j < 3; j++) {
                if (v_id == m_indices[j]) {
                    return j;
                }
            }
            return -1;
        }
    };
 
    TriMesh() {}
    virtual ~TriMesh() {}
 
    void init(size_t n_vertices, const std::vector<std::array<size_t, 3>>& tris);
 
    void init(const MatrixXi& F);
 
    std::vector<Tuple> get_vertices() const;
 
    std::vector<Tuple> get_edges() const;
 
    std::vector<Tuple> get_faces() const;
 
    Tuple tuple_from_edge(size_t vid1, size_t vid2, size_t fid) const;
 
    Tuple tuple_from_vids(size_t vid0, size_t vid1, size_t vid2) const;
 
    simplex::Vertex simplex_from_vertex(const Tuple& t) const;
    simplex::Edge simplex_from_edge(const Tuple& t) const;
    simplex::Face simplex_from_face(const Tuple& t) const;
    simplex::Face simplex_from_face(const size_t fid) const;
 
    Tuple tuple_from_simplex(const simplex::Face& s) const;
    // Tuple tuple_from_simplex(const simplex::Edge& s) const;
 
    simplex::SimplexCollection simplex_incident_triangles(const simplex::Vertex& v) const;
    simplex::SimplexCollection simplex_incident_triangles(const simplex::Edge& e) const;
    simplex::SimplexCollection simplex_link_vertices(const simplex::Vertex& v) const;
    simplex::SimplexCollection simplex_link_vertices(const simplex::Edge& e) const;
    simplex::SimplexCollection simplex_link_edges(const simplex::Vertex& v) const;
 
    template <typename T>
    using vector = std::vector<T>;
 
public:
    AbstractAttributeContainer* p_vertex_attrs = nullptr;
    AbstractAttributeContainer* p_edge_attrs = nullptr;
    AbstractAttributeContainer* p_face_attrs = nullptr;
 
public:
    void set_preallocation_factor(double factor)
    {
        if (factor >= 1.0) m_preallocation_factor = factor;
    }
    double preallocation_factor() const { return m_preallocation_factor; }
 
    // Atomically reserve `n` contiguous fresh triangle/vertex slots. Returns the
    // first index of the block, or -1 if that would exceed the preallocated
    // capacity (the caller must then abort the operation before mutating).
    long request_tri_slots(size_t n);
    long request_vert_slots(size_t n);
 
private:
    size_t reserved_capacity(size_t live_count) const
    {
        const size_t floor = 64;
        double c = std::ceil(m_preallocation_factor * static_cast<double>(live_count));
        size_t capacity = static_cast<size_t>(c);
        if (capacity < live_count) capacity = live_count;
        return capacity < floor ? floor : capacity;
    }
 
    vector<VertexConnectivity> m_vertex_connectivity;
    vector<TriangleConnectivity> m_tri_connectivity;
    std::atomic_long current_vert_size;
    std::atomic_long current_tri_size;
    double m_preallocation_factor = 6.0;
    bool m_use_link_condition = true;
 
protected:
    void vertex_fan_components(
        size_t vid,
        std::vector<size_t>& component_of,
        std::vector<size_t>& representatives) const;
 
private:
    size_t get_next_empty_slot_t();
    size_t get_next_empty_slot_v();
 
public:
    virtual bool invariants(const std::vector<Tuple>&) { return true; }
    virtual bool split_edge_before(const Tuple& t) { return true; }
    virtual bool split_edge_after(const Tuple& t) { return true; }
 
    virtual bool collapse_edge_before(const Tuple& t)
    {
        if (!m_use_link_condition) {
            return true;
        }
        return check_link_condition(t);
    }
    virtual bool collapse_edge_after(const Tuple& t) { return true; }
    virtual bool swap_edge_after(const Tuple& t) { return true; }
    virtual bool swap_edge_before(const Tuple& t);
    virtual bool smooth_before(const Tuple& t) { return true; }
    virtual bool smooth_after(const Tuple& t) { return true; }
 
    virtual bool split_face_before(const Tuple& t) { return true; }
    virtual bool split_face_after(const Tuple& t) { return true; }
 
public:
    size_t tri_capacity() const { return current_tri_size; }
    size_t vert_capacity() const { return current_vert_size; }
 
    static constexpr int EDGES_PER_CELL = 3;
    size_t cell_capacity() const { return tri_capacity(); }
    Tuple tuple_from_cell(size_t cid) const { return tuple_from_tri(cid); }
    void consolidate_mesh();
 
    void remove_tris_by_ids(const std::vector<size_t>& fids)
    {
        for (const size_t fid : fids) {
            m_tri_connectivity[fid].m_is_removed = true;
            for (int j = 0; j < 3; j++) {
                vector_erase(m_vertex_connectivity[m_tri_connectivity[fid][j]].m_conn_tris, fid);
            }
        }
        for (auto& v : m_vertex_connectivity) {
            if (v.m_is_removed) continue;
            if (v.m_conn_tris.empty()) v.m_is_removed = true;
        }
    }
    Tuple switch_vertex(const Tuple& t) const { return t.switch_vertex(*this); }
    Tuple switch_edge(const Tuple& t) const { return t.switch_edge(*this); }
    std::optional<Tuple> switch_face(const Tuple& t) const { return t.switch_face(*this); }
 
    bool check_link_condition(const Tuple& t) const;
 
    void set_use_link_condition(bool use_it) { m_use_link_condition = use_it; }
    bool use_link_condition() const { return m_use_link_condition; }
 
    bool check_mesh_connectivity_validity() const;
    bool check_edge_manifold() const;
 
    size_t edge_valence(const TriMesh::Tuple& t) const;
 
    bool is_boundary_edge(const TriMesh::Tuple& t) const;
 
    bool is_manifold_edge(const TriMesh::Tuple& t) const;
 
    size_t vertex_component_count(const size_t vid) const;
    size_t vertex_component_count(const TriMesh::Tuple& t) const
    {
        return vertex_component_count(t.vid(*this));
    }
 
    bool is_manifold_vertex(const size_t vid) const { return vertex_component_count(vid) <= 1; }
 
    std::optional<Tuple> switch_component(const TriMesh::Tuple& t) const;
 
#ifdef WMTK_DEBUG_BRUTE_FORCE_OPS
    std::string debug_canonical_form() const;
 
    std::string debug_reference_collapse(size_t v_removed, size_t v_kept) const;
 
    std::string debug_reference_swap(size_t v0, size_t v1, size_t fa, size_t fb) const;
 
    std::string debug_reference_split_face(
        size_t fid,
        size_t v0,
        size_t v1,
        size_t v2,
        size_t new_v,
        size_t tri_cap) const;
 
    std::string debug_reference_split(size_t v0, size_t v1, size_t new_v, size_t tri_cap) const;
#endif
 
    bool is_boundary_vertex(const TriMesh::Tuple& t) const
    {
        auto ve = get_one_ring_edges_for_vertex(t);
        for (auto e : ve)
            if (is_boundary_edge(e)) return true;
        return false;
    }
 
    bool split_edge(const Tuple& t, std::vector<Tuple>& new_t);
 
    virtual bool collapse_edge(const Tuple& t, std::vector<Tuple>& new_t);
 
    void collapse_edge_conn(
        const Tuple& loc0,
        std::vector<Tuple>& new_tris,
        Tuple& return_t,
        size_t& new_vid,
        std::vector<std::pair<size_t, TriangleConnectivity>>& old_tris,
        std::vector<std::pair<size_t, VertexConnectivity>>& old_vertices,
        std::vector<std::pair<size_t, size_t>>& same_edge_vid_fid,
        std::vector<size_t>& n12_intersect_fids);
 
    void collapse_edge_rollback(
        size_t& new_vid,
        std::vector<std::pair<size_t, TriangleConnectivity>>& old_tris,
        std::vector<std::pair<size_t, VertexConnectivity>>& old_vertices,
        std::vector<std::pair<size_t, size_t>>& same_edge_vid_fid,
        std::vector<size_t>& n12_intersect_fids);
 
 
    bool swap_edge(const Tuple& t, std::vector<Tuple>& new_t);
 
    bool smooth_vertex(const Tuple& t);
 
    bool split_face(const Tuple& t, std::vector<Tuple>& new_t);
 
    size_t get_valence_for_vertex(const Tuple& t) const
    {
        return m_vertex_connectivity[t.vid(*this)].m_conn_tris.size();
    }
 
    size_t vertex_valence(const size_t vid) const
    {
        return m_vertex_connectivity[vid].m_conn_tris.size();
    }
 
    std::vector<Tuple> get_one_ring_tris_for_vertex(const Tuple& t) const;
    const std::vector<size_t>& get_one_ring_fids_for_vertex(const Tuple& t) const;
    const std::vector<size_t>& get_one_ring_fids_for_vertex(const size_t vid) const;
    std::vector<size_t> get_one_ring_vids_for_vertex_duplicate(const size_t& t) const;
    void get_one_ring_vids_for_vertex_duplicate(const size_t& t, std::vector<size_t>& one_ring)
        const;
 
    std::vector<size_t> get_incident_fids_for_edge(const Tuple& t) const;
    std::vector<size_t> get_incident_fids_for_edge(const size_t vid0, const size_t vid1) const;
 
    std::vector<Tuple> get_one_ring_edges_for_vertex(const Tuple& t) const;
    std::vector<Tuple> get_one_ring_edges_for_vertex(const size_t vid) const;
 
    std::array<Tuple, 3> oriented_tri_vertices(const Tuple& t) const;
 
    std::array<size_t, 3> oriented_tri_vids(const Tuple& t) const;
    std::array<size_t, 3> oriented_tri_vids(const size_t i) const;
 
    std::array<Tuple, 2> get_edge_vertices(const Tuple& t) const;
    std::array<size_t, 2> get_edge_vids(const Tuple& t) const;
 
    Tuple tuple_from_tri(size_t fid) const
    {
        if (fid >= m_tri_connectivity.size() || m_tri_connectivity[fid].m_is_removed)
            return Tuple();
        auto vid = m_tri_connectivity[fid][0];
        return Tuple(vid, 1, fid, *this);
    }
    Tuple tuple_from_vertex(size_t vid) const
    {
        if (vid >= m_vertex_connectivity.size() || m_vertex_connectivity[vid].m_is_removed ||
            m_vertex_connectivity[vid].m_conn_tris.empty()) {
            return Tuple();
        }
        auto fid = m_vertex_connectivity[vid][0];
        auto eid = m_tri_connectivity[fid].find((int)vid);
        return Tuple(vid, (eid + 1) % 3, fid, *this);
    }
    Tuple tuple_from_edge(size_t fid, size_t local_eid) const
    {
        auto vid = m_tri_connectivity[fid][(local_eid + 1) % 3];
        return Tuple(vid, local_eid, fid, *this);
    }
 
    std::tuple<Tuple, size_t> tuple_from_edge(const std::array<size_t, 2>& vids) const;
    std::optional<std::tuple<Tuple, size_t>> try_tuple_from_edge(
        const std::array<size_t, 2>& vids) const;
 
public:
    void start_protect_attributes()
    {
        if (p_vertex_attrs) p_vertex_attrs->begin_protect();
        if (p_edge_attrs) p_edge_attrs->begin_protect();
        if (p_face_attrs) p_face_attrs->begin_protect();
    }
    void release_protect_attributes()
    {
        if (p_vertex_attrs) p_vertex_attrs->end_protect();
        if (p_edge_attrs) p_edge_attrs->end_protect();
        if (p_face_attrs) p_face_attrs->end_protect();
    }
    void rollback_protected_attributes()
    {
        if (p_vertex_attrs) p_vertex_attrs->rollback();
        if (p_edge_attrs) p_edge_attrs->rollback();
        if (p_face_attrs) p_face_attrs->rollback();
    }
 
    // Moved code from concurrent TriMesh
 
public:
    using VertexMutex = wmtk::threading::VertexMutex;
 
private:
    std::vector<VertexMutex> m_vertex_mutex;
 
    bool try_set_vertex_mutex(const Tuple& v, int threadid)
    {
        bool got = m_vertex_mutex[v.vid(*this)].trylock();
        if (got) m_vertex_mutex[v.vid(*this)].set_owner(threadid);
        return got;
    }
    bool try_set_vertex_mutex(size_t vid, int threadid)
    {
        bool got = m_vertex_mutex[vid].trylock();
        if (got) m_vertex_mutex[vid].set_owner(threadid);
        return got;
    }
 
    void unlock_vertex_mutex(const Tuple& v) { m_vertex_mutex[v.vid(*this)].unlock(); }
    void unlock_vertex_mutex(size_t vid) { m_vertex_mutex[vid].unlock(); }
 
    struct RingLockScratch
    {
        std::vector<uint32_t> stamp;
        uint32_t epoch = 0;
        std::vector<size_t> frontier;
        std::vector<size_t> next;
        std::vector<size_t> one_ring;
    };
    wmtk::threading::enumerable_thread_specific<RingLockScratch> m_ring_lock_scratch;
 
    bool lock_vertex_ball(const size_t* seeds, size_t n_seeds, int threadid, int n, size_t mark);
 
protected:
    void resize_mutex(size_t v)
    {
        if (m_vertex_mutex.size() < v) m_vertex_mutex.resize(v);
    }
 
public:
    wmtk::threading::enumerable_thread_specific<std::vector<size_t>> mutex_release_stack;
 
    int release_vertex_mutex_in_stack();
    int release_vertex_mutex_to(size_t mark);
 
    bool try_set_vertex_mutex_n_ring(const Tuple& v, int threadid, int n);
    bool try_set_vertex_mutex_n_ring(size_t vid, int threadid, int n);
    bool try_set_edge_mutex_n_ring(const Tuple& e, int threadid, int n);
 
    bool try_set_vertex_mutex_two_ring(const Tuple& v, int threadid);
    bool try_set_edge_mutex_two_ring(const Tuple& e, int threadid);
    bool try_set_vertex_mutex_one_ring(const Tuple& v, int threadid);
    bool try_set_face_mutex_one_ring(const Tuple& f, int threadid);
 
    void for_each_face(const std::function<void(const Tuple&)>&);
    void for_each_edge(const std::function<void(const Tuple&)>&);
    void for_each_vertex(const std::function<void(const Tuple&)>&);
    int NUM_THREADS = 0;
 
public:
    // substructure functionality
 
    virtual bool vertex_is_on_surface(const size_t vid) const { return false; }
 
    virtual bool edge_is_on_surface(const std::array<size_t, 2>& vids) const { return false; }
 
    simplex::SimplexCollection get_surface_edges_for_vertex(const size_t vid) const;
 
    size_t get_order_of_edge(const std::array<size_t, 2>& vids) const;
 
    size_t get_order_of_vertex(const size_t vid) const;
 
    bool substructure_link_condition(const Tuple& e_tuple) const;
};
 
} // namespace wmtk