1  using ARMeilleure.IntermediateRepresentation;
  2  using System;
  3  using System.Collections.Generic;
  4  using static ARMeilleure.IntermediateRepresentation.Operand.Factory;
  5  using static ARMeilleure.IntermediateRepresentation.Operation.Factory;
  6  
  7  namespace ARMeilleure.CodeGen.RegisterAllocators
  8  {
  9      class CopyResolver
 10      {
 11          private class ParallelCopy
 12          {
 13              private readonly struct Copy
 14              {
 15                  public Register Dest { get; }
 16                  public Register Source { get; }
 17  
 18                  public OperandType Type { get; }
 19  
 20                  public Copy(Register dest, Register source, OperandType type)
 21                  {
 22                      Dest = dest;
 23                      Source = source;
 24                      Type = type;
 25                  }
 26              }
 27  
 28              private readonly List<Copy> _copies;
 29  
 30              public int Count => _copies.Count;
 31  
 32              public ParallelCopy()
 33              {
 34                  _copies = new List<Copy>();
 35              }
 36  
 37              public void AddCopy(Register dest, Register source, OperandType type)
 38              {
 39                  _copies.Add(new Copy(dest, source, type));
 40              }
 41  
 42              public void Sequence(List<Operation> sequence)
 43              {
 44                  Dictionary<Register, Register> locations = new();
 45                  Dictionary<Register, Register> sources = new();
 46  
 47                  Dictionary<Register, OperandType> types = new();
 48  
 49                  Queue<Register> pendingQueue = new();
 50                  Queue<Register> readyQueue = new();
 51  
 52                  foreach (Copy copy in _copies)
 53                  {
 54                      locations[copy.Source] = copy.Source;
 55                      sources[copy.Dest] = copy.Source;
 56                      types[copy.Dest] = copy.Type;
 57  
 58                      pendingQueue.Enqueue(copy.Dest);
 59                  }
 60  
 61                  foreach (Copy copy in _copies)
 62                  {
 63                      // If the destination is not used anywhere, we can assign it immediately.
 64                      if (!locations.ContainsKey(copy.Dest))
 65                      {
 66                          readyQueue.Enqueue(copy.Dest);
 67                      }
 68                  }
 69  
 70                  while (pendingQueue.TryDequeue(out Register current))
 71                  {
 72                      Register copyDest;
 73                      Register origSource;
 74                      Register copySource;
 75  
 76                      while (readyQueue.TryDequeue(out copyDest))
 77                      {
 78                          origSource = sources[copyDest];
 79                          copySource = locations[origSource];
 80  
 81                          OperandType type = types[copyDest];
 82  
 83                          EmitCopy(sequence, GetRegister(copyDest, type), GetRegister(copySource, type));
 84  
 85                          locations[origSource] = copyDest;
 86  
 87                          if (origSource == copySource && sources.ContainsKey(origSource))
 88                          {
 89                              readyQueue.Enqueue(origSource);
 90                          }
 91                      }
 92  
 93                      copyDest = current;
 94                      origSource = sources[copyDest];
 95                      copySource = locations[origSource];
 96  
 97                      if (copyDest != copySource)
 98                      {
 99                          OperandType type = types[copyDest];
100  
101                          type = type.IsInteger() ? OperandType.I64 : OperandType.V128;
102  
103                          EmitXorSwap(sequence, GetRegister(copyDest, type), GetRegister(copySource, type));
104  
105                          locations[origSource] = copyDest;
106  
107                          Register swapOther = copySource;
108  
109                          if (copyDest != locations[sources[copySource]])
110                          {
111                              // Find the other swap destination register.
112                              // To do that, we search all the pending registers, and pick
113                              // the one where the copy source register is equal to the
114                              // current destination register being processed (copyDest).
115                              foreach (Register pending in pendingQueue)
116                              {
117                                  // Is this a copy of pending <- copyDest?
118                                  if (copyDest == locations[sources[pending]])
119                                  {
120                                      swapOther = pending;
121  
122                                      break;
123                                  }
124                              }
125                          }
126  
127                          // The value that was previously at "copyDest" now lives on
128                          // "copySource" thanks to the swap, now we need to update the
129                          // location for the next copy that is supposed to copy the value
130                          // that used to live on "copyDest".
131                          locations[sources[swapOther]] = copySource;
132                      }
133                  }
134              }
135  
136              private static void EmitCopy(List<Operation> sequence, Operand x, Operand y)
137              {
138                  sequence.Add(Operation(Instruction.Copy, x, y));
139              }
140  
141              private static void EmitXorSwap(List<Operation> sequence, Operand x, Operand y)
142              {
143                  sequence.Add(Operation(Instruction.BitwiseExclusiveOr, x, x, y));
144                  sequence.Add(Operation(Instruction.BitwiseExclusiveOr, y, y, x));
145                  sequence.Add(Operation(Instruction.BitwiseExclusiveOr, x, x, y));
146              }
147          }
148  
149          private Queue<Operation> _fillQueue = null;
150          private Queue<Operation> _spillQueue = null;
151          private ParallelCopy _parallelCopy = null;
152  
153          public bool HasCopy { get; private set; }
154  
155          public void AddSplit(LiveInterval left, LiveInterval right)
156          {
157              if (left.Local != right.Local)
158              {
159                  throw new ArgumentException("Intervals of different variables are not allowed.");
160              }
161  
162              OperandType type = left.Local.Type;
163  
164              if (left.IsSpilled && !right.IsSpilled)
165              {
166                  // Move from the stack to a register.
167                  AddSplitFill(left, right, type);
168              }
169              else if (!left.IsSpilled && right.IsSpilled)
170              {
171                  // Move from a register to the stack.
172                  AddSplitSpill(left, right, type);
173              }
174              else if (!left.IsSpilled && !right.IsSpilled && left.Register != right.Register)
175              {
176                  // Move from one register to another.
177                  AddSplitCopy(left, right, type);
178              }
179              else if (left.SpillOffset != right.SpillOffset)
180              {
181                  // This would be the stack-to-stack move case, but this is not supported.
182                  throw new ArgumentException("Both intervals were spilled.");
183              }
184          }
185  
186          private void AddSplitFill(LiveInterval left, LiveInterval right, OperandType type)
187          {
188              _fillQueue ??= new Queue<Operation>();
189  
190              Operand register = GetRegister(right.Register, type);
191              Operand offset = Const(left.SpillOffset);
192  
193              _fillQueue.Enqueue(Operation(Instruction.Fill, register, offset));
194  
195              HasCopy = true;
196          }
197  
198          private void AddSplitSpill(LiveInterval left, LiveInterval right, OperandType type)
199          {
200              _spillQueue ??= new Queue<Operation>();
201  
202              Operand offset = Const(right.SpillOffset);
203              Operand register = GetRegister(left.Register, type);
204  
205              _spillQueue.Enqueue(Operation(Instruction.Spill, default, offset, register));
206  
207              HasCopy = true;
208          }
209  
210          private void AddSplitCopy(LiveInterval left, LiveInterval right, OperandType type)
211          {
212              _parallelCopy ??= new ParallelCopy();
213  
214              _parallelCopy.AddCopy(right.Register, left.Register, type);
215  
216              HasCopy = true;
217          }
218  
219          public Operation[] Sequence()
220          {
221              List<Operation> sequence = new();
222  
223              if (_spillQueue != null)
224              {
225                  while (_spillQueue.TryDequeue(out Operation spillOp))
226                  {
227                      sequence.Add(spillOp);
228                  }
229              }
230  
231              _parallelCopy?.Sequence(sequence);
232  
233              if (_fillQueue != null)
234              {
235                  while (_fillQueue.TryDequeue(out Operation fillOp))
236                  {
237                      sequence.Add(fillOp);
238                  }
239              }
240  
241              return sequence.ToArray();
242          }
243  
244          private static Operand GetRegister(Register reg, OperandType type)
245          {
246              return Register(reg.Index, reg.Type, type);
247          }
248      }
249  }