revng-revng/lib/Canonicalize/PrettyIntFormattingPass.cpp

//
// This file is distributed under the MIT License. See LICENSE.md for details.
//

#include <array>
#include <optional>
#include <utility>
#include <vector>

#include "llvm/IR/Attributes.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/Pass.h"

#include "revng/Model/Binary.h"
#include "revng/Model/IRHelpers.h"
#include "revng/Model/LoadModelPass.h"
#include "revng/Support/OpaqueFunctionsPool.h"

#include "revng-c/Support/FunctionTags.h"
#include "revng-c/Support/IRHelpers.h"
#include "revng-c/Support/ModelHelpers.h"

enum class IntFormatting : uint32_t {
  NONE, // no formatting
  HEX,
  CHAR,
  BOOL,
  NULLPTR
};

struct FormatInt {
  IntFormatting Formatting;
  llvm::Use *Use;
};

static std::optional<FormatInt>
getIntFormat(llvm::Instruction &I, llvm::Use &U, const model::Binary &Model);

struct PrettyIntFormatting : public llvm::FunctionPass {
public:
  static char ID;

  PrettyIntFormatting() : llvm::FunctionPass(ID) {}

  bool runOnFunction(llvm::Function &F) override;

  void getAnalysisUsage(llvm::AnalysisUsage &AU) const override {
    AU.setPreservesCFG();
    AU.addRequired<LoadModelWrapperPass>();
  }
};

bool PrettyIntFormatting::runOnFunction(llvm::Function &F) {

  if (not FunctionTags::TagsSet::from(&F).contains(FunctionTags::Isolated))
    return false;

  const model::Binary
    &Model = *getAnalysis<LoadModelWrapperPass>().get().getReadOnlyModel();

  OpaqueFunctionsPool<llvm::Type *> HexIntegerPool(F.getParent(), false);
  initHexPrintPool(HexIntegerPool);

  OpaqueFunctionsPool<llvm::Type *> CharIntegerPool(F.getParent(), false);
  initCharPrintPool(CharIntegerPool);

  OpaqueFunctionsPool<llvm::Type *> BoolIntegerPool(F.getParent(), false);
  initBoolPrintPool(BoolIntegerPool);

  OpaqueFunctionsPool<llvm::Type *> NullPtrPool(F.getParent(), false);
  initNullPtrPrintPool(NullPtrPool);

  std::vector<FormatInt> IntsToBeFormatted;

  for (llvm::Instruction &I : llvm::instructions(F)) {
    for (llvm::Use &U : I.operands()) {
      if (auto formatting = getIntFormat(I, U, Model); formatting) {
        IntsToBeFormatted.push_back(*formatting);
      }
    }
  }

  llvm::IRBuilder<> Builder(F.getContext());
  for (const auto &[Format, Operand] : IntsToBeFormatted) {
    auto *Val = llvm::cast<llvm::ConstantInt>(Operand->get());
    llvm::Type *IntType = Val->getType();

    auto PrettyFunction = [&, Format = Format]() -> llvm::Function * {
      switch (Format) {
      case IntFormatting::HEX:
        return HexIntegerPool.get(IntType, IntType, { IntType }, "print_hex");
      case IntFormatting::CHAR:
        return CharIntegerPool.get(IntType, IntType, { IntType }, "print_char");
      case IntFormatting::BOOL:
        return BoolIntegerPool.get(IntType, IntType, { IntType }, "print_bool");
      case IntFormatting::NULLPTR:
        return NullPtrPool.get(IntType, IntType, { IntType }, "print_nullptr");
      case IntFormatting::NONE:
      default:
        return nullptr;
      }
      return nullptr;
    }();

    if (PrettyFunction) {
      Builder.SetInsertPoint(llvm::cast<llvm::Instruction>(Operand->getUser()));
      llvm::Value *Call = Builder.CreateCall(PrettyFunction, { Val });
      Operand->set(Call);
    }
  }

  return true;
}

std::optional<FormatInt>
getIntFormat(llvm::Instruction &I, llvm::Use &U, const model::Binary &Model) {
  auto &Context = I.getContext();

  // We cannot print properly characters when they are part of switch
  // instruction, because cases in LLVM switch cannot have variables inside.
  if (I.getOpcode() == llvm::Instruction::Switch) {
    return std::nullopt;
  }

  // Some intrinsic calls require ConstantInt as an argument so we are not able
  // to pass there any decorated value.
  if (auto *Intrinsic = llvm::dyn_cast<llvm::IntrinsicInst>(&I)) {
    if (Intrinsic->getIntrinsicID() == llvm::Intrinsic::abs) {
      return std::nullopt;
    }
  }

  // We want to print ints in hex format when they are left operand of shifts or
  // operands of and/or/xor instructions.
  llvm::ConstantInt *IntConstant = dyn_cast<llvm::ConstantInt>(U);
  if (not IntConstant)
    return std::nullopt;

  if (I.getOpcode() == llvm::Instruction::Shl
      || I.getOpcode() == llvm::Instruction::AShr
      || I.getOpcode() == llvm::Instruction::LShr) {
    if (U.getOperandNo() == 0) {
      return FormatInt{ IntFormatting::HEX, &U };
    }
  } else if (I.getOpcode() == llvm::Instruction::And
             || I.getOpcode() == llvm::Instruction::Or
             || I.getOpcode() == llvm::Instruction::Xor) {
    return FormatInt{ IntFormatting::HEX, &U };
  } else if (auto *Call = getCallToTagged(&I, FunctionTags::ModelCast)) {
    // If it's a ModelCast casting a zero constanto to a pointer, then we
    // decorate the constant so that it's printed as NULL.
    if (IntConstant->isZero()) {
      model::QualifiedType
        Type = deserializeFromLLVMString(Call->getArgOperand(0), Model);
      if (Type.isPointer())
        return FormatInt{ IntFormatting::NULLPTR, &U };
    }
  }

  if (U->getType() == llvm::IntegerType::getInt8Ty(Context)) {
    return FormatInt{ IntFormatting::CHAR, &U };
  }

  if (U->getType() == llvm::IntegerType::getInt1Ty(Context)) {
    return FormatInt{ IntFormatting::BOOL, &U };
  }

  return std::nullopt;
}

char PrettyIntFormatting::ID = 0;

llvm::RegisterPass<PrettyIntFormatting> X("pretty-int-formatting",
                                          "Wraps integers with decorator "
                                          "functions which informs backend "
                                          "about literal type that should be "
                                          "used",
                                          false,
                                          false);