LLVM Tutor Note

wget -O - https://apt.llvm.org/llvm-snapshot.gpg.key | sudo apt-key add -
sudo apt-add-repository "deb http://apt.llvm.org/jammy/ llvm-toolchain-jammy-17 main"
sudo apt update
sudo apt install -y llvm-17 llvm-17-dev llvm-17-tools clang-17

export LLVM_DIR=/usr/lib/llvm-17

cd /path/to/llvm-tutor
mkdir build && cd build
cmake -DLT_LLVM_INSTALL_DIR=$LLVM_DIR ../HelloWorld/
make

clang-17 -O1 -S -emit-llvm ../inputs/input_for_hello.c -o input_for_hello.ll

int foo(int a) {
	return a * 2;
}

int bar(int a, int b) {
	return (a + foo(b) * 2);
}

int fez(int a, int b, int c) {
	return (a + bar(a, b) * 2 + c * 3);
}

int main(int argc, char *argv[]) {
	int a = 123;
	int ret = 0;

	ret += foo(a);
	ret += bar(a, ret);
	ret += fez(a, ret, 123);

	return ret;
}

; ModuleID = '../inputs/input_for_hello.c'
source_filename = "../inputs/input_for_hello.c"
target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
target triple = "x86_64-pc-linux-gnu"

; Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none) uwtable
define dso_local i32 @foo(i32 noundef %0) local_unnamed_addr #0 {
  %2 = shl nsw i32 %0, 1
  ret i32 %2
}

; Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none) uwtable
define dso_local i32 @bar(i32 noundef %0, i32 noundef %1) local_unnamed_addr #0 {
  %3 = shl i32 %1, 2
  %4 = add nsw i32 %3, %0
  ret i32 %4
}

; Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none) uwtable
define dso_local i32 @fez(i32 noundef %0, i32 noundef %1, i32 noundef %2) local_unnamed_addr #0 {
  %4 = shl i32 %1, 3
  %5 = mul nsw i32 %2, 3
  %6 = mul i32 %0, 3
  %7 = add i32 %6, %4
  %8 = add nsw i32 %7, %5
  ret i32 %8
}

; Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none) uwtable
define dso_local i32 @main(i32 noundef %0, ptr nocapture noundef readnone %1) local_unnamed_addr #0 {
  ret i32 12915
}

attributes #0 = { mustprogress nofree norecurse nosync nounwind willreturn memory(none) uwtable "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cmov,+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }

!llvm.module.flags = !{!0, !1, !2, !3}
!llvm.ident = !{!4}

!0 = !{i32 1, !"wchar_size", i32 4}
!1 = !{i32 8, !"PIC Level", i32 2}
!2 = !{i32 7, !"PIE Level", i32 2}
!3 = !{i32 7, !"uwtable", i32 2}
!4 = !{!"Ubuntu clang version 17.0.6 (++20231209124227+6009708b4367-1~exp1~20231209124336.77)"}

$ opt-17 -load-pass-plugin ./libHelloWorld.so -passes=hello-world -disable-output input_for_hello.ll
(llvm-tutor) Hello from: foo
(llvm-tutor)   number of arguments: 1
(llvm-tutor) Hello from: bar
(llvm-tutor)   number of arguments: 2
(llvm-tutor) Hello from: fez
(llvm-tutor)   number of arguments: 3
(llvm-tutor) Hello from: main
(llvm-tutor)   number of arguments: 2

struct HelloWorld : PassInfoMixin<HelloWorld> {
	// Main entry point, takes IR unit to run the pass on (&F) and the
	// corresponding pass manager (to be queried if need be)
	PreservedAnalyses run(Function &F, FunctionAnalysisManager &) {
		visitor(F);
		return PreservedAnalyses::all();
	}

	// Without isRequired returning true, this pass will be skipped for functions
	// decorated with the optnone LLVM attribute. Note that clang -O0 decorates
	// all functions with optnone.
	static bool isRequired() { return true; }
};

void visitor(Function &F) {
	errs() << "(llvm-tutor) Hello from: "<< F.getName() << "\n";
	errs() << "(llvm-tutor)   number of arguments: " << F.arg_size() << "\n";
}

llvm::PassPluginLibraryInfo getHelloWorldPluginInfo() {
  return {LLVM_PLUGIN_API_VERSION, "HelloWorld", LLVM_VERSION_STRING,
          [](PassBuilder &PB) {
            PB.registerPipelineParsingCallback(
                [](StringRef Name, FunctionPassManager &FPM,
                   ArrayRef<PassBuilder::PipelineElement>) {
                  if (Name == "hello-world") {
                    FPM.addPass(HelloWorld());
                    return true;
                  }
                  return false;
                });
          }};
}

// This is the core interface for pass plugins. It guarantees that 'opt' will
// be able to recognize HelloWorld when added to the pass pipeline on the
// command line, i.e. via '-passes=hello-world'
extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
llvmGetPassPluginInfo() {
  return getHelloWorldPluginInfo();
}

/usr/bin/c++ -fPIC -shared -Wl,-soname,libHelloWorld.so -o libHelloWorld.so CMakeFiles/HelloWorld.dir/HelloWorld.cpp.o

rm -r build
mkdir build && cd build
cmake -DLT_LLVM_INSTALL_DIR=$LLVM_DIR ../
make

sudo apt install -y libedit-dev
sudo apt install -y libzstd-dev

clang-17 -emit-llvm -c ../inputs/input_for_cc.c -o input_for_cc.bc
opt-17 -load-pass-plugin ./lib/libOpcodeCounter.so -passes="print<opcode-counter>" -disable-output input_for_cc.bc

void foo() { }
void bar() {foo(); }
void fez() {bar(); }

int main() {
	foo();
	bar();
	fez();

	int ii = 0;
	for (ii = 0; ii < 10; ii++)
		foo();

	return 0;
}

opt-17 -load-pass-plugin ./lib/libOpcodeCounter.so --passes='default<O1>' -disable-output input_for_cc.bc

