Shader source mapping for node functions (#137)

MaterialGraphProject/Assets/UnityShaderEditor/Editor/Data/Graphs/MaterialGraphAsset.cs

 using System;
 using UnityEngine;
-#if UNITY_EDITOR
-
-#endif
 
 namespace UnityEditor.ShaderGraph
 {

MaterialGraphProject/Assets/UnityShaderEditor/Editor/Data/LightweightPipeline/LightWeightPBRSubShader.cs

 
         private static string GetShaderPassFromTemplate(string template, PBRMasterNode masterNode, Pass pass, GenerationMode mode, SurfaceMaterialOptions materialOptions)
         {
+            var builder = new ShaderStringBuilder();
+            builder.IncreaseIndent();
+            builder.IncreaseIndent();
-            var functionRegistry = new FunctionRegistry(2);
+            var functionRegistry = new FunctionRegistry(builder);
             var surfaceInputs = new ShaderGenerator();
 
             var shaderProperties = new PropertyCollector();
                 usedSlots);
 
             var graph = new ShaderGenerator();
-            graph.AddShaderChunk(functionRegistry.ToString(), false);
+            graph.AddShaderChunk(builder.ToString(), false);
             graph.AddShaderChunk(vertexInputs.GetShaderString(2), false);
             graph.AddShaderChunk(surfaceInputs.GetShaderString(2), false);
             graph.AddShaderChunk(surfaceDescriptionStruct.GetShaderString(2), false);

MaterialGraphProject/Assets/UnityShaderEditor/Editor/Data/LightweightPipeline/LightWeightUnlitSubShader.cs

 
         private static string GetShaderPassFromTemplate(string template, UnlitMasterNode masterNode, Pass pass, GenerationMode mode)
         {
+            var builder = new ShaderStringBuilder();
+            builder.IncreaseIndent();
+            builder.IncreaseIndent();
-            var functionRegistry = new FunctionRegistry(2);
+            var functionRegistry = new FunctionRegistry(builder);
             var surfaceInputs = new ShaderGenerator();
 
             var shaderProperties = new PropertyCollector();
                 usedSlots);
 
             var graph = new ShaderGenerator();
-            graph.AddShaderChunk(functionRegistry.ToString(), false);
+            graph.AddShaderChunk(builder.ToString(), false);
             graph.AddShaderChunk(vertexInputs.GetShaderString(2), false);
             graph.AddShaderChunk(surfaceInputs.GetShaderString(2), false);
             graph.AddShaderChunk(surfaceDescriptionStruct.GetShaderString(2), false);

MaterialGraphProject/Assets/UnityShaderEditor/Editor/Data/Nodes/CodeFunctionNode.cs

             registry.ProvideFunction(GetFunctionName(), s =>
             {
                 s.AppendLine(GetFunctionHeader());
-                s.AppendLines(GetFunctionBody(GetFunctionToConvert()).Trim('\r', '\n', '\t', ' '));
+                var functionBody = GetFunctionBody(GetFunctionToConvert());
+                var lines = functionBody.Trim('\r', '\n', '\t', ' ');
+                s.AppendLines(lines);
             });
         }
 

MaterialGraphProject/Assets/UnityShaderEditor/Editor/Data/Util/FunctionRegistry.cs

 {
     public class FunctionRegistry
     {
-        Dictionary<string, string> m_Functions = new Dictionary<string, string>();
-        ShaderStringBuilder m_StringBuilder = new ShaderStringBuilder();
-        bool m_ValidationEnabled = false;
+        Dictionary<string, string> m_Sources = new Dictionary<string, string>();
+        bool m_Validate = false;
+        ShaderStringBuilder m_Builder;
+
+        public FunctionRegistry(ShaderStringBuilder builder)
+        {
+            m_Builder = builder;
+        }
-        public FunctionRegistry(int indentLevel = 0)
+        internal ShaderStringBuilder builder
-            for (var i = 0; i < indentLevel; i++)
-                m_StringBuilder.IncreaseIndent();
+            get { return m_Builder; }
-        public bool ProvideFunction(string name, Action<ShaderStringBuilder> generator)
+        public void ProvideFunction(string name, Action<ShaderStringBuilder> generator)
-            string functionSource = string.Empty;
-            if (m_ValidationEnabled)
+            string existingSource;
+            if (m_Sources.TryGetValue(name, out existingSource))
-                var ssb = new ShaderStringBuilder();
-                generator(ssb);
-                functionSource = ssb.ToString();
+                if (m_Validate)
+                {
+                    var startIndex = builder.length;
+                    generator(builder);
+                    var length = builder.length - startIndex;
+                    var source = builder.ToString(startIndex, length);
+                    builder.length -= length;
+                    if (source != existingSource)
+                        Debug.LogErrorFormat(@"Function `{0}` has varying implementations:{1}{1}{2}{1}{1}{3}", name, Environment.NewLine, source, existingSource);
+                }
+                return;
-
-            string existingFunctionSource;
-            if (m_Functions.TryGetValue(name, out existingFunctionSource))
-                if (m_ValidationEnabled && functionSource != existingFunctionSource)
-                    Debug.LogErrorFormat(@"Function `{0}` has varying implementations:{1}{1}{2}{1}{1}{3}", name, Environment.NewLine, functionSource, existingFunctionSource);
-                return false;
+                builder.AppendNewLine();
+                var startIndex = builder.length;
+                generator(builder);
+                var length = builder.length - startIndex;
+                var source = m_Validate ? builder.ToString(startIndex, length) : string.Empty;
+                m_Sources.Add(name, source);
-            generator(m_StringBuilder);
-            m_Functions.Add(name, functionSource);
-            m_StringBuilder.AppendNewLine();
-            return true;
-        }
-
-        public override string ToString()
-        {
-            return m_StringBuilder.ToString();
         }
     }
 }

