Mesh.h

#pragma once
#include "Engine/Primitives/Primitive.h"
#include <span>
#include <tbb/spin_mutex.h>
#include <unordered_set>
 
namespace enzo::geo {
class Mesh;
class FaceNormalHandle;
class PointNormalHandle;
class VertexNormalHandle;
 
inline constexpr double kDefaultCuspAngle = 60;
 
class Mesh : public Primitive
{
  public:
    Mesh(std::string_view path = "/mesh");
    Mesh(const Mesh& other);
    Mesh& operator=(const Mesh& rhs);
    ~Mesh() override = default;
 
    PrimType getType() const override { return PrimType::MESH; }
    std::shared_ptr<Primitive> clone() const override { return std::make_shared<Mesh>(*this); }
    TransformClass transformType() const override
    {
        return TransformClass::POINT | TransformClass::PRIMITIVE;
    }
    void applyTransform(
        const Matrix4& mat,
        TransformClass transformClass = TransformClass::POINT
    ) override;
    void applyTransform(
        const Transform& transform,
        TransformClass transformClass = TransformClass::POINT
    )
    {
        applyTransform(transform.getMatrix(), transformClass);
    }
    bool canMerge() const override { return true; }
    void merge(std::shared_ptr<Primitive> other) override;
    bool hasPoints() const override { return true; }
 
    // TODO: benchmark addFace vs addFaces to quantify the speedup
    Offset addFace(const std::vector<Offset>& pointOffsets, bool closed = true);
    std::vector<Offset> addFaces(
        std::span<const Offset> pointOffsetsFlat,
        std::span<const Offset> vertexCounts,
        bool closed = true
    );
    Offset addPoint(const Vector3& pos);
    std::vector<Offset> addPoints(std::span<const Vector3> positions);
    std::vector<Offset>
    duplicatePoints(std::span<const Offset> srcPointOffsets, bool copyAttributes = true);
 
    void deleteFaces(const std::vector<Offset>& faceOffsets, bool andPoints = true);
    void deleteAllFaces(bool andPoints = true);
    void deletePoints(const std::vector<Offset>& pointOffsets) override
    {
        deletePoints(pointOffsets, false);
    }
    void deletePoints(const std::vector<Offset>& pointOffsets, bool andFaces);
    void deleteVertices(const std::vector<Offset>& vertOffsets);
    bool isValidFace(Offset offset) const;
    bool isValidVertex(Offset offset) const;
    bool isValidPoint(Offset offset) const override;
 
    void defragment() override;
 
    void merge(Mesh& other);
 
    std::unordered_set<Offset>::const_iterator soloPointsBegin() const;
    std::unordered_set<Offset>::const_iterator soloPointsEnd() const;
 
    void setPointPos(const Offset offset, const Vector3& pos);
    std::span<const Offset> getFaceStartVertices() const;
    Vector3 getPosFromVert(Offset vertexOffset) const;
    Vector3 getPointPos(Offset pointOffset) const;
    unsigned int getFaceVertCount(Offset faceOffset) const
    {
        return vertexCountFaceHandle_[faceOffset];
    }
    unsigned int getFacePointCount(Offset faceOffset) const;
    Offset getVertexFace(Offset vertexOffset) const;
 
    Offset getVertexPoint(Offset vertexOffset) const
    {
        return pointOffsetVertexHandle_.getValue(vertexOffset);
    }
 
    std::span<const intT> getFacePoints(Offset faceOffset) const
    {
        const Offset start = getFaceStartVertices()[faceOffset];
        const unsigned int count = getFaceVertCount(faceOffset);
        return pointOffsetVertexHandle_.getSpan().subspan(start, count);
    }
 
    std::span<const intT> vertexPointSpan() const { return pointOffsetVertexHandle_.getSpan(); }
 
    std::span<const Vector3> pointPosSpan() const { return posPointHandle_.getSpan(); }
 
    Offset getNumFaces() const { return getElementCount(attr::AttributeOwner::FACE); }
    Offset getNumVerts() const { return getElementCount(attr::AttributeOwner::VERTEX); }
    Offset getNumSoloPoints() const;
 
    // Face Iterator
    struct FaceOffsets
    {
        FaceOffsets(const Mesh& mesh) : mesh_(mesh) {}
        struct Iterator
        {
            using iterator_category = std::forward_iterator_tag;
            using difference_type = std::ptrdiff_t;
            using value_type = Offset;
 
            explicit Iterator(Offset current) : curOffset_(current) {}
            value_type operator*() const { return curOffset_; }
            Iterator& operator++()
            {
                ++curOffset_;
                return *this;
            }
            Iterator operator++(int)
            {
                Iterator tmp = *this;
                ++(*this);
                return tmp;
            }
            friend bool operator==(const Iterator& a, const Iterator& b)
            {
                return a.curOffset_ == b.curOffset_;
            }
            friend bool operator!=(const Iterator& a, const Iterator& b)
            {
                return a.curOffset_ != b.curOffset_;
            }
 
          private:
            Offset curOffset_ = 0;
        };
        Iterator begin() const { return Iterator(0); }
        Iterator end() const { return Iterator(mesh_.getNumFaces()); }
 
        std::vector<Offset> toVector() const
        {
            std::vector<Offset> offsets;
            offsets.reserve(mesh_.getNumFaces());
            for (const Offset faceOffset : *this)
                offsets.push_back(faceOffset);
            return offsets;
        }
 
