USD ShaderAlgo : Support Arnold-USD's convention for array connections

contrib/IECoreUSD/src/IECoreUSD/ShaderAlgo.cpp

@@ -49,6 +49,8 @@
 #include "boost/algorithm/string/replace.hpp"
 #include "boost/pointer_cast.hpp"
 
+#include <regex>
+
 #if PXR_VERSION < 2102
 #define IsContainer IsNodeGraph
 #endif
@@ -98,6 +100,25 @@ void readAdditionalLightParameters( const pxr::UsdPrim &prim, IECore::CompoundDa
 #endif
 }
 
+const std::regex g_arrayIndexFromUSDRegex( ":i([0-9]+)$" );
+const std::string g_arrayIndexFromUSDFormat( "[$1]" );
+IECore::InternedString fromUSDParameterName( const pxr::TfToken &usdName )
+{
+	// USD doesn't support connections to array indices. So Arnold-USD emulates
+	// them using its own `parameter:i<N>`syntax - see https://github.com/Autodesk/arnold-usd/pull/381.
+	// We convert these to the regular `parameter[N]` syntax during loading.
+	return std::regex_replace( usdName.GetString(), g_arrayIndexFromUSDRegex, g_arrayIndexFromUSDFormat );
+}
+
+const std::regex g_arrayIndexFromCortexRegex( "\\[([0-9]+)\\]$" );
+const std::string g_arrayIndexFromCortexFormat( ":i$1" );
+pxr::TfToken toUSDParameterName( IECore::InternedString cortexName )
+{
+	return pxr::TfToken(
+		std::regex_replace( cortexName.string(), g_arrayIndexFromCortexRegex, g_arrayIndexFromCortexFormat )
+	);
+}
+
 IECoreScene::ShaderNetwork::Parameter readShaderNetworkWalk( const pxr::SdfPath &anchorPath, const pxr::UsdShadeOutput &output, IECoreScene::ShaderNetwork &shaderNetwork );
 
 IECore::InternedString readShaderNetworkWalk( const pxr::SdfPath &anchorPath, const pxr::UsdShadeConnectableAPI &usdShader, IECoreScene::ShaderNetwork &shaderNetwork )
@@ -146,7 +167,7 @@ IECore::InternedString readShaderNetworkWalk( const pxr::SdfPath &anchorPath, co
 					anchorPath, usdSource.GetOutput( usdSourceName ), shaderNetwork
 				);
 				connections.push_back( {
-					sourceHandle, { handle, IECore::InternedString( i.GetBaseName().GetString() ) }
+					sourceHandle, { handle, fromUSDParameterName( i.GetBaseName() ) }
 				} );
 			}
 			else
@@ -160,7 +181,7 @@ IECore::InternedString readShaderNetworkWalk( const pxr::SdfPath &anchorPath, co
 
 		if( IECore::DataPtr d = IECoreUSD::DataAlgo::fromUSD( pxr::UsdAttribute( valueAttribute ) ) )
 		{
-			parameters[ i.GetBaseName().GetString() ] = d;
+			parameters[fromUSDParameterName( i.GetBaseName() )] = d;
 		}
 	}
 
