[mlir] Add lowering for abi.encodeWithSelector

Signed-off-by: Vladimir Radosavljevic <vr@matterlabs.dev>

Author: vladimirradosavljevic

mlir/include/mlir/Dialect/Sol/SolOps.td

@@ -346,12 +346,25 @@ def Sol_SliceOp : Sol_Op<"slice"> {
   }];
 }
 
-def Sol_EncodeOp : Sol_Op<"encode"> {
+def Sol_Bytes4SelectorType : Type<
+  CPred<"::llvm::isa<::mlir::sol::BytesType>($_self) && "
+        "::llvm::cast<::mlir::sol::BytesType>($_self).getSize() == 4">,
+  "!sol.bytes<4>">;
+
+def Sol_EncodeOp : Sol_Op<"encode", [AttrSizedOperandSegments]> {
   // When $packed is set, values are encoded without per-element padding.
-  let arguments = (ins Variadic<AnyType>:$ins, UnitAttr:$packed);
+  // When $selector is set, it is prepended as the first 4 bytes while $ins are
+  // ABI-encoded as a tuple.
+  let arguments = (ins Variadic<AnyType>:$ins,
+                       Optional<Sol_Bytes4SelectorType>:$selector,
+                       UnitAttr:$packed);
   // FIXME: AnyType -> Sol_StringType causes the asm printer to miss the prefix!
   let results = (outs AnyType:$res);
-  let assemblyFormat = "($ins^ `:` type($ins))? `:` type($res) attr-dict";
+  let assemblyFormat = [{
+    (`selector` `(` $selector^ `)`)? ($ins^)?
+    `:` type($selector) type($ins)
+    `:` type($res) attr-dict
+  }];
 }
 
 def Sol_DecodeOp : Sol_Op<"decode"> {

mlir/lib/Conversion/SolToStandard/SolToYul.cpp

@@ -2086,17 +2086,32 @@ struct EncodeOpLowering : public OpConversionPattern<sol::EncodeOp> {
     if (op.getPacked()) {
       Value dataStart =
           r.create<arith::AddIOp>(loc, freePtr, bExt.genI256Const(32));
-      Value dataEnd = evmB.genABIPackedEncoding(
-          op.getOperandTypes(), adaptor.getOperands(), dataStart);
+      Value dataEnd = evmB.genABIPackedEncoding(op.getIns().getType(),
+                                                adaptor.getIns(), dataStart);
       Value dataSize = r.create<arith::SubIOp>(loc, dataEnd, dataStart);
       r.create<yul::MStoreOp>(loc, freePtr, dataSize);
       Value allocationSize = r.create<arith::SubIOp>(loc, dataEnd, freePtr);
       evmB.genFreePtrUpd(freePtr, allocationSize);
+    } else if (op.getSelector()) {
+      assert(adaptor.getSelector() && "selector operand is required");
+
+      Value selectorAddr =
+          r.create<arith::AddIOp>(loc, freePtr, bExt.genI256Const(32));
+      r.create<yul::MStoreOp>(loc, selectorAddr, adaptor.getSelector());
+
+      Value tupleStart =
+          r.create<arith::AddIOp>(loc, selectorAddr, bExt.genI256Const(4));
+      Value tupleEnd = evmB.genABITupleEncoding(op.getIns().getType(),
+                                                adaptor.getIns(), tupleStart);
+      Value dataSize = r.create<arith::SubIOp>(loc, tupleEnd, selectorAddr);
+      r.create<yul::MStoreOp>(loc, freePtr, dataSize);
+      Value allocationSize = r.create<arith::SubIOp>(loc, tupleEnd, freePtr);
+      evmB.genFreePtrUpd(freePtr, allocationSize);
     } else {
       Value tupleStart =
           r.create<arith::AddIOp>(loc, freePtr, bExt.genI256Const(32));
-      Value tupleEnd = evmB.genABITupleEncoding(
-          op.getOperandTypes(), adaptor.getOperands(), tupleStart);
+      Value tupleEnd = evmB.genABITupleEncoding(op.getIns().getType(),
+                                                adaptor.getIns(), tupleStart);
       Value tupleSize = r.create<arith::SubIOp>(loc, tupleEnd, tupleStart);
       r.create<yul::MStoreOp>(loc, freePtr, tupleSize);
       Value allocationSize = r.create<arith::SubIOp>(loc, tupleEnd, freePtr);