[mlir] Lower logical not, bitwise not and ternary expressions

Author: abinavpp

libsolidity/codegen/mlir/SolidityToMLIR.cpp

@@ -351,6 +351,9 @@ class SolidityToMLIRPass {
   /// Returns the mlir expression for the binary operation.
   mlir::Value genExpr(BinaryOperation const &binOp);
 
+  /// Returns the mlir expressions for the conditional (ternary) operation.
+  mlir::SmallVector<mlir::Value> genExprs(Conditional const &cond);
+
   /// Returns the mlir expression for the call.
   mlir::SmallVector<mlir::Value> genExprs(FunctionCall const &call);
 
@@ -369,8 +372,10 @@ class SolidityToMLIRPass {
   /// to the corresponding mlir type of `resTy`.
   mlir::Value genRValExpr(Expression const &expr,
                           std::optional<mlir::Type> resTy = std::nullopt);
-  mlir::Value genRValExpr(mlir::Value val, mlir::Location loc);
-  mlir::SmallVector<mlir::Value> genRValExprs(Expression const &expr);
+  mlir::Value genRValExpr(mlir::Value val, mlir::Location loc,
+                          std::optional<mlir::Type> resTy = std::nullopt);
+  mlir::SmallVector<mlir::Value> genRValExprs(Expression const &expr,
+                                              mlir::TypeRange resTys = {});
 
   /// Generates an ir that assigns `rhs` to `lhs`.
   void genAssign(mlir::Value lhs, mlir::Value rhs, mlir::Location loc);
@@ -849,6 +854,31 @@ mlir::Value SolidityToMLIRPass::genExpr(UnaryOperation const &unaryOp) {
         b.create<mlir::sol::ConstantOp>(loc, b.getIntegerAttr(mlirTy, 0));
     return genBinExpr(Token::Sub, zero, expr, loc);
   }
+  // Logical not
+  case Token::Not: {
+    mlir::Value expr = genRValExpr(unaryOp.subExpression());
+    mlir::Value zero = b.create<mlir::sol::ConstantOp>(
+        loc, b.getIntegerAttr(expr.getType(), 0));
+    return b.create<mlir::sol::CmpOp>(loc, mlir::sol::CmpPredicate::eq, expr,
+                                      zero);
+  }
+  // Bitwise not (~x == x ^ -1)
+  case Token::BitNot: {
+    mlir::Value expr = genRValExpr(unaryOp.subExpression());
+
+    // Convert bytes to int if needed.
+    if (auto bytesTy = mlir::dyn_cast<mlir::sol::BytesType>(expr.getType())) {
+      mlir::Type intTy =
+          b.getIntegerType(bytesTy.getSize() * 8, /*isSigned=*/false);
+      expr = genCast(expr, intTy);
+    }
+
+    auto intTy = mlir::cast<mlir::IntegerType>(expr.getType());
+    mlir::Value allOnes = b.create<mlir::sol::ConstantOp>(
+        loc,
+        b.getIntegerAttr(intTy, llvm::APInt::getAllOnes(intTy.getWidth())));
+    return b.create<mlir::sol::XorOp>(loc, expr, allOnes);
+  }
   default:
     break;
   }
@@ -910,6 +940,38 @@ mlir::Value SolidityToMLIRPass::genExpr(BinaryOperation const &binOp) {
   return genBinExpr(binOp.getOperator(), lhs, rhs, loc);
 }
 
+mlir::SmallVector<mlir::Value>
+SolidityToMLIRPass::genExprs(Conditional const &cond) {
+  mlir::Location loc = getLoc(cond);
+  mlir::Value condVal = genRValExpr(cond.condition());
+
+  // Get result types - could be single type or tuple.
+  mlir::SmallVector<mlir::Type> resTys;
+  if (TupleType const *tupleTy =
+          dynamic_cast<TupleType const *>(cond.annotation().type)) {
+    for (const Type *astTy : tupleTy->components())
+      resTys.push_back(getType(astTy));
+  } else {
+    resTys.push_back(getType(cond.annotation().type));
+  }
+
+  auto ifOp =
+      b.create<mlir::scf::IfOp>(loc, resTys, condVal, /*withElse=*/true);
+  mlir::OpBuilder::InsertionGuard guard(b);
+
+  // True branch
+  b.setInsertionPointToStart(&ifOp.getThenRegion().front());
+  b.create<mlir::scf::YieldOp>(loc,
+                               genRValExprs(cond.trueExpression(), resTys));
+
+  // False branch
+  b.setInsertionPointToStart(&ifOp.getElseRegion().front());
+  b.create<mlir::scf::YieldOp>(loc,
+                               genRValExprs(cond.falseExpression(), resTys));
+
+  return ifOp.getResults();
+}
+
 mlir::Value SolidityToMLIRPass::genExpr(IndexAccess const &idxAcc) {
   mlir::Location loc = getLoc(idxAcc);
 
@@ -1935,6 +1997,13 @@ mlir::Value SolidityToMLIRPass::genLValExpr(Expression const &expr) {
     return {};
   }
 
+  // Conditional (ternary operator)
+  if (const auto *cond = dynamic_cast<Conditional const *>(&expr)) {
+    mlir::SmallVector<mlir::Value> exprs = genExprs(*cond);
+    assert(exprs.size() == 1);
+    return exprs[0];
+  }
+
   llvm_unreachable("NYI");
 }
 
@@ -1948,38 +2017,46 @@ SolidityToMLIRPass::genLValExprs(Expression const &expr) {
   if (const auto *call = dynamic_cast<FunctionCall const *>(&expr))
     return genExprs(*call);
 
+  // Conditional (ternary)
+  if (const auto *cond = dynamic_cast<Conditional const *>(&expr))
+    return genExprs(*cond);
+
   mlir::SmallVector<mlir::Value, 1> vals;
   vals.push_back(genLValExpr(expr));
   return vals;
 }
 
-mlir::Value SolidityToMLIRPass::genRValExpr(mlir::Value val,
-                                            mlir::Location loc) {
+mlir::Value SolidityToMLIRPass::genRValExpr(mlir::Value val, mlir::Location loc,
+                                            std::optional<mlir::Type> resTy) {
   if (mlir::isa<mlir::sol::PointerType>(val.getType()))
-    return b.create<mlir::sol::LoadOp>(loc, val);
+    val = b.create<mlir::sol::LoadOp>(loc, val);
+  if (resTy)
+    return genCast(val, *resTy);
   return val;
 }
 
 mlir::Value SolidityToMLIRPass::genRValExpr(Expression const &expr,
                                             std::optional<mlir::Type> resTy) {
   mlir::Value lVal = genLValExpr(expr);
   assert(lVal);
-
-  mlir::Value val = genRValExpr(lVal, getLoc(expr));
-  // Generate cast (optional).
-  if (resTy)
-    return genCast(val, *resTy);
-  return val;
+  return genRValExpr(lVal, getLoc(expr), resTy);
 }
 
 mlir::SmallVector<mlir::Value>
-SolidityToMLIRPass::genRValExprs(Expression const &expr) {
+SolidityToMLIRPass::genRValExprs(Expression const &expr,
+                                 mlir::TypeRange resTys) {
   mlir::SmallVector<mlir::Value> lVals = genLValExprs(expr);
   assert(!lVals.empty());
+  assert(resTys.empty() || lVals.size() == resTys.size());
 
   mlir::SmallVector<mlir::Value, 2> rVals;
-  for (mlir::Value lVal : lVals)
-    rVals.push_back(genRValExpr(lVal, getLoc(expr)));
+  if (resTys.empty()) {
+    for (mlir::Value lVal : lVals)
+      rVals.push_back(genRValExpr(lVal, getLoc(expr)));
+  } else {
+    for (auto [lVal, resTy] : llvm::zip(lVals, resTys))
+      rVals.push_back(genRValExpr(lVal, getLoc(expr), resTy));
+  }
 
   return rVals;
 }
@@ -1992,22 +2069,23 @@ void SolidityToMLIRPass::lower(
     VariableDeclarationStatement const &varDeclStmt) {
   mlir::Location loc = getLoc(varDeclStmt);
 
+  mlir::SmallVector<mlir::Type> varTys;
+  for (auto const &varDeclPtr : varDeclStmt.declarations())
+    varTys.push_back(getType(varDeclPtr->type(), /*indirectFn=*/true));
+
   mlir::SmallVector<mlir::Value> initExprs(varDeclStmt.declarations().size());
   if (Expression const *initExpr = varDeclStmt.initialValue())
-    initExprs = genRValExprs(*initExpr);
-
-  for (auto [varDeclPtr, initExpr] :
-       llvm::zip(varDeclStmt.declarations(), initExprs)) {
-    VariableDeclaration const &varDecl = *varDeclPtr;
+    initExprs = genRValExprs(*initExpr, varTys);
 
-    mlir::Type varTy = getType(varDecl.type(), /*indirectFn=*/true);
+  for (auto [varDeclPtr, varTy, initExpr] :
+       llvm::zip(varDeclStmt.declarations(), varTys, initExprs)) {
     mlir::Type allocTy = mlir::sol::PointerType::get(
         b.getContext(), varTy, mlir::sol::DataLocation::Stack);
 
     auto addr = b.create<mlir::sol::AllocaOp>(loc, allocTy);
-    trackLocalVarAddr(varDecl, addr);
+    trackLocalVarAddr(*varDeclPtr, addr);
     if (initExpr)
-      b.create<mlir::sol::StoreOp>(loc, genCast(initExpr, varTy), addr);
+      b.create<mlir::sol::StoreOp>(loc, initExpr, addr);
     else
       genZeroedVal(addr);
   }
@@ -2038,22 +2116,17 @@ void SolidityToMLIRPass::lower(PlaceholderStatement const &placeholder) {
 }
 
 void SolidityToMLIRPass::lower(Return const &ret) {
-  TypePointers fnResTys;
+  mlir::SmallVector<mlir::Type> fnResTys;
   for (ASTPointer<VariableDeclaration> const &retParam :
        ret.annotation().function->returnParameters())
-    fnResTys.push_back(retParam->type());
+    fnResTys.push_back(getType(retParam->type()));
 
   Expression const *astExpr = ret.expression();
-  if (astExpr) {
-    mlir::SmallVector<mlir::Value> exprs = genRValExprs(*astExpr);
-    mlir::SmallVector<mlir::Value> castedExprs;
-    for (auto [expr, dstTy] : llvm::zip(exprs, fnResTys)) {
-      castedExprs.push_back(genCast(expr, getType(dstTy)));
-    }
-    b.create<mlir::sol::ReturnOp>(getLoc(ret), castedExprs);
-  } else {
+  if (astExpr)
+    b.create<mlir::sol::ReturnOp>(getLoc(ret),
+                                  genRValExprs(*astExpr, fnResTys));
+  else
     b.create<mlir::sol::ReturnOp>(getLoc(ret));
-  }
   b.setInsertionPointToStart(b.createBlock(b.getBlock()->getParent()));
 }
 

test/libsolidity/semanticTests/mlir/arith.sol

@@ -77,6 +77,14 @@ contract C {
     return (a & b, a | b, a ^ b);
   }
 
+  function bnot(uint a) public returns (uint) {
+    return ~a;
+  }
+
+  function bnotBytes(bytes4 a) public pure returns (bytes4) {
+    return ~a;
+  }
+
   function addmod8(uint8 a, uint8 b, uint8 m) public returns (uint256) {
     return addmod(a, b, m);
   }
@@ -123,6 +131,8 @@ contract C {
 // shr8(uint8,uint256): 1, 8 -> 0
 // bit(int256,int256): 6, 3 -> 2, 7, 5
 // bit8(int8,int8): 6, 3 -> 2, 7, 5
+// bnot(uint256): 0 -> 0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff
+// bnotBytes(bytes4): left(0x12345678) -> left(0xedcba987)
 // addmod8(uint8,uint8,uint8): 42, 25, 0 -> FAILURE, hex"4e487b71", 0x12
 // addmod8(uint8,uint8,uint8): 42, 25, 24 -> 19
 // addmod8(uint8,uint8,uint8): 200, 100, 7 -> 6

test/libsolidity/semanticTests/mlir/cf.sol

@@ -1,5 +1,7 @@
 contract C {
-  function f(uint a, uint b) public returns (uint) {
+  uint public counter;
+
+  function for_brk(uint a, uint b) public returns (uint) {
     uint r = 1;
     for (uint i = 0; i < a; ++i) {
       if (i == b)
@@ -9,7 +11,7 @@ contract C {
     return r;
   }
 
-  function g(uint a) public returns (uint) {
+  function while_cont(uint a) public returns (uint) {
     uint r = 1;
     do {
       r = 2;
@@ -26,10 +28,48 @@ contract C {
     }
     return r;
   }
+
+  function tern(bool c, uint a, uint b) public pure returns (uint) {
+    return c ? a : b;
+  }
+
+  function tern_bytes(bool c, bytes4 a, bytes4 b) public pure returns (bytes4) {
+    return c ? a : b;
+  }
+
+  function tern_cast(bool c, uint8 a, uint256 b) public pure returns (uint256) {
+    return c ? a : b;
+  }
+
+  function tern_short_circuit(bool c) public returns (uint, uint) {
+    counter = 0;
+    uint result = c ? ++counter : ++counter;
+    return (result, counter);
+  }
+
+  function tern_tuple(bool c, uint a, uint b, uint x, uint y) public pure returns (uint, uint) {
+    return c ? (a, b) : (x, y);
+  }
+
+  function tern_const(bool c) public pure returns (uint) {
+    return c ? 1 : 2;
+  }
 }
 
 // ====
 // compileViaMlir: true
 // ----
-// f(uint256,uint256): 20, 10 -> 1024
-// g(uint256): 10 -> 22
+// for_brk(uint256,uint256): 20, 10 -> 1024
+// while_cont(uint256): 10 -> 22
+// tern(bool,uint256,uint256): true, 10, 20 -> 10
+// tern(bool,uint256,uint256): false, 10, 20 -> 20
+// tern_bytes(bool,bytes4,bytes4): true, left(0x12345678), left(0xaabbccdd) -> left(0x12345678)
+// tern_bytes(bool,bytes4,bytes4): false, left(0x12345678), left(0xaabbccdd) -> left(0xaabbccdd)
+// tern_cast(bool,uint8,uint256): true, 42, 1000 -> 42
+// tern_cast(bool,uint8,uint256): false, 42, 1000 -> 1000
+// tern_short_circuit(bool): true -> 1, 1
+// tern_short_circuit(bool): false -> 1, 1
+// tern_tuple(bool,uint256,uint256,uint256,uint256): true, 10, 20, 30, 40 -> 10, 20
+// tern_tuple(bool,uint256,uint256,uint256,uint256): false, 10, 20, 30, 40 -> 30, 40
+// tern_const(bool): true -> 1
+// tern_const(bool): false -> 2

test/libsolidity/semanticTests/mlir/logical.sol

@@ -23,6 +23,10 @@ contract C {
     x = 0;
     return (a || s(b), x);
   }
+
+  function j(bool a) public pure returns (bool) {
+    return !a;
+  }
 }
 
 // ====
@@ -42,3 +46,5 @@ contract C {
 // i(bool,bool): 1, 0 -> true, 0
 // i(bool,bool): 0, 1 -> true, 1
 // i(bool,bool): 0, 0 -> false, 1
+// j(bool): 0 -> true
+// j(bool): 1 -> false