MaterialGraphProject/Assets/UnityShaderEditor/Editor/Data/Util/GraphUtil.cs

 using System.Linq;
 using UnityEditor.Graphing;
 using UnityEditor.Graphing.Util;
+using UnityEngine;
 
 namespace UnityEditor.ShaderGraph
 {
             outputList.Add(node);
         }
 
-        public static string GetShader(this AbstractMaterialGraph graph, AbstractMaterialNode node, GenerationMode mode, string name, out List<PropertyCollector.TextureInfo> configuredTextures, out PreviewMode previewMode, out FloatShaderProperty outputIdProperty, Dictionary<Guid, int> ids = null)
+        public static GenerationResults GetShader(this AbstractMaterialGraph graph, AbstractMaterialNode node, GenerationMode mode, string name)
+            var results = new GenerationResults();
             bool isUber = node == null;
 
             var vertexInputs = new ShaderGenerator();
-            var functionRegistry = new FunctionRegistry(2);
+            var functionBuilder = new ShaderStringBuilder();
+            var functionRegistry = new FunctionRegistry(functionBuilder);
             var surfaceInputs = new ShaderGenerator();
 
             surfaceInputs.AddShaderChunk("struct SurfaceInputs{", false);
             if (requirements.requiresScreenPosition)
                 surfaceInputs.AddShaderChunk(String.Format("float4 {0};", ShaderGeneratorNames.ScreenPosition), false);
 
-            previewMode = PreviewMode.Preview3D;
+            results.previewMode = PreviewMode.Preview3D;
             if (!isUber)
             {
                 foreach (var pNode in activeNodeList.OfType<AbstractMaterialNode>())
-                        previewMode = PreviewMode.Preview3D;
+                        results.previewMode = PreviewMode.Preview3D;
                         break;
                     }
                 }
             GenerateSurfaceDescriptionStruct(surfaceDescriptionStruct, slots, !isUber);
 
             var shaderProperties = new PropertyCollector();
