Unity Custom PBR shader node

CustomPBR.cs

// Custom PBR Sub Shader

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using UnityEditor.Graphing;
using UnityEditor.ShaderGraph;

namespace UnityEditor.Experimental.Rendering.LightweightPipeline
{
    [Serializable]
    [FormerName("UnityEditor.ShaderGraph.CustomPBR")]
    public class CustomPBR : IPBRSubShader
    {
        struct Pass
        {
            public string Name;
            public List<int> VertexShaderSlots;
            public List<int> PixelShaderSlots;
        }

        Pass m_ForwardPassMetallic = new Pass
        {
            Name = "LightweightForward",
            PixelShaderSlots = new List<int>
            {
                PBRMasterNode.AlbedoSlotId,
                PBRMasterNode.NormalSlotId,
                PBRMasterNode.EmissionSlotId,
                PBRMasterNode.MetallicSlotId,
                PBRMasterNode.SmoothnessSlotId,
                PBRMasterNode.OcclusionSlotId,
                PBRMasterNode.AlphaSlotId,
                PBRMasterNode.AlphaThresholdSlotId
            }
        };

        Pass m_ForwardPassSpecular = new Pass()
        {
            Name = "LightweightForward",
            PixelShaderSlots = new List<int>()
            {
                PBRMasterNode.AlbedoSlotId,
                PBRMasterNode.NormalSlotId,
                PBRMasterNode.EmissionSlotId,
                PBRMasterNode.SpecularSlotId,
                PBRMasterNode.SmoothnessSlotId,
                PBRMasterNode.OcclusionSlotId,
                PBRMasterNode.AlphaSlotId,
                PBRMasterNode.AlphaThresholdSlotId
            }
        };

