LLVM ModulePass开发需求：逐次删除两个测试函数仅保留单个调用

核心目标 #

你需要的是一个可以选择性保留test1/test2/test3其中一个函数的调用与定义，同时移除另外两个测试函数的ModulePass。我们可以通过命令行参数控制要保留的目标函数，这样每次运行Pass时就能生成仅包含单个测试函数调用的模块。

完整代码实现 #

#include "llvm/IR/Module.h"
#include "llvm/Pass.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Instructions.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/raw_ostream.h"
#include <algorithm>
#include <vector>

using namespace llvm;

// 命令行参数：指定要保留的测试函数名（默认保留test1）
static cl::opt<std::string> KeepTestFunc("keep-test-func", cl::desc("Specify the test function to keep"), cl::init("test1"));

namespace {
struct KeepSingleTestFuncPass : public ModulePass {
  static char ID;
  KeepSingleTestFuncPass() : ModulePass(ID) {}

  bool runOnModule(Module &M) override {
    // 1. 定义目标测试函数名集合，用于后续匹配
    std::vector<std::string> targetTests = {"test1", "test2", "test3"};
    std::string keepFunc = KeepTestFunc;

    // 校验输入的函数名是否合法
    if (std::find(targetTests.begin(), targetTests.end(), keepFunc) == targetTests.end()) {
      errs() << "Error: Invalid test function name. Must be test1/test2/test3\n";
      return false;
    }

    // 2. 定位main函数
    Function *MainFunc = M.getFunction("main");
    if (!MainFunc) {
      errs() << "Error: No main function found in the module\n";
      return false;
    }

    // 3. 遍历main函数，标记需要移除的非保留测试函数调用
    std::vector<CallInst*> callsToRemove;
    for (BasicBlock &BB : *MainFunc) {
      for (Instruction &I : BB) {
        if (CallInst *CI = dyn_cast<CallInst>(&I)) {
          Function *Callee = CI->getCalledFunction();
          if (Callee && std::find(targetTests.begin(), targetTests.end(), Callee->getName()) != targetTests.end()) {
            if (Callee->getName() != keepFunc) {
              callsToRemove.push_back(CI);
            }
          }
        }
      }
    }

    // 统一移除标记的调用指令（避免遍历中删除导致迭代器失效）
    for (CallInst *CI : callsToRemove) {
      CI->eraseFromParent();
    }

    // 4. 遍历模块函数，标记需要删除的非保留测试函数定义
    std::vector<Function*> funcsToRemove;
    for (Module::iterator F = M.begin(), E = M.end(); F != E; ++F) {
      std::string funcName = F->getName();
      if (std::find(targetTests.begin(), targetTests.end(), funcName) != targetTests.end() && funcName != keepFunc) {
        funcsToRemove.push_back(&*F);
      }
    }

    // 统一删除标记的函数定义
    for (Function *F : funcsToRemove) {
      F->eraseFromParent();
    }

    return true; // 标记模块已被修改
  }
};
} // namespace

char KeepSingleTestFuncPass::ID = 0;
static RegisterPass<KeepSingleTestFuncPass> X("keep-single-test", "Keep only one test function call and definition", false, false);

关键步骤说明 #

• 命令行参数控制：通过-keep-test-func参数指定保留的函数，比如运行opt -load ./libKeepSingleTestFuncPass.so -keep-single-test -keep-test-func=test2 input.ll -o output.ll，就能生成仅保留test2调用与定义的IR文件。
• 避免迭代器失效：不管是处理调用指令还是函数定义，都先把要删除的对象收集到容器中，再统一删除，防止遍历过程中直接修改容器导致迭代器异常。
• 合法性校验：先检查输入的保留函数名是否在目标集合内，避免无效操作输出错误信息。

使用提示 #

• 编译Pass时要确保CMakeLists.txt正确链接LLVM相关依赖库；
• 若测试函数存在其他依赖（比如被main之外的函数调用），需要额外处理这些依赖关系，当前代码默认只有main调用测试函数。

内容的提问来源于stack exchange，提问作者mikasa