[mlir] Impl address(<contract>).code lowering (incl extcodeopy)

Author: abinavpp

libsolidity/codegen/mlir/Sol/SolOps.cpp

@@ -290,6 +290,17 @@ void EmitOp::print(OpAsmPrinter &p) {
     p << " : " << getArgs().getTypes();
 }
 
+//===----------------------------------------------------------------------===//
+// CodeOp
+//===----------------------------------------------------------------------===//
+
+void CodeOp::build(OpBuilder &odsBuilder, OperationState &odsState,
+                   Value contAddr) {
+  Type resTy =
+      sol::StringType::get(odsBuilder.getContext(), sol::DataLocation::Memory);
+  build(odsBuilder, odsState, resTy, contAddr);
+}
+
 //===----------------------------------------------------------------------===//
 // IfOp
 //===----------------------------------------------------------------------===//

libsolidity/codegen/mlir/Sol/SolOps.td

@@ -634,12 +634,27 @@ def Sol_NewOp : Sol_Op<"new", [AttrSizedOperandSegments]> {
                    Optional<UI256>:$salt,
                    Variadic<AnyType>:$ctorArgs);
   let results = (outs Sol_AddressType:$out);
+
   let assemblyFormat = [{
     $objName `value` `=` $val (`salt` `=` $salt^)?
     `ctor` `(` ($ctorArgs^ `:` type($ctorArgs))? `)` attr-dict `:` type($out)
   }];
 }
 
+def Sol_CodeOp : Sol_Op<"code"> {
+  let arguments = (ins Sol_AddressType:$contAddr);
+  // FIXME: AnyType -> Sol_StringType causes the asm printer to miss the prefix!
+  let results = (outs AnyType:$out);
+
+  let builders = [
+    OpBuilder<(ins "Value":$contAddr)>
+  ];
+
+  let assemblyFormat = [{
+    $contAddr attr-dict `:` type($contAddr) `->` type($out)
+  }];
+}
+
 def Sol_ReturnOp : Sol_Op<"return", [Pure, Terminator]> {
   // TODO: Verifier should check if a func op is an ancestor. ParentOneOf
   // doesn't work for this.

libsolidity/codegen/mlir/SolToStandardPass.cpp

@@ -159,6 +159,7 @@ struct ConvertSolToStandard
         sol::DecodeOp,
         sol::ExtCallOp,
         sol::NewOp,
+        sol::CodeOp,
         sol::EmitOp,
         sol::RequireOp,
         sol::ConvCastOp,

libsolidity/codegen/mlir/SolidityToMLIR.cpp

@@ -472,6 +472,9 @@ mlir::Type SolidityToMLIRPass::getType(Type const *ty) {
 
     return b.getFunctionType(inTys, outTys);
   }
+  case Type::Category::Contract:
+    // FIXME: 256 -> 160
+    return b.getIntegerType(256, /*isSigned=*/false);
   default:
     break;
   }
@@ -752,10 +755,15 @@ mlir::Value SolidityToMLIRPass::genExpr(MemberAccess const &memberAcc) {
                                       /*numBits=*/64, loc);
     return b.create<mlir::sol::GepOp>(loc, genRValExpr(memberAcc.expression()),
                                       memberIdx);
-    break;
   }
   default:
     break;
+  case Type::Category::Address: {
+    if (memberName == "code") {
+      return b.create<mlir::sol::CodeOp>(loc,
+                                         genRValExpr(memberAcc.expression()));
+    }
+  }
   }
 
   llvm_unreachable("NYI");

libsolidity/codegen/mlir/Target/EVM/SolToYul.cpp

@@ -1283,6 +1283,28 @@ struct NewOpLowering : public OpConversionPattern<sol::NewOp> {
   }
 };
 