        private static string GetShaderPassFromTemplate(string template, PBRMasterNode masterNode, Pass pass, GenerationMode mode, SurfaceMaterialOptions materialOptions)
        {
            var builder = new ShaderStringBuilder();
            builder.IncreaseIndent();
            builder.IncreaseIndent();
            var vertexInputs = new ShaderGenerator();
            var surfaceVertexShader = new ShaderGenerator();
            var surfaceDescriptionFunction = new ShaderGenerator();
            var surfaceDescriptionStruct = new ShaderGenerator();
            var functionRegistry = new FunctionRegistry(builder);
            var surfaceInputs = new ShaderGenerator();

            var shaderProperties = new PropertyCollector();

            surfaceInputs.AddShaderChunk("struct SurfaceInputs{", false);
            surfaceInputs.Indent();

            var activeNodeList = ListPool<INode>.Get();
            NodeUtils.DepthFirstCollectNodesFromNode(activeNodeList, masterNode, NodeUtils.IncludeSelf.Include, pass.PixelShaderSlots);

            var requirements = ShaderGraphRequirements.FromNodes(activeNodeList);

            var modelRequiements = ShaderGraphRequirements.none;
            modelRequiements.requiresNormal |= NeededCoordinateSpace.World;
            modelRequiements.requiresTangent |= NeededCoordinateSpace.World;
            modelRequiements.requiresBitangent |= NeededCoordinateSpace.World;
            modelRequiements.requiresPosition |= NeededCoordinateSpace.World;
            modelRequiements.requiresViewDir |= NeededCoordinateSpace.World;
            modelRequiements.requiresMeshUVs.Add(UVChannel.UV1);

            GraphUtil.GenerateApplicationVertexInputs(requirements.Union(modelRequiements), vertexInputs);
            ShaderGenerator.GenerateSpaceTranslationSurfaceInputs(requirements.requiresNormal, InterpolatorType.Normal, surfaceInputs);
            ShaderGenerator.GenerateSpaceTranslationSurfaceInputs(requirements.requiresTangent, InterpolatorType.Tangent, surfaceInputs);
            ShaderGenerator.GenerateSpaceTranslationSurfaceInputs(requirements.requiresBitangent, InterpolatorType.BiTangent, surfaceInputs);
            ShaderGenerator.GenerateSpaceTranslationSurfaceInputs(requirements.requiresViewDir, InterpolatorType.ViewDirection, surfaceInputs);
            ShaderGenerator.GenerateSpaceTranslationSurfaceInputs(requirements.requiresPosition, InterpolatorType.Position, surfaceInputs);

            if (requirements.requiresVertexColor)
                surfaceInputs.AddShaderChunk(string.Format("float4 {0};", ShaderGeneratorNames.VertexColor), false);

            if (requirements.requiresScreenPosition)
                surfaceInputs.AddShaderChunk(string.Format("float4 {0};", ShaderGeneratorNames.ScreenPosition), false);

            foreach (var channel in requirements.requiresMeshUVs.Distinct())
                surfaceInputs.AddShaderChunk(string.Format("half4 {0};", channel.GetUVName()), false);

            surfaceInputs.Deindent();
            surfaceInputs.AddShaderChunk("};", false);

            surfaceVertexShader.AddShaderChunk("GraphVertexInput PopulateVertexData(GraphVertexInput v){", false);
            surfaceVertexShader.Indent();
            surfaceVertexShader.AddShaderChunk("return v;", false);
            surfaceVertexShader.Deindent();
            surfaceVertexShader.AddShaderChunk("}", false);

            var slots = new List<MaterialSlot>();
            foreach (var id in pass.PixelShaderSlots)
                slots.Add(masterNode.FindSlot<MaterialSlot>(id));
            GraphUtil.GenerateSurfaceDescriptionStruct(surfaceDescriptionStruct, slots, true);

            var usedSlots = new List<MaterialSlot>();
            foreach (var id in pass.PixelShaderSlots)
                usedSlots.Add(masterNode.FindSlot<MaterialSlot>(id));

            GraphUtil.GenerateSurfaceDescription(
                activeNodeList,
                masterNode,
                masterNode.owner as AbstractMaterialGraph,
                surfaceDescriptionFunction,
                functionRegistry,
                shaderProperties,
                requirements,
                mode,
                "PopulateSurfaceData",
                "SurfaceDescription",
                null,
                usedSlots);

            var graph = new ShaderGenerator();
            graph.AddShaderChunk(shaderProperties.GetPropertiesDeclaration(2), false);
            graph.AddShaderChunk(surfaceInputs.GetShaderString(2), false);
            graph.AddShaderChunk(builder.ToString(), false);
            graph.AddShaderChunk(vertexInputs.GetShaderString(2), false);
            graph.AddShaderChunk(surfaceDescriptionStruct.GetShaderString(2), false);
            graph.AddShaderChunk(surfaceVertexShader.GetShaderString(2), false);
            graph.AddShaderChunk(surfaceDescriptionFunction.GetShaderString(2), false);

            var blendingVisitor = new ShaderGenerator();
            var cullingVisitor = new ShaderGenerator();
            var zTestVisitor = new ShaderGenerator();
            var zWriteVisitor = new ShaderGenerator();

            materialOptions.GetBlend(blendingVisitor);
            materialOptions.GetCull(cullingVisitor);
            materialOptions.GetDepthTest(zTestVisitor);
            materialOptions.GetDepthWrite(zWriteVisitor);

            var interpolators = new ShaderGenerator();
            var localVertexShader = new ShaderGenerator();
            var localPixelShader = new ShaderGenerator();
            var localSurfaceInputs = new ShaderGenerator();
            var surfaceOutputRemap = new ShaderGenerator();

            ShaderGenerator.GenerateStandardTransforms(
                3,
                10,
                interpolators,
                localVertexShader,
                localPixelShader,
                localSurfaceInputs,
                requirements,
                modelRequiements,
                CoordinateSpace.World);

            ShaderGenerator defines = new ShaderGenerator();

            if (masterNode.IsSlotConnected(PBRMasterNode.NormalSlotId))
                defines.AddShaderChunk("#define _NORMALMAP 1", true);

            if (masterNode.model == PBRMasterNode.Model.Specular)
                defines.AddShaderChunk("#define _SPECULAR_SETUP 1", true);

            if (masterNode.IsSlotConnected(PBRMasterNode.AlphaThresholdSlotId))
                defines.AddShaderChunk("#define _AlphaClip 1", true);

            if (masterNode.surfaceType == SurfaceType.Transparent && masterNode.alphaMode == AlphaMode.Premultiply)
                defines.AddShaderChunk("#define _ALPHAPREMULTIPLY_ON 1", true);

            var templateLocation = GetTemplatePath(template);

            foreach (var slot in usedSlots)
            {
                surfaceOutputRemap.AddShaderChunk(string.Format("{0} = surf.{0};", slot.shaderOutputName), true);
            }

            if (!File.Exists(templateLocation))
                return string.Empty;

            var subShaderTemplate = File.ReadAllText(templateLocation);
            var resultPass = subShaderTemplate.Replace("${Defines}", defines.GetShaderString(3));
            resultPass = resultPass.Replace("${Graph}", graph.GetShaderString(3));
            resultPass = resultPass.Replace("${Interpolators}", interpolators.GetShaderString(3));
            resultPass = resultPass.Replace("${VertexShader}", localVertexShader.GetShaderString(3));
            resultPass = resultPass.Replace("${LocalPixelShader}", localPixelShader.GetShaderString(3));
            resultPass = resultPass.Replace("${SurfaceInputs}", localSurfaceInputs.GetShaderString(3));
            resultPass = resultPass.Replace("${SurfaceOutputRemap}", surfaceOutputRemap.GetShaderString(3));

            resultPass = resultPass.Replace("${Tags}", string.Empty);
            resultPass = resultPass.Replace("${Blending}", blendingVisitor.GetShaderString(2));
            resultPass = resultPass.Replace("${Culling}", cullingVisitor.GetShaderString(2));
            resultPass = resultPass.Replace("${ZTest}", zTestVisitor.GetShaderString(2));
            resultPass = resultPass.Replace("${ZWrite}", zWriteVisitor.GetShaderString(2));
            return resultPass;
        }

