A game about forced loneliness, made by TACStudios
1using System;
2using System.Collections.Generic;
3using System.IO;
4using System.Linq;
5using UnityEditor.Graphing;
6using UnityEditor.ShaderGraph.Internal;
7using UnityEngine;
8using Pool = UnityEngine.Pool;
9
10namespace UnityEditor.ShaderGraph
11{
12 internal static class CustomInterpolatorUtils
13 {
14 // We need to be able to adapt CustomInterpolatorNode's output if the feature couldn't work per pass.
15 // there isn't really a good way to get global information about generation state during a node's generation.
16 // TODO: If we jobify our generator calls, switch these to TLS.
17 internal static bool generatorSkipFlag = false;
18
19 // For node only generation, there isn't a breadcumb of information we get in a few of the AbstractMaterialNode
20 // code generating functions. Instead of completely refactoring for this case, we've got a global flag so that
21 // CustomInterpolatorNode can redirect NODE preview graph evaluation from it's CIB's input port directly--
22 // This is necessary because node previews don't interpolate anything. So to get any previews at all we need to reroute.
23 internal static bool generatorNodeOnly = false;
24
25 // Used by preview manager to find what custom interpolator nodes need rerouting for node previews.
26 internal static IEnumerable<CustomInterpolatorNode> GetCustomBlockNodeDependents(BlockNode bnode)
27 {
28 return bnode?.owner?.GetNodes<CustomInterpolatorNode>().Where(cin => cin.e_targetBlockNode == bnode).ToList()
29 ?? new List<CustomInterpolatorNode>();
30 }
31 }
32
33 internal class CustomInterpSubGen
34 {
35 #region descriptor
36
37 // Common splicing locations or concepts. These may or may not exist in client's template code.
38 [GenerationAPI]
39 internal static class Splice
40 {
41 internal static string k_splicePreInclude => "CustomInterpolatorPreInclude";
42 internal static string k_splicePrePacking => "CustomInterpolatorPrePacking";
43 internal static string k_splicePreSurface => "CustomInterpolatorPreSurface";
44 internal static string k_splicePreVertex => "CustomInterpolatorPreVertex";
45 internal static string k_spliceCopyToSDI => "CustomInterpolatorCopyToSDI";
46 }
47
48 // Describes where/what/how custom interpolator behavior can be achieved through splicing and defines.
49 // Generally speaking, this may require a mix of changes to client template and includes.
50 [GenerationAPI]
51 internal struct Descriptor
52 {
53 internal string src, dst; // for function or block. For macro block src is start of the macro and dst is end of the macro.
54 internal string name; // for struct or function.
55 internal string define; // defined for client code to indicate we're live.
56 internal string splice; // splice location, prefer use something from the list.
57 internal string preprocessor;
58 internal bool hasMacro;
59
60 internal bool isBlock => src != null && dst != null && name == null && splice != null && !hasMacro;
61 internal bool isMacroBlock => src != null && dst != null && name == null && splice != null && hasMacro;
62 internal bool isStruct => src == null && dst == null && name != null && splice != null;
63 internal bool isFunc => src != null && dst != null && name != null && splice != null;
64 internal bool isDefine => define != null && splice != null && src == null && dst == null & name == null;
65 internal bool isValid => isDefine || isBlock || isStruct || isFunc || isMacroBlock;
66 internal bool hasPreprocessor => !String.IsNullOrEmpty(preprocessor);
67
68 internal static Descriptor MakeFunc(string splice, string name, string dstType, string srcType, string define = "", string preprocessor = "") => new Descriptor { splice = splice, name = name, dst = dstType, src = srcType, define = define, preprocessor = preprocessor };
69 internal static Descriptor MakeStruct(string splice, string name, string define = "", string preprocessor = "") => new Descriptor { splice = splice, name = name, define = define, preprocessor = preprocessor };
70 internal static Descriptor MakeBlock(string splice, string dst, string src, string preprocessor = "") => new Descriptor { splice = splice, dst = dst, src = src, preprocessor = preprocessor };
71 internal static Descriptor MakeMacroBlock(string splice, string startMacro, string endMacro, string preprocessor = "") => new Descriptor { splice = splice, dst = endMacro, src = startMacro, preprocessor = preprocessor, hasMacro = true };
72 internal static Descriptor MakeDefine(string splice, string define, string preprocessor = "") => new Descriptor { splice = splice, define = define, preprocessor = preprocessor };
73 }
74
75 [GenerationAPI]
76 internal class Collection : IEnumerable<Collection.Item>
77 {
78 public class Item
79 {
80 public Descriptor descriptor { get; }
81 public Item(Descriptor descriptor) { this.descriptor = descriptor; }
82 }
83 readonly List<Collection.Item> m_Items;
84 public Collection() { m_Items = new List<Collection.Item>(); }
85 public Collection Add(Collection structs) { foreach (Collection.Item item in structs) m_Items.Add(item); return this; }
86 public Collection Add(Descriptor descriptor) { m_Items.Add(new Collection.Item(descriptor)); return this; }
87 public IEnumerator<Item> GetEnumerator() { return m_Items.GetEnumerator(); }
88 System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() { return GetEnumerator(); }
89 }
90 #endregion
91
92 private List<BlockNode> customBlockNodes;
93 private List<FieldDescriptor> customFieldDescriptors;
94 private bool isNodePreview;
95 private Dictionary<string, ShaderStringBuilder> spliceCommandBuffer;
96
97 internal CustomInterpSubGen(bool isNodePreview)
98 {
99 this.isNodePreview = isNodePreview;
100 customBlockNodes = new List<BlockNode>();
101 customFieldDescriptors = new List<FieldDescriptor>();
102 spliceCommandBuffer = new Dictionary<String, ShaderStringBuilder>();
103 }
104
105 #region GeneratorEntryPoints
106
107
108 // This entry point handles adding our upstream antecedents to the generator's list of active nodes.
109 // Custom Interpolator Nodes have no way of expressing that their Custom Interpolator Block is a dependent within existing generator code.
110 internal void ProcessExistingStackData(List<AbstractMaterialNode> vertexNodes, List<MaterialSlot> vertexSlots, List<AbstractMaterialNode> pixelNodes, IActiveFieldsSet activeFields)
111 {
112 if (CustomInterpolatorUtils.generatorSkipFlag)
113 return;
114
115 bool needsGraphFeature = false;
116
117 // departing from current generation code, we will select what to generate based on some graph analysis.
118 foreach (var cin in pixelNodes.OfType<CustomInterpolatorNode>().ToList())
119 {
120 // The CustomBlockNode's subtree.
121 var anties = GetAntecedents(cin.e_targetBlockNode);
122 // cin contains an inlined value, so there is nothing to do.
123 if (anties == null)
124 {
125 continue;
126 }
127 else if (isNodePreview)
128 {
129 foreach (var ant in anties)
130 {
131 // sorted insertion, based on dependencies already present in pixelNodes (an issue because we're faking for the preview).
132 if (!pixelNodes.Contains(ant))
133 InsertAntecedent(pixelNodes, ant);
134 }
135 }
136 else // it's a full compile and cin isn't inlined, so do all the things.
137 {
138 if (!customBlockNodes.Contains(cin.e_targetBlockNode))
139 {
140 activeFields.AddAll(cin.e_targetBlockNode.descriptor); // add the BlockFieldDescriptor for VertexDescription
141 customBlockNodes.Add(cin.e_targetBlockNode);
142 }
143
144 foreach (var ant in anties)
145 {
146 if (!vertexNodes.Contains(ant))
147 InsertAntecedent(vertexNodes, ant);
148 }
149
150 if (!vertexNodes.Contains(cin.e_targetBlockNode))
151 vertexNodes.Add(cin.e_targetBlockNode);
152 if (!vertexSlots.Contains(cin.e_targetBlockNode.FindSlot<MaterialSlot>(0)))
153 vertexSlots.Add(cin.e_targetBlockNode.FindSlot<MaterialSlot>(0));
154
155 needsGraphFeature = true;
156 }
157 }
158 // if a target has allowed custom interpolators, it should expect that the vertex feature can be forced on.
159 if (needsGraphFeature)
160 activeFields.AddAll(Fields.GraphVertex);
161 }
162
163 internal void AddCustomInterpolant(FieldDescriptor field)
164 {
165 customFieldDescriptors.Add(field);
166 }
167
168 // This entry point is to inject custom interpolator fields into the appropriate structs for struct generation.
169 internal List<StructDescriptor> CopyModifyExistingPassStructs(IEnumerable<StructDescriptor> passStructs, IActiveFieldsSet activeFields)
170 {
171 if (CustomInterpolatorUtils.generatorSkipFlag)
172 return passStructs.ToList();
173
174 var newPassStructs = new List<StructDescriptor>();
175
176 // StructDescriptor is (kind-of) immutable, so we need to do some copy/modify shenanigans to make this work.
177 foreach (var ps in passStructs)
178 {
179 if (ps.populateWithCustomInterpolators)
180 {
181 var agg = new List<FieldDescriptor>();
182 foreach (var cib in customBlockNodes)
183 {
184 var fd = new FieldDescriptor(ps.name, cib.customName, "", ShaderValueTypeFrom((int)cib.customWidth), subscriptOptions: StructFieldOptions.Generated);
185
186 agg.Add(fd);
187 activeFields.AddAll(fd);
188 }
189 foreach (var field in customFieldDescriptors)
190 {
191 // Don't add duplicate fields
192 if (ps.fields.Any((f) => (f.name == field.name)))
193 continue;
194 agg.Add(field);
195 }
196 newPassStructs.Add(new StructDescriptor { name = ps.name, packFields = ps.packFields, fields = ps.fields.Union(agg).ToArray() });
197 }
198 else
199 {
200 newPassStructs.Add(ps);
201 }
202 }
203
204 foreach (var cid in customBlockNodes.Select(bn => bn.descriptor))
205 activeFields.AddAll(cid);
206
207 return newPassStructs;
208 }
209
210 // Custom Interpolator descriptors indicate how and where code should be generated.
211 // At this entry point, we can process the descriptors on a provided pass and generate
212 // the corresponding splices.
213 internal void ProcessDescriptors(IEnumerable<Descriptor> descriptors)
214 {
215 if (CustomInterpolatorUtils.generatorSkipFlag)
216 return;
217
218 ShaderStringBuilder builder = new ShaderStringBuilder();
219 foreach (var desc in descriptors)
220 {
221 builder.Clear();
222 if (!desc.isValid)
223 continue;
224
225 if (desc.hasPreprocessor)
226 builder.AppendLine($"#ifdef {desc.preprocessor}");
227
228 if (desc.isBlock) GenCopyBlock(desc.dst, desc.src, builder);
229 else if (desc.isMacroBlock) GenCopyMacroBlock(desc.src, desc.dst, builder);
230 else if (desc.isFunc) GenCopyFunc(desc.name, desc.dst, desc.src, builder, desc.define);
231 else if (desc.isStruct) GenStruct(desc.name, builder, desc.define);
232 else if (desc.isDefine) builder.AppendLine($"#define {desc.define}");
233
234 if (desc.hasPreprocessor)
235 builder.AppendLine("#endif");
236
237 if (!spliceCommandBuffer.ContainsKey(desc.splice))
238 spliceCommandBuffer.Add(desc.splice, new ShaderStringBuilder());
239
240 spliceCommandBuffer[desc.splice].Concat(builder);
241 }
242 }
243
244 // add our splices to the generator's dictionary.
245 internal void AppendToSpliceCommands(Dictionary<string, string> spliceCommands)
246 {
247 if (CustomInterpolatorUtils.generatorSkipFlag)
248 return;
249
250 foreach (var spliceKV in spliceCommandBuffer)
251 spliceCommands.Add(spliceKV.Key, spliceKV.Value.ToCodeBlock());
252 }
253
254 #endregion
255
256 #region helpers
257 private void GenStruct(string structName, ShaderStringBuilder builder, string makeDefine = "")
258 {
259 builder.AppendLine($"struct {structName}");
260 builder.AppendLine("{");
261 using (builder.IndentScope())
262 {
263 foreach (var bn in customBlockNodes)
264 {
265 builder.AppendLine($"float{(int)bn.customWidth} {bn.customName};");
266 }
267 foreach (var field in customFieldDescriptors)
268 builder.AppendLine($"float{(int)field.vectorCount} {field.name};");
269 }
270 builder.AppendLine("};");
271 if (makeDefine != null && makeDefine != "")
272 builder.AppendLine($"#define {makeDefine}");
273
274 builder.AppendNewLine();
275 }
276
277 private void GenCopyBlock(string dst, string src, ShaderStringBuilder builder)
278 {
279 foreach (var bnode in customBlockNodes)
280 builder.AppendLine($"{dst}.{bnode.customName} = {src}.{bnode.customName};");
281 foreach (var field in customFieldDescriptors)
282 builder.AppendLine($"{dst}.{field.name} = {src}.{field.name};");
283 }
284
285 private void GenCopyMacroBlock(string startMacro, string endMacro, ShaderStringBuilder builder)
286 {
287 foreach (var bnode in customBlockNodes)
288 builder.AppendLine($"{startMacro}{bnode.customName}{endMacro};");
289 foreach (var field in customFieldDescriptors)
290 builder.AppendLine($"{startMacro}{field.name}{endMacro};");
291 }
292
293 private void GenCopyFunc(string funcName, string dstType, string srcType, ShaderStringBuilder builder, string makeDefine = "")
294 {
295 builder.AppendLine($"{dstType} {funcName}(inout {dstType} output, {srcType} input)");
296 using (builder.BlockScope())
297 {
298 GenCopyBlock("output", "input", builder);
299 builder.AppendLine("return output;");
300 }
301 if (makeDefine != null && makeDefine != "")
302 builder.AppendLine($"#define {makeDefine}");
303 }
304
305 private static HashSet<AbstractMaterialNode> GetAntecedents(BlockNode blockNode)
306 {
307 if (blockNode != null && blockNode.isCustomBlock && blockNode.isActive && blockNode.GetInputNodeFromSlot(0) != null)
308 {
309 var results = new HashSet<AbstractMaterialNode>();
310 NodeUtils.DepthFirstCollectNodesFromNode(results, blockNode, NodeUtils.IncludeSelf.Exclude);
311 return results != null && results.Count == 0 ? null : results;
312 }
313 return null;
314 }
315
316 private static void InsertAntecedent(List<AbstractMaterialNode> nodes, AbstractMaterialNode node)
317 {
318 var upstream = node.GetInputSlots<MaterialSlot>().Where(slot => slot.isConnected).Select(slot => node.GetInputNodeFromSlot(slot.id));
319 int safeIdx = nodes.FindLastIndex(n => upstream.Contains(n)) + 1;
320 nodes.Insert(safeIdx, node);
321 }
322
323 private static ShaderValueType ShaderValueTypeFrom(int width)
324 {
325 switch (width)
326 {
327 case 1: return ShaderValueType.Float;
328 case 2: return ShaderValueType.Float2;
329 case 3: return ShaderValueType.Float3;
330 default: return ShaderValueType.Float4;
331 }
332 }
333
334 #endregion
335 }
336}