llvm::PassPluginLibraryInfo getOpcodeCounterPluginInfo() {
  return {
    LLVM_PLUGIN_API_VERSION, "OpcodeCounter", LLVM_VERSION_STRING,
        [](PassBuilder &PB) {
			// #1 REGISTRATION FOR "opt -passes=print<opcode-counter>"
			// Register OpcodeCounterPrinter so that it can be used when
			// specifying pass pipelines with `-passes=`.
			PB.registerPipelineParsingCallback(
				[&](StringRef Name, FunctionPassManager &FPM,
					ArrayRef<PassBuilder::PipelineElement>) {
					if (Name == "print<opcode-counter>") {
					FPM.addPass(OpcodeCounterPrinter(llvm::errs()));
					return true;
					}
					return false;
				});
			// #2 REGISTRATION FOR "-O{1|2|3|s}"
			// Register OpcodeCounterPrinter as a step of an existing pipeline.
			// The insertion point is specified by using the
			// 'registerVectorizerStartEPCallback' callback. To be more precise,
			// using this callback means that OpcodeCounterPrinter will be called
			// whenever the vectoriser is used (i.e. when using '-O{1|2|3|s}'.
			PB.registerVectorizerStartEPCallback(
				[](llvm::FunctionPassManager &PM,
					llvm::OptimizationLevel Level) {
					PM.addPass(OpcodeCounterPrinter(llvm::errs()));
				});
			// #3 REGISTRATION FOR "FAM.getResult<OpcodeCounter>(Func)"
			// Register OpcodeCounter as an analysis pass. This is required so that
			// OpcodeCounterPrinter (or any other pass) can request the results
			// of OpcodeCounter.
			PB.registerAnalysisRegistrationCallback(
				[](FunctionAnalysisManager &FAM) {
					FAM.registerPass([&] { return OpcodeCounter(); });
				});
		}
	};
}

extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
llvmGetPassPluginInfo() {
	return getOpcodeCounterPluginInfo();
}

OpcodeCounter::Result OpcodeCounter::run(llvm::Function &Func,
                                         llvm::FunctionAnalysisManager &) {
	return generateOpcodeMap(Func);
}

PreservedAnalyses OpcodeCounterPrinter::run(Function &Func,
                                            FunctionAnalysisManager &FAM) {
	auto &OpcodeMap = FAM.getResult<OpcodeCounter>(Func);

	// In the legacy PM, the following string is printed automatically by the
	// pass manager. For the sake of consistency, we're adding this here so that
	// it's also printed when using the new PM.
	OS << "Printing analysis 'OpcodeCounter Pass' for function '"
		<< Func.getName() << "':\n";

	printOpcodeCounterResult(OS, OpcodeMap);
	return PreservedAnalyses::all();
}

OpcodeCounter::Result OpcodeCounter::generateOpcodeMap(llvm::Function &Func) {
	OpcodeCounter::Result OpcodeMap;

	for (auto &BB : Func) {
		for (auto &Inst : BB) {
		StringRef Name = Inst.getOpcodeName();

		if (OpcodeMap.find(Name) == OpcodeMap.end()) {
			OpcodeMap[Inst.getOpcodeName()] = 1;
		} else {
			OpcodeMap[Inst.getOpcodeName()]++;
		}
		}
	}

	return OpcodeMap;
}

printf("(llvm-tutor) Hello from: %s\n(llvm-tutor)   number of arguments: %d\n", FuncName, FuncNumArgs)

clang-17 -O0 -emit-llvm -c ../inputs/input_for_hello.c -o input_for_hello.bc
opt-17 -load-pass-plugin ./lib/libInjectFuncCall.so --passes="inject-func-call" input_for_hello.bc -o instrumented.bin

$ lli-17 instrumented.bin
(llvm-tutor) Hello from: main
(llvm-tutor)   number of arguments: 2
(llvm-tutor) Hello from: foo
(llvm-tutor)   number of arguments: 1
(llvm-tutor) Hello from: bar
(llvm-tutor)   number of arguments: 2
(llvm-tutor) Hello from: foo
(llvm-tutor)   number of arguments: 1
(llvm-tutor) Hello from: fez
(llvm-tutor)   number of arguments: 3
(llvm-tutor) Hello from: bar
(llvm-tutor)   number of arguments: 2
(llvm-tutor) Hello from: foo
(llvm-tutor)   number of arguments: 1

$ diff input_for_hello.ll instrumented.bin.ll
1c1
- ; ModuleID = 'input_for_hello.bc'
---
+ ; ModuleID = './instrumented.bin'
5a6,11
+ @PrintfFormatStr = global [68 x i8] c"(llvm-tutor) Hello from: %s\0A(llvm-tutor)   number of arguments: %d\0A\00"
+ @0 = private unnamed_addr constant [4 x i8] c"foo\00", align 1
+ @1 = private unnamed_addr constant [4 x i8] c"bar\00", align 1
+ @2 = private unnamed_addr constant [4 x i8] c"fez\00", align 1
+ @3 = private unnamed_addr constant [5 x i8] c"main\00", align 1
+
8,12c14,19
-   %2 = alloca i32, align 4
-   store i32 %0, ptr %2, align 4
-   %3 = load i32, ptr %2, align 4
-   %4 = mul nsw i32 %3, 2
-   ret i32 %4
---
+   %2 = call i32 (ptr, ...) @printf(ptr @PrintfFormatStr, ptr @0, i32 1)
+   %3 = alloca i32, align 4
+   store i32 %0, ptr %3, align 4
+   %4 = load i32, ptr %3, align 4
+   %5 = mul nsw i32 %4, 2
+   ret i32 %5
...

PreservedAnalyses InjectFuncCall::run(llvm::Module &M,
                                       llvm::ModuleAnalysisManager &) {
	bool Changed =  runOnModule(M);

	return (Changed ? llvm::PreservedAnalyses::none()
					: llvm::PreservedAnalyses::all());
}

