USDShader : Shader type namespacing and assignment improvements

Changes.md

@@ -5,6 +5,10 @@ Improvements
 ------------
 
 - Arnold : Added support for Int64Data and UInt64Data custom attributes, allowing USD's `instanceId` to be used as a custom attribute in the Instancer node. Warnings are emitted if values are out of range for Arnold's 32 bit ints.
+- USDShader :
+  - A namespace prefix is now set for the shader type from the USD source type, except for `USD` and `glslfx` built-ins which correspond to USDLux lights and USDPreviewSurface built-ins
+  - Added a way to register namespace overrides from USD shader source types if the resulting USD source type is undesirable
+  - The correct shader assignment for displacement and volume shaders are now set for shaders which have a context of `displacement` or `volume` set
 
 Fixes
 -----

include/GafferUSD/USDShader.h

@@ -56,6 +56,10 @@ class GAFFERUSD_API USDShader : public GafferScene::Shader
 
 		void loadShader( const std::string &shaderName, bool keepExistingValues = false ) override;
 
+		/// Set a namespace remap based on the sourceType of the registered USD shader
+		/// eg. `USD` remaps to no namespace or `arnold` into `ai`.
+		static bool registerShaderNameSpace( const IECore::InternedString sourceType, const IECore::InternedString nameSpace );
+
 	protected :
 
 		IECore::ConstCompoundObjectPtr attributes( const Gaffer::Plug *output ) const override;

python/GafferUSDTest/USDShaderTest.py

@@ -258,5 +258,59 @@ def testUsdPreviewSurfaceAssignment( self ) :
 		self.assertEqual( shaderAssignment["out"].attributes( "/sphere" ).keys(), [ "displacement" ] )
 		self.assertIsInstance( shaderAssignment["out"].attributes( "/sphere" )["displacement"], IECoreScene.ShaderNetwork )
 
+	def testMtlxSurfaceAssignment( self ) :
+
+		sphere = GafferScene.Sphere()
+
+		shader = GafferUSD.USDShader()
+		shader.loadShader( "ND_surface" )
+
+		sphereFilter = GafferScene.PathFilter()
+		sphereFilter["paths"].setValue( IECore.StringVectorData( [ "/sphere" ] ) )
+
+		shaderAssignment = GafferScene.ShaderAssignment()
+		shaderAssignment["in"].setInput( sphere["out"] )
+		shaderAssignment["filter"].setInput( sphereFilter["out"] )
+		shaderAssignment["shader"].setInput( shader["out"]["out"] )
+
+		self.assertEqual( shaderAssignment["out"].attributes( "/sphere" ).keys(), [ "mtlx:surface" ] )
+		self.assertIsInstance( shaderAssignment["out"].attributes( "/sphere" )["mtlx:surface"], IECoreScene.ShaderNetwork )
+
+	def testMtlxDisplacementAssignment( self ) :
+
+		sphere = GafferScene.Sphere()
+
+		shader = GafferUSD.USDShader()
+		shader.loadShader( "ND_displacement_float" )
+
+		sphereFilter = GafferScene.PathFilter()
+		sphereFilter["paths"].setValue( IECore.StringVectorData( [ "/sphere" ] ) )
+
+		shaderAssignment = GafferScene.ShaderAssignment()
+		shaderAssignment["in"].setInput( sphere["out"] )
+		shaderAssignment["filter"].setInput( sphereFilter["out"] )
+		shaderAssignment["shader"].setInput( shader["out"]["out"] )
+
+		self.assertEqual( shaderAssignment["out"].attributes( "/sphere" ).keys(), [ "mtlx:displacement" ] )
+		self.assertIsInstance( shaderAssignment["out"].attributes( "/sphere" )["mtlx:displacement"], IECoreScene.ShaderNetwork )
+
+	def testMtlxVolumeAssignment( self ) :
+
+		sphere = GafferScene.Sphere()
+
+		shader = GafferUSD.USDShader()
+		shader.loadShader( "ND_volume" )
+
+		sphereFilter = GafferScene.PathFilter()
+		sphereFilter["paths"].setValue( IECore.StringVectorData( [ "/sphere" ] ) )
+
+		shaderAssignment = GafferScene.ShaderAssignment()
+		shaderAssignment["in"].setInput( sphere["out"] )
+		shaderAssignment["filter"].setInput( sphereFilter["out"] )
+		shaderAssignment["shader"].setInput( shader["out"]["out"] )
+
+		self.assertEqual( shaderAssignment["out"].attributes( "/sphere" ).keys(), [ "mtlx:volume" ] )
+		self.assertIsInstance( shaderAssignment["out"].attributes( "/sphere" )["mtlx:volume"], IECoreScene.ShaderNetwork )
+
 if __name__ == "__main__":
 	unittest.main()

