path: root/willow/include/willow/IR/Instruction.h

#ifndef WILLOW_INCLUDE_IR_INSTRUCTION_H
#define WILLOW_INCLUDE_IR_INSTRUCTION_H

#include <willow/ADT/IList.h>
#include <willow/IR/Location.h>
#include <willow/IR/Types.h>
#include <willow/IR/Value.h>
#include <willow/Util/LogicalResult.h>

#include <cassert>
#include <cstddef>
#include <optional>
#include <ostream>
#include <utility>
#include <vector>

namespace willow {

class Value;
class BasicBlock;
class Function;

/// Helper structure for passing around the successors of a block
///
/// A terminator can induce at most two succesors, so no need to allocate.
class Successors {
  friend class Instruction;
  std::array<BasicBlock *, 2> data;
  uint8_t n;

  Successors() : data{}, n{0} {}
  explicit Successors(BasicBlock *a) : data{a}, n{1} {}
  Successors(BasicBlock *a, BasicBlock *b) : data{a, b}, n{2} {}

public:
  auto begin() { return data.data(); }
  auto begin() const { return data.data(); }
  auto end() { return data.data() + n; }
  auto end() const { return data.data() + n; }
  size_t size() const { return n; }

  BasicBlock *operator[](std::size_t i) const { return data[i]; }
};

/// Defines an IR instruction.
class Instruction : public Value, public IListTrait<Instruction> {
public:
  enum class Opcode {
    Add,  ///< Int addition
    Mul,  ///< Int multiplication
    Sub,  ///< Int subtraction
    Div,  ///< Int division
    Mod,  ///< Int modulo
    Shl,  ///< Int shift left
    Shr,  ///< Int shift right
    Ashl, ///< Int shift left arithmetic
    Ashr, ///< Int shift right arithmetic

    Eq, ///< %0 == %1
    Lt, ///< %0 <  %1
    Gt, ///< %0 >  %1
    Le, ///< %0 <= %1
    Ge, ///< %0 >= %1

    And, ///< logical and
    Or,  ///< logical or
    Not, ///< logical not

    Jmp,  ///< goto %0
    Br,   ///< goto (%0 ? %1 : %2)
    Call, ///< call %0 ...
    Ret,  ///< ret val?

    Phi,   ///< phi ^label1 %val1 ^label2 %val2 ...
    Alloca ///< %a : ptr<type> = alloca type
  };

private:
  BasicBlock *parent = nullptr;
  Opcode op;

  std::optional<Location> loc;
  std::vector<Value *> operands;

protected:
  void setOperand(std::size_t index, Value *operand) {
    assert(index < operands.size() && "Operand index out of bounds");
    assert(operand && "Operand cannot be null");
    Value *old = operands[index];
    if (old == operand)
      return;

    old->delUse(this);
    operands[index] = operand;
    operand->addUse(this);
  }

public:
  /// \param op Opcode for this instruction.
  /// \param type Type of the result of this instruction.
  /// \param loc Source location of this instruction.
  /// \param prev Previous instruction in the block.
  /// \param next Next instruction in the block.
  Instruction(Opcode op, Type type, std::optional<Location> loc = std::nullopt,
              Instruction *prev = nullptr, Instruction *next = nullptr)
      : Value(ValueKind::Instruction, type),
        IListTrait<Instruction>{prev, next}, op(op), loc(loc) {}

  ~Instruction() override {
    assert(!IListTrait<Instruction>::isLinked() &&
           "Instruction is destroyed before it is unlinked");
    assert(getUses().empty() && "Removing an instruction that still has uses");
    // TODO prob need a use wrapper, for op %1 %1
    for (Value *operand : operands)
      operand->delUse(this);
  }

  bool hasParent() const { return parent != nullptr; }
  BasicBlock *getParent() { return parent; }
  const BasicBlock *getParent() const { return parent; }
  void setParent(BasicBlock *block) { parent = block; }

  std::optional<Location> getLoc() const { return loc; }

  static bool isTerminatorOp(Opcode op);
  Successors succs();

  bool isTerminator() const { return isTerminatorOp(op); }

  Opcode opcode() const { return op; }
  std::size_t getNumOperands() const { return operands.size(); }

  Value *getOperand(std::size_t index) {
    assert(index < operands.size() && "Operand index out of bounds");
    return operands[index];
  }

  const Value *getOperand(std::size_t index) const {
    assert(index < operands.size() && "Operand index out of bounds");
    return operands[index];
  }

  std::vector<Value *> &getOperands() { return operands; }
  const std::vector<Value *> &getOperands() const { return operands; }

  void addOperand(Value *operand) {
    assert(operand && "Operand cannot be null");
    operands.push_back(operand);
    operand->addUse(this);
  }

  static constexpr std::string_view opcodeName(Opcode op);
};

constexpr std::string_view Instruction::opcodeName(Instruction::Opcode op) {
  using enum Instruction::Opcode;
  switch (op) {
  case Add:
    return "add";
  case Mul:
    return "mul";
  case Sub:
    return "sub";
  case Div:
    return "div";
  case Mod:
    return "mod";
  case Shl:
    return "shl";
  case Shr:
    return "shr";
  case Ashl:
    return "ashl";
  case Ashr:
    return "ashr";
  case Eq:
    return "eq";
  case Lt:
    return "lt";
  case Gt:
    return "gt";
  case Le:
    return "le";
  case Ge:
    return "ge";
  case And:
    return "and";
  case Or:
    return "or";
  case Not:
    return "not";
  case Jmp:
    return "jmp";
  case Br:
    return "br";
  case Call:
    return "call";
  case Ret:
    return "ret";
  case Phi:
    return "phi";
  case Alloca:
    return "alloca";
  }

  std::unreachable();
}

} // namespace willow

template <>
struct std::formatter<willow::Instruction::Opcode> {
  constexpr auto parse(std::format_parse_context &ctx) { return ctx.begin(); }

  constexpr auto format(const willow::Instruction::Opcode &op,
                        std::format_context &ctx) const {
    return std::format_to(ctx.out(), "{}", willow::Instruction::opcodeName(op));
  }
};

inline std::ostream &operator<<(std::ostream &os,
                                const willow::Instruction::Opcode op) {
  return os << willow::Instruction::opcodeName(op);
}

inline std::ostream &operator<<(std::ostream &os,
                                const willow::Instruction &inst) {
  auto vty = inst.getType();

  os << vty << " " << willow::Instruction::opcodeName(inst.opcode());

  for (auto *operand : inst.getOperands())
    os << " " << operand;

  return os;
}

#endif // WILLOW_INCLUDE_IR_INSTRUCTION_H