bool InjectFuncCall::runOnModule(Module &M) {
	bool InsertedAtLeastOnePrintf = false;

	auto &CTX = M.getContext();
	PointerType *PrintfArgTy = PointerType::getUnqual(Type::getInt8Ty(CTX));

	// STEP 1: Inject the declaration of printf
	// ----------------------------------------
	// Create (or _get_ in cases where it's already available) the following
	// declaration in the IR module:
	//    declare i32 @printf(i8*, ...)
	// It corresponds to the following C declaration:
	//    int printf(char *, ...)
	FunctionType *PrintfTy = FunctionType::get(
		IntegerType::getInt32Ty(CTX),
		PrintfArgTy,
		/*IsVarArgs=*/true);

	FunctionCallee Printf = M.getOrInsertFunction("printf", PrintfTy);

	// Set attributes as per inferLibFuncAttributes in BuildLibCalls.cpp
	Function *PrintfF = dyn_cast<Function>(Printf.getCallee());
	PrintfF->setDoesNotThrow();
	PrintfF->addParamAttr(0, Attribute::NoCapture);
	PrintfF->addParamAttr(0, Attribute::ReadOnly);


	// STEP 2: Inject a global variable that will hold the printf format string
	// ------------------------------------------------------------------------
	llvm::Constant *PrintfFormatStr = llvm::ConstantDataArray::getString(
		CTX, "(llvm-tutor) Hello from: %s\n(llvm-tutor)   number of arguments: %d\n");

	Constant *PrintfFormatStrVar =
		M.getOrInsertGlobal("PrintfFormatStr", PrintfFormatStr->getType());
	dyn_cast<GlobalVariable>(PrintfFormatStrVar)->setInitializer(PrintfFormatStr);

	// STEP 3: For each function in the module, inject a call to printf
	// ----------------------------------------------------------------
	for (auto &F : M) {
		if (F.isDeclaration())
		continue;

		// Get an IR builder. Sets the insertion point to the top of the function
		IRBuilder<> Builder(&*F.getEntryBlock().getFirstInsertionPt());

		// Inject a global variable that contains the function name
		auto FuncName = Builder.CreateGlobalStringPtr(F.getName());

		// Printf requires i8*, but PrintfFormatStrVar is an array: [n x i8]. Add
		// a cast: [n x i8] -> i8*
		llvm::Value *FormatStrPtr =
			Builder.CreatePointerCast(PrintfFormatStrVar, PrintfArgTy, "formatStr");

		// The following is visible only if you pass -debug on the command line
		// *and* you have an assert build.
		LLVM_DEBUG(dbgs() << " Injecting call to printf inside " << F.getName()
						<< "\n");

		// Finally, inject a call to printf
		Builder.CreateCall(
			Printf, {FormatStrPtr, FuncName, Builder.getInt32(F.arg_size())});

		InsertedAtLeastOnePrintf = true;
	}

	return InsertedAtLeastOnePrintf;
}

for (ii = 0; ii < 10; ii++)
	foo();

opt-17 -load-pass-plugin ./lib/libStaticCallCounter.so -passes="print<static-cc>" -disable-output input_for_cc.bc

llvm::PassPluginLibraryInfo getStaticCallCounterPluginInfo() {
	return {LLVM_PLUGIN_API_VERSION, "static-cc", LLVM_VERSION_STRING,
			[](PassBuilder &PB) {
				// #1 REGISTRATION FOR "opt -passes=print<static-cc>"
				PB.registerPipelineParsingCallback(
					[&](StringRef Name, ModulePassManager &MPM,
						ArrayRef<PassBuilder::PipelineElement>) {
					if (Name == "print<static-cc>") {
						MPM.addPass(StaticCallCounterPrinter(llvm::errs()));
						return true;
					}
					return false;
					});
				// #2 REGISTRATION FOR "MAM.getResult<StaticCallCounter>(Module)"
				PB.registerAnalysisRegistrationCallback(
					[](ModuleAnalysisManager &MAM) {
					MAM.registerPass([&] { return StaticCallCounter(); });
					});
			}};
};

StaticCallCounter::Result StaticCallCounter::runOnModule(Module &M) {
	llvm::MapVector<const llvm::Function *, unsigned> Res;

	for (auto &Func : M) {
		for (auto &BB : Func) {
		for (auto &Ins : BB) {

			// If this is a call instruction then CB will be not null.
			auto *CB = dyn_cast<CallBase>(&Ins);
			if (nullptr == CB) {
			continue;
			}

			// If CB is a direct function call then DirectInvoc will be not null.
			auto DirectInvoc = CB->getCalledFunction();
			if (nullptr == DirectInvoc) {
			continue;
			}

			// We have a direct function call - update the count for the function
			// being called.
			auto CallCount = Res.find(DirectInvoc);
			if (Res.end() == CallCount) {
			CallCount = Res.insert(std::make_pair(DirectInvoc, 0)).first;
			}
			++CallCount->second;
		}
		}
	}

	return Res;
}

StaticCallCounter::Result
StaticCallCounter::run(llvm::Module &M, llvm::ModuleAnalysisManager &) {
	return runOnModule(M);
}

opt-17 -load-pass-plugin ./lib/libDynamicCallCounter.so -passes="dynamic-cc" input_for_cc.bc -o instrumented.bin
lli-17 instrumented.bin

bool Instrumented = false;

// Function name <--> IR variable that holds the call counter
llvm::StringMap<Constant *> CallCounterMap;
// Function name <--> IR variable that holds the function name
llvm::StringMap<Constant *> FuncNameMap;

auto &CTX = M.getContext();

// STEP 1: For each function in the module, inject a call-counting code
// --------------------------------------------------------------------
for (auto &F : M) {
if (F.isDeclaration())
  continue;

// Get an IR builder. Sets the insertion point to the top of the function
IRBuilder<> Builder(&*F.getEntryBlock().getFirstInsertionPt());

// Create a global variable to count the calls to this function
std::string CounterName = "CounterFor_" + std::string(F.getName());
Constant *Var = CreateGlobalCounter(M, CounterName);
CallCounterMap[F.getName()] = Var;

// Create a global variable to hold the name of this function
auto FuncName = Builder.CreateGlobalStringPtr(F.getName());
FuncNameMap[F.getName()] = FuncName;

// Inject instruction to increment the call count each time this function
// executes
LoadInst *Load2 = Builder.CreateLoad(IntegerType::getInt32Ty(CTX), Var);
Value *Inc2 = Builder.CreateAdd(Builder.getInt32(1), Load2);
Builder.CreateStore(Inc2, Var);

// The following is visible only if you pass -debug on the command line
// *and* you have an assert build.
LLVM_DEBUG(dbgs() << " Instrumented: " << F.getName() << "\n");

Instrumented = true;
}

