[ReuseIR] implement checks for closure ops (#26)

reuse-mlir/include/ReuseIR/IR/ReuseIROps.td

@@ -189,38 +189,41 @@ def DestroyOp : ReuseIR_Op<"destroy"> {
   let hasVerifier = 0;
 }
 
-def NewClosureOp : ReuseIR_Op<"new_closure"> {
+def ClosureNewOp : ReuseIR_Op<"closure.new"> {
   let summary = "Create a new closure";
   let description = [{
-    `reuse_ir.new_closure` creates a new closure object. 
+    `reuse_ir.closure.new` creates a new closure object. 
     The closure can either be a value or wrapped in a rc pointer.
     The closure shall not capture mutable references (for now).
   }];
   let results = (outs Res<AnyTypeOf<[ReuseIR_RcTypeOf<[ReuseIR_ClosureType]>, ReuseIR_ClosureType]>, 
-    "Closure Object", [MemAlloc<DefaultResource>]>:$result);
+    "Closure Object", [MemAlloc<DefaultResource>]>:$closure);
   let regions = (region MinSizedRegion<1>:$body);
   let assemblyFormat = [{
-    $body `:` type($result) attr-dict
+    $body `:` type($closure) attr-dict
+  }];
+  let extraClassDeclaration = [{
+    ::mlir::reuse_ir::ClosureType getClosureType();
   }];
   // TODO: add verifier
-  let hasVerifier = 0;
+  let hasVerifier = 1;
 }
 
 def ClosureYieldOp : ReuseIR_Op<"closure.yield", [ReturnLike, Terminator,
-    ParentOneOf<["::mlir::reuse_ir::NewClosureOp"]>]> {
+    ParentOneOf<["::mlir::reuse_ir::ClosureNewOp"]>]> {
   let summary = "Yield closure return value";
   let description = [{
     `closure.yield` yields the return value of the closure.
     This serves as the terminator of the closure body.
   }];
 
-  let arguments = (ins AnyType:$value);
+  let arguments = (ins Optional<AnyType>:$value);
 
   let assemblyFormat = [{
-    $value `:` type($value) attr-dict
+   ($value^ `:` type($value))? attr-dict
   }];
 
   // TODO: add verifier
-  let hasVerifier = 0;
+  let hasVerifier = 1;
 }
 #endif // REUSE_IR_OPS

reuse-mlir/include/ReuseIR/IR/ReuseIRTypes.td

@@ -333,11 +333,11 @@ def ReuseIR_ClosureType : ReuseIR_Type<"Closure", "closure",
 
   let parameters = (ins
     ArrayRefParameter<"::mlir::Type">:$inputTypes,
-    "::mlir::Type":$outputType
+    OptionalParameter<"::mlir::Type">:$outputType
   );
 
   let assemblyFormat = [{
-    `<` `(` $inputTypes `)` `->` $outputType `>`
+    `<` `(` $inputTypes `)` (`->` $outputType^)? `>`
   }];
 }
 

reuse-mlir/src/cxx/IR/ReuseIROps.cpp

