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 a list of
annotations is exposed.

Signed-off-by: Anton Johansson <anjo@rev.ng>
---
 .../include/FunctionAnnotation.h              | 54 ++++++++++++
 .../helper-to-tcg/include/PrepareForOptPass.h |  7 +-
 .../PrepareForOptPass/PrepareForOptPass.cpp   | 87 +++++++++++++++++++
 .../helper-to-tcg/pipeline/Pipeline.cpp       |  3 +-
 4 files changed, 149 insertions(+), 2 deletions(-)
 create mode 100644 subprojects/helper-to-tcg/include/FunctionAnnotation.h

diff --git a/subprojects/helper-to-tcg/include/FunctionAnnotation.h b/subprojects/helper-to-tcg/include/FunctionAnnotation.h
new file mode 100644
index 0000000000..b562f7c892
--- /dev/null
+++ b/subprojects/helper-to-tcg/include/FunctionAnnotation.h
@@ -0,0 +1,54 @@
+//
+//  Copyright(c) 2024 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 <stdint.h>
+
+namespace llvm
+{
+class Function;
+}
+
+// Different kind of function annotations which control the behaviour
+// of helper-to-tcg.
+enum class AnnotationKind : uint8_t {
+    // Function should be translated
+    HelperToTcg,
+    // Declares a list of arguments as immediates
+    Immediate,
+    // Declares a list of arguments as vectors, represented by offsets into
+    // the CPU state
+    PtrToOffset,
+};
+
+// Annotation data which may be attached to a function
+struct Annotation {
+    // Indices of function arguments the annotation applies to, only
+    // used for AnnotationKind::[Immediate|PtrToOffset].
+    llvm::SmallVector<uint8_t, 4> ArgIndices;
+    AnnotationKind Kind;
+};
+
+// Map from Function * to a list of struct Annotation.  std::map is used here
+// which allocates for each mapped pair due to the value being large
+// (at least 48*3 bits).  If ArgIndices were to be stored out-of-band this could
+// be reduced, and DenseMap would be more appropriate.
+using AnnotationVectorTy = llvm::SmallVector<Annotation, 3>;
+using AnnotationMapTy = llvm::DenseMap<llvm::Function *, AnnotationVectorTy>;
diff --git a/subprojects/helper-to-tcg/include/PrepareForOptPass.h b/subprojects/helper-to-tcg/include/PrepareForOptPass.h
index d74618613f..5f9c059b97 100644
--- a/subprojects/helper-to-tcg/include/PrepareForOptPass.h
+++ b/subprojects/helper-to-tcg/include/PrepareForOptPass.h
@@ -17,6 +17,7 @@
 
 #pragma once
 
+#include "FunctionAnnotation.h"
 #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/passes/PrepareForOptPass/PrepareForOptPass.cpp b/subprojects/helper-to-tcg/passes/PrepareForOptPass/PrepareForOptPass.cpp
index 0a018494fe..9f1d4df102 100644
--- a/subprojects/helper-to-tcg/passes/PrepareForOptPass/PrepareForOptPass.cpp
+++ b/subprojects/helper-to-tcg/passes/PrepareForOptPass/PrepareForOptPass.cpp
@@ -16,10 +16,97 @@
 //
 
 #include <PrepareForOptPass.h>
+#include <Error.h>
+
+#include <llvm/IR/Constants.h>
+#include <llvm/IR/Function.h>
+#include <llvm/IR/Instruction.h>
+#include <llvm/IR/Module.h>
 
 using namespace llvm;
 
+static Expected<Annotation> parseAnnotationStr(StringRef Str,
+                                               uint32_t num_function_args)
+{
+    Annotation Ann;
+
+    Str = Str.trim();
+
+    if (Str.consume_front("helper-to-tcg")) {
+        Ann.Kind = AnnotationKind::HelperToTcg;
+        // Early return, no additional info to parse from annotation string
+        return Ann;
+    } else if (Str.consume_front("immediate")) {
+        Ann.Kind = AnnotationKind::Immediate;
+    } else if (Str.consume_front("ptr-to-offset")) {
+        Ann.Kind = AnnotationKind::PtrToOffset;
+    } else {
+        return mkError("Unknown annotation");
+    }
+
+    // Parse comma separated list of argument indices
+
+    if (!Str.consume_front(":")) {
+        return mkError("Expected \":\"");
+    }
+
+    Str = Str.ltrim(' ');
+    do {
+        Str = Str.ltrim(' ');
+        uint32_t i = 0;
+        Str.consumeInteger(10, i);
+        if (i >= num_function_args) {
+            return mkError("Annotation has out of bounds argument index");
+        }
+        Ann.ArgIndices.push_back(i);
+    } while (Str.consume_front(","));
+
+    return Ann;
+}
+
+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);
+        Constant *UseOfF = Struct->getOperand(0);
+        if (isa<UndefValue>(UseOfF)) {
+            continue;
+        }
+        auto *F = cast<Function>(UseOfF->getOperand(0));
+        auto *AnnVar =
+            cast<GlobalVariable>(Struct->getOperand(1)->getOperand(0));
+        auto *AnnData = cast<ConstantDataArray>(AnnVar->getOperand(0));
+
+        StringRef AnnStr = AnnData->getAsString();
+        AnnStr = AnnStr.substr(0, AnnStr.size() - 1);
+        Expected<Annotation> Ann = parseAnnotationStr(AnnStr, F->arg_size());
+        if (!Ann) {
+            dbgs() << "Failed to parse annotation: \"" << Ann.takeError()
+                   << "\" for function " << F->getName() << "\n";
+            continue;
+        }
+        ResultAnnotations[F].push_back(*Ann);
+    }
+}
+
 PreservedAnalyses PrepareForOptPass::run(Module &M, ModuleAnalysisManager &MAM)
 {
+    collectAnnotations(M, ResultAnnotations);
     return PreservedAnalyses::none();
 }
diff --git a/subprojects/helper-to-tcg/pipeline/Pipeline.cpp b/subprojects/helper-to-tcg/pipeline/Pipeline.cpp
index fad335f4a9..3b9493bc73 100644
--- a/subprojects/helper-to-tcg/pipeline/Pipeline.cpp
+++ b/subprojects/helper-to-tcg/pipeline/Pipeline.cpp
@@ -174,7 +174,8 @@ int main(int argc, char **argv)
         MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
     }
 
-    MPM.addPass(PrepareForOptPass());
+    AnnotationMapTy Annotations;
+    MPM.addPass(PrepareForOptPass(Annotations));
 
     {
         FunctionPassManager FPM;
-- 
2.45.2