        static string GetTemplatePath(string templateName)
        {
            var dirs = Directory.GetDirectories("./Assets/Scripts/Editor");
            foreach (var path in dirs)
            {
                var templatePath = Path.Combine(path, templateName);
                if (File.Exists(templatePath))
                    return templatePath;
            }
            throw new FileNotFoundException(string.Format(@"Cannot find a template with name ""{0}"".", templateName));
        }

        public string GetSubshader(IMasterNode inMasterNode, GenerationMode mode)
        {
            var masterNode = inMasterNode as PBRMasterNode;
            var subShader = new ShaderGenerator();
            subShader.AddShaderChunk("SubShader", true);
            subShader.AddShaderChunk("{", true);
            subShader.Indent();
            subShader.AddShaderChunk("Tags{ \"RenderPipeline\" = \"LightweightPipeline\"}", true);

            var materialOptions = ShaderGenerator.GetMaterialOptions(masterNode.surfaceType, masterNode.alphaMode, masterNode.twoSided.isOn);
            var tagsVisitor = new ShaderGenerator();
            materialOptions.GetTags(tagsVisitor);
            subShader.AddShaderChunk(tagsVisitor.GetShaderString(0), true);

            subShader.AddShaderChunk(
                GetShaderPassFromTemplate(
                    "lightweightPBRForwardPass.template",
                    masterNode,
                    masterNode.model == PBRMasterNode.Model.Metallic ? m_ForwardPassMetallic : m_ForwardPassSpecular,
                    mode,
                    materialOptions),
                true);

            var extraPassesTemplateLocation = GetTemplatePath("lightweightPBRExtraPasses.template");
            if (File.Exists(extraPassesTemplateLocation))
            {
                var extraPassesTemplate = File.ReadAllText(extraPassesTemplateLocation);
                extraPassesTemplate = extraPassesTemplate.Replace("${Culling}", materialOptions.cullMode.ToString());
                subShader.AddShaderChunk(extraPassesTemplate, true);
            }

            subShader.Deindent();
            subShader.AddShaderChunk("}", true);

            return subShader.GetShaderString(0);
        }
    }
}

// Editor: 
public class CreateCustomPBRShaderGraph : EndNameEditAction
{
    [MenuItem("Assets/Create/Shader/Custom PBR Graph", false, 208)]
    public static void CreateMaterialGraph()
    {
        ProjectWindowUtil.StartNameEditingIfProjectWindowExists(0, CreateInstance<CreateCustomPBRShaderGraph>(),
            "New Shader Graph.ShaderGraph", null, null);
    }

    public override void Action(int instanceId, string pathName, string resourceFile)
    {
        var graph = new MaterialGraph();
        graph.AddNode(new PBRMasterNode());
        File.WriteAllText(pathName, EditorJsonUtility.ToJson(graph));
        AssetDatabase.Refresh();
    }
}

