1  using Ryujinx.Common;
  2  using Ryujinx.Memory;
  3  using Ryujinx.Memory.Tracking;
  4  using System;
  5  using System.Collections.Generic;
  6  using System.Diagnostics;
  7  
  8  namespace Ryujinx.Cpu.Jit.HostTracked
  9  {
 10      class AddressSpacePartitioned : IDisposable
 11      {
 12          private const int PartitionBits = 25;
 13          private const ulong PartitionSize = 1UL << PartitionBits;
 14  
 15          private readonly MemoryBlock _backingMemory;
 16          private readonly List<AddressSpacePartition> _partitions;
 17          private readonly AddressSpacePartitionAllocator _asAllocator;
 18          private readonly Action<ulong, IntPtr, ulong> _updatePtCallback;
 19          private readonly bool _useProtectionMirrors;
 20  
 21          public AddressSpacePartitioned(MemoryTracking tracking, MemoryBlock backingMemory, NativePageTable nativePageTable, bool useProtectionMirrors)
 22          {
 23              _backingMemory = backingMemory;
 24              _partitions = new();
 25              _asAllocator = new(tracking, nativePageTable.Read, _partitions);
 26              _updatePtCallback = nativePageTable.Update;
 27              _useProtectionMirrors = useProtectionMirrors;
 28          }
 29  
 30          public void Map(ulong va, ulong pa, ulong size)
 31          {
 32              ulong endVa = va + size;
 33  
 34              lock (_partitions)
 35              {
 36                  EnsurePartitionsLocked(va, size);
 37  
 38                  while (va < endVa)
 39                  {
 40                      int partitionIndex = FindPartitionIndexLocked(va);
 41                      AddressSpacePartition partition = _partitions[partitionIndex];
 42  
 43                      (ulong clampedVa, ulong clampedEndVa) = ClampRange(partition, va, endVa);
 44  
 45                      partition.Map(clampedVa, pa, clampedEndVa - clampedVa);
 46  
 47                      ulong currentSize = clampedEndVa - clampedVa;
 48  
 49                      va += currentSize;
 50                      pa += currentSize;
 51  
 52                      InsertOrRemoveBridgeIfNeeded(partitionIndex);
 53                  }
 54              }
 55          }
 56  
 57          public void Unmap(ulong va, ulong size)
 58          {
 59              ulong endVa = va + size;
 60  
 61              while (va < endVa)
 62              {
 63                  AddressSpacePartition partition;
 64  
 65                  lock (_partitions)
 66                  {
 67                      int partitionIndex = FindPartitionIndexLocked(va);
 68                      if (partitionIndex < 0)
 69                      {
 70                          va += PartitionSize - (va & (PartitionSize - 1));
 71  
 72                          continue;
 73                      }
 74  
 75                      partition = _partitions[partitionIndex];
 76  
 77                      (ulong clampedVa, ulong clampedEndVa) = ClampRange(partition, va, endVa);
 78  
 79                      partition.Unmap(clampedVa, clampedEndVa - clampedVa);
 80  
 81                      va += clampedEndVa - clampedVa;
 82  
 83                      InsertOrRemoveBridgeIfNeeded(partitionIndex);
 84  
 85                      if (partition.IsEmpty())
 86                      {
 87                          _partitions.Remove(partition);
 88                          partition.Dispose();
 89                      }
 90                  }
 91              }
 92          }
 93  
 94          public void Reprotect(ulong va, ulong size, MemoryPermission protection)
 95          {
 96              ulong endVa = va + size;
 97  
 98              lock (_partitions)
 99              {
100                  while (va < endVa)
101                  {
102                      AddressSpacePartition partition = FindPartitionWithIndex(va, out int partitionIndex);
103  
104                      if (partition == null)
105                      {
106                          va += PartitionSize - (va & (PartitionSize - 1));
107  
108                          continue;
109                      }
110  
111                      (ulong clampedVa, ulong clampedEndVa) = ClampRange(partition, va, endVa);
112  
113                      if (_useProtectionMirrors)
114                      {
115                          partition.Reprotect(clampedVa, clampedEndVa - clampedVa, protection, this, _updatePtCallback);
116                      }
117                      else
118                      {
119                          partition.ReprotectAligned(clampedVa, clampedEndVa - clampedVa, protection);
120  
121                          if (clampedVa == partition.Address &&
122                              partitionIndex > 0 &&
123                              _partitions[partitionIndex - 1].EndAddress == partition.Address)
124                          {
125                              _partitions[partitionIndex - 1].ReprotectBridge(protection);
126                          }
127                      }
128  
129                      va += clampedEndVa - clampedVa;
130                  }
131              }
132          }
133  
134          public PrivateRange GetPrivateAllocation(ulong va)
135          {
136              AddressSpacePartition partition = FindPartition(va);
137  
138              if (partition == null)
139              {
140                  return PrivateRange.Empty;
141              }
142  
143              return partition.GetPrivateAllocation(va);
144          }
145  
146          public PrivateRange GetFirstPrivateAllocation(ulong va, ulong size, out ulong nextVa)
147          {
148              AddressSpacePartition partition = FindPartition(va);
149  
150              if (partition == null)
151              {
152                  nextVa = (va & ~(PartitionSize - 1)) + PartitionSize;
153  
154                  return PrivateRange.Empty;
155              }
156  
157              return partition.GetFirstPrivateAllocation(va, size, out nextVa);
158          }
159  
160          public bool HasAnyPrivateAllocation(ulong va, ulong size, out PrivateRange range)
161          {
162              range = PrivateRange.Empty;
163  
164              ulong startVa = va;
165              ulong endVa = va + size;
166  
167              while (va < endVa)
168              {
169                  AddressSpacePartition partition = FindPartition(va);
170  
171                  if (partition == null)
172                  {
173                      va += PartitionSize - (va & (PartitionSize - 1));
174  
175                      continue;
176                  }
177  
178                  (ulong clampedVa, ulong clampedEndVa) = ClampRange(partition, va, endVa);
179  
180                  if (partition.HasPrivateAllocation(clampedVa, clampedEndVa - clampedVa, startVa, size, ref range))
181                  {
182                      return true;
183                  }
184  
185                  va += clampedEndVa - clampedVa;
186              }
187  
188              return false;
189          }
190  
191          private void InsertOrRemoveBridgeIfNeeded(int partitionIndex)
192          {
193              if (partitionIndex > 0)
194              {
195                  if (_partitions[partitionIndex - 1].EndAddress == _partitions[partitionIndex].Address)
196                  {
197                      _partitions[partitionIndex - 1].InsertBridgeAtEnd(_partitions[partitionIndex], _useProtectionMirrors);
198                  }
199                  else
200                  {
201                      _partitions[partitionIndex - 1].InsertBridgeAtEnd(null, _useProtectionMirrors);
202                  }
203              }
204  
205              if (partitionIndex + 1 < _partitions.Count && _partitions[partitionIndex].EndAddress == _partitions[partitionIndex + 1].Address)
206              {
207                  _partitions[partitionIndex].InsertBridgeAtEnd(_partitions[partitionIndex + 1], _useProtectionMirrors);
208              }
209              else
210              {
211                  _partitions[partitionIndex].InsertBridgeAtEnd(null, _useProtectionMirrors);
212              }
213          }
214  
215          public IntPtr GetPointer(ulong va, ulong size)
216          {
217              AddressSpacePartition partition = FindPartition(va);
218  
219              return partition.GetPointer(va, size);
220          }
221  
222          private static (ulong, ulong) ClampRange(AddressSpacePartition partition, ulong va, ulong endVa)
223          {
224              if (va < partition.Address)
225              {
226                  va = partition.Address;
227              }
228  
229              if (endVa > partition.EndAddress)
230              {
231                  endVa = partition.EndAddress;
232              }
233  
234              return (va, endVa);
235          }
236  
237          private AddressSpacePartition FindPartition(ulong va)
238          {
239              lock (_partitions)
240              {
241                  int index = FindPartitionIndexLocked(va);
242                  if (index >= 0)
243                  {
244                      return _partitions[index];
245                  }
246              }
247  
248              return null;
249          }
250  
251          private AddressSpacePartition FindPartitionWithIndex(ulong va, out int index)
252          {
253              lock (_partitions)
254              {
255                  index = FindPartitionIndexLocked(va);
256                  if (index >= 0)
257                  {
258                      return _partitions[index];
259                  }
260              }
261  
262              return null;
263          }
264  
265          private int FindPartitionIndexLocked(ulong va)
266          {
267              int left = 0;
268              int middle;
269              int right = _partitions.Count - 1;
270  
271              while (left <= right)
272              {
273                  middle = left + ((right - left) >> 1);
274  
275                  AddressSpacePartition partition = _partitions[middle];
276  
277                  if (partition.Address <= va && partition.EndAddress > va)
278                  {
279                      return middle;
280                  }
281  
282                  if (partition.Address >= va)
283                  {
284                      right = middle - 1;
285                  }
286                  else
287                  {
288                      left = middle + 1;
289                  }
290              }
291  
292              return -1;
293          }
294  
295          private void EnsurePartitionsLocked(ulong va, ulong size)
296          {
297              ulong endVa = BitUtils.AlignUp(va + size, PartitionSize);
298              va = BitUtils.AlignDown(va, PartitionSize);
299  
300              for (int i = 0; i < _partitions.Count && va < endVa; i++)
301              {
302                  AddressSpacePartition partition = _partitions[i];
303  
304                  if (partition.Address <= va && partition.EndAddress > va)
305                  {
306                      if (partition.EndAddress >= endVa)
307                      {
308                          // Fully mapped already.
309                          va = endVa;
310  
311                          break;
312                      }
313  
314                      ulong gapSize;
315  
316                      if (i + 1 < _partitions.Count)
317                      {
318                          AddressSpacePartition nextPartition = _partitions[i + 1];
319  
320                          if (partition.EndAddress == nextPartition.Address)
321                          {
322                              va = partition.EndAddress;
323  
324                              continue;
325                          }
326  
327                          gapSize = Math.Min(endVa, nextPartition.Address) - partition.EndAddress;
328                      }
329                      else
330                      {
331                          gapSize = endVa - partition.EndAddress;
332                      }
333  
334                      _partitions.Insert(i + 1, CreateAsPartition(partition.EndAddress, gapSize));
335                      va = partition.EndAddress + gapSize;
336                      i++;
337                  }
338                  else if (partition.EndAddress > va)
339                  {
340                      Debug.Assert(partition.Address > va);
341  
342                      ulong gapSize;
343  
344                      if (partition.Address < endVa)
345                      {
346                          gapSize = partition.Address - va;
347                      }
348                      else
349                      {
350                          gapSize = endVa - va;
351                      }
352  
353                      _partitions.Insert(i, CreateAsPartition(va, gapSize));
354                      va = Math.Min(partition.EndAddress, endVa);
355                      i++;
356                  }
357              }
358  
359              if (va < endVa)
360              {
361                  _partitions.Add(CreateAsPartition(va, endVa - va));
362              }
363  
364              ValidatePartitionList();
365          }
366  
367          [Conditional("DEBUG")]
368          private void ValidatePartitionList()
369          {
370              for (int i = 1; i < _partitions.Count; i++)
371              {
372                  Debug.Assert(_partitions[i].Address > _partitions[i - 1].Address);
373                  Debug.Assert(_partitions[i].EndAddress > _partitions[i - 1].EndAddress);
374              }
375          }
376  
377          private AddressSpacePartition CreateAsPartition(ulong va, ulong size)
378          {
379              return new(CreateAsPartitionAllocation(va, size), _backingMemory, va, size);
380          }
381  
382          public AddressSpacePartitionAllocation CreateAsPartitionAllocation(ulong va, ulong size)
383          {
384              return _asAllocator.Allocate(va, size + MemoryBlock.GetPageSize());
385          }
386  
387          protected virtual void Dispose(bool disposing)
388          {
389              if (disposing)
390              {
391                  foreach (AddressSpacePartition partition in _partitions)
392                  {
393                      partition.Dispose();
394                  }
395  
396                  _partitions.Clear();
397                  _asAllocator.Dispose();
398              }
399          }
400  
401          public void Dispose()
402          {
403              Dispose(disposing: true);
404              GC.SuppressFinalize(this);
405          }
406      }
407  }