[EVM] Improve Operation implementation

Author: akiramenai

llvm/lib/Target/EVM/EVMStackDebug.cpp

@@ -19,11 +19,6 @@
 
 using namespace llvm;
 
-template <class... Ts> struct Overload : Ts... {
-  using Ts::operator()...;
-};
-template <class... Ts> Overload(Ts...) -> Overload<Ts...>;
-
 std::string llvm::stackToString(const Stack &S) {
   std::string Result("[ ");
   for (const auto *Slot : S)
@@ -55,32 +50,14 @@ void StackLayoutPrinter::operator()() {
 void StackLayoutPrinter::printBlock(MachineBasicBlock const &Block) {
   OS << "Block" << getBlockId(Block) << " [\n";
   OS << stackToString(Layout.getMBBEntryLayout(&Block)) << "\n";
-  for (auto const &Operation : StackModel.getOperations(&Block)) {
+  for (auto const &Op : StackModel.getOperations(&Block)) {
     OS << "\n";
-    Stack EntryLayout = Layout.getOperationEntryLayout(&Operation);
+    Stack EntryLayout = Layout.getOperationEntryLayout(&Op);
     OS << stackToString(EntryLayout) << "\n";
-    std::visit(Overload{[&](FunctionCall const &Call) {
-                          const MachineOperand *Callee =
-                              Call.MI->explicit_uses().begin();
-                          OS << Callee->getGlobal()->getName();
-                        },
-                        [&](BuiltinCall const &Call) {
-                          OS << getInstName(Call.MI);
-                        },
-                        [&](Assignment const &Assignment) {
-                          OS << "Assignment(";
-                          for (const auto *Var : Assignment.Variables)
-                            OS << printReg(Var->getReg(), nullptr, 0, nullptr)
-                               << ", ";
-                          OS << ")";
-                        }},
-               Operation.Operation);
-    OS << "\n";
-
-    assert(Operation.Input.size() <= EntryLayout.size());
-    for (size_t i = 0; i < Operation.Input.size(); ++i)
-      EntryLayout.pop_back();
-    EntryLayout.append(Operation.Output);
+    OS << Op.toString() << "\n";
+    assert(Op.getInput().size() <= EntryLayout.size());
+    EntryLayout.resize(EntryLayout.size() - Op.getInput().size());
+    EntryLayout.append(Op.getOutput());
     OS << stackToString(EntryLayout) << "\n";
   }
   OS << "\n";

llvm/lib/Target/EVM/EVMStackLayoutGenerator.cpp

@@ -261,10 +261,9 @@ std::unique_ptr<EVMStackLayout> EVMStackLayoutGenerator::run() {
 }
 
 Stack EVMStackLayoutGenerator::propagateStackThroughOperation(
-    Stack ExitStack, const Operation &Operation,
-    bool AggressiveStackCompression) {
+    Stack ExitStack, const Operation &Op, bool AggressiveStackCompression) {
   // Enable aggressive stack compression for recursive calls.
-  if (std::holds_alternative<FunctionCall>(Operation.Operation))
+  if (Op.isFunctionCall())
     // TODO: compress stack for recursive functions.
     AggressiveStackCompression = false;
 
@@ -277,25 +276,25 @@ Stack EVMStackLayoutGenerator::propagateStackThroughOperation(
   // operation outputs (and not to be generated on the fly), s.t. shuffling the
   // 'IdealStack + Operation.output' to ExitLayout is cheap.
   Stack IdealStack =
-      createIdealLayout(Operation.Output, ExitStack, generateSlotOnTheFly);
+      createIdealLayout(Op.getOutput(), ExitStack, generateSlotOnTheFly);
 
   // Make sure the resulting previous slots do not overlap with any assignmed
   // variables.
-  if (auto const *Assign = std::get_if<Assignment>(&Operation.Operation))
+  if (Op.isAssignment())
     for (auto *StackSlot : IdealStack)
       if (const auto *VarSlot = dyn_cast<VariableSlot>(StackSlot))
-        assert(!is_contained(Assign->Variables, VarSlot));
+        assert(!is_contained(Op.getOutput(), VarSlot));
 
   // Since stack+Operation.output can be easily shuffled to ExitLayout, the
   // desired layout before the operation is stack+Operation.input;
-  IdealStack.append(Operation.Input);
+  IdealStack.append(Op.getInput());
 
   // Store the exact desired operation entry layout. The stored layout will be
   // recreated by the code transform before executing the operation. However,
   // this recreation can produce slots that can be freely generated or are
   // duplicated, i.e. we can compress the stack afterwards without causing
   // problems for code generation later.
-  OperationEntryLayoutMap[&Operation] = IdealStack;
+  OperationEntryLayoutMap[&Op] = IdealStack;
 
   // Remove anything from the stack top that can be freely generated or dupped
   // from deeper on the stack.
@@ -751,10 +750,10 @@ void EVMStackLayoutGenerator::addJunksToStackBottom(
     EntryTmp.append(MBBEntryLayoutMap.at(MBB));
     MBBEntryLayoutMap[MBB] = std::move(EntryTmp);
 
-    for (const Operation &Operation : StackModel.getOperations(MBB)) {
+    for (const Operation &Op : StackModel.getOperations(MBB)) {
       Stack OpEntryTmp(NumJunk, EVMStackModel::getJunkSlot());
-      OpEntryTmp.append(OperationEntryLayoutMap.at(&Operation));
-      OperationEntryLayoutMap[&Operation] = std::move(OpEntryTmp);
+      OpEntryTmp.append(OperationEntryLayoutMap.at(&Op));
+      OperationEntryLayoutMap[&Op] = std::move(OpEntryTmp);
     }
 
     Stack ExitTmp(NumJunk, EVMStackModel::getJunkSlot());

llvm/lib/Target/EVM/EVMStackModel.cpp

@@ -29,8 +29,7 @@ static const Function *getCalledFunction(const MachineInstr &MI) {
   }
   return nullptr;
 }
-// TODO: make it static once Operation gets rid of std::variant.
-std::string llvm::getInstName(const MachineInstr *MI) {
+static std::string getInstName(const MachineInstr *MI) {
   const MachineFunction *MF = MI->getParent()->getParent();
   const TargetInstrInfo *TII = MF->getSubtarget().getInstrInfo();
   return TII->getName(MI->getOpcode()).str();
@@ -48,6 +47,24 @@ std::string TemporarySlot::toString() const {
   OS << "TMP[" << getInstName(MI) << ", " << std::to_string(Index) + "]";
   return std::string(S.str());
 }
+std::string Operation::toString() const {
+  if (isFunctionCall()) {
+    const MachineOperand *Callee = MI->explicit_uses().begin();
+    return Callee->getGlobal()->getName().str();
+  }
+  if (isBuiltinCall())
+    return getInstName(MI);
+
+  assert(isAssignment());
+  SmallString<128> S;
+  raw_svector_ostream OS(S);
+  OS << "Assignment(";
+  for (const auto *S : Output)
+    OS << printReg(cast<VariableSlot>(S)->getReg(), nullptr, 0, nullptr)
+       << ", ";
+  OS << ")";
+  return std::string(S);
+}
 
 EVMStackModel::EVMStackModel(MachineFunction &MF, const LiveIntervals &LIS)
     : MF(MF), LIS(LIS) {
@@ -128,22 +145,18 @@ void EVMStackModel::createOperation(MachineInstr &MI,
     return;
   case EVM::FCALL: {
     Stack Input;
-    bool IsNoReturn = false;
     for (const MachineOperand &MO : MI.operands()) {
       if (MO.isGlobal()) {
-        const auto *Func = dyn_cast<Function>(MO.getGlobal());
-        assert(Func);
-        IsNoReturn = Func->hasFnAttribute(Attribute::NoReturn);
-        if (!IsNoReturn)
+        const auto *Func = cast<Function>(MO.getGlobal());
+        if (!Func->hasFnAttribute(Attribute::NoReturn))
           Input.push_back(getFunctionCallReturnLabelSlot(&MI));
         break;
       }
     }
     const Stack &Tmp = getInstrInput(MI);
     Input.insert(Input.end(), Tmp.begin(), Tmp.end());
-    size_t NumArgs = Input.size() - (IsNoReturn ? 0 : 1);
-    Ops.emplace_back(Operation{std::move(Input), getInstrOutput(MI),
-                               FunctionCall{&MI, NumArgs}});
+    Ops.emplace_back(Operation::FunctionCall, std::move(Input),
+                     getInstrOutput(MI), &MI);
   } break;
   case EVM::RET:
   case EVM::JUMP:
@@ -165,21 +178,19 @@ void EVMStackModel::createOperation(MachineInstr &MI,
       return;
   } break;
   default: {
-    Ops.emplace_back(
-        Operation{getInstrInput(MI), getInstrOutput(MI), BuiltinCall{&MI}});
+    Ops.emplace_back(Operation::BuiltinCall, getInstrInput(MI),
+                     getInstrOutput(MI), &MI);
   } break;
   }
 
   // Create CFG::Assignment object for the MI.
   Stack Input, Output;
-  SmallVector<VariableSlot *> Variables;
   switch (MI.getOpcode()) {
   case EVM::CONST_I256: {
     const Register DefReg = MI.getOperand(0).getReg();
     const APInt Imm = MI.getOperand(1).getCImm()->getValue();
     Input.push_back(getLiteralSlot(std::move(Imm)));
     Output.push_back(getVariableSlot(DefReg));
-    Variables.push_back(getVariableSlot(DefReg));
   } break;
   case EVM::DATASIZE:
   case EVM::DATAOFFSET:
@@ -188,15 +199,13 @@ void EVMStackModel::createOperation(MachineInstr &MI,
     MCSymbol *Sym = MI.getOperand(1).getMCSymbol();
     Input.push_back(getSymbolSlot(Sym, &MI));
     Output.push_back(getVariableSlot(DefReg));
-    Variables.push_back(getVariableSlot(DefReg));
   } break;
   case EVM::COPY_I256: {
     // Copy instruction corresponds to the assignment operator, so
     // we do not need to create intermediate TmpSlots.
     Input = getInstrInput(MI);
     const Register DefReg = MI.getOperand(0).getReg();
     Output.push_back(getVariableSlot(DefReg));
-    Variables.push_back(getVariableSlot(DefReg));
   } break;
   default: {
     unsigned ArgsNumber = 0;
@@ -205,15 +214,14 @@ void EVMStackModel::createOperation(MachineInstr &MI,
       const Register Reg = MO.getReg();
       Input.push_back(getTemporarySlot(&MI, ArgsNumber++));
       Output.push_back(getVariableSlot(Reg));
-      Variables.push_back(getVariableSlot(Reg));
     }
   } break;
   }
   // We don't need an assignment part of the instructions that do not write
   // results.
   if (!Input.empty() || !Output.empty())
-    Ops.emplace_back(Operation{std::move(Input), std::move(Output),
-                               Assignment{std::move(Variables)}});
+    Ops.emplace_back(Operation::Assignment, std::move(Input), std::move(Output),
+                     &MI);
 }
 
 Stack EVMStackModel::getReturnArguments(const MachineInstr &MI) const {

llvm/lib/Target/EVM/EVMStackModel.h

@@ -36,8 +36,6 @@ namespace llvm {
 class MachineFunction;
 class MachineBasicBlock;
 
-std::string getInstName(const MachineInstr *MI);
-
 class StackSlot {
 public:
   enum SlotKind {
@@ -203,28 +201,33 @@ class JunkSlot final : public StackSlot {
 /// The stack top is the last element of the vector.
 using Stack = SmallVector<StackSlot *>;
 
-struct BuiltinCall {
+class Operation {
+public:
+  enum OpType { BuiltinCall, FunctionCall, Assignment };
+
+private:
+  OpType Type;
+  // Stack slots this operation expects at the top of the stack and consumes.
+  Stack Input;
+  // Stack slots this operation leaves on the stack as output.
+  Stack Output;
+  // The emulated machine instruction.
   MachineInstr *MI = nullptr;
-};
 
-struct FunctionCall {
-  const MachineInstr *MI;
-  size_t NumArguments = 0;
-};
+public:
+  Operation(OpType Type, Stack Input, Stack Output, MachineInstr *MI)
+      : Type(Type), Input(std::move(Input)), Output(std::move(Output)), MI(MI) {
+  }
 
-struct Assignment {
-  /// The variables being assigned to also occur as 'Output' in the
-  /// 'Operation' containing the assignment, but are also stored here for
-  /// convenience.
-  SmallVector<VariableSlot *> Variables;
-};
+  const Stack &getInput() const { return Input; }
+  const Stack &getOutput() const { return Output; }
+  MachineInstr *getMachineInstr() const { return MI; }
 
-struct Operation {
-  /// Stack slots this operation expects at the top of the stack and consumes.
-  Stack Input;
-  /// Stack slots this operation leaves on the stack as output.
-  Stack Output;
-  std::variant<FunctionCall, BuiltinCall, Assignment> Operation;
+  bool isBuiltinCall() const { return Type == BuiltinCall; }
+  bool isFunctionCall() const { return Type == FunctionCall; }
+  bool isAssignment() const { return Type == Assignment; }
+
+  std::string toString() const;
 };
 
 class EVMStackModel {