using namespace llvm;

#define DEBUG_TYPE "undobool"

PreservedAnalyses clspv::UndoBoolPass::run(Module &M, ModuleAnalysisManager &) {
  PreservedAnalyses PA;

  SmallVector<Instruction *, 8> WorkList;
  for (Function &F : M) {
    for (BasicBlock &BB : F) {
      for (Instruction &I : BB) {
        switch (I.getOpcode()) {
        default:
          // Skip this instruction as it did not match
          continue;
        case Instruction::ZExt:
          // We are looking for an instruction that produces an i8 from an i1.
          if (I.getOperand(0)->getType()->isIntegerTy(1) &&
              I.getType()->isIntegerTy(8)) {
            WorkList.push_back(&I);
          }
          break;
        }
      }
    }
  }

  for (Instruction *I : WorkList) {
    SmallVector<Instruction *, 1> ToReplace;

    // Walk the users of the instruction looking for a trunc
    for (User *U : I->users()) {
      if (Instruction *OI = dyn_cast<Instruction>(U)) {
        // We are looking for a trunc instruction that produces an i1 from i8.
        if ((Instruction::Trunc == OI->getOpcode()) &&
            OI->getOperand(0)->getType()->isIntegerTy(8) &&
            OI->getType()->isIntegerTy(1)) {
          ToReplace.push_back(OI);
        }
      }
    }

    for (Instruction *OI : ToReplace) {
      OI->replaceAllUsesWith(I->getOperand(0));
      OI->eraseFromParent();
    }

    // Delete trivially dead instructions.
    if (I->use_empty()) {
      I->eraseFromParent();
    }
  }

  return PA;
}