@@ -6,6 +6,7 @@
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "llvm/Support/ErrorHandling.h"
+#include <algorithm>
 
 namespace mlir {
 namespace REUSE_IR_DECL_SCOPE {
@@ -96,6 +97,45 @@ mlir::reuse_ir::LogicalResult LoadOp::verify() {
   return mlir::reuse_ir::success();
 }
 
+// ClosureNewOp
+ClosureType ClosureNewOp::getClosureType() {
+  return llvm::TypeSwitch<mlir::Type, ClosureType>(getClosure().getType())
+      .Case<RcType>(
+          [](const RcType &ty) { return cast<ClosureType>(ty.getPointee()); })
+      .Default([](const mlir::Type &ty) { return cast<ClosureType>(ty); });
+}
+
+mlir::reuse_ir::LogicalResult ClosureNewOp::verify() {
+  Region *region = &getRegion();
+  ClosureType closureTy = getClosureType();
+  if (region->getArguments().size() != closureTy.getInputTypes().size())
+    return emitOpError("the number of arguments in the region must match the "
+                       "number of input types in the closure type");
+  if (std::any_of(region->getArguments().begin(), region->getArguments().end(),
+                  [&](BlockArgument arg) {
+                    return arg.getType() !=
+                           closureTy.getInputTypes()[arg.getArgNumber()];
+                  }))
+    return emitOpError("the types of arguments in the region must match the "
+                       "input types in the closure type");
+  return mlir::reuse_ir::success();
+}
+
+// ClosureYieldOp
+mlir::reuse_ir::LogicalResult ClosureYieldOp::verify() {
+  ClosureNewOp op = getParentOp();
+  ClosureType closureTy = op.getClosureType();
+  if (getValue() && !closureTy.getOutputType())
+    return emitOpError("cannot yield a value in a closure without output");
+  if (!getValue() && closureTy.getOutputType())
+    return emitOpError("must yield a value in a closure with output");
+  if (getValue().getType() != closureTy.getOutputType())
+    return emitOpError("expected to yield a value of ")
+           << closureTy.getOutputType() << ", but " << getValue().getType()
+           << " is found instead";
+  return mlir::reuse_ir::success();
+}
+
 template <StringLiteral Literal> struct ParseKeywordAsUnitAttr {
   OptionalParseResult operator()(OpAsmParser &parser, UnitAttr &attr) {
     if (parser.parseOptionalKeyword(Literal).succeeded())

reuse-mlir/test/integration/basic/closure_syntax.mlir

@@ -3,7 +3,7 @@
 module @test {
     // CHECK: func.func @closure_test() -> !reuse_ir.rc<!reuse_ir.closure<(i32, i32) -> i32>, nonatomic, nonfreezing>
     func.func @closure_test() -> !reuse_ir.rc<!closure, nonatomic, nonfreezing> {
-        %1 = reuse_ir.new_closure {
+        %1 = reuse_ir.closure.new {
             ^bb(%arg0: i32, %arg1: i32):
                 %2 = arith.addi %arg0, %arg1 : i32
                 reuse_ir.closure.yield %2 : i32

reuse-mlir/test/integration/basic/invalid_closure_args_number.mlir

@@ -0,0 +1,12 @@
+// RUN: %not %reuse-opt %s 2>&1 | %FileCheck %s
+!closure = !reuse_ir.closure<(i32, i32) -> i32>
+module @test {
+    func.func @closure_test() -> !reuse_ir.rc<!closure, nonatomic, nonfreezing> {
+        // CHECK: error: 'reuse_ir.closure.new' op the number of arguments in the region must match the number of input types in the closure type
+        %1 = reuse_ir.closure.new {
+            ^bb(%arg0: i32):
+                reuse_ir.closure.yield %arg0 : i32
+        } : !reuse_ir.rc<!closure, nonatomic, nonfreezing>
+        return %1 : !reuse_ir.rc<!closure, nonatomic, nonfreezing>
+    }
+}

reuse-mlir/test/integration/basic/invalid_closure_args_type.mlir

@@ -0,0 +1,13 @@
+// RUN: %not %reuse-opt %s 2>&1 | %FileCheck %s
+!closure = !reuse_ir.closure<(index) -> index>
+module @test {
+    func.func @closure_test() -> !reuse_ir.rc<!closure, nonatomic, nonfreezing> {
+        // CHECK: error: 'reuse_ir.closure.new' op the types of arguments in the region must match the input types in the closure type
+        %1 = reuse_ir.closure.new {
+            ^bb(%arg0: i32):
+                %1 = arith.index_castui %arg0 : i32 to index
+                reuse_ir.closure.yield %1 : index
+        } : !reuse_ir.rc<!closure, nonatomic, nonfreezing>
+        return %1 : !reuse_ir.rc<!closure, nonatomic, nonfreezing>
+    }
+}

reuse-mlir/test/integration/basic/invalid_closure_nonnull_return.mlir

@@ -0,0 +1,13 @@
+// RUN: %not %reuse-opt %s 2>&1 | %FileCheck %s
+!closure = !reuse_ir.closure<(i32, i32)>
+module @test {
+    func.func @closure_test() -> !reuse_ir.rc<!closure, nonatomic, nonfreezing> {
+        // CHECK: error: 'reuse_ir.closure.yield' op cannot yield a value in a closure without output
+        %1 = reuse_ir.closure.new {
+            ^bb(%arg0: i32, %arg1: i32):
+                %2 = arith.addi %arg0, %arg1 : i32
+                reuse_ir.closure.yield %2 : i32
+        } : !reuse_ir.rc<!closure, nonatomic, nonfreezing>
+        return %1 : !reuse_ir.rc<!closure, nonatomic, nonfreezing>
+    }
+}

reuse-mlir/test/integration/basic/invalid_closure_null_return.mlir

@@ -0,0 +1,13 @@
+// RUN: %not %reuse-opt %s 2>&1 | %FileCheck %s
+!closure = !reuse_ir.closure<(i32, i32) -> i32>
+module @test {
+    func.func @closure_test() -> !reuse_ir.rc<!closure, nonatomic, nonfreezing> {
+        // CHECK: error: 'reuse_ir.closure.yield' op must yield a value in a closure with output
+        %1 = reuse_ir.closure.new {
+            ^bb(%arg0: i32, %arg1: i32):
+                %2 = arith.addi %arg0, %arg1 : i32
+                reuse_ir.closure.yield
+        } : !reuse_ir.rc<!closure, nonatomic, nonfreezing>
+        return %1 : !reuse_ir.rc<!closure, nonatomic, nonfreezing>
+    }
+}

reuse-mlir/test/integration/basic/invalid_closure_yield_type.mlir

@@ -0,0 +1,14 @@
+// RUN: %not %reuse-opt %s 2>&1 | %FileCheck %s
+!closure = !reuse_ir.closure<(i32, i32) -> i32>
+module @test {
+    func.func @closure_test() -> !reuse_ir.rc<!closure, nonatomic, nonfreezing> {
+        // CHECK: error: 'reuse_ir.closure.yield' op expected to yield a value of 'i32', but 'index' is found instead
+        %1 = reuse_ir.closure.new {
+            ^bb(%arg0: i32, %arg1: i32):
+                %2 = arith.addi %arg0, %arg1 : i32
+                %3 = arith.index_castui %2 : i32 to index
+                reuse_ir.closure.yield %3 : index
+        } : !reuse_ir.rc<!closure, nonatomic, nonfreezing>
+        return %1 : !reuse_ir.rc<!closure, nonatomic, nonfreezing>
+    }
+}