From ad1606d6aae284898904ceb0a8c0f4f4249295b4 Mon Sep 17 00:00:00 2001
From: Mark Tucker <mtucker@sidefx.com>
Date: Fri, 24 Jan 2025 11:44:01 -0500
Subject: [PATCH] Implement a specialization of
 UsdImagingDataSourceAttribute<SdfAssetPath> that properly resolves UDIM
 relative paths.

---
 .../usdImaging/dataSourceAttribute.cpp        | 79 +++++++++++++++++++
 .../usdImaging/dataSourceAttribute.h          |  6 +-
 2 files changed, 83 insertions(+), 2 deletions(-)

diff --git a/pxr/usdImaging/usdImaging/dataSourceAttribute.cpp b/pxr/usdImaging/usdImaging/dataSourceAttribute.cpp
index 20b617f298..2bb97091ad 100644
--- a/pxr/usdImaging/usdImaging/dataSourceAttribute.cpp
+++ b/pxr/usdImaging/usdImaging/dataSourceAttribute.cpp
@@ -8,10 +8,78 @@
 
 #include "pxr/imaging/hd/dataSourceLocator.h"
 
+#include "pxr/usd/usdShade/udimUtils.h"
+
 PXR_NAMESPACE_OPEN_SCOPE
 
 namespace {
 
+// We need to find the first layer that changes the value
+// of the parameter so that we anchor relative paths to that.
+static
+SdfLayerHandle
+_FindLayerHandle(const UsdAttribute& attr, const UsdTimeCode& time) {
+    for (const auto& spec: attr.GetPropertyStack(time)) {
+        if (spec->HasDefaultValue() ||
+            spec->GetLayer()->GetNumTimeSamplesForPath(
+                spec->GetPath()) > 0) {
+            return spec->GetLayer();
+        }
+    }
+    return TfNullPtr;
+}
+
+class UsdImagingDataSourceAssetPathAttribute :
+    public UsdImagingDataSourceAttribute<SdfAssetPath>
+{
+public:
+    HD_DECLARE_DATASOURCE(UsdImagingDataSourceAssetPathAttribute);
+
+    /// Returns the extracted SdfAssetPath value of the attribute at
+    /// \p shutterOffset, with proper handling for UDIM paths.
+    SdfAssetPath
+    GetTypedValue(HdSampledDataSource::Time shutterOffset) override
+    {
+        // Zero-initialization for numerical types.
+        SdfAssetPath result{};
+        UsdTimeCode time = _stageGlobals.GetTime();
+        if (time.IsNumeric()) {
+            time = UsdTimeCode(time.GetValue() + shutterOffset);
+        }
+        _usdAttrQuery.Get<SdfAssetPath>(&result, time);
+        if (UsdShadeUdimUtils::IsUdimIdentifier(result.GetAssetPath())) {
+            const std::string resolvedPath =
+                UsdShadeUdimUtils::ResolveUdimPath(result.GetAssetPath(),
+                    _FindLayerHandle(_usdAttrQuery.GetAttribute(), time));
+            if (!resolvedPath.empty()) {
+                result = SdfAssetPath(result.GetAssetPath(), resolvedPath);
+            }
+        }
+        return result;
+    }
+
+protected:
+    UsdImagingDataSourceAssetPathAttribute(
+            const UsdAttribute &usdAttr,
+            const UsdImagingDataSourceStageGlobals &stageGlobals,
+            const SdfPath &sceneIndexPath = SdfPath::EmptyPath(),
+            const HdDataSourceLocator &timeVaryingFlagLocator =
+                    HdDataSourceLocator::EmptyLocator())
+        : UsdImagingDataSourceAttribute<SdfAssetPath>(usdAttr,
+            stageGlobals, sceneIndexPath, timeVaryingFlagLocator)
+    { }
+
+    UsdImagingDataSourceAssetPathAttribute(
+            const UsdAttributeQuery &usdAttrQuery,
+            const UsdImagingDataSourceStageGlobals &stageGlobals,
+            const SdfPath &sceneIndexPath = SdfPath::EmptyPath(),
+            const HdDataSourceLocator &timeVaryingFlagLocator =
+                    HdDataSourceLocator::EmptyLocator())
+        : UsdImagingDataSourceAttribute<SdfAssetPath>(usdAttrQuery,
+            stageGlobals, sceneIndexPath, timeVaryingFlagLocator)
+    { }
+};
+
 typedef HdSampledDataSourceHandle (*_DataSourceFactory)(
         const UsdAttributeQuery &usdAttrQuery,
         const UsdImagingDataSourceStageGlobals &stageGlobals,
@@ -32,6 +100,17 @@ HdSampledDataSourceHandle _FactoryImpl(
             usdAttrQuery, stageGlobals, sceneIndexPath, timeVaryingFlagLocator);
 }
 
+template <>
+HdSampledDataSourceHandle _FactoryImpl<SdfAssetPath>(
+        const UsdAttributeQuery &usdAttrQuery,
+        const UsdImagingDataSourceStageGlobals &stageGlobals,
+        const SdfPath &sceneIndexPath,
+        const HdDataSourceLocator &timeVaryingFlagLocator)
+{
+    return UsdImagingDataSourceAssetPathAttribute::New(
+            usdAttrQuery, stageGlobals, sceneIndexPath, timeVaryingFlagLocator);
+}
+
 static _FactoryMap _CreateFactoryMap()
 {
     _FactoryMap map;
diff --git a/pxr/usdImaging/usdImaging/dataSourceAttribute.h b/pxr/usdImaging/usdImaging/dataSourceAttribute.h
index 112ae89819..fa1052ab46 100644
--- a/pxr/usdImaging/usdImaging/dataSourceAttribute.h
+++ b/pxr/usdImaging/usdImaging/dataSourceAttribute.h
@@ -106,7 +106,10 @@ class UsdImagingDataSourceAttribute : public HdTypedSampledDataSource<T>
         return outSampleTimes->size() > 1;
     }
 
-private:
+protected:
+    /// Constructors and member data are protected instead of private
+    /// for use by a specialized subclass for accessing SdfAssetPath
+    /// attributes which need special handling for UDIMs.
 
     /// Constructs a new UsdImagingDataSourceAttribute for the given \p usdAttr
     ///
@@ -137,7 +140,6 @@ class UsdImagingDataSourceAttribute : public HdTypedSampledDataSource<T>
             const HdDataSourceLocator &timeVaryingFlagLocator =
                     HdDataSourceLocator::EmptyLocator());
 
-private:
     UsdAttributeQuery _usdAttrQuery;
     const UsdImagingDataSourceStageGlobals &_stageGlobals;
 };