Demo.cpp

#include "llvm/ADT/Statistic.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/Pass.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/raw_ostream.h"
#include <bits/stdc++.h>
#include <map>
#include <vector>
using namespace llvm;
using namespace std;

struct variable_values {
  bool TOP;
  bool BOTTOM;
  APInt constant_value;

  variable_values() : TOP(false), BOTTOM(false) {}
};

namespace {
struct constant_p : public FunctionPass {
  static char ID;
  constant_p() : FunctionPass(ID) {}
  map<StringRef, map<StringRef, variable_values>>
      global_map; // This is a map per Basic Block. Key contains the name of the
                  // Basic Block and the value contains the map of the variables
                  // and their value.
  queue<BasicBlock *> BB_to_execute; // This queue is used for maintaining the
                                     // execution order of the Basic Blocks

  queue<StringRef> worklist;

  bool runOnFunction(Function &F) override {
    /* We Push the entry Basic Block first */
    errs() << "Name of the function : " << F.getName() << "\n";
    worklist.push(F.getName());
    errs() << "Main was inserted"
           << "\n";

    BasicBlock *startBB = &F.getEntryBlock();
    BB_to_execute.push(startBB);

    while (!BB_to_execute.empty()) {
      BasicBlock *front_BB =
          BB_to_execute
              .front(); /* We fetch the address of the front BB to process */
      BB_to_execute
          .pop(); /* Then we remove that BB, since we're anyway processing it */

      /* We then push all the successor Basic Blocks of the current Basic Block
       */
      for (BasicBlock *Succ : successors(front_BB)) {
        if (Succ != BB_to_execute.back()) {
          BB_to_execute.push(Succ);
        }
      }

      map<StringRef, variable_values>
          variable_state; /* This Map contains the name of the Variable and the
                             Value contained in it */
      StringRef current_BB_name =
          front_BB->getName(); /* This holds the name of the current BB */

      /*Here, we get the set of all pred basic block */
      queue<BasicBlock *> pred_of_current_block;
      for (BasicBlock *pred : predecessors(front_BB)) {
        pred_of_current_block.push(pred);
      }
      if (current_BB_name != "entry" and pred_of_current_block.size() == 1) {
        /* If the current BB has only one Predecessor BB, then we simply
         * propagate the variable values of the Pred BB to the current basic
         * block */
        variable_state = global_map[pred_of_current_block.front()->getName()];
      } else if (current_BB_name != "entry" and
                 pred_of_current_block.size() ==
                     2) /* Since the current BB has 2 predecessor, so we need to
                           merge the values of the variables present in both the
                           Basic Block */
      {

        for (auto it =
                 global_map[pred_of_current_block.front()->getName()].begin();
             it != global_map[pred_of_current_block.front()->getName()].end();
             ++it) {
          for (auto iit =
                   global_map[pred_of_current_block.back()->getName()].begin();
               iit != global_map[pred_of_current_block.back()->getName()].end();
               ++iit) {
            if ((iit->first == it->first)) /* If the Key is same */
            {

              if ((iit->second.constant_value == it->second.constant_value) and
                  (iit->second.TOP == false) and
                  (iit->second.BOTTOM == false) and
                  (it->second.TOP == false) and
                  (it->second.BOTTOM == false)) // c1 meet c1 = c1
              {
                variable_state[iit->first].constant_value =
                    iit->second.constant_value;
              } else if ((iit->second.constant_value !=
                          it->second.constant_value) and
                         (iit->second.TOP == false) and
                         (iit->second.BOTTOM == false) and
                         (it->second.TOP == false) and
                         (it->second.BOTTOM == false)) // c1 meet c2 = BOTTOM
              {
                variable_state[iit->first].BOTTOM = true;
              } else if (iit->second.BOTTOM == true or
                         it->second.BOTTOM == true) {
                variable_state[iit->first].BOTTOM = true;
              } else if (iit->second.TOP == true) {
                variable_state[iit->first] = it->second;
              } else if (it->second.TOP == true) {
                variable_state[it->first] = iit->second;
              } else {
              }
              break;
            } else /* If the keys are not same */
            {
              variable_state[iit->first] = iit->second;
            }
          }
        }

        for (auto it =
                 global_map[pred_of_current_block.back()->getName()].begin();
             it != global_map[pred_of_current_block.back()->getName()].end();
             ++it) {

          for (auto iit =
                   global_map[pred_of_current_block.front()->getName()].begin();
               iit !=
               global_map[pred_of_current_block.front()->getName()].end();
               ++iit) {
            if ((iit->first == it->first)) // if the key is same
            {

              if ((iit->second.constant_value == it->second.constant_value) and
                  (iit->second.TOP == false) and
                  (iit->second.BOTTOM == false) and
                  (it->second.TOP == false) and
                  (it->second.BOTTOM == false)) // c1 meet c1 = c1
              {
                variable_state[iit->first].constant_value =
                    iit->second.constant_value;
              } else if ((iit->second.constant_value !=
                          it->second.constant_value) and
                         (iit->second.TOP == false) and
                         (iit->second.BOTTOM == false) and
                         (it->second.TOP == false) and
                         (it->second.BOTTOM == false)) // c1 meet c2 = BOTTOM
              {
                variable_state[iit->first].BOTTOM = true;
              } else if (iit->second.BOTTOM == true or
                         it->second.BOTTOM == true) {
                variable_state[iit->first].BOTTOM = true;
              } else if (iit->second.TOP == true) {
                variable_state[iit->first] = it->second;
              } else if (it->second.TOP == true) {
                variable_state[it->first] = iit->second;
              } else {
              }
              break;
            } else // if the keys are not same
            {
              variable_state[iit->first] = iit->second;
            }
          }
        }
      }

      /* We execute all the instructions in the current Basic Block sequentially
       */
      for (Instruction &I : *front_BB) {
        errs() << I << "\n";

        /* Processing Store Instructions */
        if (StoreInst *Store = dyn_cast<StoreInst>(&I)) {
          Value *StoredValue =
              Store->getValueOperand(); // gets the Operand value
          auto *Var = dyn_cast<AllocaInst>(Store->getPointerOperand());

          /* If the Store operation contains a Constant value then this code
           * "if" will be executed */
          if (ConstantInt *ConstantValue = dyn_cast<ConstantInt>(StoredValue)) {

            /* Here, we store the variable and the value assigned to it*/
            variable_state[Var->getName()].constant_value =
                ConstantValue->getValue();
            variable_state[Var->getName()].TOP = false;
            variable_state[Var->getName()].BOTTOM = false;

            errs() << I << " --> %" << Var->getName() << " = "
                   << variable_state[Var->getName()].constant_value << "\n";
          } else {

            errs() << I << " --> %" << Var->getName() << " = "
                   << variable_state[StoredValue->getName()].constant_value
                   << "\n";
            variable_state[Var->getName()] =
                variable_state[StoredValue->getName()];
          }
        }

        /* Processing Alloca Instructions */
        else if (AllocaInst *Alloca = dyn_cast<AllocaInst>(&I)) {
          variable_state[Alloca->getName()].TOP =
              true; // For all alloca variables, we intialize it to TOP
          errs() << I << " --> %" << Alloca->getName() << " = TOP"
                 << "\n";
        }

        /* Processing Load instructions */
        else if (LoadInst *Load = dyn_cast<LoadInst>(&I)) {
          Value *LoadedValue = Load->getPointerOperand();
          auto *Var = dyn_cast<AllocaInst>(LoadedValue);
          Value *destination = &I;
          Value *ptrOperand = Load->getPointerOperand();

          variable_state[I.getNameOrAsOperand()] =
              variable_state[LoadedValue->getName()]; /* Here, we copy the
                                                         value of the source
                                                         variable to the
                                                         destination variable
                                                       */
          if (variable_state[I.getNameOrAsOperand()].TOP) {
            errs() << I << " --> " << I.getNameOrAsOperand() << " = TOP, "
                   << LoadedValue->getName() << " = TOP"
                   << "\n";
          } else if (variable_state[I.getNameOrAsOperand()].BOTTOM) {
            errs() << I << " --> " << I.getNameOrAsOperand() << " = BOTTOM, "
                   << LoadedValue->getName() << " = BOTTOM"
                   << "\n";
          } else {
            errs() << I << " --> " << I.getNameOrAsOperand() << " = "
                   << variable_state[LoadedValue->getName()].constant_value
                   << ", " << LoadedValue->getName() << " = "
                   << variable_state[LoadedValue->getName()].constant_value
                   << "\n";
          }
        }

        /* Processing the Call Instructions */
        else if (CallInst *Call = dyn_cast<CallInst>(&I)) {
          Function *calledFunction = Call->getCalledFunction();
          if (calledFunction->getName() ==
              "__isoc99_scanf") { // If "scanf" is called, then we assign
                                  // BOTTOM
            // to the variable
            variable_state[(Call->getArgOperand(1))->getName()].BOTTOM = true;
            variable_state[(Call->getArgOperand(1))->getName()].TOP = false;

            errs() << I << " --> %" << (Call->getArgOperand(1))->getName()
                   << " = BOTTOM"
                   << "\n";
          } else if (calledFunction->getName() ==
                     "printf") { // If "printf" is called, then no changes
                                 // will
            // be done
            errs() << I << "\n";
          } else {
            StringRef FunctionName = calledFunction->getName();
            errs() << "Call to the Function: " << FunctionName << "\n";
          }
        }

        /* Processing Return Instructions */
        else if (ReturnInst *RetIns = dyn_cast<ReturnInst>(&I)) {
          errs() << I << "\n";
          global_map[current_BB_name] = variable_state;
        }

        /* Processing Arithmetic Instructions */
        else if (auto *binaryOp = dyn_cast<BinaryOperator>(&I)) {
          if (binaryOp->getOpcode() == Instruction::Add) {
            Value *operand1 = binaryOp->getOperand(1);
            if (ConstantInt *op1 = dyn_cast<ConstantInt>(operand1)) {
              APInt result =
                  op1->getValue() + variable_state[(binaryOp->getOperandUse(0))
                                                       ->getNameOrAsOperand()]
                                        .constant_value;

              variable_state[I.getName()].constant_value = result;
              variable_state[I.getName()].TOP = false;
              variable_state[I.getName()].BOTTOM = false;

              errs() << I << " --> %" << I.getName() << " = "
                     << variable_state[I.getName()].constant_value << "\n";
            } else {
              llvm::APInt operand1 = variable_state[(binaryOp->getOperandUse(0))
                                                        ->getNameOrAsOperand()]
                                         .constant_value;
              llvm::APInt operand2 = variable_state[(binaryOp->getOperandUse(1))
                                                        ->getNameOrAsOperand()]
                                         .constant_value;

              // Ensure operands have the same bit width
              if (operand1.getBitWidth() == operand2.getBitWidth()) {
                unsigned int maxBitWidth =
                    std::max(operand1.getBitWidth(), operand2.getBitWidth());

                // Extend both operands to have the maximum bit width
                operand1 = operand1.zextOrTrunc(maxBitWidth);
                operand2 = operand2.zextOrTrunc(maxBitWidth);

                APInt result_final = operand1 + operand2;
                errs() << I << " --> %" << I.getName() << " = " << result_final
                       << "\n";
                variable_state[I.getName()].constant_value = result_final;
                variable_state[I.getName()].TOP = false;
                variable_state[I.getName()].BOTTOM = false;
              }
            }
          } else if (binaryOp->getOpcode() ==
                     Instruction::Sub) // Handling Sub Instruction
          {

            if (ConstantInt *op1 =
                    dyn_cast<ConstantInt>(binaryOp->getOperand(1))) {
              APInt result = variable_state[(binaryOp->getOperandUse(0))
                                                ->getNameOrAsOperand()]
                                 .constant_value -
                             op1->getValue();

              variable_state[I.getName()].constant_value = result;
              variable_state[I.getName()].TOP = false;
              variable_state[I.getName()].BOTTOM = false;

              errs() << I << " --> %" << I.getName() << " = "
                     << variable_state[I.getName()].constant_value << "\n";
            }
          }
        }

        /* Processing Comparison Instruction */
        else if (CmpInst *cmpinst = dyn_cast<CmpInst>(&I)) {
          variable_state[I.getName()].BOTTOM = true;
          errs() << I << " --> %" << I.getName() << " = BOTTOM"
                 << "\n";
        }

        /* Processing Branch Instruction */
        else if (BranchInst *BrInst = dyn_cast<BranchInst>(&I)) {
          global_map[current_BB_name] = variable_state;
          if (BrInst->isConditional()) {
            if (variable_state[(BrInst->getOperand(0))->getName()].BOTTOM) {
              errs() << I << " --> %" << (BrInst->getOperand(0))->getName()
                     << " = BOTTOM"
                     << "\n";
            } else if (variable_state[(BrInst->getOperand(0))->getName()].TOP) {
              errs() << I << " --> %" << (BrInst->getOperand(0))->getName()
                     << " = TOP"
                     << "\n";
            } else {
              errs() << I << " --> %" << (BrInst->getOperand(0))->getName()
                     << " = "
                     << variable_state[(BrInst->getOperand(0))->getName()]
                            .constant_value
                     << "\n";
            }
          } else {
            errs() << I << "\n";
          }
        }
      }
    }

    return false;
  }
};
} // namespace

char constant_p::ID = 0;
static RegisterPass<constant_p> X("Demoold",
                                  "Constant Propagation Pass for Assignment",
                                  false /* Only looks at CFG */,
                                  false /* Analysis Pass */);