-            outputIdProperty = new FloatShaderProperty
+            results.outputIdProperty = new FloatShaderProperty
             {
                 displayName = "OutputId",
                 generatePropertyBlock = false,
-                shaderProperties.AddShaderProperty(outputIdProperty);
+                shaderProperties.AddShaderProperty(results.outputIdProperty);
 
             GenerateSurfaceDescription(
                 activeNodeList,
                 shaderProperties,
                 requirements,
                 mode,
-                outputIdProperty: outputIdProperty,
-                ids: ids);
+                outputIdProperty: results.outputIdProperty,
+                ids: results.ids);
-            var finalShader = new ShaderGenerator();
-            finalShader.AddShaderChunk(String.Format(@"Shader ""{0}""", name), false);
-            finalShader.AddShaderChunk("{", false);
-            finalShader.Indent();
-
-            finalShader.AddShaderChunk("Properties", false);
-            finalShader.AddShaderChunk("{", false);
-            finalShader.Indent();
-            finalShader.AddShaderChunk(shaderProperties.GetPropertiesBlock(2), false);
-            finalShader.Deindent();
-            finalShader.AddShaderChunk("}", false);
-
-            finalShader.AddShaderChunk("CGINCLUDE", false);
-            finalShader.AddShaderChunk("#include \"UnityCG.cginc\"", false);
-            finalShader.AddShaderChunk(functionRegistry.ToString(), false);
-            finalShader.AddShaderChunk(vertexInputs.GetShaderString(2), false);
-            finalShader.AddShaderChunk(surfaceInputs.GetShaderString(2), false);
-            finalShader.AddShaderChunk(surfaceDescriptionStruct.GetShaderString(2), false);
-            finalShader.AddShaderChunk(shaderProperties.GetPropertiesDeclaration(2), false);
-            finalShader.AddShaderChunk(vertexShader.GetShaderString(2), false);
-            finalShader.AddShaderChunk(surfaceDescriptionFunction.GetShaderString(2), false);
-            finalShader.AddShaderChunk("ENDCG", false);
+            var finalBuilder = new ShaderStringBuilder();
+            finalBuilder.AppendLine(@"Shader ""{0}""", name);
+            using (finalBuilder.BlockScope())
+            {
+                finalBuilder.AppendLine("Properties");
+                using (finalBuilder.BlockScope())
+                {
+                    finalBuilder.AppendLines(shaderProperties.GetPropertiesBlock(0));
+                }
-            finalShader.AddShaderChunk(ShaderGenerator.GetPreviewSubShader(node, requirements), false);
+                finalBuilder.AppendLine(@"CGINCLUDE");
+                finalBuilder.AppendLine(@"#include ""UnityCG.cginc""");
+                finalBuilder.Concat(functionBuilder);
+                finalBuilder.AppendLines(vertexInputs.GetShaderString(0));
+                finalBuilder.AppendLines(surfaceInputs.GetShaderString(0));
+                finalBuilder.AppendLines(surfaceDescriptionStruct.GetShaderString(0));
+                finalBuilder.AppendLines(shaderProperties.GetPropertiesDeclaration(0));
+                finalBuilder.AppendLines(vertexShader.GetShaderString(0));
+                finalBuilder.AppendLines(surfaceDescriptionFunction.GetShaderString(0));
+                finalBuilder.AppendLine(@"ENDCG");
-            ListPool<INode>.Release(activeNodeList);
+                finalBuilder.AppendLines(ShaderGenerator.GetPreviewSubShader(node, requirements));
+                ListPool<INode>.Release(activeNodeList);
+            }
-            finalShader.Deindent();
-            finalShader.AddShaderChunk("}", false);
-            configuredTextures = shaderProperties.GetConfiguredTexutres();
-            return finalShader.GetShaderString(0);
+            results.configuredTextures = shaderProperties.GetConfiguredTexutres();
+            ShaderSourceMap sourceMap;
+            results.shader = finalBuilder.ToString(out sourceMap);
+            results.sourceMap = sourceMap;
+            return results;
         }
 
         internal static void GenerateSurfaceDescriptionStruct(ShaderGenerator surfaceDescriptionStruct, List<MaterialSlot> slots, bool isMaster)
             foreach (var activeNode in activeNodeList.OfType<AbstractMaterialNode>())
             {
                 if (activeNode is IGeneratesFunction)
+                {
+                    functionRegistry.builder.currentNode = activeNode;
+                }
                 if (activeNode is IGeneratesBodyCode)
                     (activeNode as IGeneratesBodyCode).GenerateNodeCode(surfaceDescriptionFunction, mode);
                 if (masterNode == null && activeNode.hasPreview)
 
                 activeNode.CollectShaderProperties(shaderProperties, mode);
             }
