ImplicitSurfaceNetwork-xj/primitive_process/src/base/primitive.cpp

#include <data/data_center.hpp>
#include <data/data_type.hpp>
#include <base/subface.hpp>
#include <base/primitive.hpp>

dynamic_bitset_mp<> primitive::judge_sign_by_subface_sign(const stl_vector_mp<dynamic_bitset_mp<>>& subface_signs) const
{
    // NOTE: This overload assumes subface_signs[i] is already in the same order as
    // get_subfaces()[i] (i.e. local init order, not baked order). Only use when
    // signs have NOT been reordered by bake_blobtree.
    auto subfaces = get_subfaces();

    // Apply mark for i=0 as well (was incorrectly skipped before)
    dynamic_bitset_mp<> res = subfaces[0].get_mark() ? ~subface_signs[0] : subface_signs[0];
    for (size_t i = 1; i < subfaces.size(); ++i) {
        if (!subfaces[i].get_mark())
            res &= subface_signs[i];
        else
            res &= ~subface_signs[i];
    }

    return res;
}

dynamic_bitset_mp<> primitive::judge_sign_by_subface_sign(const stl_vector_mp<dynamic_bitset_mp<>>& subface_signs,
                                                          const stl_vector_mp<uint32_t>&            subface_indices,
                                                          const stl_vector_mp<bool>&                invert_flags) const
{
    // subface_indices[i] and invert_flags[i] are both in BAKED order (as stored in
    // baked_blobtree_t::subfaces_of_primitives / invert_flags_of_primitives).
    // invert_flags[i] is the mark that corresponds to baked subface subface_indices[i],
    // populated by bake_blobtree via pointer-based lookup so the order is always correct
    // regardless of how baking reorders subfaces relative to initialize() order.
    assert(subface_indices.size() == invert_flags.size());

    dynamic_bitset_mp<> res = invert_flags[0] ? ~subface_signs[subface_indices[0]] : subface_signs[subface_indices[0]];
    for (size_t i = 1; i < subface_indices.size(); ++i) {
        if (!invert_flags[i])
            res &= subface_signs[subface_indices[i]];
        else
            res &= ~subface_signs[subface_indices[i]];
    }

    return res;
}

void primitive::initialize(const stl_vector_mp<internal::paired_model_matrix_ptr_t>& matrices,
                           const stl_vector_mp<bool>&                                reversed_flags,
                           const stl_vector_mp<const void*>&                         geometry,
                           const aabb_t&                                             aabb)
{
    auto                       subfaces      = get_subfaces();
    auto                       subface_types = get_subface_types();
    // 创建 surface（geometry 所有权转移到 surface）
    stl_vector_mp<const void*> raw_geometry  = geometry.empty() ? get_subface_geometries() : geometry;

    stl_vector_mp<const void*> canonical_geometry;
    canonical_geometry.reserve(raw_geometry.size());

    for (size_t i = 0; i < subface_types.size(); ++i) {
        const void* canonical_ptr = nullptr;

        if (raw_geometry[i] != nullptr) {
            switch (subface_types[i]) {
                case surface_type::plane: {
                    canonical_ptr = data_center->acquire_plane_geometry(raw_geometry[i]);
                    std::cout << "[PRIMITIVE_INIT] Acquired plane geometry for subface " << i << "\n";
                    std::cout << "[PRIMITIVE_INIT] Plane geometry descriptor: local_aabb_min=("
                              << static_cast<const plane_descriptor_t*>(canonical_ptr)->local_aabb_min.x << ", "
                              << static_cast<const plane_descriptor_t*>(canonical_ptr)->local_aabb_min.y << ")\n";
                    break;
                }
                case surface_type::cylinder: {
                    // 将原始 geometry 注册到 data_center，获取规范化指针
                    canonical_ptr = data_center->acquire_cylinder_geometry(raw_geometry[i]);
                    break;
                }
                case surface_type::sphere: {
                    // Sphere 不依赖任何外部 geometry descriptor，
                    // 其形状完全由 paired_model_matrix 编码，canonical_ptr 保持 nullptr。
                    break;
                }
                case surface_type::extrude_polyline_side: {
                    canonical_ptr = data_center->acquire_extrude_polyline_geometry(raw_geometry[i]);
                    break;
                }
                case surface_type::extrude_helixline_side: {
                    canonical_ptr = data_center->acquire_extrude_helixline_geometry(raw_geometry[i]);
                    break;
                }
                default:                  canonical_ptr = nullptr; break;
            }
        }
        canonical_geometry.push_back(canonical_ptr);
    }
    for (size_t i = 0; i < subfaces.size(); ++i) {
        data_center->require_surface(subface_types[i], *matrices[i], canonical_geometry[i], subfaces[i]);
        subfaces[i].set_mark(static_cast<size_t>(reversed_flags[i]));
    }
    m_aabb = aabb;
}

