RayCaster.h

/*
 * Copyright 2017 Adobe. All rights reserved.
 * This file is licensed to you under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License. You may obtain a copy
 * of the License at http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under
 * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
 * OF ANY KIND, either express or implied. See the License for the specific language
 * governing permissions and limitations under the License.
 */
#pragma once
 
#include <lagrange/SurfaceMesh.h>
#include <lagrange/raycasting/api.h>
#include <lagrange/scene/SimpleScene.h>
#include <lagrange/utils/BitField.h>
#include <lagrange/utils/function_ref.h>
#include <lagrange/utils/value_ptr.h>
 
#include <Eigen/Core>
 
#include <cstdint>
#include <functional>
#include <limits>
#include <memory>
#include <optional>
#include <variant>
 
namespace lagrange::raycasting {
 
namespace internal {
struct RayCasterOBBAccess;
} // namespace internal
 

struct HitBase
{
    uint32_t mesh_index = invalid<uint32_t>();

    uint32_t instance_index = invalid<uint32_t>();

    uint32_t facet_index = invalid<uint32_t>();

    Eigen::Vector2f barycentric_coord = Eigen::Vector2f::Zero();

    Eigen::Vector3f position = Eigen::Vector3f::Zero();
};

struct ClosestPointHit : public HitBase
{
    float distance = std::numeric_limits<float>::infinity();
};

struct RayHit : public HitBase
{
    float ray_depth = 0;

    Eigen::Vector3f normal = Eigen::Vector3f::Zero();
};

template <size_t N>
struct HitBaseN
{
    uint32_t valid_mask = 0;

    Eigen::Vector<uint32_t, N> mesh_indices =
        Eigen::Vector<uint32_t, N>::Constant(invalid<uint32_t>());

    Eigen::Vector<uint32_t, N> instance_indices =
        Eigen::Vector<uint32_t, N>::Constant(invalid<uint32_t>());

    Eigen::Vector<uint32_t, N> facet_indices =
        Eigen::Vector<uint32_t, N>::Constant(invalid<uint32_t>());

    Eigen::Matrix<float, 2, N> barycentric_coords = Eigen::Matrix<float, 2, N>::Zero();

    Eigen::Matrix<float, 3, N> positions = Eigen::Matrix<float, 3, N>::Zero();
 
    bool is_valid(size_t i) const { return (valid_mask & (1u << i)) != 0; }
};

template <size_t N>
struct ClosestPointHitN : public HitBaseN<N>
{
    Eigen::Vector<float, N> distances =
        Eigen::Vector<float, N>::Constant(std::numeric_limits<float>::infinity());
};

template <size_t N>
struct RayHitN : public HitBaseN<N>
{
    Eigen::Vector<float, N> ray_depths = Eigen::Vector<float, N>::Zero();

    Eigen::Matrix<float, 3, N> normals = Eigen::Matrix<float, 3, N>::Zero();
};

enum class SceneFlags {
    None = 0,

    Dynamic = 1 << 0,

    Compact = 1 << 1,

    Robust = 1 << 2,

    Filter = 1 << 3,
};

enum class BuildQuality {
    Low = 0,

    Medium = 1,

    High = 2,
};

class RayCaster
{
public:
    using Pointf = Eigen::Vector3f;

    using Directionf = Eigen::Vector3f;

    template <typename Scalar>
    using Affine = Eigen::Transform<Scalar, 3, Eigen::Affine>;

    using Hit = RayHit;

    using Point4f = Eigen::Matrix<float, 4, 3, Eigen::RowMajor>;
    using Direction4f = Eigen::Matrix<float, 4, 3, Eigen::RowMajor>;
    using Float4 = Eigen::Vector<float, 4>;
    using Mask4 = Eigen::Vector<bool, 4>;
    using RayHit4 = RayHitN<4>;
    using ClosestPointHit4 = ClosestPointHitN<4>;

    using Point8f = Eigen::Matrix<float, 8, 3, Eigen::RowMajor>;
    using Direction8f = Eigen::Matrix<float, 8, 3, Eigen::RowMajor>;
    using Float8 = Eigen::Vector<float, 8>;
    using Mask8 = Eigen::Vector<bool, 8>;
    using RayHit8 = RayHitN<8>;
    using ClosestPointHit8 = ClosestPointHitN<8>;

    using Point16f = Eigen::Matrix<float, 16, 3, Eigen::RowMajor>;
    using Direction16f = Eigen::Matrix<float, 16, 3, Eigen::RowMajor>;
    using Float16 = Eigen::Vector<float, 16>;
    using Mask16 = Eigen::Vector<bool, 16>;
    using RayHit16 = RayHitN<16>;
    using ClosestPointHit16 = ClosestPointHitN<16>;
 
public:

    explicit RayCaster(
        BitField<SceneFlags> scene_flags = SceneFlags::Robust,
        BuildQuality build_quality = BuildQuality::Medium);

    ~RayCaster();

    RayCaster(RayCaster&& other) noexcept;

    RayCaster& operator=(RayCaster&& other) noexcept;

    RayCaster(const RayCaster&) = delete;

    RayCaster& operator=(const RayCaster&) = delete;



    template <typename Scalar, typename Index>
    uint32_t add_mesh(
        SurfaceMesh<Scalar, Index> mesh,
        const std::optional<Affine<Scalar>>& transform = Affine<Scalar>::Identity());

    uint32_t add_instance(uint32_t mesh_index, const Eigen::Affine3f& transform);

