ImageTracker: Load AOT images

For now, a simple directory structure with filename based mapping is
used:
LOCALAPPDATA/cache/<main image id>/{<dll 1 image id>, <dll 2 image id>}

All cache blobs are loaded once at the start for simplicity, this will
also be useful in the future as metadata validation for settings etc can
be done just once.

Source/Windows/Common/ImageTracker.cpp

@@ -167,6 +167,12 @@ FEXCore::ExecutableFileSectionInfo ImageTracker::HandleImageMap(std::string_view
         Writer->AppendSetMainExecutable(ImageInfo->Info);
         CTX.SetCodeMapWriter(std::move(Writer));
       }
+      LoadAOTImages(*ImageInfo);
+    }
+
+    auto AOTImage = AOTImages.find(ID);
+    if (AOTImage != AOTImages.end()) {
+      CTX.GetCodeCache().LoadData(nullptr, AOTImage->second.Data, ImageInfo->SectionInfo);
     }
   }
 
@@ -209,4 +215,79 @@ int ImageTracker::OpenCodeMapFile() {
   }
   return _sopen(ActiveCodeMapPath.c_str(), O_CREAT | O_TRUNC | O_WRONLY | O_APPEND, _SH_DENYRW, 0644);
 }
+
+void ImageTracker::LoadAOTImages(MappedImageInfo& ImageInfo) {
+  const auto AnsiPath =
+    fmt::format("\\??\\{}cache\\{}", FEX::Config::GetCacheDirectory(), FEXCore::CodeMap::GetBaseFilename(ImageInfo.Info, false));
+
+  // Iterate over all files in the given executable's cache directory, mapping them into memory for future use.
+  // Each cache file name matches the unique ID (as returned by FEXCore::CodeMap::GetBaseFilename) of the image it corresponds to.
+  ScopedUnicodeString NtPath(AnsiPath.c_str());
+
+  OBJECT_ATTRIBUTES DirAttr;
+  InitializeObjectAttributes(&DirAttr, &*NtPath, OBJ_CASE_INSENSITIVE, NULL, NULL);
+
+  ScopedHandle DirHandle;
+  IO_STATUS_BLOCK IOSB;
+  if (!NT_SUCCESS(NtOpenFile(&*DirHandle, FILE_LIST_DIRECTORY | SYNCHRONIZE, &DirAttr, &IOSB, FILE_SHARE_READ | FILE_SHARE_WRITE,
+                             FILE_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT))) {
+    return;
+  }
+
+  std::array<uint8_t, 0x1000> DirBuffer;
+  bool FirstScan = true;
+  auto QueryDir = [&]() {
+    NTSTATUS Status = NtQueryDirectoryFile(*DirHandle, nullptr, nullptr, nullptr, &IOSB, DirBuffer.data(), DirBuffer.size(),
+                                           FileBothDirectoryInformation, FALSE, nullptr, FirstScan);
+    if (FirstScan) {
+      FirstScan = false;
+    }
+    return NT_SUCCESS(Status);
+  };
+
+  while (QueryDir()) {
+    auto* Info = reinterpret_cast<PFILE_BOTH_DIRECTORY_INFORMATION>(DirBuffer.data());
+
+    while (true) {
+      UNICODE_STRING CurrentFileName;
+      CurrentFileName.Buffer = Info->FileName;
+      CurrentFileName.Length = static_cast<USHORT>(Info->FileNameLength);
+      CurrentFileName.MaximumLength = CurrentFileName.Length;
+
+      bool Skip = (Info->FileAttributes & FILE_ATTRIBUTE_DIRECTORY) || (Info->FileNameLength == 2 && Info->FileName[0] == '.') ||
+                  (Info->FileNameLength == 4 && Info->FileName[0] == '.' && Info->FileName[1] == '.');
+
+      if (!Skip) {
+        OBJECT_ATTRIBUTES FileAttr;
+        InitializeObjectAttributes(&FileAttr, &CurrentFileName, OBJ_CASE_INSENSITIVE, *DirHandle, nullptr);
+
+        ScopedHandle FileHandle;
+        if (NT_SUCCESS(NtOpenFile(&*FileHandle, GENERIC_READ | SYNCHRONIZE, &FileAttr, &IOSB, FILE_SHARE_READ, FILE_SYNCHRONOUS_IO_NONALERT))) {
+          ScopedHandle SectionHandle;
+          if (NT_SUCCESS(NtCreateSection(&*SectionHandle, SECTION_MAP_EXECUTE | SECTION_MAP_READ, nullptr, nullptr, PAGE_EXECUTE_READ,
+                                         SEC_COMMIT, *FileHandle))) {
+            void* LoadAddress = nullptr;
+            SIZE_T MappedSize = 0;
+            if (NT_SUCCESS(NtMapViewOfSection(*SectionHandle, NtCurrentProcess(), &LoadAddress, 0, 0, nullptr, &MappedSize, ViewUnmap,
+                                              MEM_RESERVE | MEM_TOP_DOWN, PAGE_EXECUTE_READ))) {
+              fextl::string UniqueId;
+              ULONG AnsiLength = 0;
+              RtlUnicodeToMultiByteSize(&AnsiLength, Info->FileName, Info->FileNameLength);
+              UniqueId.resize(AnsiLength);
+              RtlUnicodeToMultiByteN(UniqueId.data(), AnsiLength, NULL, Info->FileName, Info->FileNameLength);
+
+              AOTImages[UniqueId] = {.Data = static_cast<std::byte*>(LoadAddress)};
+              LogMan::Msg::IFmt("Loaded cache: {}", UniqueId);
+            }
+          }
+        }
+      }
+
+      if (Info->NextEntryOffset == 0) {
+        break;
+      }
+      Info = reinterpret_cast<PFILE_BOTH_DIRECTORY_INFORMATION>(reinterpret_cast<uint8_t*>(Info) + Info->NextEntryOffset);
+    }
+  }
+}
 } // namespace FEX::Windows

