[PATCH v2 14/50] helper-to-tcg: PrepareForOptPass, map annotations

In the LLVM IR module, function annotations are stored in one big global
array of strings.  Traverse this array and parse the data into a format
more useful for future passes.  A map between Functions * and an
`Annotations` structure is exposed.

Signed-off-by: Anton Johansson <[email protected]>
---
 .../include/FunctionAnnotation.hpp            | 104 ++++++++++++++++++
 .../include/PrepareForOptPass.hpp             |   7 +-
 subprojects/helper-to-tcg/src/Pipeline.cpp    |   6 +-
 .../PrepareForOptPass/PrepareForOptPass.cpp   |  94 ++++++++++++++++
 4 files changed, 209 insertions(+), 2 deletions(-)
 create mode 100644 subprojects/helper-to-tcg/include/FunctionAnnotation.hpp

diff --git a/subprojects/helper-to-tcg/include/FunctionAnnotation.hpp b/subprojects/helper-to-tcg/include/FunctionAnnotation.hpp
new file mode 100644
index 0000000000..398dd53ef5
--- /dev/null
+++ b/subprojects/helper-to-tcg/include/FunctionAnnotation.hpp
@@ -0,0 +1,104 @@
+//
+//  Copyright(c) 2026 rev.ng Labs Srl. All Rights Reserved.
+//
+//  This program is free software; you can redistribute it and/or modify
+//  it under the terms of the GNU General Public License as published by
+//  the Free Software Foundation; either version 2 of the License, or
+//  (at your option) any later version.
+//
+//  This program is distributed in the hope that it will be useful,
+//  but WITHOUT ANY WARRANTY; without even the implied warranty of
+//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+//  GNU General Public License for more details.
+//
+//  You should have received a copy of the GNU General Public License
+//  along with this program; if not, see <http://www.gnu.org/licenses/>.
+//
+
+#pragma once
+
+#include <llvm/ADT/DenseMap.h>
+#include <llvm/ADT/SmallVector.h>
+#include <llvm/Support/Format.h>
+#include <llvm/Support/raw_ostream.h>
+#include <stdint.h>
+
+namespace llvm {
+class Function;
+}
+
+// Different kind of function annotations which control the behaviour
+// of helper-to-tcg.
+enum class ArgumentAnnotation : uint8_t {
+    // Declares a list of arguments as immediates
+    Immediate = 1,
+    // Declares a list of arguments as vectors, represented by offsets into
+    // the CPU state
+    PtrToOffset = 2,
+};
+
+// Different kind of function annotations which control the behaviour
+// of helper-to-tcg.
+enum class FunctionAnnotation : uint8_t {
+    // Function should be translated
+    HelperToTcg = 1,
+    // Return value of function is an immediate
+    ReturnsImmediate = 2,
+};
+
+// Annotation data which may be attached to a function
+class Annotations {
+    // 8-bit flag for each argument in a function, fields defined by
+    // `ArgumentAnnotions`.
+    llvm::SmallVector<uint8_t, 4> ArgumentAnnotations;
+    // Flag of function annotations, fields defined by `FunctionsAnnotations`.
+    uint8_t FunctionAnnotations = 0;
+
+  public:
+    inline uint8_t getArgFlag(size_t Index) const {
+        if (Index >= ArgumentAnnotations.size()) {
+            return 0;
+        }
+        return ArgumentAnnotations[Index];
+    }
+
+    // Getters and setters for annotations flags.
+
+    inline void set(FunctionAnnotation FA) {
+        FunctionAnnotations |= (uint8_t)FA;
+    }
+
+    inline void set(size_t Index, ArgumentAnnotation AA) {
+        if (Index >= ArgumentAnnotations.size()) {
+            // Resizing will default initialize any new elements.
+            ArgumentAnnotations.resize(Index + 1);
+        }
+        ArgumentAnnotations[Index] |= (uint8_t)AA;
+    }
+
+    inline bool isSet(FunctionAnnotation FA) const {
+        return (FunctionAnnotations & (uint8_t)FA) != 0;
+    }
+
+    inline bool isSet(size_t Index, ArgumentAnnotation AA) const {
+        uint8_t Flag = getArgFlag(Index);
+        return (Flag & (uint8_t)AA) != 0;
+    }
+
+    // Pretty printing debug information
+    inline void dump(llvm::raw_ostream &Out) const {
+        Out << "Annotations:\n";
+        Out << "  Function: " << llvm::format_hex(FunctionAnnotations, 4)
+            << "\n";
+        for (size_t I = 0; I < ArgumentAnnotations.size(); ++I) {
+            const uint8_t Flag = ArgumentAnnotations[I];
+            Out << "  Argument[" << I << "]: " << llvm::format_hex(Flag, 4)
+                << "\n";
+        }
+    }
+};
+
+// Mapping from functions to annotations, this is the main structure to be used
+// by other parts of the codebase when referencing annotations, filled out by
+// `PrepareForOptPass`.
+using AnnotationMapTy = llvm::DenseMap<llvm::Function *, Annotations>;
diff --git a/subprojects/helper-to-tcg/include/PrepareForOptPass.hpp b/subprojects/helper-to-tcg/include/PrepareForOptPass.hpp
index 2b3694c536..e007243578 100644
--- a/subprojects/helper-to-tcg/include/PrepareForOptPass.hpp
+++ b/subprojects/helper-to-tcg/include/PrepareForOptPass.hpp
@@ -17,6 +17,7 @@
 
 #pragma once
 
+#include "FunctionAnnotation.hpp"
 #include <llvm/IR/PassManager.h>
 
 //