    template <typename Scalar, typename Index>
    void add_scene(scene::SimpleScene<Scalar, Index, 3> simple_scene);

    void commit_updates();



    template <typename Scalar, typename Index>
    void update_mesh(uint32_t mesh_index, const SurfaceMesh<Scalar, Index>& mesh);

    template <typename Scalar, typename Index>
    void update_vertices(uint32_t mesh_index, const SurfaceMesh<Scalar, Index>& mesh);

    template <typename Scalar>
    void update_vertices(uint32_t mesh_index, span<const Scalar> vertices);

    Eigen::Affine3f get_transform(uint32_t mesh_index, uint32_t instance_index) const;

    void update_transform(
        uint32_t mesh_index,
        uint32_t instance_index,
        const Eigen::Affine3f& transform);

    bool get_visibility(uint32_t mesh_index, uint32_t instance_index);

    void update_visibility(uint32_t mesh_index, uint32_t instance_index, bool visible);



    auto get_intersection_filter(uint32_t mesh_index) const
        -> std::function<bool(uint32_t instance_index, uint32_t facet_index)>;

    void set_intersection_filter(
        uint32_t mesh_index,
        std::function<bool(uint32_t instance_index, uint32_t facet_index)>&& filter);

    auto get_occlusion_filter(uint32_t mesh_index) const
        -> std::function<bool(uint32_t instance_index, uint32_t facet_index)>;

    void set_occlusion_filter(
        uint32_t mesh_index,
        std::function<bool(uint32_t instance_index, uint32_t facet_index)>&& filter);



    std::optional<ClosestPointHit> closest_point(const Pointf& query_point) const;

    ClosestPointHit4 closest_point4(
        const Point4f& query_points,
        const std::variant<Mask4, size_t>& active) const;

    ClosestPointHit8 closest_point8(
        const Point8f& query_points,
        const std::variant<Mask8, size_t>& active) const;

    ClosestPointHit16 closest_point16(
        const Point16f& query_points,
        const std::variant<Mask16, size_t>& active) const;



    std::optional<ClosestPointHit> closest_vertex(const Pointf& query_point) const;

    ClosestPointHit4 closest_vertex4(
        const Point4f& query_points,
        const std::variant<Mask4, size_t>& active) const;

    ClosestPointHit8 closest_vertex8(
        const Point8f& query_points,
        const std::variant<Mask8, size_t>& active) const;

    ClosestPointHit16 closest_vertex16(
        const Point16f& query_points,
        const std::variant<Mask16, size_t>& active) const;

 


    std::optional<RayHit> cast(
        const Pointf& origin,
        const Directionf& direction,
        float tmin = 0,
        float tmax = std::numeric_limits<float>::infinity()) const;

    bool occluded(
        const Pointf& origin,
        const Directionf& direction,
        float tmin = 0,
        float tmax = std::numeric_limits<float>::infinity()) const;



    RayHit4 cast4(
        const Point4f& origins,
        const Direction4f& directions,
        const std::variant<Mask4, size_t>& active = size_t(4),
        const Float4& tmin = Float4::Zero(),
        const Float4& tmax = Float4::Constant(std::numeric_limits<float>::infinity())) const;

    uint32_t occluded4(
        const Point4f& origins,
        const Direction4f& directions,
        const std::variant<Mask4, size_t>& active = size_t(4),
        const Float4& tmin = Float4::Zero(),
        const Float4& tmax = Float4::Constant(std::numeric_limits<float>::infinity())) const;



    RayHit8 cast8(
        const Point8f& origins,
        const Direction8f& directions,
        const std::variant<Mask8, size_t>& active = size_t(8),
        const Float8& tmin = Float8::Zero(),
        const Float8& tmax = Float8::Constant(std::numeric_limits<float>::infinity())) const;

    uint32_t occluded8(
        const Point8f& origins,
        const Direction8f& directions,
        const std::variant<Mask8, size_t>& active = size_t(8),
        const Float8& tmin = Float8::Zero(),
        const Float8& tmax = Float8::Constant(std::numeric_limits<float>::infinity())) const;



    RayHit16 cast16(
        const Point16f& origins,
        const Direction16f& directions,
        const std::variant<Mask16, size_t>& active = size_t(16),
        const Float16& tmin = Float16::Zero(),
        const Float16& tmax = Float16::Constant(std::numeric_limits<float>::infinity())) const;

    uint32_t occluded16(
        const Point16f& origins,
        const Direction16f& directions,
        const std::variant<Mask16, size_t>& active = size_t(16),
        const Float16& tmin = Float16::Zero(),
        const Float16& tmax = Float16::Constant(std::numeric_limits<float>::infinity())) const;

 
private:
    friend struct internal::RayCasterOBBAccess;

    struct OrientedBox
    {
        Eigen::Vector3f center = Eigen::Vector3f::Zero();
        Eigen::Matrix3f axes = Eigen::Matrix3f::Identity(); 
        Eigen::Vector3f half_extents = Eigen::Vector3f::Zero();
    };
 
    void overlap_obb_internal(
        const OrientedBox& obb,
        function_ref<bool(uint32_t mesh_index, uint32_t instance_index, uint32_t facet_index)>
            callback) const;

    void overlap_obb16_internal(
        span<const OrientedBox> obbs,
        std::variant<Mask16, size_t> active,
        function_ref<
            bool(uint32_t lane, uint32_t mesh_index, uint32_t instance_index, uint32_t facet_index)>
            callback) const;
 
    struct Impl;
    value_ptr<Impl> m_impl;
};
 
} // namespace lagrange::raycasting