ShaderTemplate.template

Pass
{
    Name "StandardLit"
	Tags{"LightMode" = "LightweightForward"}
	${Tags}
	${Blending}
	${Culling}
	${ZTest}
	${ZWrite}

	HLSLPROGRAM
    // Required to compile gles 2.0 with standard srp library
    #pragma prefer_hlslcc gles
	#pragma exclude_renderers d3d11_9x
	#pragma target 2.0

	// -------------------------------------
	// Lightweight Pipeline keywords
	#pragma multi_compile _ _ADDITIONAL_LIGHTS
	#pragma multi_compile _ _VERTEX_LIGHTS
	#pragma multi_compile _ _MIXED_LIGHTING_SUBTRACTIVE
	#pragma multi_compile _ _SHADOWS_ENABLED
    #pragma multi_compile _ _LOCAL_SHADOWS_ENABLED
    #pragma multi_compile _ _SHADOWS_SOFT

	// -------------------------------------
	// Unity defined keywords
	#pragma multi_compile _ DIRLIGHTMAP_COMBINED
	#pragma multi_compile _ LIGHTMAP_ON
	#pragma multi_compile_fog

	//--------------------------------------
	// GPU Instancing
	#pragma multi_compile_instancing

    #pragma vertex vert
	#pragma fragment frag

	${Defines}

	#include "LWRP/ShaderLibrary/Core.hlsl"
	#include "LWRP/ShaderLibrary/Lighting.hlsl"
	#include "CoreRP/ShaderLibrary/Color.hlsl"
	#include "CoreRP/ShaderLibrary/UnityInstancing.hlsl"
	#include "ShaderGraphLibrary/Functions.hlsl"

	${Graph}

	struct GraphVertexOutput
    {
        float4 clipPos                : SV_POSITION;
        DECLARE_LIGHTMAP_OR_SH(lightmapUV, vertexSH, 0);
		half4 fogFactorAndVertexLight : TEXCOORD1; // x: fogFactor, yzw: vertex light
    	float4 shadowCoord            : TEXCOORD2;
        ${Interpolators}
        UNITY_VERTEX_INPUT_INSTANCE_ID
    };

    GraphVertexOutput vert (GraphVertexInput v)
	{
	    v = PopulateVertexData(v);

        GraphVertexOutput o = (GraphVertexOutput)0;

        UNITY_SETUP_INSTANCE_ID(v);
    	UNITY_TRANSFER_INSTANCE_ID(v, o);

        ${VertexShader}

		float3 lwWNormal = TransformObjectToWorldNormal(v.normal);
		float3 lwWorldPos = TransformObjectToWorld(v.vertex.xyz);
		float4 clipPos = TransformWorldToHClip(lwWorldPos);

 		// We either sample GI from lightmap or SH.
	    // Lightmap UV and vertex SH coefficients use the same interpolator ("float2 lightmapUV" for lightmap or "half3 vertexSH" for SH)
        // see DECLARE_LIGHTMAP_OR_SH macro.
	    // The following funcions initialize the correct variable with correct data
	    OUTPUT_LIGHTMAP_UV(v.texcoord1, unity_LightmapST, o.lightmapUV);
	    OUTPUT_SH(lwWNormal, o.vertexSH);

	    half3 vertexLight = VertexLighting(lwWorldPos, lwWNormal);
	    half fogFactor = ComputeFogFactor(clipPos.z);
	    o.fogFactorAndVertexLight = half4(fogFactor, vertexLight);
	    o.clipPos = clipPos;

	    #ifdef _SHADOWS_ENABLED
		#if SHADOWS_SCREEN
			o.shadowCoord = ComputeShadowCoord(clipPos);
		#else
			o.shadowCoord = TransformWorldToShadowCoord(lwWorldPos);
		#endif
		#endif
		
		return o;
	}

	half3 CustomLightingPhysicallyBased(BRDFData brdfData, half3 lightColor, half3 lightDirectionWS, half lightAttenuation, half3 normalWS, half3 viewDirectionWS)
	{
		half3 halfDir = Unity_SafeNormalize(lightDirectionWS + viewDirectionWS); 
		half NdotL = smoothstep(0, 0.05, saturate(dot(normalWS, lightDirectionWS)));
		
		half3 radiance = lightColor * (lightAttenuation * NdotL);
		return DirectBDRF(brdfData, normalWS, lightDirectionWS, viewDirectionWS) * radiance;
	}

	half3 CustomLightingPhysicallyBased(BRDFData brdfData, Light light, half3 normalWS, half3 viewDirectionWS)
	{
		return CustomLightingPhysicallyBased(brdfData, light.color, light.direction, light.attenuation, normalWS, viewDirectionWS);
	}

	half4 CustomLightweightFragmentPBR(InputData inputData, half3 albedo, half metallic, half3 specular, half smoothness, half occlusion, half3 emission, half alpha)
	{
		BRDFData brdfData;
		InitializeBRDFData(albedo, metallic, specular, smoothness, alpha, brdfData);

		Light mainLight = GetMainLight();
		mainLight.attenuation = MainLightRealtimeShadowAttenuation(inputData.shadowCoord);

		MixRealtimeAndBakedGI(mainLight, inputData.normalWS, inputData.bakedGI, half4(0, 0, 0, 0));
		half3 color = GlobalIllumination(brdfData, inputData.bakedGI, occlusion, inputData.normalWS, inputData.viewDirectionWS);
		color += CustomLightingPhysicallyBased(brdfData, mainLight, inputData.normalWS, inputData.viewDirectionWS);

	#ifdef _ADDITIONAL_LIGHTS
		int pixelLightCount = GetPixelLightCount();
		for (int i = 0; i < pixelLightCount; ++i)
		{
			Light light = GetLight(i, inputData.positionWS);
			light.attenuation *= LocalLightRealtimeShadowAttenuation(light.index, inputData.positionWS);
			color += CustomLightingPhysicallyBased(brdfData, light, inputData.normalWS, inputData.viewDirectionWS);
		}
	#endif

		color += inputData.vertexLighting * brdfData.diffuse;
		color += emission;
		return half4(color, alpha);
	}

	half4 frag (GraphVertexOutput IN) : SV_Target
    {
    	UNITY_SETUP_INSTANCE_ID(IN);

    	${LocalPixelShader}

        SurfaceInputs surfaceInput = (SurfaceInputs)0;
        ${SurfaceInputs}

        SurfaceDescription surf = PopulateSurfaceData(surfaceInput);

		float3 Albedo = float3(0.5, 0.5, 0.5);
		float3 Specular = float3(0, 0, 0);
		float Metallic = 1;
		float3 Normal = float3(0, 0, 1);
		float3 Emission = 0;
		float Smoothness = 0.5;
		float Occlusion = 1;
		float Alpha = 1;
		float AlphaClipThreshold = 0;

        ${SurfaceOutputRemap}

		InputData inputData;
		inputData.positionWS = WorldSpacePosition;

#ifdef _NORMALMAP
	    inputData.normalWS = TangentToWorldNormal(Normal, WorldSpaceTangent, WorldSpaceBiTangent, WorldSpaceNormal);
#else
    #if !SHADER_HINT_NICE_QUALITY
        inputData.normalWS = WorldSpaceNormal;
    #else
	    inputData.normalWS = normalize(WorldSpaceNormal);
    #endif
#endif

#if !SHADER_HINT_NICE_QUALITY
	    // viewDirection should be normalized here, but we avoid doing it as it's close enough and we save some ALU.
	    inputData.viewDirectionWS = WorldSpaceViewDirection;
#else
	    inputData.viewDirectionWS = normalize(WorldSpaceViewDirection);
#endif

	    inputData.shadowCoord = IN.shadowCoord;

	    inputData.fogCoord = IN.fogFactorAndVertexLight.x;
	    inputData.vertexLighting = IN.fogFactorAndVertexLight.yzw;
	    inputData.bakedGI = SAMPLE_GI(IN.lightmapUV, IN.vertexSH, inputData.normalWS);

		half4 color = CustomLightweightFragmentPBR(
			inputData, 
			Albedo, 
			Metallic, 
			Specular, 
			Smoothness, 
			Occlusion, 
			Emission, 
			Alpha);

		// Computes fog factor per-vertex
    	ApplyFog(color.rgb, IN.fogFactorAndVertexLight.x);

#if _AlphaClip
		clip(Alpha - AlphaClipThreshold);
#endif
		return color;
    }

	ENDHLSL
}