hydra-rs 0.0.1

#include "hydra_bridge.h"

#include <pxr/base/gf/frustum.h>
#include <pxr/base/gf/half.h>
#include <pxr/base/gf/matrix4d.h>
#include <pxr/base/gf/vec3d.h>
#include <pxr/base/gf/vec4d.h>
#include <pxr/base/gf/vec4f.h>
#include <pxr/base/vt/value.h>
#include <pxr/imaging/glf/simpleLight.h>
#include <pxr/imaging/glf/simpleMaterial.h>
#include <pxr/imaging/hd/aov.h>
#include <pxr/imaging/hd/driver.h>
#include <pxr/imaging/hd/renderBuffer.h>
#include <pxr/imaging/hd/rendererPluginRegistry.h>
#include <pxr/imaging/hd/types.h>
#include <pxr/imaging/hdx/tokens.h>
#include <pxr/imaging/hf/pluginDesc.h>
#include <pxr/imaging/hgi/hgi.h>
#include <pxr/imaging/hgi/tokens.h>
#include <pxr/usd/usd/timeCode.h>
#include <pxr/usdImaging/usdImagingGL/engine.h>
#include <pxr/usdImaging/usdImagingGL/renderParams.h>

#include <algorithm>
#include <cstring>
#include <functional>
#include <stdexcept>