// Stop here if there are no function definitions in this module
if (false == Instrumented)
return Instrumented;

// STEP 2: Inject the declaration of printf
// ----------------------------------------
// Create (or _get_ in cases where it's already available) the following
// declaration in the IR module:
//    declare i32 @printf(i8*, ...)
// It corresponds to the following C declaration:
//    int printf(char *, ...)
PointerType *PrintfArgTy = PointerType::getUnqual(Type::getInt8Ty(CTX));
FunctionType *PrintfTy = FunctionType::get(IntegerType::getInt32Ty(CTX), PrintfArgTy,
						/*IsVarArgs=*/true);

FunctionCallee Printf = M.getOrInsertFunction("printf", PrintfTy);

// Set attributes as per inferLibFuncAttributes in BuildLibCalls.cpp
Function *PrintfF = dyn_cast<Function>(Printf.getCallee());
PrintfF->setDoesNotThrow();
PrintfF->addParamAttr(0, Attribute::NoCapture);
PrintfF->addParamAttr(0, Attribute::ReadOnly);

// STEP 3: Inject a global variable that will hold the printf format string
// ------------------------------------------------------------------------
llvm::Constant *ResultFormatStr =
	llvm::ConstantDataArray::getString(CTX, "%-20s %-10lu\n");

Constant *ResultFormatStrVar =
  M.getOrInsertGlobal("ResultFormatStrIR", ResultFormatStr->getType());
dyn_cast<GlobalVariable>(ResultFormatStrVar)->setInitializer(ResultFormatStr);

std::string out = "";
out += "=================================================\n";
out += "LLVM-TUTOR: dynamic analysis results\n";
out += "=================================================\n";
out += "NAME                 #N DIRECT CALLS\n";
out += "-------------------------------------------------\n";

llvm::Constant *ResultHeaderStr =
	llvm::ConstantDataArray::getString(CTX, out.c_str());

Constant *ResultHeaderStrVar =
	M.getOrInsertGlobal("ResultHeaderStrIR", ResultHeaderStr->getType());
dyn_cast<GlobalVariable>(ResultHeaderStrVar)->setInitializer(ResultHeaderStr);

// STEP 4: Define a printf wrapper that will print the results
// -----------------------------------------------------------
// Define `printf_wrapper` that will print the results stored in FuncNameMap
// and CallCounterMap.  It is equivalent to the following C++ function:
// ```
//    void printf_wrapper() {
//      for (auto &item : Functions)
//        printf("llvm-tutor): Function %s was called %d times. \n",
//        item.name, item.count);
//    }
// ```
// (item.name comes from FuncNameMap, item.count comes from
// CallCounterMap)
FunctionType *PrintfWrapperTy =
	FunctionType::get(llvm::Type::getVoidTy(CTX), {},
						/*IsVarArgs=*/false);
Function *PrintfWrapperF = dyn_cast<Function>(
	M.getOrInsertFunction("printf_wrapper", PrintfWrapperTy).getCallee());

// Create the entry basic block for printf_wrapper ...
llvm::BasicBlock *RetBlock =
	llvm::BasicBlock::Create(CTX, "enter", PrintfWrapperF);
IRBuilder<> Builder(RetBlock);

// ... and start inserting calls to printf
// (printf requires i8*, so cast the input strings accordingly)
llvm::Value *ResultHeaderStrPtr =
	Builder.CreatePointerCast(ResultHeaderStrVar, PrintfArgTy);
llvm::Value *ResultFormatStrPtr =
	Builder.CreatePointerCast(ResultFormatStrVar, PrintfArgTy);

Builder.CreateCall(Printf, {ResultHeaderStrPtr});

LoadInst *LoadCounter;
for (auto &item : CallCounterMap) {
LoadCounter = Builder.CreateLoad(IntegerType::getInt32Ty(CTX), item.second);
// LoadCounter = Builder.CreateLoad(item.second);
Builder.CreateCall(
	Printf, {ResultFormatStrPtr, FuncNameMap[item.first()], LoadCounter});
}

// Finally, insert return instruction
Builder.CreateRetVoid();

// STEP 5: Call `printf_wrapper` at the very end of this module
// ------------------------------------------------------------
appendToGlobalDtors(M, PrintfWrapperF, /*Priority=*/0);

return true;

llc-17 -filetype=asm instrumented.bin -o instrumented.s

.section	.fini_array.0,"aw",@fini_array
.p2align	3, 0x90
.quad	printf_wrapper

clang-17 -emit-llvm -S ../inputs/input_for_mba_sub.c -o input_for_sub.ll
opt-17 -load-pass-plugin ./lib/libMBASub.so -passes="mba-sub" -S input_for_sub.ll -o out.ll

define dso_local i32 @main(i32 noundef %0, ptr noundef %1) #0 {
  ...
  %30 = sub nsw i32 %28, %29
  ...
  %33 = sub nsw i32 %31, %32
  ...
  %36 = sub nsw i32 %34, %35
  ret i32 %36
}

define dso_local i32 @main(i32 noundef %0, ptr noundef %1) #0 {
  ...
  %30 = xor i32 %29, -1
  %31 = add i32 %28, %30
  %32 = add i32 %31, 1
  ...
  %35 = xor i32 %34, -1
  %36 = add i32 %33, %35
  %37 = add i32 %36, 1
  ...
  %40 = xor i32 %39, -1
  %41 = add i32 %38, %40
  %42 = add i32 %41, 1
  ret i32 %42
}