void primitive::destroy()
{
    auto subfaces      = get_subfaces();
    auto subface_types = get_subface_types();
    assert(subfaces.front().get_ptr() != nullptr);

    for (size_t i = 0; i < subfaces.size(); ++i) {
        auto        get_geom_func = internal::get_geometry_accessor_ptr(subface_types[i]);
        const void* geometry_ptr  = get_geom_func(make_pointer_wrapper(subfaces[i].get_ptr()));

        data_center->release_surface(subface_types[i], subfaces[i]);
        subfaces[i].set_ptr(nullptr);

        // 释放descriptor
        if (geometry_ptr) {
            switch (subface_types[i]) {
                case surface_type::plane: {
                    data_center->release_plane_geometry(geometry_ptr);
                    break;
                }
                case surface_type::cylinder: {
                    data_center->release_cylinder_geometry(geometry_ptr);
                    break;
                }
                case surface_type::sphere: {
                    // Sphere 没有 descriptor，geometry_ptr 始终为 nullptr，此分支不会执行。
                    break;
                }
                case surface_type::extrude_polyline_side: {
                    data_center->release_extrude_polyline_geometry(geometry_ptr);
                    break;
                }
                case surface_type::extrude_helixline_side: {
                    data_center->release_extrude_helixline_geometry(geometry_ptr);
                    break;
                }
                default: break;
            }
        }
    }
}

void primitive::apply_transform(internal::transform_type type, Eigen::Vector4d param)
{
    auto subfaces      = get_subfaces();
    auto subface_types = get_subface_types();

    // get a copy of old model matrices
    stl_vector_mp<internal::paired_model_matrix> new_model_matrices{};
    new_model_matrices.reserve(subfaces.size());
    for (size_t i = 0; i < subfaces.size(); ++i) {
        new_model_matrices.emplace_back(static_cast<internal::paired_model_matrix>(*subfaces[i]));
    }

    stl_vector_mp<const void*> geometry_ptrs;
    geometry_ptrs.reserve(subfaces.size());

    for (size_t i = 0; i < subfaces.size(); ++i) {
        // Plane and sphere both carry no descriptor: geometry pointer is always nullptr.
        if (subface_types[i] == surface_type::plane || subface_types[i] == surface_type::sphere) {
            geometry_ptrs.push_back(nullptr);
        } else {
            auto        get_geom_func = internal::get_geometry_accessor_ptr(subface_types[i]);
            const void* geom_ptr      = get_geom_func(make_pointer_wrapper(subfaces[i].get_ptr()));
            geometry_ptrs.push_back(geom_ptr);
        }
    }

    // make affine transform matrix from type and param
    // and apply transform to model matrices by the way, since this is faster than applying affine matrix
    Eigen::Affine3d affine = Eigen::Affine3d::Identity();
    switch (type) {
        case internal::transform_type::scale: {
            affine.linear() = param.head<3>().asDiagonal();
            Eigen::Affine3d affine_inv = affine.inverse();
            for (auto& model_matrix : new_model_matrices) {
                model_matrix.local_to_world = affine * model_matrix.local_to_world;
                model_matrix.world_to_local = model_matrix.world_to_local * affine_inv;
            }
        } break;
        case internal::transform_type::rotation: {
            affine.linear() = Eigen::Quaterniond(param).toRotationMatrix();
            Eigen::Affine3d affine_inv = affine.inverse();
            for (auto& model_matrix : new_model_matrices) {
                model_matrix.local_to_world = affine * model_matrix.local_to_world;
                model_matrix.world_to_local = model_matrix.world_to_local * affine_inv;
            }
        } break;
        case internal::transform_type::translation: {
            Eigen::Translation3d translation(param.head<3>());
            affine = Eigen::Affine3d(translation);
            for (auto& model_matrix : new_model_matrices) {
                model_matrix.local_to_world = translation * model_matrix.local_to_world;
                model_matrix.world_to_local = model_matrix.world_to_local * translation.inverse();
            }
        } break;
        default: throw std::invalid_argument("Invalid transform type");
    }

    stl_vector_mp<size_t> saved_marks;
    saved_marks.reserve(subfaces.size());
    for (size_t i = 0; i < subfaces.size(); ++i) { saved_marks.push_back(subfaces[i].get_mark()); }

    for (size_t i = 0; i < subfaces.size(); ++i) {
        data_center->require_surface(subface_types[i], new_model_matrices[i], geometry_ptrs[i], subfaces[i]);
        // Restore the mark that was set during initialize() via reversed_flags.
        subfaces[i].set_mark(saved_marks[i]);
    }
    recompute_aabb();
}