+            functionRegistry.builder.currentNode = null;
 
             if (masterNode != null)
             {
             surfaceDescriptionFunction.AddShaderChunk("}", false);
         }
 
-        public static string GetPreviewShader(this AbstractMaterialGraph graph, AbstractMaterialNode node, out PreviewMode previewMode)
+        public static GenerationResults GetPreviewShader(this AbstractMaterialGraph graph, AbstractMaterialNode node)
-            List<PropertyCollector.TextureInfo> configuredTextures;
-            FloatShaderProperty outputIdProperty;
-            return graph.GetShader(node, GenerationMode.Preview, String.Format("hidden/preview/{0}", node.GetVariableNameForNode()), out configuredTextures, out previewMode, out outputIdProperty);
+            return graph.GetShader(node, GenerationMode.Preview, String.Format("hidden/preview/{0}", node.GetVariableNameForNode()));
-        public static string GetUberPreviewShader(this AbstractMaterialGraph graph, Dictionary<Guid, int> ids, out FloatShaderProperty outputIdProperty)
+        public static GenerationResults GetUberPreviewShader(this AbstractMaterialGraph graph)
-            List<PropertyCollector.TextureInfo> configuredTextures;
-            PreviewMode previewMode;
-            return graph.GetShader(null, GenerationMode.Preview, "hidden/preview", out configuredTextures, out previewMode, out outputIdProperty, ids);
+            return graph.GetShader(null, GenerationMode.Preview, "hidden/preview");
         }
 
         static Dictionary<SerializationHelper.TypeSerializationInfo, SerializationHelper.TypeSerializationInfo> s_LegacyTypeRemapping;

MaterialGraphProject/Assets/UnityShaderEditor/Editor/Data/Util/ShaderStringBuilder.cs

 using System;
 using System.Collections.Generic;
 using System.Text;
-using System.Text.RegularExpressions;
+using UnityEditor.Graphing;
+    struct ShaderStringMapping
+    {
+        public INode node { get; set; }
+        public int startIndex { get; set; }
+        public int count { get; set; }
+    }
+
     public class ShaderStringBuilder : IDisposable
     {
         enum ScopeType
         StringBuilder m_StringBuilder;
         Stack<ScopeType> m_ScopeStack;
         int m_IndentationLevel;
+        ShaderStringMapping m_CurrentMapping;
+        List<ShaderStringMapping> m_Mappings;
+
+        internal INode currentNode
+        {
+            get { return m_CurrentMapping.node; }
+            set
+            {
+                m_CurrentMapping.count = m_StringBuilder.Length - m_CurrentMapping.startIndex;
+                if (m_CurrentMapping.count > 0)
+                    m_Mappings.Add(m_CurrentMapping);
+                m_CurrentMapping.node = value;
+                m_CurrentMapping.startIndex = m_StringBuilder.Length;
+                m_CurrentMapping.count = 0;
+            }
+        }
+
+            m_Mappings = new List<ShaderStringMapping>();
+            m_CurrentMapping = new ShaderStringMapping();
         }
 
         public void AppendNewLine()
 
         public void AppendLine(string value)
         {
-            AppendIndentation();
-            m_StringBuilder.Append(value);
+            if (!string.IsNullOrEmpty(value))
+            {
+                AppendIndentation();
+                m_StringBuilder.Append(value);
+            }
             AppendNewLine();
         }
 
 
         public void AppendLines(string lines)
         {
-            foreach (var line in Regex.Split(lines, Environment.NewLine))
-                AppendLine(line);
+            if (string.IsNullOrEmpty(lines))
+                return;
+            var splitLines = lines.Split('\n');
+            var lineCount = splitLines.Length;
+            var lastLine = splitLines[lineCount - 1];
+            if (string.IsNullOrEmpty(lastLine) || lastLine == "\r")
+                lineCount--;
+            for (var i = 0; i < lineCount; i++)
+                AppendLine(splitLines[i].Trim('\r'));
         }
 
         public void Append(string value)
             }
         }
 