bool MBASub::runOnBasicBlock(BasicBlock &BB) {
	bool Changed = false;

	// Loop over all instructions in the block. Replacing instructions requires
	// iterators, hence a for-range loop wouldn't be suitable.
	for (auto Inst = BB.begin(), IE = BB.end(); Inst != IE; ++Inst) {

		// Skip non-binary (e.g. unary or compare) instruction.
		auto *BinOp = dyn_cast<BinaryOperator>(Inst);
		if (!BinOp)
		continue;

		/// Skip instructions other than integer sub.
		unsigned Opcode = BinOp->getOpcode();
		if (Opcode != Instruction::Sub || !BinOp->getType()->isIntegerTy())
		continue;

		// A uniform API for creating instructions and inserting
		// them into basic blocks.
		IRBuilder<> Builder(BinOp);

		// Create an instruction representing (a + ~b) + 1
		Instruction *NewValue = BinaryOperator::CreateAdd(
			Builder.CreateAdd(BinOp->getOperand(0),
							Builder.CreateNot(BinOp->getOperand(1))),
			ConstantInt::get(BinOp->getType(), 1));

		// The following is visible only if you pass -debug on the command line
		// *and* you have an assert build.
		LLVM_DEBUG(dbgs() << *BinOp << " -> " << *NewValue << "\n");

		// Replace `(a - b)` (original instructions) with `(a + ~b) + 1`
		// (the new instruction)
		ReplaceInstWithInst(&BB, Inst, NewValue);
		Changed = true;

		// Update the statistics
		++SubstCount;
	}
	return Changed;
}

clang-17 -O1 -emit-llvm -S ../inputs/input_for_mba.c -o input_for_mba.ll
opt-17 -load-pass-plugin ./lib/libMBAAdd.so -passes="mba-add" -S input_for_mba.ll -o out.ll

$ diff input_for_mba.ll out.ll
1c1
- ; ModuleID = '../inputs/input_for_mba.c'
---
+ ; ModuleID = 'input_for_mba.ll'
8,11c8,32
-   %5 = add i8 %1, %0
-   %6 = add i8 %5, %2
-   %7 = add i8 %6, %3
-   ret i8 %7
---
+   %5 = and i8 %1, %0
+   %6 = mul i8 2, %5
+   %7 = xor i8 %1, %0
+   %8 = add i8 %7, %6
+   %9 = mul i8 39, %8
+   %10 = add i8 23, %9
+   %11 = mul i8 -105, %10
+   %12 = add i8 111, %11
+   %13 = and i8 %12, %2
+   %14 = mul i8 2, %13
+   %15 = xor i8 %12, %2
+   %16 = add i8 %15, %14
+   %17 = mul i8 39, %16
+   %18 = add i8 23, %17
+   %19 = mul i8 -105, %18
+   %20 = add i8 111, %19
+   %21 = and i8 %20, %3
+   %22 = mul i8 2, %21
+   %23 = xor i8 %20, %3
+   %24 = add i8 %23, %22
+   %25 = mul i8 39, %24
+   %26 = add i8 23, %25
+   %27 = mul i8 -105, %26
+   %28 = add i8 111, %27
+   ret i8 %28

bool MBAAdd::runOnBasicBlock(BasicBlock &BB) {
	bool Changed = false;
	
	// Loop over all instructions in the block. Replacing instructions requires
	// iterators, hence a for-range loop wouldn't be suitable
	for (auto Inst = BB.begin(), IE = BB.end(); Inst != IE; ++Inst) {
		// Skip non-binary (e.g. unary or compare) instructions
		auto *BinOp = dyn_cast<BinaryOperator>(Inst);
		if (!BinOp)
		continue;

		// Skip instructions other than add
		if (BinOp->getOpcode() != Instruction::Add)
		continue;

		// Skip if the result is not 8-bit wide (this implies that the operands are
		// also 8-bit wide)
		if (!BinOp->getType()->isIntegerTy() ||
			!(BinOp->getType()->getIntegerBitWidth() == 8))
		continue;

		// A uniform API for creating instructions and inserting
		// them into basic blocks
		IRBuilder<> Builder(BinOp);

		// Constants used in building the instruction for substitution
		auto Val39 = ConstantInt::get(BinOp->getType(), 39);
		auto Val151 = ConstantInt::get(BinOp->getType(), 151);
		auto Val23 = ConstantInt::get(BinOp->getType(), 23);
		auto Val2 = ConstantInt::get(BinOp->getType(), 2);
		auto Val111 = ConstantInt::get(BinOp->getType(), 111);

		// Build an instruction representing `(((a ^ b) + 2 * (a & b)) * 39 + 23) *
		// 151 + 111`
		Instruction *NewInst =
			// E = e5 + 111
			BinaryOperator::CreateAdd(
				Val111,
				// e5 = e4 * 151
				Builder.CreateMul(
					Val151,
					// e4 = e2 + 23
					Builder.CreateAdd(
						Val23,
						// e3 = e2 * 39
						Builder.CreateMul(
							Val39,
							// e2 = e0 + e1
							Builder.CreateAdd(
								// e0 = a ^ b
								Builder.CreateXor(BinOp->getOperand(0),
												BinOp->getOperand(1)),
								// e1 = 2 * (a & b)
								Builder.CreateMul(
									Val2, Builder.CreateAnd(BinOp->getOperand(0),
															BinOp->getOperand(1))))
						) // e3 = e2 * 39
					) // e4 = e2 + 23
				) // e5 = e4 * 151
			); // E = e5 + 111

		// The following is visible only if you pass -debug on the command line
		// *and* you have an assert build.
		LLVM_DEBUG(dbgs() << *BinOp << " -> " << *NewInst << "\n");

		// Replace `(a + b)` (original instructions) with `(((a ^ b) + 2 * (a & b))
		// * 39 + 23) * 151 + 111` (the new instruction)
		ReplaceInstWithInst(&BB, Inst, NewInst);
		Changed = true;

		// Update the statistics
		++SubstCount;
	}
	return Changed;
}

