rust-usd 0.0.4

#include "usd_bridge.h"

#include "rust-usd/src/lib.rs.h"

#include <pxr/base/gf/matrix4d.h>
#include <pxr/base/gf/vec2f.h>
#include <pxr/base/gf/vec3f.h>
#include <pxr/base/plug/registry.h>
#include <pxr/base/tf/token.h>
#include <pxr/base/tf/type.h>
#include <pxr/base/vt/array.h>
#include <pxr/usd/ar/defaultResolver.h>
#include <pxr/usd/ar/defineResolver.h>
#include <pxr/usd/ar/resolvedPath.h>
#include <pxr/usd/ar/resolver.h>
#include <pxr/usd/sdf/assetPath.h>
#include <pxr/usd/sdf/layer.h>
#include <pxr/usd/sdf/path.h>
#include <pxr/usd/usd/primRange.h>
#include <pxr/usd/usd/timeCode.h>
#include <pxr/usd/usdGeom/primvarsAPI.h>
#include <pxr/usd/usdGeom/tokens.h>
#include <pxr/usd/usdGeom/xformCache.h>
#include <pxr/usd/usdShade/material.h>
#include <pxr/usd/usdShade/materialBindingAPI.h>
#include <pxr/usd/usdShade/shader.h>

#include <algorithm>
#include <cstring>
#include <mutex>
#include <stdexcept>
#include <string>

namespace rust_usd {

namespace {

void flatten_vec3(const pxr::VtArray<pxr::GfVec3f>& src, std::vector<float>& dst) {
    dst.reserve(src.size() * 3);
    for (const auto& v : src) {
        dst.push_back(v[0]);
        dst.push_back(v[1]);
        dst.push_back(v[2]);
    }
}

void flatten_vec2(const pxr::VtArray<pxr::GfVec2f>& src, std::vector<float>& dst) {
    dst.reserve(src.size() * 2);
    for (const auto& v : src) {
        dst.push_back(v[0]);
        dst.push_back(v[1]);
    }
}

std::string asset_path_to_string(const pxr::SdfAssetPath& asset) {
    const std::string& resolved = asset.GetResolvedPath();
    return resolved.empty() ? asset.GetAssetPath() : resolved;
}

// Registered URI schemes (e.g. {"forge", "s3"}) the Rust resolver claims.
// Anything not matching falls through to ArDefaultResolver.
std::mutex& schemes_mutex() {
    static std::mutex m;
    return m;
}

std::vector<std::string>& registered_schemes() {
    static std::vector<std::string> v;
    return v;
}

bool path_uses_registered_scheme(const std::string& path) {
    auto colon = path.find(':');
    if (colon == std::string::npos) return false;
    if (path.size() < colon + 3) return false;
    if (path[colon + 1] != '/' || path[colon + 2] != '/') return false;

    std::lock_guard<std::mutex> lock(schemes_mutex());
    const auto& schemes = registered_schemes();
    for (const auto& s : schemes) {
        if (s.size() == colon &&
            std::memcmp(path.data(), s.data(), colon) == 0) {
            return true;
        }
    }
    return false;
}

}  // namespace

// ---------------------------------------------------------------------------
// RustUriResolver
//
// Inherits ArDefaultResolver's filesystem behavior; intercepts URI paths
// whose scheme is in the registered set and bounces them through the
// Rust-side AssetResolver trait. Once the URI is translated to a real
// filesystem path, every other resolver method (OpenAsset, timestamps, etc.)
// just works via the inherited implementation.
// ---------------------------------------------------------------------------
class RustUriResolver : public pxr::ArDefaultResolver {
public:
    RustUriResolver() = default;
    ~RustUriResolver() override = default;

protected:
    std::string _CreateIdentifier(
        const std::string& assetPath,
        const pxr::ArResolvedPath& anchorAssetPath) const override
    {
        if (path_uses_registered_scheme(assetPath)) return assetPath;
        return pxr::ArDefaultResolver::_CreateIdentifier(assetPath, anchorAssetPath);
    }

