1  using System.Collections.Generic;
  2  using System.Diagnostics;
  3  using System.IO;
  4  using static Spv.Specification;
  5  
  6  namespace Spv.Generator
  7  {
  8      public partial class Module
  9      {
 10          // TODO: Register on SPIR-V registry.
 11          private const int GeneratorId = 0;
 12  
 13          private readonly uint _version;
 14  
 15          private uint _bound;
 16  
 17          // Follow spec order here while keeping it as simple as possible.
 18          private readonly List<Capability> _capabilities;
 19          private readonly List<string> _extensions;
 20          private readonly Dictionary<DeterministicStringKey, Instruction> _extInstImports;
 21          private AddressingModel _addressingModel;
 22          private MemoryModel _memoryModel;
 23  
 24          private readonly List<Instruction> _entrypoints;
 25          private readonly List<Instruction> _executionModes;
 26          private readonly List<Instruction> _debug;
 27          private readonly List<Instruction> _annotations;
 28  
 29          // In the declaration block.
 30          private readonly Dictionary<TypeDeclarationKey, Instruction> _typeDeclarations;
 31          private readonly List<Instruction> _typeDeclarationsList;
 32          // In the declaration block.
 33          private readonly List<Instruction> _globals;
 34          // In the declaration block.
 35          private readonly Dictionary<ConstantKey, Instruction> _constants;
 36          // In the declaration block, for function that aren't defined in the module.
 37          private readonly List<Instruction> _functionsDeclarations;
 38  
 39          private readonly List<Instruction> _functionsDefinitions;
 40  
 41          private readonly GeneratorPool<Instruction> _instPool;
 42          private readonly GeneratorPool<LiteralInteger> _integerPool;
 43  
 44          public Module(uint version, GeneratorPool<Instruction> instPool = null, GeneratorPool<LiteralInteger> integerPool = null)
 45          {
 46              _version = version;
 47              _bound = 1;
 48              _capabilities = new List<Capability>();
 49              _extensions = new List<string>();
 50              _extInstImports = new Dictionary<DeterministicStringKey, Instruction>();
 51              _addressingModel = AddressingModel.Logical;
 52              _memoryModel = MemoryModel.Simple;
 53              _entrypoints = new List<Instruction>();
 54              _executionModes = new List<Instruction>();
 55              _debug = new List<Instruction>();
 56              _annotations = new List<Instruction>();
 57              _typeDeclarations = new Dictionary<TypeDeclarationKey, Instruction>();
 58              _typeDeclarationsList = new List<Instruction>();
 59              _constants = new Dictionary<ConstantKey, Instruction>();
 60              _globals = new List<Instruction>();
 61              _functionsDeclarations = new List<Instruction>();
 62              _functionsDefinitions = new List<Instruction>();
 63  
 64              _instPool = instPool ?? new GeneratorPool<Instruction>();
 65              _integerPool = integerPool ?? new GeneratorPool<LiteralInteger>();
 66  
 67              LiteralInteger.RegisterPool(_integerPool);
 68          }
 69  
 70          private uint GetNewId()
 71          {
 72              return _bound++;
 73          }
 74  
 75          public void AddCapability(Capability capability)
 76          {
 77              _capabilities.Add(capability);
 78          }
 79  
 80          public void AddExtension(string extension)
 81          {
 82              _extensions.Add(extension);
 83          }
 84  
 85          public Instruction NewInstruction(Op opcode, uint id = Instruction.InvalidId, Instruction resultType = null)
 86          {
 87              var result = _instPool.Allocate();
 88              result.Set(opcode, id, resultType);
 89  
 90              return result;
 91          }
 92  
 93          public Instruction AddExtInstImport(string import)
 94          {
 95              var key = new DeterministicStringKey(import);
 96  
 97              if (_extInstImports.TryGetValue(key, out Instruction extInstImport))
 98              {
 99                  // Update the duplicate instance to use the good id so it ends up being encoded correctly.
100                  return extInstImport;
101              }
102  
103              Instruction instruction = NewInstruction(Op.OpExtInstImport);
104              instruction.AddOperand(import);
105  
106              instruction.SetId(GetNewId());
107  
108              _extInstImports.Add(key, instruction);
109  
110              return instruction;
111          }
112  
113          private void AddTypeDeclaration(Instruction instruction, bool forceIdAllocation)
114          {
115              var key = new TypeDeclarationKey(instruction);
116  
117              if (!forceIdAllocation)
118              {
119                  if (_typeDeclarations.TryGetValue(key, out Instruction typeDeclaration))
120                  {
121                      // Update the duplicate instance to use the good id so it ends up being encoded correctly.
122  
123                      instruction.SetId(typeDeclaration.Id);
124  
125                      return;
126                  }
127              }
128  
129              instruction.SetId(GetNewId());
130  
131              _typeDeclarations[key] = instruction;
132              _typeDeclarationsList.Add(instruction);
133          }
134  
135          public void AddEntryPoint(ExecutionModel executionModel, Instruction function, string name, params Instruction[] interfaces)
136          {
137              Debug.Assert(function.Opcode == Op.OpFunction);
138  
139              Instruction entryPoint = NewInstruction(Op.OpEntryPoint);
140  
141              entryPoint.AddOperand(executionModel);
142              entryPoint.AddOperand(function);
143              entryPoint.AddOperand(name);
144              entryPoint.AddOperand(interfaces);
145  
146              _entrypoints.Add(entryPoint);
147          }
148  
149          public void AddExecutionMode(Instruction function, ExecutionMode mode, params IOperand[] parameters)
150          {
151              Debug.Assert(function.Opcode == Op.OpFunction);
152  
153              Instruction executionModeInstruction = NewInstruction(Op.OpExecutionMode);
154  
155              executionModeInstruction.AddOperand(function);
156              executionModeInstruction.AddOperand(mode);
157              executionModeInstruction.AddOperand(parameters);
158  
159              _executionModes.Add(executionModeInstruction);
160          }
161  
162          private void AddToFunctionDefinitions(Instruction instruction)
163          {
164              Debug.Assert(instruction.Opcode != Op.OpTypeInt);
165              _functionsDefinitions.Add(instruction);
166          }
167  
168          private void AddAnnotation(Instruction annotation)
169          {
170              _annotations.Add(annotation);
171          }
172  
173          private void AddDebug(Instruction debug)
174          {
175              _debug.Add(debug);
176          }
177  
178          public void AddLabel(Instruction label)
179          {
180              Debug.Assert(label.Opcode == Op.OpLabel);
181  
182              label.SetId(GetNewId());
183  
184              AddToFunctionDefinitions(label);
185          }
186  
187          public void AddLocalVariable(Instruction variable)
188          {
189              // TODO: Ensure it has the local modifier.
190              Debug.Assert(variable.Opcode == Op.OpVariable);
191  
192              variable.SetId(GetNewId());
193  
194              AddToFunctionDefinitions(variable);
195          }
196  
197          public void AddGlobalVariable(Instruction variable)
198          {
199              // TODO: Ensure it has the global modifier.
200              // TODO: All constants opcodes (OpSpecXXX and the rest of the OpConstantXXX).
201              Debug.Assert(variable.Opcode == Op.OpVariable);
202  
203              variable.SetId(GetNewId());
204  
205              _globals.Add(variable);
206          }
207  
208          private void AddConstant(Instruction constant)
209          {
210              Debug.Assert(constant.Opcode == Op.OpConstant ||
211                           constant.Opcode == Op.OpConstantFalse ||
212                           constant.Opcode == Op.OpConstantTrue ||
213                           constant.Opcode == Op.OpConstantNull ||
214                           constant.Opcode == Op.OpConstantComposite);
215  
216              var key = new ConstantKey(constant);
217  
218              if (_constants.TryGetValue(key, out Instruction global))
219              {
220                  // Update the duplicate instance to use the good id so it ends up being encoded correctly.
221                  constant.SetId(global.Id);
222  
223                  return;
224              }
225  
226              constant.SetId(GetNewId());
227  
228              _constants.Add(key, constant);
229          }
230  
231          public Instruction ExtInst(Instruction resultType, Instruction set, LiteralInteger instruction, params IOperand[] parameters)
232          {
233              Instruction result = NewInstruction(Op.OpExtInst, GetNewId(), resultType);
234  
235              result.AddOperand(set);
236              result.AddOperand(instruction);
237              result.AddOperand(parameters);
238              AddToFunctionDefinitions(result);
239  
240              return result;
241          }
242  
243          public void SetMemoryModel(AddressingModel addressingModel, MemoryModel memoryModel)
244          {
245              _addressingModel = addressingModel;
246              _memoryModel = memoryModel;
247          }
248  
249          // TODO: Find a way to make the auto generate one used.
250          public Instruction OpenClPrintf(Instruction resultType, Instruction format, params Instruction[] additionalarguments)
251          {
252              Instruction result = NewInstruction(Op.OpExtInst, GetNewId(), resultType);
253  
254              result.AddOperand(AddExtInstImport("OpenCL.std"));
255              result.AddOperand((LiteralInteger)184);
256              result.AddOperand(format);
257              result.AddOperand(additionalarguments);
258              AddToFunctionDefinitions(result);
259  
260              return result;
261          }
262  
263          public byte[] Generate()
264          {
265              // Estimate the size needed for the generated code, to avoid expanding the MemoryStream.
266              int sizeEstimate = 1024 + _functionsDefinitions.Count * 32;
267  
268              using MemoryStream stream = new(sizeEstimate);
269  
270              BinaryWriter writer = new(stream, System.Text.Encoding.ASCII);
271  
272              // Header
273              writer.Write(MagicNumber);
274              writer.Write(_version);
275              writer.Write(GeneratorId);
276              writer.Write(_bound);
277              writer.Write(0u);
278  
279              // 1.
280              foreach (Capability capability in _capabilities)
281              {
282                  Instruction capabilityInstruction = NewInstruction(Op.OpCapability);
283  
284                  capabilityInstruction.AddOperand(capability);
285                  capabilityInstruction.Write(writer);
286              }
287  
288              // 2.
289              foreach (string extension in _extensions)
290              {
291                  Instruction extensionInstruction = NewInstruction(Op.OpExtension);
292  
293                  extensionInstruction.AddOperand(extension);
294                  extensionInstruction.Write(writer);
295              }
296  
297              // 3.
298              foreach (Instruction extInstImport in _extInstImports.Values)
299              {
300                  extInstImport.Write(writer);
301              }
302  
303              // 4.
304              Instruction memoryModelInstruction = NewInstruction(Op.OpMemoryModel);
305              memoryModelInstruction.AddOperand(_addressingModel);
306              memoryModelInstruction.AddOperand(_memoryModel);
307              memoryModelInstruction.Write(writer);
308  
309              // 5.
310              foreach (Instruction entrypoint in _entrypoints)
311              {
312                  entrypoint.Write(writer);
313              }
314  
315              // 6.
316              foreach (Instruction executionMode in _executionModes)
317              {
318                  executionMode.Write(writer);
319              }
320  
321              // 7.
322              // TODO: Order debug information correctly.
323              foreach (Instruction debug in _debug)
324              {
325                  debug.Write(writer);
326              }
327  
328              // 8.
329              foreach (Instruction annotation in _annotations)
330              {
331                  annotation.Write(writer);
332              }
333  
334              // Ensure that everything is in the right order in the declarations section.
335              List<Instruction> declarations = new();
336              declarations.AddRange(_typeDeclarationsList);
337              declarations.AddRange(_globals);
338              declarations.AddRange(_constants.Values);
339              declarations.Sort((Instruction x, Instruction y) => x.Id.CompareTo(y.Id));
340  
341              // 9.
342              foreach (Instruction declaration in declarations)
343              {
344                  declaration.Write(writer);
345              }
346  
347              // 10.
348              foreach (Instruction functionDeclaration in _functionsDeclarations)
349              {
350                  functionDeclaration.Write(writer);
351              }
352  
353              // 11.
354              foreach (Instruction functionDefinition in _functionsDefinitions)
355              {
356                  functionDefinition.Write(writer);
357              }
358  
359              _instPool.Clear();
360              _integerPool.Clear();
361  
362              LiteralInteger.UnregisterPool();
363  
364              return stream.ToArray();
365          }
366      }
367  }