@@ -244,7 +265,7 @@ pxr::UsdShadeOutput IECoreUSD::ShaderAlgo::writeShaderNetwork( const IECoreScene
 		for( const auto &p : expandedParameters->readable() )
 		{
 			pxr::UsdShadeInput input = usdShader.CreateInput(
-				pxr::TfToken( p.first.string() ),
+				toUSDParameterName( p.first ),
 				DataAlgo::valueTypeName( p.second.get() )
 			);
 			input.Set( DataAlgo::toUSD( p.second.get() ) );
@@ -278,7 +299,7 @@ pxr::UsdShadeOutput IECoreUSD::ShaderAlgo::writeShaderNetwork( const IECoreScene
 			pxr::UsdShadeInput dest = usdShader.GetInput( pxr::TfToken( c.destination.name.string() ) );
 			if( ! dest.GetPrim().IsValid() )
 			{
-				dest = usdShader.CreateInput( pxr::TfToken( c.destination.name.string() ), pxr::SdfValueTypeNames->Token );
+				dest = usdShader.CreateInput( toUSDParameterName( c.destination.name ), pxr::SdfValueTypeNames->Token );
 			}
 
 			pxr::UsdShadeShader sourceUsdShader = pxr::UsdShadeShader::Get( shaderContainer.GetStage(), shaderContainer.GetPath().AppendChild( pxr::TfToken( pxr::TfMakeValidIdentifier( c.source.shader.string() ) ) ) );

contrib/IECoreUSD/test/IECoreUSD/USDSceneTest.py

@@ -3271,5 +3271,34 @@ def testPointInstancerPrimvars( self ) :
 		self.assertEqual( points["myColor"].interpolation, IECoreScene.PrimitiveVariable.Interpolation.Vertex )
 		self.assertEqual( points["myColor"].indices, None )
 
+	def testArnoldArrayInputs( self ) :
+
+		def assertExpectedArrayInputs( network ) :
+
+			inputs = network.inputConnections( "rampRGB" )
+			self.assertEqual( len( inputs ), 2 )
+			self.assertEqual( inputs[0], ( ( "noise", "out" ), ( "rampRGB", "color[0]" ) ) )
+			self.assertEqual( inputs[1], ( ( "flat", "out" ), ( "rampRGB", "color[1]" ) ) )
+
+		# Load original USD out of USD-Arnold.
+
+		scene = IECoreScene.SceneInterface.create(
+			os.path.join( os.path.dirname( __file__ ), "data", "arnoldArrayInputs.usda" ),
+			IECore.IndexedIO.OpenMode.Read
+		)
+		network = scene.child( "sphere" ).readAttribute( "ai:surface", 0 )
+
+		assertExpectedArrayInputs( network )
+
+		# Write our own USD from that data, to check we can round-trip it.
+
+		fileName = os.path.join( self.temporaryDirectory(), "arnoldArrayInputsRewritten.usda" )
+		scene = IECoreScene.SceneInterface.create( fileName, IECore.IndexedIO.OpenMode.Write )
+		scene.createChild( "sphere" ).writeAttribute( "ai:surface", network, 0 )
+
+		del scene
+		scene = IECoreScene.SceneInterface.create( fileName, IECore.IndexedIO.OpenMode.Read )
+		assertExpectedArrayInputs( scene.child( "sphere" ).readAttribute( "ai:surface", 0 ) )
+
 if __name__ == "__main__":
 	unittest.main()

contrib/IECoreUSD/test/IECoreUSD/data/arnoldArrayInputs.usda

@@ -0,0 +1,47 @@
+#usda 1.0
+
+def Sphere "sphere"
+{
+    rel material:binding = </rampSurface>
+}
+
+def Material "rampSurface"
+{
+    token outputs:arnold:surface.connect = </rampSurface/standardSurface.outputs:surface>
+
+    def Shader "standardSurface"
+    {
+        uniform token info:id = "arnold:standard_surface"
+        color3f inputs:base_color = (0.8, 0.8, 0.8)
+        prepend color3f inputs:base_color.connect = </rampSurface/rampRGB.outputs:out>
+        float inputs:indirect_specular = 0
+        token outputs:surface
+    }
+
+    def Shader "rampRGB"
+    {
+        uniform token info:id = "arnold:ramp_rgb"
+        color3f[] inputs:color = [(0, 0, 0), (0.5, 0.5, 0.5)]
+        prepend color3f inputs:color:i0.connect = </rampSurface/noise.outputs:out>
+        prepend color3f inputs:color:i1.connect = </rampSurface/flat.outputs:out>
+        int[] inputs:interpolation = [1, 1]
+        float[] inputs:position = [0, 1]
+        token inputs:type = "v"
+        color3f outputs:out
+    }
+
+    def Shader "noise"
+    {
+        uniform token info:id = "arnold:noise"
+        color3f inputs:color1 = (1, 0, 0)
+        color3f inputs:color2 = (0, 1, 0)
+        color3f outputs:out
+    }
+
+    def Shader "flat"
+    {
+        uniform token info:id = "arnold:flat"
+        color3f inputs:color = (0, 0, 1)
+        color3f outputs:out
+    }
+}