Example of an optimization pass

#include "llvm/Analysis/ConstantFolding.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PassManager.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/Passes/PassPlugin.h"
#include "llvm/Support/raw_ostream.h"

#include <set>

using namespace llvm;

namespace {
struct DTigerOpts : public PassInfoMixin<DTigerOpts> {
  // Return the function __not if it is used in the module
  Function *find_not_function(Module &M) {
    for (auto &fun : M.getFunctionList())
      if (fun.getName() == "__not") {
        errs() << "Found function __not\n";
        return &fun;
      }
    return nullptr;
  }

  // Given a use of function __not, this function checks if the use
  // corresponds to a call with a constant integer parameter, such as
  // __not(1). In that case, it directly evaluates the expected result
  // and removes the call.
  bool replace_constant_call(User *use, Module &M, const DataLayout &DL) {
    // Retrieve the CallSite of the use
    CallInst *call = dyn_cast<CallInst>(use);
    // Check if it is a direct call instruction
    if (!call)
      return false;
    errs() << "Found call to function __not\n";

    // Since the typechecker has run, we assume that there is exactly one
    // argument on the call to __not
    assert(call->arg_begin() != nullptr);
    Value *arg = call->arg_begin()->get();

    // If the argument is not a constant integer skip to next call
    ConstantInt *c;
    if (!(c = dyn_cast<ConstantInt>(arg)))
      return false;

    // Evaluate the constant call result
    char result = (c->equalsInt(0) ? 1 : 0);

    // Record all uses of the call result
    std::set<Instruction *> WorkList;
    for (auto u : call->users()) {
      if (auto i = dyn_cast<Instruction>(u)) {
        WorkList.insert(i);
      }
    }

    // Replace the call by the constant result
    call->replaceAllUsesWith(
        ConstantInt::getSigned(Type::getInt32Ty(M.getContext()), result));
    call->eraseFromParent();

    // Constant propagate to all users of the call
    while (!WorkList.empty()) {
      Instruction *i = *WorkList.begin();
      WorkList.erase(WorkList.begin());
      if (Constant *c = ConstantFoldInstruction(i, DL)) {
        for (auto u : i->users())
          WorkList.insert(cast<Instruction>(u));
        i->replaceAllUsesWith(c);
        i->eraseFromParent();
      }
    }

    return true;
  }

  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM) {
    const DataLayout &DL = M.getDataLayout();

    bool modified = false, did_replace;

    // Check if the function __not is used in the module
    Function *not_fun = find_not_function(M);
    if (not_fun == nullptr)
      return PreservedAnalyses::all();

    // If __not is referenced, check every use for constant calls
    // that we could optimize
    do {
      errs() << "Entering replace loop\n";
      did_replace = false;
      SmallVector<User *, 8> Users(not_fun->users());
      for (const auto user : Users) {
        if (replace_constant_call(user, M, DL)) {
          did_replace = true;
          errs() << "Found and replaced constant call\n";
        }
      }
      modified |= did_replace;
      // If we did replace some calls to __not; do a further pass
      // to enable successive folds such as not(not(0)) -> 0
    } while (did_replace);

    return modified ? PreservedAnalyses::none() : PreservedAnalyses::all();
  }
};
} // namespace

extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
llvmGetPassPluginInfo() {
  return {LLVM_PLUGIN_API_VERSION, "dtigeropts", LLVM_VERSION_STRING,
          [](PassBuilder &PB) {
            PB.registerPipelineParsingCallback(
                [](StringRef Name, ModulePassManager &MPM,
                   ArrayRef<PassBuilder::PipelineElement>) {
                  if (Name == "dtigeropts") {
                    MPM.addPass(DTigerOpts());
                    return true;
                  }
                  return false;
                });
          }};
}