// 重新计算 Primitive 的 AABB
void primitive::recompute_aabb()
{
    auto subfaces      = get_subfaces();
    auto subface_types = get_subface_types();

    // 重置为空 AABB
    m_aabb.setEmpty();
    bool has_valid_aabb = false;

    for (size_t i = 0; i < subfaces.size(); ++i) {
        aabb_t local_aabb = aabb_t();
        if (subface_types[i] == surface_type::extrude_polyline_side) {
            // Extrude faces own their own AABB computation.
            auto* face = static_cast<internal::extrude_polyline_side_face_t*>(subfaces[i].get_ptr());
            local_aabb.extend(face->get_aabb()); // 通过访问器读取
        } else if (subface_types[i] == surface_type::extrude_helixline_side) {
            auto* face = static_cast<internal::extrude_helixline_side_face_t*>(subfaces[i].get_ptr());
            local_aabb.extend(face->get_aabb());
        } else {
            switch (subface_types[i]) {
                case surface_type::sphere: {
                    constexpr double aabb_epsilon = 1e-9;
                    local_aabb                    = aabb_t(internal::k_aabb_unit.min().array() - aabb_epsilon,
                                        internal::k_aabb_unit.max().array() + aabb_epsilon);
                    break;
                }
                case surface_type::cylinder: {
                    if (subfaces[i].get_ptr() != nullptr) {
                        auto        get_geom_func = internal::get_geometry_accessor_ptr(surface_type::cylinder);
                        const void* geom_raw      = get_geom_func(make_pointer_wrapper(subfaces[i].get_ptr()));
                        const auto* desc          = static_cast<const cylinder_descriptor_t*>(geom_raw);
                        if (desc) {
                            const double r = desc->radius;
                            const double h = desc->height; // use the explicit height field, not offset magnitude
                            local_aabb     = aabb_t(Eigen::Vector3d(-r, -r, 0.0), Eigen::Vector3d(r, r, h));
                        }
                    }
                    break;
                }
                case surface_type::plane:
                default:                  continue;
            }

            if (local_aabb.isEmpty()) { continue; }
            if (!local_aabb.min().allFinite() || !local_aabb.max().allFinite()) {
                std::cerr << "[RECOMPUTE-AABB-WARNING] Invalid local AABB - skipping\n";
                continue;
            }
        }
        // 变换到世界坐标系
        aabb_t world_aabb = local_aabb;
        world_aabb.transform(subfaces[i]->local_to_world);
        if (!world_aabb.min().allFinite() || !world_aabb.max().allFinite()) {
            std::cerr << "[RECOMPUTE-AABB-WARNING] Invalid world AABB after transform\n";
            continue;
        }
        // 合并到总 AABB
        if (!has_valid_aabb) {
            m_aabb         = world_aabb;
            has_valid_aabb = true;
        } else {
            m_aabb.extend(world_aabb);
        }
    }
    if (!has_valid_aabb) {
        std::cerr << "[RECOMPUTE-AABB-WARNING] No valid AABB computed for primitive\n";
        {
            m_aabb = internal::k_aabb_unit;
        }
    }
}

