Make an early pass over all functions in the input module and filter out
functions with:
  1. Invalid return type;
  2. No helper-to-tcg annotation.

Signed-off-by: Anton Johansson <anjo@rev.ng>
---
 .../helper-to-tcg/include/PrepareForOptPass.h |  7 +-
 .../PrepareForOptPass/PrepareForOptPass.cpp   | 93 +++++++++++++++++++
 .../helper-to-tcg/pipeline/Pipeline.cpp       |  7 +-
 3 files changed, 104 insertions(+), 3 deletions(-)

diff --git a/subprojects/helper-to-tcg/include/PrepareForOptPass.h b/subprojects/helper-to-tcg/include/PrepareForOptPass.h
index 5f9c059b97..8615625f09 100644
--- a/subprojects/helper-to-tcg/include/PrepareForOptPass.h
+++ b/subprojects/helper-to-tcg/include/PrepareForOptPass.h
@@ -29,9 +29,12 @@
 
 class PrepareForOptPass : public llvm::PassInfoMixin<PrepareForOptPass> {
     AnnotationMapTy &ResultAnnotations;
+    bool TranslateAllHelpers;
 public:
-    PrepareForOptPass(AnnotationMapTy &ResultAnnotations)
-        : ResultAnnotations(ResultAnnotations)
+    PrepareForOptPass(AnnotationMapTy &ResultAnnotations,
+                      bool TranslateAllHelpers)
+        : ResultAnnotations(ResultAnnotations),
+          TranslateAllHelpers(TranslateAllHelpers)
     {
     }
     llvm::PreservedAnalyses run(llvm::Module &M,
diff --git a/subprojects/helper-to-tcg/passes/PrepareForOptPass/PrepareForOptPass.cpp b/subprojects/helper-to-tcg/passes/PrepareForOptPass/PrepareForOptPass.cpp
index 9f1d4df102..22509008c8 100644
--- a/subprojects/helper-to-tcg/passes/PrepareForOptPass/PrepareForOptPass.cpp
+++ b/subprojects/helper-to-tcg/passes/PrepareForOptPass/PrepareForOptPass.cpp
@@ -17,12 +17,17 @@
 
 #include <PrepareForOptPass.h>
 #include <Error.h>
+#include <FunctionAnnotation.h>
 
 #include <llvm/IR/Constants.h>
 #include <llvm/IR/Function.h>
 #include <llvm/IR/Instruction.h>
+#include <llvm/IR/Instructions.h>
 #include <llvm/IR/Module.h>
 
+#include <queue>
+#include <set>
+
 using namespace llvm;
 
 static Expected<Annotation> parseAnnotationStr(StringRef Str,
@@ -105,8 +110,96 @@ static void collectAnnotations(Module &M, AnnotationMapTy &ResultAnnotations)
     }
 }
 
+inline bool hasValidReturnTy(const Module &M, const Function *F)
+{
+    Type *RetTy = F->getReturnType();
+    return RetTy == Type::getVoidTy(F->getContext()) ||
+           RetTy == Type::getInt8Ty(M.getContext()) ||
+           RetTy == Type::getInt16Ty(M.getContext()) ||
+           RetTy == Type::getInt32Ty(M.getContext()) ||
+           RetTy == Type::getInt64Ty(M.getContext());
+}
+
+// Functions that should be removed:
+//   - No helper-to-tcg annotation (if TranslateAllHelpers == false);
+//   - Invalid (non-integer/void) return type
+static bool shouldRemoveFunction(const Module &M, const Function &F,
+                                 const AnnotationMapTy &AnnotationMap,
+                                 bool TranslateAllHelpers)
+{
+    if (F.isDeclaration()) {
+        return false;
+    }
+
+    if (!hasValidReturnTy(M, &F)) {
+        return true;
+    }
+
+    auto hasCorrectAnnotation = [](const Annotation &Ann) {
+        return Ann.Kind == AnnotationKind::HelperToTcg;
+    };
+
+    std::queue<const Function *> Worklist;
+    std::set<const Function *> Visited;
+    Worklist.push(&F);
+    while (!Worklist.empty()) {
+        const Function *F = Worklist.front();
+        Worklist.pop();
+        if (F->isDeclaration() or Visited.find(F) != Visited.end()) {
+            continue;
+        }
+        Visited.insert(F);
+
+        // Check for llvm-to-tcg annotation
+        if (TranslateAllHelpers and F->getName().startswith("helper_")) {
+            return false;
+        } else {
+            auto It = AnnotationMap.find(F);
+            if (It != AnnotationMap.end()) {
+                const auto &AnnotationVec = It->second;
+                auto Res = find_if(AnnotationVec, hasCorrectAnnotation);
+                if (Res != AnnotationVec.end()) {
+                    return false;
+                }
+            }
+        }
+
+        // Push functions that call F to the worklist, this way we retain
+        // functions that are being called by functions with the llvm-to-tcg
+        // annotation.
+        for (const User *U : F->users()) {
+            auto Call = dyn_cast<CallInst>(U);
+            if (!Call) {
+                continue;
+            }
+            const Function *ParentF = Call->getParent()->getParent();
+            Worklist.push(ParentF);
+        }
+    }
+
+    return true;
+}
+
+static void cullUnusedFunctions(Module &M, AnnotationMapTy &Annotations,
+                                bool TranslateAllHelpers)
+{
+    SmallVector<Function *, 16> FunctionsToRemove;
+    for (auto &F : M) {
+        if (shouldRemoveFunction(M, F, Annotations, TranslateAllHelpers)) {
+            FunctionsToRemove.push_back(&F);
+        }
+    }
+
+    for (Function *F : FunctionsToRemove) {
+        Annotations.erase(F);
+        F->setComdat(nullptr);
+        F->deleteBody();
+    }
+}
+
 PreservedAnalyses PrepareForOptPass::run(Module &M, ModuleAnalysisManager &MAM)
 {
     collectAnnotations(M, ResultAnnotations);
+    cullUnusedFunctions(M, ResultAnnotations, TranslateAllHelpers);
     return PreservedAnalyses::none();
 }
diff --git a/subprojects/helper-to-tcg/pipeline/Pipeline.cpp b/subprojects/helper-to-tcg/pipeline/Pipeline.cpp
index 3b9493bc73..dde3641ab3 100644
--- a/subprojects/helper-to-tcg/pipeline/Pipeline.cpp
+++ b/subprojects/helper-to-tcg/pipeline/Pipeline.cpp
@@ -47,6 +47,11 @@ cl::OptionCategory Cat("helper-to-tcg Options");
 cl::opt<std::string> InputFile(cl::Positional, cl::desc("[input LLVM module]"),
                                cl::cat(Cat));
 
+// Options for PrepareForOptPass
+cl::opt<bool> TranslateAllHelpers(
+    "translate-all-helpers", cl::init(false),
+    cl::desc("Translate all functions starting with helper_*"), cl::cat(Cat));
+
 // Define a TargetTransformInfo (TTI) subclass, this allows for overriding
 // common per-llvm-target information expected by other LLVM passes, such
 // as the width of the largest scalar/vector registers.  Needed for consistent
@@ -175,7 +180,7 @@ int main(int argc, char **argv)
     }
 
     AnnotationMapTy Annotations;
-    MPM.addPass(PrepareForOptPass(Annotations));
+    MPM.addPass(PrepareForOptPass(Annotations, TranslateAllHelpers));
 
     {
         FunctionPassManager FPM;
-- 
2.45.2