+        public void Concat(ShaderStringBuilder other)
+        {
+            // First re-add all the mappings from `other`, such that their mappings are transformed.
+            foreach (var mapping in other.m_Mappings)
+            {
+                currentNode = mapping.node;
+
+                // Use `AppendLines` to indent according to the current indentation.
+                AppendLines(other.ToString(mapping.startIndex, mapping.count));
+            }
+            currentNode = other.currentNode;
+            AppendLines(other.ToString(other.m_CurrentMapping.startIndex, other.length - other.m_CurrentMapping.startIndex));
+        }
+
+        }
+
+        public string ToString(out ShaderSourceMap sourceMap)
+        {
+            m_CurrentMapping.count = m_StringBuilder.Length - m_CurrentMapping.startIndex;
+            if (m_CurrentMapping.count > 0)
+                m_Mappings.Add(m_CurrentMapping);
+            var source = m_StringBuilder.ToString();
+            sourceMap = new ShaderSourceMap(source, m_Mappings);
+            m_Mappings.RemoveAt(m_Mappings.Count - 1);
+            m_CurrentMapping.count = 0;
+            return source;
+        }
+
+        public string ToString(int startIndex, int length)
+        {
+            return m_StringBuilder.ToString(startIndex, length);
+        }
+
+        internal void Clear()
+        {
+            m_StringBuilder.Length = 0;
+        }
+
+        internal int length
+        {
+            get { return m_StringBuilder.Length; }
+            set { m_StringBuilder.Length = value; }
         }
     }
 }

MaterialGraphProject/Assets/UnityShaderEditor/Editor/Drawing/PreviewManager.cs

 using System.Diagnostics;
 using System.IO;
 using System.Linq;
-using System.Runtime.InteropServices;
-using System.Text;
-using System.Threading;
-using UnityEditor.Graphing.Util;
-using UnityEditor.ShaderGraph;
 using UnityEngine.Rendering;
 using Debug = UnityEngine.Debug;
 using Object = UnityEngine.Object;
                     if (uberNodes.Count > 0)
                     {
                         m_UberShaderIds.Clear();
-                        m_UberShaderString = m_Graph.GetUberPreviewShader(m_UberShaderIds, out m_OutputIdProperty);
+                        var results = m_Graph.GetUberPreviewShader();
+                        m_UberShaderString = results.shader;
+                        m_OutputIdProperty = results.outputIdProperty;
+                        m_UberShaderIds = results.ids;
                         ShaderUtil.UpdateShaderAsset(m_UberShader, m_UberShaderString);
                         File.WriteAllText(Application.dataPath + "/../UberShader.shader", (m_UberShaderString ?? "null").Replace("UnityEngine.MaterialGraph", "Generated"));
                         bool uberShaderHasError = false;
                             message.AppendLine(@"Preview shader for graph has {0} error{1}:", errors.Length, errors.Length != 1 ? "s" : "");
                             foreach (var error in errors)
                             {
-                                message.AppendLine("{0} at line {1} (on {2})", error.message, error.line, error.platform);
+                                INode node = null;
+                                try
+                                {
+                                    node = results.sourceMap.FindNode(error.line);
+                                }
+                                catch (Exception)
+                                {
+                                    Debug.LogWarning("ERROR");
+                                    continue;
+                                }
+                                message.AppendLine("{0} in {3} at line {1} (on {2})", error.message, error.line, error.platform, node != null ? string.Format("node {0} ({1})", node.name, node.guid) : "graph");
+                                message.AppendNewLine();
                             }
                             Debug.LogWarning(message.ToString());
                             ShaderUtil.ClearShaderErrors(m_UberShader);
             }
             else
             {
-                PreviewMode mode;
-                if (node is IMasterNode)
+                var masterNode = node as IMasterNode;
+                if (masterNode != null)
-                    shaderData.shaderString = ((IMasterNode)node).GetShader(GenerationMode.Preview, node.name, out configuredTextures);
+                    shaderData.shaderString = masterNode.GetShader(GenerationMode.Preview, node.name, out configuredTextures);
-                    shaderData.shaderString = m_Graph.GetPreviewShader(node, out mode);
+                    shaderData.shaderString = m_Graph.GetPreviewShader(node).shader;
             }
 
             File.WriteAllText(Application.dataPath + "/../GeneratedShader.shader", (shaderData.shaderString ?? "null").Replace("UnityEngine.MaterialGraph", "Generated"));

