From 2e046be90e51df4978a72fb1665493ed3f57882e Mon Sep 17 00:00:00 2001
From: Rahul Joshi <jurahul@google.com>
Date: Thu, 16 Jul 2020 07:35:46 -0700
Subject: [PATCH] [flang] Adopt NoRegionArguments (WhereOp) and ParentOneOf
 (ResultOp) traits

Differential Revision: https://reviews.llvm.org/D83522
---
 .../include/flang/Optimizer/Dialect/FIROps.td |  6 ++--
 flang/lib/Optimizer/Dialect/FIROps.cpp        | 29 ++++---------------
 2 files changed, 10 insertions(+), 25 deletions(-)

diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index 58bf38aa8a4b..0bc543882a26 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -1853,7 +1853,9 @@ def fir_LenParamIndexOp : fir_OneResultOp<"len_param_index", [NoSideEffect]> {
 // Fortran loops
 //===----------------------------------------------------------------------===//
 
-def fir_ResultOp : fir_Op<"result", [NoSideEffect, ReturnLike, Terminator]> {
+def fir_ResultOp : fir_Op<"result",
+    [NoSideEffect, ReturnLike, Terminator,
+     ParentOneOf<["WhereOp", "LoopOp", "IterWhileOp"]>]> {
   let summary = "special terminator for use in fir region operations";
 
   let description = [{
@@ -1970,7 +1972,7 @@ def fir_LoopOp : region_Op<"do_loop",
   }];
 }
 
-def fir_WhereOp : region_Op<"if"> {
+def fir_WhereOp : region_Op<"if", [NoRegionArguments]> {
   let summary = "if-then-else conditional operation";
   let description = [{
     Used to conditionally execute operations. This operation is the FIR
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 44310d6e0691..36334167184d 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -968,19 +968,12 @@ static mlir::LogicalResult verify(fir::ResultOp op) {
   auto results = parentOp->getResults();
   auto operands = op.getOperands();
 
-  if (isa<fir::WhereOp>(parentOp) || isa<fir::LoopOp>(parentOp) ||
-      isa<fir::IterWhileOp>(parentOp)) {
-    if (parentOp->getNumResults() != op.getNumOperands())
-      return op.emitOpError() << "parent of result must have same arity";
-    for (auto e : llvm::zip(results, operands)) {
-      if (std::get<0>(e).getType() != std::get<1>(e).getType())
-        return op.emitOpError()
-               << "types mismatch between result op and its parent";
-    }
-  } else {
-    return op.emitOpError()
-           << "result only terminates if, do_loop, or iterate_while regions";
-  }
+  if (parentOp->getNumResults() != op.getNumOperands())
+    return op.emitOpError() << "parent of result must have same arity";
+  for (auto e : llvm::zip(results, operands))
+    if (std::get<0>(e).getType() != std::get<1>(e).getType())
+      return op.emitOpError()
+             << "types mismatch between result op and its parent";
   return success();
 }
 
@@ -1452,16 +1445,6 @@ static mlir::ParseResult parseWhereOp(OpAsmParser &parser,
 }
 
 static LogicalResult verify(fir::WhereOp op) {
-  // Verify that the entry of each child region does not have arguments.
-  for (auto &region : op.getOperation()->getRegions()) {
-    if (region.empty())
-      continue;
-
-    for (auto &b : region)
-      if (b.getNumArguments() != 0)
-        return op.emitOpError(
-            "requires that child entry blocks have no arguments");
-  }
   if (op.getNumResults() != 0 && op.otherRegion().empty())
     return op.emitOpError("must have an else block if defining values");