    std::string _CreateIdentifierForNewAsset(
        const std::string& assetPath,
        const pxr::ArResolvedPath& anchorAssetPath) const override
    {
        if (path_uses_registered_scheme(assetPath)) return assetPath;
        return pxr::ArDefaultResolver::_CreateIdentifierForNewAsset(assetPath, anchorAssetPath);
    }

    pxr::ArResolvedPath _Resolve(const std::string& assetPath) const override
    {
        if (path_uses_registered_scheme(assetPath)) {
            rust::String resolved = uri_resolve(rust::Str(assetPath));
            if (resolved.size() == 0) return pxr::ArResolvedPath();
            return pxr::ArResolvedPath(std::string(resolved));
        }
        return pxr::ArDefaultResolver::_Resolve(assetPath);
    }

    pxr::ArResolvedPath _ResolveForNewAsset(const std::string& assetPath) const override
    {
        return pxr::ArDefaultResolver::_ResolveForNewAsset(assetPath);
    }
};

}  // namespace rust_usd

PXR_NAMESPACE_USING_DIRECTIVE
AR_DEFINE_RESOLVER(rust_usd::RustUriResolver, pxr::ArDefaultResolver)

namespace rust_usd {

void install_preferred_resolver() {
    static std::once_flag once;
    std::call_once(once, [] {
        pxr::TfType type = pxr::TfType::Find<RustUriResolver>();
        if (!type.IsUnknown()) {
            pxr::ArSetPreferredResolver(type.GetTypeName());
        }
    });
}

void register_plugin_directory(rust::Str dir) {
    std::string s(dir.data(), dir.size());
    pxr::PlugRegistry::GetInstance().RegisterPlugins(s);
}

void register_uri_scheme(rust::Str scheme) {
    std::lock_guard<std::mutex> lock(schemes_mutex());
    registered_schemes().push_back(std::string(scheme.data(), scheme.size()));
}

void clear_uri_schemes() {
    std::lock_guard<std::mutex> lock(schemes_mutex());
    registered_schemes().clear();
}

std::unique_ptr<Stage> open_stage(rust::Str path) {
    return open_stage_with_load(path, /*load_all=*/true);
}

std::unique_ptr<Stage> open_stage_with_load(rust::Str path, bool load_all) {
    std::string path_str(path.data(), path.size());
    auto load_set = load_all
        ? pxr::UsdStage::InitialLoadSet::LoadAll
        : pxr::UsdStage::InitialLoadSet::LoadNone;
    pxr::UsdStageRefPtr stage_ptr = pxr::UsdStage::Open(path_str, load_set);
    if (!stage_ptr) {
        throw std::runtime_error("UsdStage::Open returned null for path: " + path_str);
    }
    auto wrapped = std::make_unique<Stage>();
    wrapped->stage = stage_ptr;
    return wrapped;
}

std::unique_ptr<Stage> open_for_painting(rust::Str asset_path, rust::Str edit_layer_path) {
    std::string asset(asset_path.data(), asset_path.size());
    std::string edits(edit_layer_path.data(), edit_layer_path.size());

    pxr::SdfLayerRefPtr edit_layer = pxr::SdfLayer::FindOrOpen(edits);
    if (!edit_layer) {
        edit_layer = pxr::SdfLayer::CreateNew(edits);
        if (!edit_layer) {
            throw std::runtime_error("Failed to create edit layer: " + edits);
        }
    }

    // Make sure the asset is in the edit layer's sublayer list. The edit
    // layer sits ABOVE the asset, so opinions in the edit layer override
    // anything authored in the asset, but the asset file itself stays
    // untouched.
    auto sublayers = edit_layer->GetSubLayerPaths();
    bool found = std::any_of(sublayers.begin(), sublayers.end(),
                             [&](const std::string& s) { return s == asset; });
    if (!found) {
        edit_layer->InsertSubLayerPath(asset);
    }

    pxr::UsdStageRefPtr stage_ptr = pxr::UsdStage::Open(edit_layer);
    if (!stage_ptr) {
        throw std::runtime_error("Failed to open stage from edit layer: " + edits);
    }
    // Default edit target is the root layer (= edit_layer here), so any
    // subsequent authoring lands in the edit layer.

    auto wrapped = std::make_unique<Stage>();
    wrapped->stage = stage_ptr;
    return wrapped;
}

std::unique_ptr<Prim> Stage::pseudo_root() const {
    if (!stage) {
        throw std::runtime_error("rust_usd::Stage holds a null UsdStageRefPtr");
    }
    auto wrapped = std::make_unique<Prim>();
    wrapped->prim = stage->GetPseudoRoot();
    return wrapped;
}

std::unique_ptr<Prim> Stage::prim_at_path(rust::Str sdf_path) const {
    if (!stage) return nullptr;
    std::string s(sdf_path.data(), sdf_path.size());
    pxr::SdfPath path(s);
    if (!path.IsAbsolutePath()) return nullptr;
    pxr::UsdPrim prim = stage->GetPrimAtPath(path);
    if (!prim) return nullptr;
    auto wrapped = std::make_unique<Prim>();
    wrapped->prim = prim;
    return wrapped;
}

std::unique_ptr<std::vector<Mesh>> Stage::all_meshes() const {
    auto vec = std::make_unique<std::vector<Mesh>>();
    if (!stage) return vec;
    for (const pxr::UsdPrim& prim : stage->Traverse()) {
        if (prim.IsA<pxr::UsdGeomMesh>()) {
            vec->push_back(Mesh{pxr::UsdGeomMesh(prim)});
        }
    }
    return vec;
}

void Stage::load_path(rust::Str sdf_path) const {
    if (!stage) return;
    std::string s(sdf_path.data(), sdf_path.size());
    stage->Load(pxr::SdfPath(s), pxr::UsdLoadPolicy::UsdLoadWithDescendants);
}

void Stage::unload_path(rust::Str sdf_path) const {
    if (!stage) return;
    std::string s(sdf_path.data(), sdf_path.size());
    stage->Unload(pxr::SdfPath(s));
}

void Stage::save_edit_layer() const {
    if (!stage) return;
    auto target = stage->GetEditTarget();
    auto layer = target.GetLayer();
    if (layer && !layer->IsAnonymous()) {
        layer->Save();
    }
}

rust::String Stage::edit_layer_path() const {
    if (!stage) return rust::String();
    auto target = stage->GetEditTarget();
    auto layer = target.GetLayer();
    if (!layer) return rust::String();
    return rust::String(layer->GetIdentifier());
}

rust::String Prim::path() const {
    return rust::String(prim.GetPath().GetString());
}

rust::String Prim::type_name() const {
    return rust::String(prim.GetTypeName().GetString());
}

bool Prim::is_mesh() const {
    return prim.IsA<pxr::UsdGeomMesh>();
}

std::unique_ptr<std::vector<Prim>> Prim::children() const {
    auto vec = std::make_unique<std::vector<Prim>>();
    for (const pxr::UsdPrim& child : prim.GetChildren()) {
        vec->push_back(Prim{child});
    }
    return vec;
}

std::unique_ptr<Mesh> Prim::as_mesh() const {
    if (!prim.IsA<pxr::UsdGeomMesh>()) return nullptr;
    auto wrapped = std::make_unique<Mesh>();
    wrapped->mesh = pxr::UsdGeomMesh(prim);
    return wrapped;
}

std::unique_ptr<std::vector<std::string>> Prim::variant_set_names() const {
    auto vec = std::make_unique<std::vector<std::string>>();
    if (!prim || !prim.HasVariantSets()) return vec;
    auto names = prim.GetVariantSets().GetNames();
    *vec = std::vector<std::string>(names.begin(), names.end());
    return vec;
}

std::unique_ptr<VariantSet> Prim::get_variant_set(rust::Str name) const {
    if (!prim) return nullptr;
    std::string n(name.data(), name.size());
    pxr::UsdVariantSet vs = prim.GetVariantSet(n);
    if (!vs.IsValid()) return nullptr;
    return std::make_unique<VariantSet>(std::move(vs));
}

rust::String VariantSet::name() const {
    return rust::String(variant_set.GetName());
}

std::unique_ptr<std::vector<std::string>> VariantSet::variant_names() const {
    auto vec = std::make_unique<std::vector<std::string>>();
    auto names = variant_set.GetVariantNames();
    *vec = std::vector<std::string>(names.begin(), names.end());
    return vec;
}

rust::String VariantSet::selection() const {
    return rust::String(variant_set.GetVariantSelection());
}

bool VariantSet::set_selection(rust::Str variant_name) const {
    std::string n(variant_name.data(), variant_name.size());
    return variant_set.SetVariantSelection(n);
}

void VariantSet::clear_selection() const {
    variant_set.ClearVariantSelection();
}

bool VariantSet::has_authored_selection() const {
    return variant_set.HasAuthoredVariantSelection();
}

rust::String Mesh::prim_path() const {
    return rust::String(mesh.GetPrim().GetPath().GetString());
}

rust::String Mesh::subdivision_scheme() const {
    pxr::TfToken scheme;
    if (mesh.GetSubdivisionSchemeAttr().Get(&scheme, pxr::UsdTimeCode::Default())) {
        return rust::String(scheme.GetString());
    }
    return rust::String("catmullClark");
}

std::unique_ptr<std::vector<float>> Mesh::points() const {
    auto vec = std::make_unique<std::vector<float>>();
    pxr::VtArray<pxr::GfVec3f> arr;
    if (mesh.GetPointsAttr().Get(&arr, pxr::UsdTimeCode::Default())) {
        flatten_vec3(arr, *vec);
    }
    return vec;
}

std::unique_ptr<std::vector<int32_t>> Mesh::face_vertex_counts() const {
    auto vec = std::make_unique<std::vector<int32_t>>();
    pxr::VtArray<int> arr;
    if (mesh.GetFaceVertexCountsAttr().Get(&arr, pxr::UsdTimeCode::Default())) {
        vec->reserve(arr.size());
        for (int v : arr) vec->push_back(static_cast<int32_t>(v));
    }
    return vec;
}

std::unique_ptr<std::vector<int32_t>> Mesh::face_vertex_indices() const {
    auto vec = std::make_unique<std::vector<int32_t>>();
    pxr::VtArray<int> arr;
    if (mesh.GetFaceVertexIndicesAttr().Get(&arr, pxr::UsdTimeCode::Default())) {
        vec->reserve(arr.size());
        for (int v : arr) vec->push_back(static_cast<int32_t>(v));
    }
    return vec;
}

std::unique_ptr<std::vector<float>> Mesh::normals() const {
    auto vec = std::make_unique<std::vector<float>>();
    pxr::VtArray<pxr::GfVec3f> arr;
    if (mesh.GetNormalsAttr().Get(&arr, pxr::UsdTimeCode::Default())) {
        flatten_vec3(arr, *vec);
    }
    return vec;
}

rust::String Mesh::normals_interpolation() const {
    return rust::String(mesh.GetNormalsInterpolation().GetString());
}

rust::String Mesh::orientation() const {
    pxr::TfToken tok;
    if (auto attr = mesh.GetOrientationAttr()) {
        attr.Get(&tok, pxr::UsdTimeCode::Default());
    }
    if (tok.IsEmpty()) {
        tok = pxr::UsdGeomTokens->rightHanded;
    }
    return rust::String(tok.GetString());
}

std::unique_ptr<std::vector<float>> Mesh::st() const {
    auto vec = std::make_unique<std::vector<float>>();
    auto pv = pxr::UsdGeomPrimvarsAPI(mesh.GetPrim()).GetPrimvar(pxr::TfToken("st"));
    if (!pv) return vec;
    pxr::VtArray<pxr::GfVec2f> arr;
    if (pv.Get(&arr, pxr::UsdTimeCode::Default())) {
        flatten_vec2(arr, *vec);
    }
    return vec;
}

std::unique_ptr<std::vector<int32_t>> Mesh::st_indices() const {
    auto vec = std::make_unique<std::vector<int32_t>>();
    auto pv = pxr::UsdGeomPrimvarsAPI(mesh.GetPrim()).GetPrimvar(pxr::TfToken("st"));
    if (!pv || !pv.IsIndexed()) return vec;
    pxr::VtArray<int> arr;
    if (pv.GetIndices(&arr, pxr::UsdTimeCode::Default())) {
        vec->reserve(arr.size());
        for (int v : arr) vec->push_back(static_cast<int32_t>(v));
    }
    return vec;
}

std::unique_ptr<std::vector<float>> Mesh::local_to_world() const {
    auto vec = std::make_unique<std::vector<float>>();
    vec->reserve(16);
    pxr::UsdGeomXformCache cache(pxr::UsdTimeCode::Default());
    pxr::GfMatrix4d m = cache.GetLocalToWorldTransform(mesh.GetPrim());
    for (size_t r = 0; r < 4; ++r) {
        for (size_t c = 0; c < 4; ++c) {
            vec->push_back(static_cast<float>(m[r][c]));
        }
    }
    return vec;
}

std::unique_ptr<std::vector<std::string>> Mesh::bound_texture_paths() const {
    auto vec = std::make_unique<std::vector<std::string>>();

    pxr::UsdShadeMaterialBindingAPI binding_api(mesh.GetPrim());
    pxr::UsdShadeMaterial material = binding_api.ComputeBoundMaterial();
    if (!material) return vec;

    static const pxr::TfToken kUvTexture("UsdUVTexture");
    static const pxr::TfToken kFile("inputs:file");

    for (const pxr::UsdPrim& child : material.GetPrim().GetChildren()) {
        pxr::UsdShadeShader shader(child);
        if (!shader) continue;

        pxr::TfToken id;
        if (!shader.GetIdAttr().Get(&id)) continue;
        if (id != kUvTexture) continue;

        pxr::UsdAttribute file_attr = child.GetAttribute(kFile);
        if (!file_attr) continue;

        pxr::SdfAssetPath asset;
        if (!file_attr.Get(&asset, pxr::UsdTimeCode::Default())) continue;

        vec->push_back(asset_path_to_string(asset));
    }

    return vec;
}

std::unique_ptr<std::vector<std::string>> Mesh::primvar_names() const {
    auto vec = std::make_unique<std::vector<std::string>>();
    pxr::UsdGeomPrimvarsAPI api(mesh.GetPrim());
    auto pvs = api.GetPrimvars();
    vec->reserve(pvs.size());
    for (const auto& pv : pvs) {
        vec->push_back(pv.GetPrimvarName().GetString());
    }
    return vec;
}

std::unique_ptr<Primvar> Mesh::primvar(rust::Str name) const {
    pxr::UsdGeomPrimvarsAPI api(mesh.GetPrim());
    pxr::TfToken n(std::string(name.data(), name.size()));
    auto pv = api.GetPrimvar(n);
    if (!pv) return nullptr;
    return std::make_unique<Primvar>(std::move(pv));
}

bool Mesh::create_primvar_float(
    rust::Str name,
    rust::Slice<const float> values,
    rust::Str interpolation) const
{
    pxr::UsdGeomPrimvarsAPI api(mesh.GetPrim());
    pxr::TfToken n(std::string(name.data(), name.size()));
    pxr::TfToken interp(std::string(interpolation.data(), interpolation.size()));
    auto pv = api.CreatePrimvar(n, pxr::SdfValueTypeNames->FloatArray, interp);
    if (!pv) return false;
    pxr::VtArray<float> arr(values.size());
    std::copy(values.begin(), values.end(), arr.begin());
    return pv.Set(arr, pxr::UsdTimeCode::Default());
}

bool Mesh::create_primvar_int(
    rust::Str name,
    rust::Slice<const int32_t> values,
    rust::Str interpolation) const
{
    pxr::UsdGeomPrimvarsAPI api(mesh.GetPrim());
    pxr::TfToken n(std::string(name.data(), name.size()));
    pxr::TfToken interp(std::string(interpolation.data(), interpolation.size()));
    auto pv = api.CreatePrimvar(n, pxr::SdfValueTypeNames->IntArray, interp);
    if (!pv) return false;
    pxr::VtArray<int> arr(values.size());
    std::copy(values.begin(), values.end(), arr.begin());
    return pv.Set(arr, pxr::UsdTimeCode::Default());
}

bool Mesh::create_primvar_vec2f(
    rust::Str name,
    rust::Slice<const float> values,
    rust::Str interpolation) const
{
    if (values.size() % 2 != 0) return false;
    pxr::UsdGeomPrimvarsAPI api(mesh.GetPrim());
    pxr::TfToken n(std::string(name.data(), name.size()));
    pxr::TfToken interp(std::string(interpolation.data(), interpolation.size()));
    auto pv = api.CreatePrimvar(n, pxr::SdfValueTypeNames->Float2Array, interp);
    if (!pv) return false;
    pxr::VtArray<pxr::GfVec2f> arr(values.size() / 2);
    for (size_t i = 0; i < arr.size(); ++i) {
        arr[i] = pxr::GfVec2f(values[i * 2], values[i * 2 + 1]);
    }
    return pv.Set(arr, pxr::UsdTimeCode::Default());
}

bool Mesh::create_primvar_vec3f(
    rust::Str name,
    rust::Slice<const float> values,
    rust::Str interpolation) const
{
    if (values.size() % 3 != 0) return false;
    pxr::UsdGeomPrimvarsAPI api(mesh.GetPrim());
    pxr::TfToken n(std::string(name.data(), name.size()));
    pxr::TfToken interp(std::string(interpolation.data(), interpolation.size()));
    auto pv = api.CreatePrimvar(n, pxr::SdfValueTypeNames->Float3Array, interp);
    if (!pv) return false;
    pxr::VtArray<pxr::GfVec3f> arr(values.size() / 3);
    for (size_t i = 0; i < arr.size(); ++i) {
        arr[i] = pxr::GfVec3f(values[i * 3], values[i * 3 + 1], values[i * 3 + 2]);
    }
    return pv.Set(arr, pxr::UsdTimeCode::Default());
}

bool Mesh::create_primvar_color3f(
    rust::Str name,
    rust::Slice<const float> values,
    rust::Str interpolation) const
{
    if (values.size() % 3 != 0) return false;
    pxr::UsdGeomPrimvarsAPI api(mesh.GetPrim());
    pxr::TfToken n(std::string(name.data(), name.size()));
    pxr::TfToken interp(std::string(interpolation.data(), interpolation.size()));
    auto pv = api.CreatePrimvar(n, pxr::SdfValueTypeNames->Color3fArray, interp);
    if (!pv) return false;
    pxr::VtArray<pxr::GfVec3f> arr(values.size() / 3);
    for (size_t i = 0; i < arr.size(); ++i) {
        arr[i] = pxr::GfVec3f(values[i * 3], values[i * 3 + 1], values[i * 3 + 2]);
    }
    return pv.Set(arr, pxr::UsdTimeCode::Default());
}

bool Mesh::remove_primvar(rust::Str name) const {
    pxr::UsdGeomPrimvarsAPI api(mesh.GetPrim());
    pxr::TfToken n(std::string(name.data(), name.size()));
    return api.RemovePrimvar(n);
}

rust::String Primvar::name() const {
    return rust::String(primvar.GetPrimvarName().GetString());
}

rust::String Primvar::interpolation() const {
    return rust::String(primvar.GetInterpolation().GetString());
}

rust::String Primvar::type_name() const {
    return rust::String(primvar.GetAttr().GetTypeName().GetAsToken().GetString());
}

bool Primvar::has_authored_value() const {
    return primvar.GetAttr().HasAuthoredValue();
}

bool Primvar::is_indexed() const {
    return primvar.IsIndexed();
}

std::unique_ptr<std::vector<int32_t>> Primvar::indices() const {
    auto vec = std::make_unique<std::vector<int32_t>>();
    if (!primvar.IsIndexed()) return vec;
    pxr::VtArray<int> arr;
    if (primvar.GetIndices(&arr, pxr::UsdTimeCode::Default())) {
        vec->reserve(arr.size());
        for (int v : arr) vec->push_back(static_cast<int32_t>(v));
    }
    return vec;
}

std::unique_ptr<std::vector<float>> Primvar::as_float_array() const {
    auto vec = std::make_unique<std::vector<float>>();
    pxr::VtArray<float> arr;
    if (primvar.Get(&arr, pxr::UsdTimeCode::Default())) {
        vec->insert(vec->end(), arr.begin(), arr.end());
    }
    return vec;
}

std::unique_ptr<std::vector<int32_t>> Primvar::as_int_array() const {
    auto vec = std::make_unique<std::vector<int32_t>>();
    pxr::VtArray<int> arr;
    if (primvar.Get(&arr, pxr::UsdTimeCode::Default())) {
        vec->reserve(arr.size());
        for (int v : arr) vec->push_back(static_cast<int32_t>(v));
    }
    return vec;
}

std::unique_ptr<std::vector<float>> Primvar::as_vec2f_array() const {
    auto vec = std::make_unique<std::vector<float>>();
    pxr::VtArray<pxr::GfVec2f> arr;
    if (primvar.Get(&arr, pxr::UsdTimeCode::Default())) {
        flatten_vec2(arr, *vec);
    }
    return vec;
}

std::unique_ptr<std::vector<float>> Primvar::as_vec3f_array() const {
    auto vec = std::make_unique<std::vector<float>>();
    pxr::VtArray<pxr::GfVec3f> arr;
    if (primvar.Get(&arr, pxr::UsdTimeCode::Default())) {
        flatten_vec3(arr, *vec);
    }
    return vec;
}

bool Mesh::bind_material(const Material& material) const {
    auto api = pxr::UsdShadeMaterialBindingAPI::Apply(mesh.GetPrim());
    if (!api) return false;
    return api.Bind(material.material);
}

std::unique_ptr<Prim> Stage::define_prim(rust::Str sdf_path, rust::Str type_name) const {
    if (!stage) return nullptr;
    pxr::SdfPath path(std::string(sdf_path.data(), sdf_path.size()));
    if (!path.IsAbsolutePath()) return nullptr;
    pxr::UsdPrim prim = stage->DefinePrim(
        path,
        pxr::TfToken(std::string(type_name.data(), type_name.size())));
    if (!prim) return nullptr;
    auto wrapped = std::make_unique<Prim>();
    wrapped->prim = prim;
    return wrapped;
}

std::unique_ptr<Material> Stage::create_material(rust::Str sdf_path) const {
    if (!stage) return nullptr;
    pxr::SdfPath path(std::string(sdf_path.data(), sdf_path.size()));
    if (!path.IsAbsolutePath()) return nullptr;
    auto mat = pxr::UsdShadeMaterial::Define(stage, path);
    if (!mat) return nullptr;
    return std::make_unique<Material>(std::move(mat));
}

std::unique_ptr<Material> Stage::material_at_path(rust::Str sdf_path) const {
    if (!stage) return nullptr;
    pxr::SdfPath path(std::string(sdf_path.data(), sdf_path.size()));
    if (!path.IsAbsolutePath()) return nullptr;
    pxr::UsdPrim prim = stage->GetPrimAtPath(path);
    if (!prim || !prim.IsA<pxr::UsdShadeMaterial>()) return nullptr;
    return std::make_unique<Material>(pxr::UsdShadeMaterial(prim));
}

namespace {

pxr::SdfValueTypeName resolve_type_name(const std::string& name) {
    if (name == "float")    return pxr::SdfValueTypeNames->Float;
    if (name == "float2")   return pxr::SdfValueTypeNames->Float2;
    if (name == "float3")   return pxr::SdfValueTypeNames->Float3;
    if (name == "color3f")  return pxr::SdfValueTypeNames->Color3f;
    if (name == "asset")    return pxr::SdfValueTypeNames->Asset;
    if (name == "token")    return pxr::SdfValueTypeNames->Token;
    if (name == "int")      return pxr::SdfValueTypeNames->Int;
    if (name == "string")   return pxr::SdfValueTypeNames->String;
    return pxr::SdfValueTypeName();
}

}  // namespace

rust::String Material::path() const {
    return rust::String(material.GetPrim().GetPath().GetString());
}

std::unique_ptr<Shader> Material::create_shader(rust::Str name, rust::Str shader_id) const {
    pxr::UsdPrim mat_prim = material.GetPrim();
    if (!mat_prim) return nullptr;
    auto stage = mat_prim.GetStage();
    if (!stage) return nullptr;

    std::string n(name.data(), name.size());
    pxr::SdfPath shader_path = mat_prim.GetPath().AppendChild(pxr::TfToken(n));
    pxr::UsdShadeShader shader = pxr::UsdShadeShader::Define(stage, shader_path);
    if (!shader) return nullptr;
    shader.SetShaderId(pxr::TfToken(std::string(shader_id.data(), shader_id.size())));
    return std::make_unique<Shader>(std::move(shader));
}

bool Material::connect_surface(const Shader& source) const {
    pxr::UsdShadeOutput surface = material.CreateSurfaceOutput();
    if (!surface) return false;
    pxr::UsdShadeOutput src_out = source.shader.GetOutput(pxr::TfToken("surface"));
    if (!src_out) {
        // Auto-declare a token-typed surface output if the shader didn't yet.
        src_out = source.shader.CreateOutput(pxr::TfToken("surface"),
                                             pxr::SdfValueTypeNames->Token);
    }
    return surface.ConnectToSource(src_out);
}

rust::String Shader::path() const {
    return rust::String(shader.GetPrim().GetPath().GetString());
}

rust::String Shader::shader_id() const {
    pxr::TfToken id;
    shader.GetShaderId(&id);
    return rust::String(id.GetString());
}

bool Shader::set_input_float(rust::Str name, float value) const {
    auto input = shader.CreateInput(
        pxr::TfToken(std::string(name.data(), name.size())),
        pxr::SdfValueTypeNames->Float);
    if (!input) return false;
    return input.Set(value);
}

bool Shader::set_input_color3f(rust::Str name, float r, float g, float b) const {
    auto input = shader.CreateInput(
        pxr::TfToken(std::string(name.data(), name.size())),
        pxr::SdfValueTypeNames->Color3f);
    if (!input) return false;
    return input.Set(pxr::GfVec3f(r, g, b));
}

bool Shader::set_input_asset(rust::Str name, rust::Str asset_path) const {
    auto input = shader.CreateInput(
        pxr::TfToken(std::string(name.data(), name.size())),
        pxr::SdfValueTypeNames->Asset);
    if (!input) return false;
    return input.Set(pxr::SdfAssetPath(std::string(asset_path.data(), asset_path.size())));
}

bool Shader::set_input_token(rust::Str name, rust::Str value) const {
    auto input = shader.CreateInput(
        pxr::TfToken(std::string(name.data(), name.size())),
        pxr::SdfValueTypeNames->Token);
    if (!input) return false;
    return input.Set(pxr::TfToken(std::string(value.data(), value.size())));
}

bool Shader::connect_input(rust::Str input_name, const Shader& source, rust::Str output_name) const {
    pxr::UsdShadeInput input = shader.GetInput(
        pxr::TfToken(std::string(input_name.data(), input_name.size())));
    if (!input) {
        // Need a typed input to connect; default to Color3f for the common
        // diffuseColor case. Caller can declare with set_input_color3f first
        // for other types.
        input = shader.CreateInput(
            pxr::TfToken(std::string(input_name.data(), input_name.size())),
            pxr::SdfValueTypeNames->Color3f);
    }
    pxr::UsdShadeOutput src_out = source.shader.GetOutput(
        pxr::TfToken(std::string(output_name.data(), output_name.size())));
    if (!src_out) return false;
    return input.ConnectToSource(src_out);
}

bool Shader::declare_output(rust::Str output_name, rust::Str type_name) const {
    auto type = resolve_type_name(std::string(type_name.data(), type_name.size()));
    if (!type) return false;
    auto out = shader.CreateOutput(
        pxr::TfToken(std::string(output_name.data(), output_name.size())),
        type);
    return bool(out);
}

std::unique_ptr<Prim> copy_prim(const Prim& prim) {
    auto wrapped = std::make_unique<Prim>();
    wrapped->prim = prim.prim;
    return wrapped;
}

std::unique_ptr<Mesh> copy_mesh(const Mesh& mesh) {
    auto wrapped = std::make_unique<Mesh>();
    wrapped->mesh = mesh.mesh;
    return wrapped;
}

}  // namespace rust_usd