MaterialGraphProject/Assets/UnityShaderEditor/Editor/Drawing/Views/MaterialNodeView.cs

             }
             else
             {
-                PreviewMode previewMode;
-                FloatShaderProperty outputIdProperty;
-                var shader = graph.GetShader(node, GenerationMode.ForReals, node.name, out textureInfo, out previewMode, out outputIdProperty);
-                GUIUtility.systemCopyBuffer = shader;
+                GUIUtility.systemCopyBuffer = graph.GetShader(node, GenerationMode.ForReals, node.name).shader;
             }
         }
 

MaterialGraphProject/Assets/UnityShaderEditor/Editor/Data/Util/GenerationResults.cs

+using System;
+using System.Collections.Generic;
+
+namespace UnityEditor.ShaderGraph
+{
+    public class GenerationResults
+    {
+        public string shader { get; set; }
+        public List<PropertyCollector.TextureInfo> configuredTextures;
+        public PreviewMode previewMode { get; set; }
+        public FloatShaderProperty outputIdProperty { get; set; }
+        public Dictionary<Guid, int> ids { get; set; }
+        public ShaderSourceMap sourceMap { get; set; }
+
+        public GenerationResults()
+        {
+            configuredTextures = new List<PropertyCollector.TextureInfo>();
+            ids = new Dictionary<Guid, int>();
+        }
+    }
+}

MaterialGraphProject/Assets/UnityShaderEditor/Editor/Data/Util/GenerationResults.cs.meta

+fileFormatVersion: 2
+guid: c922293e622b43f4a1cd25a1246a17dc
+timeCreated: 1513866299

MaterialGraphProject/Assets/UnityShaderEditor/Editor/Data/Util/ShaderSourceMap.cs

+using System;
+using System.Collections.Generic;
+using UnityEditor.Graphing;
+
+namespace UnityEditor.ShaderGraph
+{
+    public class ShaderSourceMap
+    {
+        // Indicates where a new node begins
+        List<int> m_LineStarts;
+        List<INode> m_Nodes;
+        int m_LineCount;
+
+        internal ShaderSourceMap(string source, List<ShaderStringMapping> mappings)
+        {
+            m_LineStarts = new List<int>();
+            m_Nodes = new List<INode>();
+
+            var line = 0;
+            var currentIndex = 0;
+            foreach (var mapping in mappings)
+            {
+                var stopIndex = mapping.startIndex + mapping.count;
+                if (currentIndex >= stopIndex)
+                    continue;
+
+                m_LineStarts.Add(line);
+                m_Nodes.Add(mapping.node);
+
+                while (currentIndex < stopIndex && currentIndex != -1)
+                {
+                    currentIndex = source.IndexOf('\n', currentIndex+1);
+                    line++;
+                }
+
+                if (currentIndex == -1)
+                    break;
+            }
+
+            m_LineCount = line;
+        }
+
+        public INode FindNode(int line)
+        {
+            if (line >= m_LineCount || line < 0)
+                return null;
+            var l = 0;
+            var r = m_LineStarts.Count - 1;
+            while (l <= r)
+            {
+                var m = (l + r) / 2;
+                var lineStart = m_LineStarts[m];
+                var lineStop = m == m_LineStarts.Count ? m_LineCount : m_LineStarts[m + 1];
+                if (line >= lineStop)
+                    l = m + 1;
+                else if (line < lineStart)
+                    r = m - 1;
+                else
+                    return m_Nodes[m];
+            }
+            throw new Exception("Something went wrong in binary search");
+        }
+    }
+}

MaterialGraphProject/Assets/UnityShaderEditor/Editor/Data/Util/ShaderSourceMap.cs.meta

+fileFormatVersion: 2
+guid: f6c222989d9946e8a85c01ba541f4b41
+timeCreated: 1513867574