@@ -27,8 +28,12 @@
 //
 
 class PrepareForOptPass : public llvm::PassInfoMixin<PrepareForOptPass> {
+    AnnotationMapTy &ResultAnnotations;
 public:
-    PrepareForOptPass() {}
+    PrepareForOptPass(AnnotationMapTy &ResultAnnotations)
+        : ResultAnnotations(ResultAnnotations)
+    {
+    }
     llvm::PreservedAnalyses run(llvm::Module &M,
                                 llvm::ModuleAnalysisManager &MAM);
 };
diff --git a/subprojects/helper-to-tcg/src/Pipeline.cpp b/subprojects/helper-to-tcg/src/Pipeline.cpp
index 59de572bf6..051611b0f3 100644
--- a/subprojects/helper-to-tcg/src/Pipeline.cpp
+++ b/subprojects/helper-to-tcg/src/Pipeline.cpp
@@ -184,7 +184,11 @@ int main(int argc, char **argv) {
         MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
     }
 
-    MPM.addPass(PrepareForOptPass());
+    // TODO: Get pass results via dependencies instead? Adds more boiler-plate
+    // but is correlct in LLVM-terms.
+
+    AnnotationMapTy Annotations;
+    MPM.addPass(PrepareForOptPass(Annotations));
 
     {
         FunctionPassManager FPM;
diff --git a/subprojects/helper-to-tcg/src/PrepareForOptPass/PrepareForOptPass.cpp b/subprojects/helper-to-tcg/src/PrepareForOptPass/PrepareForOptPass.cpp
index c15c0af6ea..1228ac952f 100644
--- a/subprojects/helper-to-tcg/src/PrepareForOptPass/PrepareForOptPass.cpp
+++ b/subprojects/helper-to-tcg/src/PrepareForOptPass/PrepareForOptPass.cpp
@@ -16,10 +16,15 @@
 //
 
 #include "PrepareForOptPass.hpp"
+#include "Error.hpp"
 
 #include <llvm/ADT/StringRef.h>
 #include <llvm/ADT/StringSet.h>
 #include <llvm/Demangle/Demangle.h>
+#include <llvm/IR/Constants.h>
+#include <llvm/IR/Function.h>
+#include <llvm/IR/Instruction.h>
+#include <llvm/IR/Module.h>
 #include <llvm/Support/Debug.h>
 
 #define DEBUG_TYPE "prepare-for-opt"
@@ -63,8 +68,97 @@ static void demangleFunctionNames(Module &M) {
     }
 }
 
+static Error parseAnnotationStr(Annotations &Ann, StringRef Str,
+                                size_t NumArgs) {
+    Str = Str.trim();
+
+    // Function annotations
+    if (Str.consume_front("helper-to-tcg")) {
+        Ann.set(FunctionAnnotation::HelperToTcg);
+        return Error::success();
+    } else if (Str.consume_front("returns-immediate")) {
+        Ann.set(FunctionAnnotation::ReturnsImmediate);
+        return Error::success();
+    }
+
+    // Argument annotations
+    ArgumentAnnotation AA;
+    if (Str.consume_front("immediate")) {
+        AA = ArgumentAnnotation::Immediate;
+    } else if (Str.consume_front("ptr-to-offset")) {
+        AA = ArgumentAnnotation::PtrToOffset;
+    } else {
+        return mkError("Unknown annotation");
+    }
+
+    // An argument annotation looks like
+    //
+    //  "immediate: 0, 1, 2",
+    //
+    // parse the comma separated list of argument indices.
+    if (!Str.consume_front(":")) {
+        return mkError("Expected \":\"");
+    }
+    Str = Str.ltrim(' ');
+    do {
+        Str = Str.ltrim(' ');
+        size_t I = 0;
+        Str.consumeInteger(10, I);
+        if (I >= NumArgs) {
+            return mkError("Annotation has out of bounds argument index");
+        }
+        Ann.set(I, AA);
+    } while (Str.consume_front(","));
+
+    return Error::success();
+}
+
+static void collectAnnotations(Module &M, AnnotationMapTy &ResultAnnotations) {
+    // cast over dyn_cast is being used here to
+    // assert that the structure of
+    //
+    //     llvm.global.annotation
+    //
+    // is what we expect.
+
+    GlobalVariable *GA = M.getGlobalVariable("llvm.global.annotations");
+    if (!GA) {
+        return;
+    }
+
+    // Get the metadata which is stored in the first op
+    auto *CA = cast<ConstantArray>(GA->getOperand(0));
+    // Loop over metadata
+    for (Value *CAOp : CA->operands()) {
+        auto *Struct = cast<ConstantStruct>(CAOp);
+        assert(Struct->getNumOperands() >= 2);
+
+        Function *F = cast<Function>(Struct->getOperand(0));
+        ConstantDataArray *AnnData =
+            cast<ConstantDataArray>(Struct->getOperand(1)->getOperand(0));
+
+        StringRef AnnStr = AnnData->getAsString();
+        AnnStr = AnnStr.substr(0, AnnStr.size() - 1);
+        Annotations Ann = ResultAnnotations[F];
+        if (auto Err = parseAnnotationStr(Ann, AnnStr, F->arg_size()); Err) {
+            errs() << "Failed to parse annotation: \"" << Err
+                   << "\" for function " << F->getName() << "\n";
+            continue;
+        }
+        ResultAnnotations[F] = Ann;
+    }
+
+    LLVM_DEBUG({
+        for (auto &P : ResultAnnotations) {
+            dbgs() << "Annotations for " << P.first->getName() << "\n";
+            P.second.dump(dbgs());
+        }
+    });
+}
+
 PreservedAnalyses PrepareForOptPass::run(Module &M,
                                          ModuleAnalysisManager &MAM) {
     demangleFunctionNames(M);
+    collectAnnotations(M, ResultAnnotations);
     return PreservedAnalyses::none();
 }
-- 
2.52.0