src/GafferUSD/USDShader.cpp

@@ -59,6 +59,7 @@
 #include "pxr/usd/usdLux/nonboundableLightBase.h"
 
 #include "boost/algorithm/string/predicate.hpp"
+#include "boost/container/flat_map.hpp"
 
 #include "fmt/format.h"
 
@@ -319,6 +320,61 @@ Plug *loadPrimDefinitionAttribute( const UsdPrimDefinition::Attribute &attribute
 
 const IECore::InternedString g_surface( "surface" );
 const IECore::InternedString g_displacement( "displacement" );
+const IECore::InternedString g_volume( "volume" );
+
+using ShaderRemap = boost::container::flat_map<IECore::InternedString, IECore::InternedString>;
+
+ShaderRemap &shaderNameSpace()
+{
+	static ShaderRemap g_shaderNameSpace;
+	return g_shaderNameSpace;
+}
+
+const std::string remapShaderType( const TfToken sourceType, const TfToken context )
+{
+	const IECore::InternedString stype( sourceType.GetString() );
+	const IECore::InternedString ctx( context.GetString() );
+	std::string nameSpace = stype.string();
+	std::string outputType = g_surface.string();
+
+	const ShaderRemap &shaderNS = shaderNameSpace();
+	ShaderRemap::const_iterator it = shaderNS.find( stype );
+	if( it != shaderNS.end() )
+	{
+		nameSpace = it->second.string();
+	}
+
+	if( ctx == g_displacement )
+	{
+		outputType = g_displacement.string();
+	}
+	else if( ctx == g_volume )
+	{
+		outputType = g_volume.string();
+	}
+
+	if( nameSpace.empty() )
+	{
+		return outputType;
+	}
+
+	return nameSpace + ":" + outputType;
+}
+
+// Ideally this one should translate to `gl` however USDPreviewSurface and friends
+// have their sourceType set to `glslfx` so we leave these without a namespace.
+const bool g_glslfxShaderNameSpaceRegistration = USDShader::registerShaderNameSpace( "glslfx", "" );
+// Prman tends to register their built-in OSL shaders as `OSL` sourceType, which might
+// need to be re-looked at if we don't want Prman built-in shaders to be used in
+// other OSL-compatible renderers or in other GafferOSL functionality.
+const bool g_oslShaderNameSpaceRegistration = USDShader::registerShaderNameSpace( "OSL", "osl" );
+
+// TODO: Put these registers into their respective plugins maybe?
+
+// We shouldn't use Arnold shaders as USD shaders anyways as they're already available in Gaffer from Arnold itself.
+const bool g_arnoldShaderNameSpaceRegistration = USDShader::registerShaderNameSpace( "arnold", "ai" );
+const bool g_prmanShaderNameSpaceRegistration = USDShader::registerShaderNameSpace( "RmanCpp", "ri" );
+
 
 } // namespace
 
@@ -386,7 +442,14 @@ void USDShader::loadShader( const std::string &shaderName, bool keepExistingValu
 	// Set name and type and delete old parameters if necessary.
 
 	namePlug()->setValue( shaderName );
-	typePlug()->setValue( "surface" );
+	if( shader )
+	{
+		typePlug()->setValue( remapShaderType( shader->GetSourceType(), shader->GetContext() ) );
+	}
+	else
+	{
+		typePlug()->setValue( "surface" );
+	}
 
 	Plug *parametersPlug = this->parametersPlug()->source();
 	Plug *outPlug = this->outPlug();
@@ -474,3 +537,9 @@ IECore::ConstCompoundObjectPtr USDShader::attributes( const Gaffer::Plug *output
 	}
 	return result;
 }
+
+bool USDShader::registerShaderNameSpace( const IECore::InternedString sourceType, const IECore::InternedString nameSpace )
+{
+	ShaderRemap &shaderNS = shaderNameSpace();
+	return shaderNS.insert( { sourceType, nameSpace } ).second;
+}