void primitive::update_geometry(const void* new_geometry_ptr, size_t subface_index)
{
    auto subfaces      = get_subfaces();
    auto subface_types = get_subface_types();

    if (subface_types[subface_index] == surface_type::plane) {
        throw std::invalid_argument("Cannot update descriptor for plane surface");
    }

    auto        get_geom_func = internal::get_geometry_accessor_ptr(subface_types[subface_index]);
    const void* current_geom  = get_geom_func(make_pointer_wrapper(subfaces[subface_index].get_ptr()));
    if (current_geom == new_geometry_ptr) { return; }
    stl_vector_mp<internal::paired_model_matrix> current_matrices;
    current_matrices.reserve(subfaces.size());
    for (size_t i = 0; i < subfaces.size(); ++i) {
        current_matrices.emplace_back(static_cast<internal::paired_model_matrix>(*subfaces[i]));
    }

    stl_vector_mp<const void*> new_geometry;
    new_geometry.reserve(subfaces.size());
    for (size_t i = 0; i < subfaces.size(); ++i) {
        if (i == subface_index) {
            //  规范化新 geometry
            const void* canonical_ptr = nullptr;

            switch (subface_types[i]) {
                case surface_type::plane: {
                    canonical_ptr = data_center->acquire_plane_geometry(new_geometry_ptr);
                    break;
                }
                case surface_type::cylinder: {
                    canonical_ptr = data_center->acquire_cylinder_geometry(new_geometry_ptr);
                    break;
                }
                case surface_type::extrude_polyline_side: {
                    canonical_ptr = data_center->acquire_extrude_polyline_geometry(new_geometry_ptr);
                    break;
                }
                case surface_type::extrude_helixline_side: {
                    canonical_ptr = data_center->acquire_extrude_helixline_geometry(new_geometry_ptr);
                    break;
                }
                default: throw std::invalid_argument("Unsupported surface type for descriptor update");
            }
            new_geometry.push_back(canonical_ptr);
        } else {
            // 保持原有 descriptor
            if (subface_types[i] == surface_type::plane || subface_types[i] == surface_type::sphere) {
                // Plane and sphere carry no descriptor.
                new_geometry.push_back(nullptr);
            } else {
                auto        get_desc_func     = internal::get_geometry_accessor_ptr(subface_types[i]);
                const void* existing_desc_ptr = get_desc_func(make_pointer_wrapper(subfaces[i].get_ptr()));
                new_geometry.push_back(existing_desc_ptr);
            }
        }
    }
    // 先释放 surface
    data_center->release_surface(subface_types[subface_index], subfaces[subface_index]);
    subfaces[subface_index].set_ptr(nullptr); // 清零，避免 require_surface 内双重释放

    if (current_geom && current_geom != new_geometry[subface_index]) {
        switch (subface_types[subface_index]) {
            case surface_type::plane: {
                data_center->release_plane_geometry(current_geom);
                break;
            }
            case surface_type::cylinder: {
                data_center->release_cylinder_geometry(current_geom);
                break;
            }
            case surface_type::extrude_polyline_side: {
                data_center->release_extrude_polyline_geometry(current_geom);
                break;
            }
            case surface_type::extrude_helixline_side: {
                data_center->release_extrude_helixline_geometry(current_geom);
                break;
            }
            default: break;
        }
    }
    const size_t saved_mark = subfaces[subface_index].get_mark();
    data_center->require_surface(subface_types[subface_index],
                                 current_matrices[subface_index],
                                 new_geometry[subface_index],
                                 subfaces[subface_index]);
    subfaces[subface_index].set_mark(saved_mark);
    // 重新计算 AABB
    recompute_aabb();
}