int foo(int a, int b, int c) {
	int result = 123 + a;

	if (a > 0) {
		int d = a * b;
		int e = b / c;
		if (d == e) {
		int f = d * e;
		result = result - 2*f;
		} else {
		int g = 987;
		result = g * c * e;
		}
	} else {
		result = 321;
	}

	return result;
}

clang-17 -O1 -emit-llvm -S ../inputs/input_for_riv.c -o input_for_riv.ll
opt-17 -load-pass-plugin ./lib/libRIV.so -passes="print<riv>" -disable-output input_for_riv.ll

define dso_local i32 @foo(i32 noundef %0, i32 noundef %1, i32 noundef %2) local_unnamed_addr #0 {
  %4 = add nsw i32 %0, 123
  %5 = icmp sgt i32 %0, 0
  br i1 %5, label %6, label %17

6:                                                ; preds = %3
  %7 = mul nsw i32 %1, %0
  %8 = sdiv i32 %1, %2
  %9 = icmp eq i32 %7, %8
  br i1 %9, label %10, label %14

10:                                               ; preds = %6
  %11 = mul i32 %7, -2
  %12 = mul i32 %11, %8
  %13 = add i32 %4, %12
  br label %17

14:                                               ; preds = %6
  %15 = mul nsw i32 %2, 987
  %16 = mul nsw i32 %15, %8
  br label %17

17:                                               ; preds = %3, %10, %14
  %18 = phi i32 [ %13, %10 ], [ %16, %14 ], [ 321, %3 ]
  ret i32 %18
}

=================================================
LLVM-TUTOR: RIV analysis results
=================================================
BB id      Reachable Ineger Values
-------------------------------------------------
BB %3
             i32 %0
             i32 %1
             i32 %2
BB %6
               %4 = add nsw i32 %0, 123
               %5 = icmp sgt i32 %0, 0
             i32 %0
             i32 %1
             i32 %2
BB %17
               %4 = add nsw i32 %0, 123
               %5 = icmp sgt i32 %0, 0
             i32 %0
             i32 %1
             i32 %2
BB %10
               %7 = mul nsw i32 %1, %0
               %8 = sdiv i32 %1, %2
               %9 = icmp eq i32 %7, %8
               %4 = add nsw i32 %0, 123
               %5 = icmp sgt i32 %0, 0
             i32 %0
             i32 %1
             i32 %2
BB %14
               %7 = mul nsw i32 %1, %0
               %8 = sdiv i32 %1, %2
               %9 = icmp eq i32 %7, %8
               %4 = add nsw i32 %0, 123
               %5 = icmp sgt i32 %0, 0
             i32 %0
             i32 %1
             i32 %2

// DominatorTree node types used in RIV. One could use auto instead, but IMO
// being verbose makes it easier to follow.
using NodeTy = DomTreeNodeBase<llvm::BasicBlock> *;
// A map that a basic block BB holds a set of pointers to values defined in BB.
using DefValMapTy = RIV::Result;

RIV::Result RIV::buildRIV(Function &F, NodeTy CFGRoot) {
	Result ResultMap;

	// Initialise a double-ended queue that will be used to traverse all BBs in F
	std::deque<NodeTy> BBsToProcess;
	BBsToProcess.push_back(CFGRoot);

	// STEP 1: For every basic block BB compute the set of integer values defined
	// in BB
	DefValMapTy DefinedValuesMap;
	for (BasicBlock &BB : F) {
		auto &Values = DefinedValuesMap[&BB];
		for (Instruction &Inst : BB)
		if (Inst.getType()->isIntegerTy())
			Values.insert(&Inst);
	}

	// STEP 2: Compute the RIVs for the entry BB. This will include global
	// variables and input arguments.
	auto &EntryBBValues = ResultMap[&F.getEntryBlock()];

	for (auto &Global : F.getParent()->globals())
		if (Global.getValueType()->isIntegerTy())
		EntryBBValues.insert(&Global);

	for (Argument &Arg : F.args())
		if (Arg.getType()->isIntegerTy())
		EntryBBValues.insert(&Arg);

	// STEP 3: Traverse the CFG for every BB in F calculate its RIVs
	while (!BBsToProcess.empty()) {
		auto *Parent = BBsToProcess.back();
		BBsToProcess.pop_back();

		// Get the values defined in Parent
		auto &ParentDefs = DefinedValuesMap[Parent->getBlock()];
		// Get the RIV set of for Parent
		// (Since RIVMap is updated on every iteration, its contents are likely to
		// be moved around when resizing. This means that we need a copy of it
		// (i.e. a reference is not sufficient).
		llvm::SmallPtrSet<llvm::Value *, 8> ParentRIVs =
			ResultMap[Parent->getBlock()];

		// Loop over all BBs that Parent dominates and update their RIV sets
		for (NodeTy Child : *Parent) {
		BBsToProcess.push_back(Child);
		auto ChildBB = Child->getBlock();

		// Add values defined in Parent to the current child's set of RIV
		ResultMap[ChildBB].insert(ParentDefs.begin(), ParentDefs.end());

		// Add Parent's set of RIVs to the current child's RIV
		ResultMap[ChildBB].insert(ParentRIVs.begin(), ParentRIVs.end());
		}
	}

	return ResultMap;
}

clang-17 -O1 -emit-llvm -S ../inputs/input_for_duplicate_bb.c -o input_for_duplicate_bb.ll
opt-17 -load-pass-plugin ./lib/libRIV.so -load-pass-plugin ./lib/libDuplicateBB.so -passes=duplicate-bb -S input_for_duplicate_bb.ll -o duplicate.ll

define dso_local i32 @foo(i32 noundef %0) local_unnamed_addr #0 {
  ret i32 1
}

define dso_local i32 @foo(i32 noundef %0) local_unnamed_addr #0 {
lt-if-then-else-0:
  %1 = icmp eq i32 %0, 0
  br i1 %1, label %lt-clone-1-0, label %lt-clone-2-0

lt-clone-1-0:                                     ; preds = %lt-if-then-else-0
  br label %lt-tail-0

lt-clone-2-0:                                     ; preds = %lt-if-then-else-0
  br label %lt-tail-0