namespace hydra_rs {

std::unique_ptr<SceneIndex> populate_from_path(rust::Str usd_path) {
    std::string p(usd_path.data(), usd_path.size());
    pxr::UsdStageRefPtr stage = pxr::UsdStage::Open(p);
    if (!stage) {
        throw std::runtime_error("UsdStage::Open returned null for: " + p);
    }

    auto si = pxr::UsdImagingStageSceneIndex::New();
    si->SetStage(stage);
    si->SetTime(pxr::UsdTimeCode::Default());

    auto wrapped = std::make_unique<SceneIndex>();
    wrapped->stage_owner = stage;
    wrapped->scene_index = si;
    return wrapped;
}

rust::String SceneIndex::stage_root() const {
    if (!stage_owner) return rust::String();
    auto root = stage_owner->GetRootLayer();
    if (!root) return rust::String();
    return rust::String(root->GetIdentifier());
}

size_t SceneIndex::prim_count() const {
    if (!scene_index) return 0;
    size_t count = 0;
    std::function<void(const pxr::SdfPath&)> walk = [&](const pxr::SdfPath& p) {
        ++count;
        for (const auto& child : scene_index->GetChildPrimPaths(p)) {
            walk(child);
        }
    };
    walk(pxr::SdfPath::AbsoluteRootPath());
    return count;
}

std::unique_ptr<std::vector<std::string>> SceneIndex::prim_paths() const {
    auto vec = std::make_unique<std::vector<std::string>>();
    if (!scene_index) return vec;
    std::function<void(const pxr::SdfPath&)> walk = [&](const pxr::SdfPath& p) {
        vec->push_back(p.GetString());
        for (const auto& child : scene_index->GetChildPrimPaths(p)) {
            walk(child);
        }
    };
    walk(pxr::SdfPath::AbsoluteRootPath());
    return vec;
}

std::unique_ptr<std::vector<std::string>> list_render_delegate_ids() {
    auto vec = std::make_unique<std::vector<std::string>>();
    auto& reg = pxr::HdRendererPluginRegistry::GetInstance();
    pxr::HfPluginDescVector descs;
    reg.GetPluginDescs(&descs);
    vec->reserve(descs.size());
    for (const auto& d : descs) {
        vec->push_back(d.id.GetString());
    }
    return vec;
}

std::unique_ptr<std::vector<uint8_t>> render_to_rgba(
    rust::Str usd_path,
    rust::Str render_delegate_id,
    uint32_t width,
    uint32_t height)
{
    if (width == 0 || height == 0) {
        throw std::runtime_error("render_to_rgba: width and height must be > 0");
    }

    std::string p(usd_path.data(), usd_path.size());
    pxr::UsdStageRefPtr stage = pxr::UsdStage::Open(p);
    if (!stage) {
        throw std::runtime_error("UsdStage::Open returned null for: " + p);
    }

    // Headless rendering with Storm requires an HGI that's actually
    // constructed up-front. The default UsdImagingGLEngine ctor lazy-inits
    // HGI on first render, and on macOS that path segfaults when there's no
    // pre-existing GL/Metal context. Constructing HgiMetal explicitly via
    // Hgi::CreatePlatformDefaultHgi and handing it to the engine through
    // HdDriver gives Storm a real Metal device to render into.
    //
    // hgi MUST outlive `engine` — declared first so it's destroyed last.
    pxr::HgiUniquePtr hgi = pxr::Hgi::CreatePlatformDefaultHgi();
    if (!hgi) {
        throw std::runtime_error(
            "Hgi::CreatePlatformDefaultHgi returned null — no usable GPU "
            "backend. On macOS, expected HgiMetal; on Linux, HgiGL/Vulkan.");
    }
    pxr::HdDriver hd_driver;
    hd_driver.name = pxr::HgiTokens->renderDriver;
    hd_driver.driver = pxr::VtValue(hgi.get());

    // High-level engine that owns the render delegate, render index, task
    // controller, and scene-index/imaging-delegate. Same API usdview and
    // usdrecord sit on top of, just fed an explicit HGI.
    pxr::UsdImagingGLEngine engine(hd_driver);

    // Headless: tell the engine to skip HdxPresentTask. Otherwise Storm
    // renders into a Metal texture and PresentTask tries to copy it to a
    // GL framebuffer via HgiInterop, which calls glGetString and segfaults
    // when there's no active GL context. We read the color AOV directly
    // through GetAovRenderBuffer below.
    engine.SetEnablePresentation(false);

    std::string delegate_str(render_delegate_id.data(), render_delegate_id.size());
    if (!delegate_str.empty()) {
        if (!engine.SetRendererPlugin(pxr::TfToken(delegate_str))) {
            throw std::runtime_error("SetRendererPlugin failed: " + delegate_str);
        }
    }

    // Free camera: 45° vertical FOV, looking at origin from (5, 5, 5).
    pxr::GfMatrix4d view_mat;
    view_mat.SetLookAt(
        pxr::GfVec3d(5.0, 5.0, 5.0),
        pxr::GfVec3d(0.0, 0.0, 0.0),
        pxr::GfVec3d(0.0, 1.0, 0.0));

    pxr::GfFrustum frustum;
    const double aspect = static_cast<double>(width) / static_cast<double>(height);
    frustum.SetPerspective(45.0, aspect, 0.1, 1000.0);
    pxr::GfMatrix4d proj_mat = frustum.ComputeProjectionMatrix();

    engine.SetCameraState(view_mat, proj_mat);
    engine.SetRenderViewport(pxr::GfVec4d(0.0, 0.0,
                                          static_cast<double>(width),
                                          static_cast<double>(height)));
    engine.SetRendererAov(pxr::HdAovTokens->color);

    // The Hydra-2 scene-index lighting path doesn't engage headlessly in this
    // configuration (engine.SetEnablePresentation(false) skips the present
    // task and at least some of the lighting setup that goes with it).
    // Hand the engine an explicit GlfSimpleLightVector via the legacy
    // SetLightingState API and Storm honors it. The same trick keeps usdview
    // and usdrecord working when there's no scene-authored light.
    pxr::GlfSimpleLight key_light;
    pxr::GfVec4f key_pos(2.5f, 4.0f, 5.0f, 1.0f);  // w=1: positional
    key_light.SetPosition(key_pos);
    key_light.SetAmbient(pxr::GfVec4f(0.0f, 0.0f, 0.0f, 1.0f));
    key_light.SetDiffuse(pxr::GfVec4f(1.0f, 1.0f, 1.0f, 1.0f));
    key_light.SetSpecular(pxr::GfVec4f(1.0f, 1.0f, 1.0f, 1.0f));

    pxr::GlfSimpleLightVector lights = {key_light};
    pxr::GlfSimpleMaterial material;
    material.SetAmbient(pxr::GfVec4f(0.2f, 0.2f, 0.2f, 1.0f));
    material.SetDiffuse(pxr::GfVec4f(1.0f, 1.0f, 1.0f, 1.0f));
    pxr::GfVec4f scene_ambient(0.1f, 0.1f, 0.1f, 1.0f);
    engine.SetLightingState(lights, material, scene_ambient);

    pxr::UsdImagingGLRenderParams params;
    params.frame = pxr::UsdTimeCode::Default();
    params.complexity = 1.0f;
    params.enableLighting = true;
    params.enableSceneLights = false;  // use the SetLightingState lights only
    params.enableSceneMaterials = true;
    params.clearColor = pxr::GfVec4f(0.1f, 0.1f, 0.15f, 1.0f);

    // Storm submits async; loop until converged. Embree-style ray tracers
    // would converge over many iterations — bound the loop to avoid infinite
    // loops on misbehaving delegates.
    engine.Render(stage->GetPseudoRoot(), params);
    int safety = 0;
    while (!engine.IsConverged() && safety < 256) {
        engine.Render(stage->GetPseudoRoot(), params);
        ++safety;
    }

    pxr::HdRenderBuffer* color = engine.GetAovRenderBuffer(pxr::HdAovTokens->color);
    if (!color) {
        throw std::runtime_error("UsdImagingGLEngine has no color render buffer");
    }
    color->Resolve();

    const uint32_t bw = color->GetWidth();
    const uint32_t bh = color->GetHeight();
    const pxr::HdFormat fmt = color->GetFormat();

    void* mapped = color->Map();
    if (!mapped) {
        throw std::runtime_error("HdRenderBuffer::Map returned null — render may have failed");
    }

    auto out = std::make_unique<std::vector<uint8_t>>();
    out->resize(static_cast<size_t>(bw) * bh * 4);

    if (fmt == pxr::HdFormatFloat32Vec4) {
        const float* src = static_cast<const float*>(mapped);
        for (size_t i = 0; i < static_cast<size_t>(bw) * bh; ++i) {
            for (int c = 0; c < 4; ++c) {
                float v = src[i * 4 + c];
                v = std::max(0.0f, std::min(1.0f, v));
                (*out)[i * 4 + c] = static_cast<uint8_t>(v * 255.0f + 0.5f);
            }
        }
    } else if (fmt == pxr::HdFormatFloat16Vec4) {
        // Storm's default color AOV is half-float on macOS Metal.
        const pxr::GfHalf* src = static_cast<const pxr::GfHalf*>(mapped);
        for (size_t i = 0; i < static_cast<size_t>(bw) * bh; ++i) {
            for (int c = 0; c < 4; ++c) {
                float v = static_cast<float>(src[i * 4 + c]);
                v = std::max(0.0f, std::min(1.0f, v));
                (*out)[i * 4 + c] = static_cast<uint8_t>(v * 255.0f + 0.5f);
            }
        }
    } else if (fmt == pxr::HdFormatUNorm8Vec4) {
        std::memcpy(out->data(), mapped, out->size());
    } else {
        color->Unmap();
        throw std::runtime_error(
            "Unsupported render buffer format from delegate (HdFormat=" +
            std::to_string(static_cast<int>(fmt)) + ")");
    }

    color->Unmap();
    return out;
}

}  // namespace hydra_rs