+struct CodeOpLowering : public OpConversionPattern<sol::CodeOp> {
+  using OpConversionPattern<sol::CodeOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(sol::CodeOp op, OpAdaptor adaptor,
+                                ConversionPatternRewriter &r) const override {
+
+    Location loc = op.getLoc();
+    solidity::mlirgen::BuilderExt bExt(r, loc);
+    evm::Builder evmB(r, loc);
+
+    auto extCodeSize = r.create<yul::ExtCodeSizeOp>(loc, adaptor.getContAddr());
+    Value alloc = evmB.genMemAlloc(op.getType(), /*zeroInit=*/false,
+                                   /*initVals=*/{}, extCodeSize);
+    auto codeAddr = evmB.genDataAddrPtr(alloc, sol::DataLocation::Memory);
+    r.create<yul::ExtCodeCopyOp>(
+        loc, adaptor.getContAddr(), /*dstOffset=*/codeAddr,
+        /*srcOffset=*/bExt.genI256Const(0), extCodeSize);
+    r.replaceOp(op, alloc);
+    return success();
+  }
+};
+
 struct TryOpLowering : public OpConversionPattern<sol::TryOp> {
   using OpConversionPattern<sol::TryOp>::OpConversionPattern;
 
@@ -2231,8 +2253,8 @@ void evm::populateAbiPats(mlir::RewritePatternSet &pats,
 }
 
 void evm::populateExtCallPat(RewritePatternSet &pats, TypeConverter &tyConv) {
-  pats.add<ExtCallOpLowering, TryOpLowering, NewOpLowering>(tyConv,
-                                                            pats.getContext());
+  pats.add<ExtCallOpLowering, TryOpLowering, NewOpLowering, CodeOpLowering>(
+      tyConv, pats.getContext());
 }
 
 void evm::populateEmitPat(RewritePatternSet &pats, TypeConverter &tyConv) {

libsolidity/codegen/mlir/Target/EVM/YulToStandard.cpp

@@ -315,6 +315,23 @@ struct ExtCodeSizeOpLowering : public OpRewritePattern<yul::ExtCodeSizeOp> {
   }
 };
 
+struct ExtCodeCopyOpLowering : public OpRewritePattern<yul::ExtCodeCopyOp> {
+  using OpRewritePattern<yul::ExtCodeCopyOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(yul::ExtCodeCopyOp op,
+                                PatternRewriter &r) const override {
+    evm::Builder evmB(r, op.getLoc());
+
+    r.replaceOpWithNewOp<LLVM::IntrCallOp>(
+        op, llvm::Intrinsic::evm_extcodecopy, /*resTy=*/Type{}, /*ins=*/
+        ValueRange{op.getAddr(),
+                   /*dst=*/evmB.genHeapPtr(op.getDst()),
+                   /*src=*/evmB.genCodePtr(op.getSrc()), op.getSize()},
+        "evm.extcodecopy");
+    return success();
+  }
+};
+
 struct CreateOpLowering : public OpRewritePattern<yul::CreateOp> {
   using OpRewritePattern<yul::CreateOp>::OpRewritePattern;
 
@@ -699,6 +716,7 @@ void evm::populateYulPats(RewritePatternSet &pats) {
       CodeSizeOpLowering,
       CodeCopyOpLowering,
       ExtCodeSizeOpLowering,
+      ExtCodeCopyOpLowering,
       CreateOpLowering,
       Create2OpLowering,
       MLoadOpLowering,

libsolidity/codegen/mlir/Yul/YulOps.td

@@ -254,6 +254,12 @@ def Yul_ExtCodeSizeOp : Yul_Op<"extcodesize", [Pure]> {
   let assemblyFormat = "$addr attr-dict";
 }
 
+def Yul_ExtCodeCopyOp : Yul_Op<"extcodecopy"> {
+  let summary = "Represents the `extcodecopy` call in yul";
+  let arguments = (ins I256:$addr, I256:$dst, I256:$src, I256:$size);
+  let assemblyFormat = "$addr `,` $dst `,` $src `,` $size attr-dict";
+}
+
 def Yul_CreateOp : Yul_Op<"create"> {
   let summary = "Represents the `create` call in yul";
   let arguments = (ins I256:$val, I256:$addr, I256:$size);

libsolidity/codegen/mlir/YulToMLIR.cpp

@@ -308,6 +308,7 @@ void YulToMLIRPass::populateBuiltinGenMap() {
   defSimpleBuiltinGenNoRet<CodeCopyOp>("codecopy");
   defSimpleBuiltinGen<CodeSizeOp>("codesize");
   defSimpleBuiltinGen<ExtCodeSizeOp>("extcodesize");
+  defSimpleBuiltinGenNoRet<ExtCodeCopyOp>("extcodecopy");
   defSimpleBuiltinGen<CreateOp>("create");
   defSimpleBuiltinGen<Create2Op>("create2");
   defSimpleBuiltinGen<SLoadOp>("sload");

test/libsolidity/semanticTests/mlir/code.sol

@@ -0,0 +1,17 @@
+contract C {
+  function f() public returns (uint) { return 42; }
+}
+
+contract D {
+  function g() public returns (uint) {
+    C c1 = new C();
+    C c2 = c1;
+    require(address(c1).code.length > 50);
+    return 1;
+  }
+}
+
+// ====
+// compileViaMlir: true
+// ----
+// g() -> 1