      private:
        const Mesh& mesh_;
    };
    // TODO walk only valid faces like Primitive::PointOffsets once deletion settles. For now it is
    // dense.
    FaceOffsets getFaces() const { return FaceOffsets(*this); }
 
    attr::AttributeHandleBool addVertexGroup(std::string name)
    {
        return addGroup(attr::AttributeOwner::VERTEX, std::move(name));
    }
    attr::AttributeHandleBool addFaceGroup(std::string name)
    {
        return addGroup(attr::AttributeOwner::FACE, std::move(name));
    }
    void addToVertexGroup(const std::string& name, const std::vector<Offset>& offsets)
    {
        addToGroup(attr::AttributeOwner::VERTEX, name, offsets);
    }
    void addToFaceGroup(const std::string& name, const std::vector<Offset>& offsets)
    {
        addToGroup(attr::AttributeOwner::FACE, name, offsets);
    }
 
    boolT isClosed(Offset faceOffset) const;
 
    void computeFaceStartVertices() const;
 
    FaceNormalHandle getFaceNormal(bool precompute = false) const;
 
    VertexNormalHandle getVertexNormal(double cuspAngle = kDefaultCuspAngle) const;
 
    PointNormalHandle getPointNormal() const;
 
    friend class FaceNormalHandle;
    friend class VertexNormalHandle;
 
  protected:
    attr::attribVector& getAttributeStore(const attr::AttributeOwner& owner) override;
    const attr::attribVector& getAttributeStore(const attr::AttributeOwner& owner) const override;
    attr::attribVector& getGroupStore(const attr::AttributeOwner& owner) override;
    const attr::attribVector& getGroupStore(const attr::AttributeOwner& owner) const override;
 
  private:
    void mergeAppend(std::shared_ptr<attr::Attribute> dst, std::shared_ptr<attr::Attribute> src);
    template <typename T>
    void mergeAppendImpl(std::shared_ptr<attr::Attribute> dst, std::shared_ptr<attr::Attribute> src)
    {
        auto dstHandle = attr::AttributeHandle<T>(dst);
        auto srcHandle = attr::AttributeHandle<T>(src);
 
        const Offset srcCount = srcHandle.getSize();
        const Offset dstCount = dstHandle.getSize();
 
        dstHandle.resize(dstCount + srcCount);
 
        for (Offset i = 0; i < srcCount; ++i)
        {
            const T value = srcHandle.getValue(i);
            const Offset dstOffset = dstCount + i;
            dstHandle.setValue(dstOffset, value);
        }
    };
 
    attr::attribVector vertexAttributes_;
    attr::attribVector faceAttributes_;
    attr::attribVector vertexGroups_;
    attr::attribVector faceGroups_;
 
    mutable std::unordered_set<Offset> soloPoints_;
    mutable bool soloPointsDirty_ = true;
    void rebuildSoloPoints() const;
    bool needsDefrag_ = false;
 
    mutable std::vector<Offset> faceStarts_;
    mutable std::vector<Offset> vertexFaces_;
 
    mutable std::atomic<bool> faceStartsDirty_{true};
    mutable tbb::spin_mutex faceStartsMutex_;
 
    // intrinsic handles
    enzo::attr::AttributeHandleInt vertexCountFaceHandle_;
    enzo::attr::AttributeHandleBool closedFaceHandle_;
    enzo::attr::AttributeHandleInt pointOffsetVertexHandle_;
    enzo::attr::AttributeHandleVector3 posPointHandle_;
    enzo::attr::AttributeHandleBool validFaceHandle_;
    enzo::attr::AttributeHandleBool validVertexHandle_;
    enzo::attr::AttributeHandleBool validPointHandle_;
};
 
class FaceNormalHandle
{
  public:
    Vector3 operator[](Offset faceOffset) const
    {
        if (cached_) return (*cached_)[faceOffset];
        if (!precomputed_.empty()) return precomputed_[faceOffset];
        return computeNormal(faceOffset);
    }
 
  private:
    friend class Mesh;
    FaceNormalHandle(const Mesh& mesh, bool precompute);
    Vector3 computeNormal(Offset faceOffset) const;
    const Mesh& mesh_;
    std::optional<attr::AttributeHandleRO<Vector3>> cached_;
    std::vector<Vector3> precomputed_;
};
 
class PointNormalHandle
{
  public:
    Vector3 operator[](Offset pointOffset) const
    {
        if (attribute_) return (*attribute_)[pointOffset];
        return computed_[pointOffset];
    }
 
  private:
    friend class Mesh;
    explicit PointNormalHandle(const Mesh& mesh);
    std::optional<attr::AttributeHandleRO<Vector3>> attribute_;
    std::vector<Vector3> computed_;
};
 
class VertexNormalHandle
{
  public:
    Vector3 operator[](Offset vertexOffset) const;
 
  private:
    friend class Mesh;
    VertexNormalHandle(const Mesh& mesh, double cuspAngle);
    const Mesh& mesh_;
    std::optional<attr::AttributeHandleRO<Vector3>> vertexAttribute_;
    std::optional<attr::AttributeHandleRO<Vector3>> pointAttribute_;
    std::vector<Vector3> computed_;
};
} // namespace enzo::geo