lt-tail-0:                                        ; preds = %lt-clone-2-0, %lt-clone-1-0
  ret i32 1
}

PreservedAnalyses DuplicateBB::run(llvm::Function &F,
                                   llvm::FunctionAnalysisManager &FAM) {
	if (!pRNG)
		pRNG = F.getParent()->createRNG("duplicate-bb");
	
	BBToSingleRIVMap Targets = findBBsToDuplicate(F, FAM.getResult<RIV>(F));

	// This map is used to keep track of the new bindings. Otherwise, the
	// information from RIV will become obsolete.
	ValueToPhiMap ReMapper;

	// Duplicate
	for (auto &BB_Ctx : Targets) {
		cloneBB(*std::get<0>(BB_Ctx), std::get<1>(BB_Ctx), ReMapper);
	}

	DuplicateBBCountStats = DuplicateBBCount;
	return (Targets.empty() ? llvm::PreservedAnalyses::all()
							: llvm::PreservedAnalyses::none());
}

DuplicateBB::BBToSingleRIVMap
DuplicateBB::findBBsToDuplicate(Function &F, const RIV::Result &RIVResult) {
	BBToSingleRIVMap BlocksToDuplicate;

	for (BasicBlock &BB : F) {
		// Basic blocks which are landing pads are used for handling exceptions.
		// That's out of scope of this pass.
		if (BB.isLandingPad())
		continue;

		// Get the set of RIVs for this block
		auto const &ReachableValues = RIVResult.lookup(&BB);
		size_t ReachableValuesCount = ReachableValues.size();

		// Are there any RIVs for this BB? We need at least one to be able to
		// duplicate this BB.
		if (0 == ReachableValuesCount) {
		LLVM_DEBUG(errs() << "No context values for this BB\n");
		continue;
		}

		// Get a random context value from the RIV set
		auto Iter = ReachableValues.begin();
		std::uniform_int_distribution<> Dist(0, ReachableValuesCount - 1);
		std::advance(Iter, Dist(*pRNG));

		if (dyn_cast<GlobalValue>(*Iter)) {
		LLVM_DEBUG(errs() << "Random context value is a global variable. "
							<< "Skipping this BB\n");
		continue;
		}

		LLVM_DEBUG(errs() << "Random context value: " << **Iter << "\n");

		// Store the binding between the current BB and the context variable that
		// will be used for the `if-then-else` construct.
		BlocksToDuplicate.emplace_back(&BB, *Iter);
	}

	return BlocksToDuplicate;
}

void DuplicateBB::cloneBB(BasicBlock &BB, Value *ContextValue,
                          ValueToPhiMap &ReMapper) {
	// Don't duplicate Phi nodes - start right after them
	Instruction *BBHead = BB.getFirstNonPHI();

	// Create the condition for 'if-then-else'
	IRBuilder<> Builder(BBHead);
	Value *Cond = Builder.CreateIsNull(
		ReMapper.count(ContextValue) ? ReMapper[ContextValue] : ContextValue);

	// Create and insert the 'if-else' blocks. At this point both blocks are
	// trivial and contain only one terminator instruction branching to BB's
	// tail, which contains all the instructions from BBHead onwards.
	Instruction *ThenTerm = nullptr;
	Instruction *ElseTerm = nullptr;
	SplitBlockAndInsertIfThenElse(Cond, &*BBHead, &ThenTerm, &ElseTerm);
	BasicBlock *Tail = ThenTerm->getSuccessor(0);

	assert(Tail == ElseTerm->getSuccessor(0) && "Inconsistent CFG");

	// Give the new basic blocks some meaningful names. This is not required, but
	// makes the output easier to read.
	std::string DuplicatedBBId = std::to_string(DuplicateBBCount);
	ThenTerm->getParent()->setName("lt-clone-1-" + DuplicatedBBId);
	ElseTerm->getParent()->setName("lt-clone-2-" + DuplicatedBBId);
	Tail->setName("lt-tail-" + DuplicatedBBId);
	ThenTerm->getParent()->getSinglePredecessor()->setName("lt-if-then-else-" +
															DuplicatedBBId);

	// Variables to keep track of the new bindings
	ValueToValueMapTy TailVMap, ThenVMap, ElseVMap;

	// The list of instructions in Tail that don't produce any values and thus
	// can be removed
	SmallVector<Instruction *, 8> ToRemove;

	// Iterate through the original basic block and clone every instruction into
	// the 'if-then' and 'else' branches. Update the bindings/uses on the fly
	// (through ThenVMap, ElseVMap, TailVMap). At this stage, all instructions
	// apart from PHI nodes, are stored in Tail.
	for (auto IIT = Tail->begin(), IE = Tail->end(); IIT != IE; ++IIT) {
		Instruction &Instr = *IIT;
		assert(!isa<PHINode>(&Instr) && "Phi nodes have already been filtered out");

		// Skip terminators - duplicating them wouldn't make sense unless we want
		// to delete Tail completely.
		if (Instr.isTerminator()) {
		RemapInstruction(&Instr, TailVMap, RF_IgnoreMissingLocals);
		continue;
		}

		// Clone the instructions.
		Instruction *ThenClone = Instr.clone(), *ElseClone = Instr.clone();

		// Operands of ThenClone still hold references to the original BB.
		// Update/remap them.
		RemapInstruction(ThenClone, ThenVMap, RF_IgnoreMissingLocals);
		ThenClone->insertBefore(ThenTerm);
		ThenVMap[&Instr] = ThenClone;

		// Operands of ElseClone still hold references to the original BB.
		// Update/remap them.
		RemapInstruction(ElseClone, ElseVMap, RF_IgnoreMissingLocals);
		ElseClone->insertBefore(ElseTerm);
		ElseVMap[&Instr] = ElseClone;

		// Instructions that don't produce values can be safely removed from Tail
		if (ThenClone->getType()->isVoidTy()) {
		ToRemove.push_back(&Instr);
		continue;
		}

		// Instruction that produce a value should not require a slot in the
		// TAIL *but* they can be used from the context, so just always
		// generate a PHI, and let further optimization do the cleaning
		PHINode *Phi = PHINode::Create(ThenClone->getType(), 2);
		Phi->addIncoming(ThenClone, ThenTerm->getParent());
		Phi->addIncoming(ElseClone, ElseTerm->getParent());
		TailVMap[&Instr] = Phi;

		ReMapper[&Instr] = Phi;

		// Instructions are modified as we go, use the iterator version of
		// ReplaceInstWithInst.
		ReplaceInstWithInst(Tail, IIT, Phi);
	}

	// Purge instructions that don't produce any value
	for (auto *I : ToRemove)
		I->eraseFromParent();

	++DuplicateBBCount;
}

