modify

Author: SapphireLinXu

lab6/llvm-pass.so.cc

@@ -1,7 +1,6 @@
 #include "llvm/Passes/PassPlugin.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/IR/IRBuilder.h"
-
 using namespace llvm;
 
 struct LLVMPass : public PassInfoMixin<LLVMPass> {
@@ -11,37 +10,53 @@ struct LLVMPass : public PassInfoMixin<LLVMPass> {
 PreservedAnalyses LLVMPass::run(Module &M, ModuleAnalysisManager &MAM) {
   LLVMContext &Ctx = M.getContext();
   IntegerType *Int32Ty = IntegerType::getInt32Ty(Ctx);
-  FunctionCallee debug_func = M.getOrInsertFunction("debug", Int32Ty);
-  ConstantInt *debug_arg = ConstantInt::get(Int32Ty, 48763);
+
+  // Get or insert debug function declaration
+  FunctionType *debugTy = FunctionType::get(Type::getVoidTy(Ctx), {Int32Ty}, false);
+  FunctionCallee debugFunc = M.getOrInsertFunction("debug", debugTy);
+  ConstantInt *debugArg = ConstantInt::get(Int32Ty, 48763);
 
   for (auto &F : M) {
-    if (F.getName() == "main") {
-      IRBuilder<> Builder(&*F.getEntryBlock().getFirstInsertionPt());
-
-      // 1. 插入 debug(48763);
-      Builder.CreateCall(debug_func, debug_arg);
-
-      // 2. 將 argv[1] = "hayaku... motohayaku!"
-      // 建立常數字串
-      Constant *StrConstant = Builder.CreateGlobalStringPtr("hayaku... motohayaku!", "hayaku_str");
-
-      // 取得參數 argc 和 argv
-      Function::arg_iterator args = F.arg_begin();
-      Value *argcArg = &*args++;
-      Value *argvArg = &*args;
-
-      // argv[1] = "hayaku... motohayaku!"
-      Value *idxs[] = {
-        ConstantInt::get(Int32Ty, 1)
-      };
-      Value *argv1Ptr = Builder.CreateInBoundsGEP(argvArg, idxs, "argv1_ptr");
-      Builder.CreateStore(StrConstant, argv1Ptr);
-
-      // 3. 將 argc 改為 48763
-      argcArg->replaceAllUsesWith(debug_arg);  // 把 argc 的所有使用都替換成 48763
+    if (F.getName() != "main" || F.isDeclaration())
+      continue;
+
+    errs() << "instrumenting main\n";
+
+    IRBuilder<> Builder(&*F.getEntryBlock().getFirstInsertionPt());
+
+    // Call debug(48763)
+    Builder.CreateCall(debugFunc, {debugArg});
+
+    // Get arguments
+    auto ArgIter = F.arg_begin();
+    Argument *argcArg = &*ArgIter++;
+    Argument *argvArg = &*ArgIter;
+
+    // Allocate space and store 48763 to replace argc
+    AllocaInst *argcAlloca = Builder.CreateAlloca(Int32Ty, nullptr, "argcVar");
+    Builder.CreateStore(debugArg, argcAlloca);
+
+    // Replace uses of argcArg with loaded value
+    for (auto &BB : F) {
+      for (auto &I : BB) {
+        for (unsigned i = 0; i < I.getNumOperands(); ++i) {
+          if (I.getOperand(i) == argcArg) {
+            // Load from new variable instead of using original argc
+            IRBuilder<> B(&I);
+            Value *loadedArgc = B.CreateLoad(Int32Ty, argcAlloca);
+            I.setOperand(i, loadedArgc);
+          }
+        }
+      }
     }
 
+    // Replace argv[1] = "hayaku... motohayaku!"
+    Value *index1 = ConstantInt::get(Int32Ty, 1);
+    Value *argv1Ptr = Builder.CreateInBoundsGEP(argvArg->getType()->getPointerElementType(), argvArg, index1);
+    Value *hayakuStr = Builder.CreateGlobalStringPtr("hayaku... motohayaku!");
+    Builder.CreateStore(hayakuStr, argv1Ptr);
   }
+
   return PreservedAnalyses::none();
 }
 
@@ -55,4 +70,3 @@ llvmGetPassPluginInfo() {
         });
     }};
 }
-