Source/Windows/Common/ImageTracker.h

@@ -55,6 +55,12 @@ private:
                     fextl::unordered_map<uint32_t, FEXCore::GuestRelocationType> Relocations);
   };
 
+  struct AOTImageInfo {
+    std::byte* Data;
+  };
+
+  void LoadAOTImages(MappedImageInfo& Info);
+
   FEXCore::Context::Context& CTX;
 
   FEX_CONFIG_OPT(ExtendedVolatileMetadataConfig, EXTENDEDVOLATILEMETADATA);
@@ -62,6 +68,7 @@ private:
 
   std::shared_mutex ImagesLock;
   std::map<uint64_t, MappedImageInfo> MappedImages;
+  std::map<fextl::string, AOTImageInfo> AOTImages;
 
   std::string ActiveCodeMapPath;
   bool IsGeneratingCache;

Source/Windows/include/winternl.h

@@ -519,6 +519,8 @@ NTSYSAPI NTSTATUS WINAPI NtMapViewOfSection(HANDLE, HANDLE, PVOID*, ULONG_PTR, S
 NTSTATUS WINAPI NtOpenKeyEx(PHANDLE, ACCESS_MASK, const OBJECT_ATTRIBUTES*, ULONG);
 NTSTATUS WINAPI NtProtectVirtualMemory(HANDLE, PVOID*, SIZE_T*, ULONG, ULONG*);
 NTSTATUS WINAPI NtQueryAttributesFile(const OBJECT_ATTRIBUTES*, FILE_BASIC_INFORMATION*);
+NTSTATUS WINAPI NtQueryDirectoryFile(HANDLE, HANDLE, PIO_APC_ROUTINE, PVOID, PIO_STATUS_BLOCK, PVOID, ULONG, FILE_INFORMATION_CLASS,
+                                     BOOLEAN, PUNICODE_STRING, BOOLEAN);
 NTSTATUS WINAPI NtQueryValueKey(HANDLE, const UNICODE_STRING*, KEY_VALUE_INFORMATION_CLASS, void*, DWORD, DWORD*);
 NTSTATUS WINAPI NtQueryVirtualMemory(HANDLE, LPCVOID, MEMORY_INFORMATION_CLASS, PVOID, SIZE_T, SIZE_T*);
 NTSTATUS WINAPI NtReadFile(HANDLE, HANDLE, PIO_APC_ROUTINE, PVOID, PIO_STATUS_BLOCK, PVOID, ULONG, PLARGE_INTEGER, PULONG);