opt-17 -load-pass-plugin ./lib/libMergeBB.so -passes=merge-bb -S duplicate.ll -o merge.ll

define dso_local i32 @foo(i32 noundef %0) local_unnamed_addr #0 {
lt-if-then-else-0:
  %1 = icmp eq i32 %0, 0
  br i1 %1, label %lt-clone-2-0, label %lt-clone-2-0

lt-clone-2-0:                                     ; preds = %lt-if-then-else-0, %lt-if-then-else-0
  br label %lt-tail-0

lt-tail-0:                                        ; preds = %lt-clone-2-0
  ret i32 1
}

opt-17 -load-pass-plugin ./lib/libRIV.so -load-pass-plugin ./lib/libDuplicateBB.so -load-pass-plugin ./lib/libMergeBB.so -passes=d
uplicate-bb,merge-bb -S input_for_duplicate_bb.ll -o merge_after_duplicate.ll

PreservedAnalyses MergeBB::run(llvm::Function &Func,
                               llvm::FunctionAnalysisManager &) {
	bool Changed = false;
	SmallPtrSet<BasicBlock *, 8> DeleteList;
	for (auto &BB : Func) {
		Changed |= mergeDuplicatedBlock(&BB, DeleteList);
	}

	for (BasicBlock *BB : DeleteList) {
		DeleteDeadBlock(BB);
	}

	return (Changed ? llvm::PreservedAnalyses::none()
					: llvm::PreservedAnalyses::all());
}

bool MergeBB::mergeDuplicatedBlock(BasicBlock *BB1,
                                   SmallPtrSet<BasicBlock *, 8> &DeleteList) {
	// Do not optimize the entry block
	if (BB1 == &BB1->getParent()->getEntryBlock())
		return false;

	// Only merge CFG edges of unconditional branch
	BranchInst *BB1Term = dyn_cast<BranchInst>(BB1->getTerminator());
	if (!(BB1Term && BB1Term->isUnconditional()))
		return false;

	// Do not optimize non-branch and non-switch CFG edges (to keep things
	// relatively simple)
	for (auto *B : predecessors(BB1))
		if (!(isa<BranchInst>(B->getTerminator()) ||
			isa<SwitchInst>(B->getTerminator())))
		return false;

	BasicBlock *BBSucc = BB1Term->getSuccessor(0);

	BasicBlock::iterator II = BBSucc->begin();
	const PHINode *PN = dyn_cast<PHINode>(II);
	Value *InValBB1 = nullptr;
	Instruction *InInstBB1 = nullptr;
	BBSucc->getFirstNonPHI();
	if (nullptr != PN) {
		// Do not optimize if multiple PHI instructions exist in the successor (to
		// keep things relatively simple)
		if (++II != BBSucc->end() && isa<PHINode>(II))
		return false;

		InValBB1 = PN->getIncomingValueForBlock(BB1);
		InInstBB1 = dyn_cast<Instruction>(InValBB1);
	}

	unsigned BB1NumInst = getNumNonDbgInstrInBB(BB1);
	for (auto *BB2 : predecessors(BBSucc)) {
		// Do not optimize the entry block
		if (BB2 == &BB2->getParent()->getEntryBlock())
		continue;

		// Only merge CFG edges of unconditional branch
		BranchInst *BB2Term = dyn_cast<BranchInst>(BB2->getTerminator());
		if (!(BB2Term && BB2Term->isUnconditional()))
		continue;

		// Do not optimize non-branch and non-switch CFG edges (to keep things
		// relatively simple)
		for (auto *B : predecessors(BB2))
		if (!(isa<BranchInst>(B->getTerminator()) ||
				isa<SwitchInst>(B->getTerminator())))
			continue;

		// Skip basic blocks that have already been marked for merging
		if (DeleteList.end() != DeleteList.find(BB2))
		continue;

		// Make sure that BB2 != BB1
		if (BB2 == BB1)
		continue;

		// BB1 and BB2 are definitely different if the number of instructions is
		// not identical
		if (BB1NumInst != getNumNonDbgInstrInBB(BB2))
		continue;

		// Control flow can be merged if incoming values to the PHI node
		// at the successor are same values or both defined in the BBs to merge.
		// For the latter case, canMergeInstructions executes further analysis.
		if (nullptr != PN) {
		Value *InValBB2 = PN->getIncomingValueForBlock(BB2);
		Instruction *InInstBB2 = dyn_cast<Instruction>(InValBB2);

		bool areValuesSimilar = (InValBB1 == InValBB2);
		bool bothValuesDefinedInParent =
			((InInstBB1 && InInstBB1->getParent() == BB1) &&
			(InInstBB2 && InInstBB2->getParent() == BB2));
		if (!areValuesSimilar && !bothValuesDefinedInParent)
			continue;
		}

		// Finally, check that all instructions in BB1 and BB2 are identical
		LockstepReverseIterator LRI(BB1, BB2);
		while (LRI.isValid() && canMergeInstructions(*LRI)) {
		--LRI;
		}

		// Valid iterator  means that a mismatch was found in middle of BB
		if (LRI.isValid())
		continue;

		// It is safe to de-duplicate - do so.
		unsigned UpdatedTargets = updateBranchTargets(BB1, BB2);
		assert(UpdatedTargets && "No branch target was updated");
		OverallNumOfUpdatedBranchTargets += UpdatedTargets;
		DeleteList.insert(BB1);
		NumDedupBBs++;

		return true;
	}

	return false;
}