From bd74adffb60ceeceaa448d1b238bf3965ad764f2 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Tue, 8 Oct 2024 20:37:57 +0100
Subject: [PATCH] Add FunctionIndexMap to track function indices.

---
 llvm/lib/YkIR/YkIRWriter.cpp | 29 +++++++++++++++++++----------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/YkIR/YkIRWriter.cpp b/llvm/lib/YkIR/YkIRWriter.cpp
index cab8f57f4e869e3..28bc8fccb3b89f5 100644
--- a/llvm/lib/YkIR/YkIRWriter.cpp
+++ b/llvm/lib/YkIR/YkIRWriter.cpp
@@ -305,7 +305,10 @@ class YkIRWriter {
   vector<llvm::Type *> Types;
   vector<llvm::Constant *> Constants;
   vector<llvm::GlobalVariable *> Globals;
-
+  // Maps a function to its index. This is not the same as the index in
+  // the module's function list because we skip cloned functions in
+  // serialisation.
+  std::unordered_map<llvm::Function *, size_t> FunctionIndexMap;
   // File paths.
   vector<string> Paths;
 
@@ -365,9 +368,13 @@ class YkIRWriter {
   }
 
   size_t functionIndex(llvm::Function *F) {
-    // FIXME: For now we assume that function indicies in LLVM IR and our IR
-    // are the same.
-    return getIndex(&M, F);
+    auto it = FunctionIndexMap.find(F);
+    if (it != FunctionIndexMap.end()) {
+      return it->second;
+    }
+    llvm::errs() << "Function not found in function index map: " << F->getName()
+                 << "\n";
+    llvm::report_fatal_error("Function not found in function index map");
   }
 
   // Serialises a null-terminated string.
@@ -1751,15 +1758,17 @@ class YkIRWriter {
     OutStreamer.emitInt8(IdxBitWidth);
 
     // num_funcs:
-    // Count non-cloned functions
-    unsigned numFuncs = 0;
-    for (const llvm::Function &F : M) {
+    // Count functions for serilaisation and populate functions map
+    int functionCount = 0;
+    for (llvm::Function &F : M) {
+      // Skip cloned functions
       if (!StringRef(F.getName()).startswith(YK_CLONE_PREFIX)) {
-        numFuncs++;
+        FunctionIndexMap[&F] = functionCount;
+        functionCount++;
       }
     }
-    // Emit the number of non-cloned functions
-    OutStreamer.emitSizeT(numFuncs);
+    // Emit the number of functions
+    OutStreamer.emitSizeT(functionCount);
     // funcs:
     for (llvm::Function &F : M) {
       // Skip cloned functions