From 641fc7007c5cc554d8edd137257cbb9aca69e5c6 Mon Sep 17 00:00:00 2001
From: River Riddle <riverriddle@google.com>
Date: Tue, 6 Aug 2019 10:33:11 -0700
Subject: [PATCH] NFC: Simplify ModuleOp by using the
 SingleBlockImplicitTerminator trait. PiperOrigin-RevId: 261944712

---
 mlir/include/mlir/IR/Module.h       |  9 +++++++--
 mlir/lib/IR/Module.cpp              | 20 ++------------------
 mlir/test/IR/invalid-module-op.mlir |  4 ++--
 3 files changed, 11 insertions(+), 22 deletions(-)

diff --git a/mlir/include/mlir/IR/Module.h b/mlir/include/mlir/IR/Module.h
index 5936ae4..0bdd5f2 100644
--- a/mlir/include/mlir/IR/Module.h
+++ b/mlir/include/mlir/IR/Module.h
@@ -25,6 +25,8 @@
 #include "mlir/IR/SymbolTable.h"
 
 namespace mlir {
+class ModuleTerminatorOp;
+
 //===----------------------------------------------------------------------===//
 // Module Operation.
 //===----------------------------------------------------------------------===//
@@ -33,8 +35,11 @@ namespace mlir {
 /// single block containing opaque operations. The region of a module is not
 /// allowed to implicitly capture global values, and all external references
 /// must use symbolic references via attributes(e.g. via a string name).
-class ModuleOp : public Op<ModuleOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
-                           OpTrait::IsIsolatedFromAbove, OpTrait::SymbolTable> {
+class ModuleOp
+    : public Op<
+          ModuleOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
+          OpTrait::IsIsolatedFromAbove, OpTrait::SymbolTable,
+          OpTrait::SingleBlockImplicitTerminator<ModuleTerminatorOp>::Impl> {
 public:
   using Op::Op;
   using Op::print;
diff --git a/mlir/lib/IR/Module.cpp b/mlir/lib/IR/Module.cpp
index ff986b8..73510a1 100644
--- a/mlir/lib/IR/Module.cpp
+++ b/mlir/lib/IR/Module.cpp
@@ -25,16 +25,8 @@ using namespace mlir;
 // Module Operation.
 //===----------------------------------------------------------------------===//
 
-// Insert `module_terminator` at the end of the region's only block if it does
-// not have a terminator already. If the region is empty, insert a new block
-// first.
-static void ensureModuleTerminator(Region &region, Builder &builder,
-                                   Location loc) {
-  impl::ensureRegionTerminator<ModuleTerminatorOp>(region, builder, loc);
-}
-
 void ModuleOp::build(Builder *builder, OperationState *result) {
-  ensureModuleTerminator(*result->addRegion(), *builder, result->location);
+  ensureTerminator(*result->addRegion(), *builder, result->location);
 }
 
 /// Construct a module from the given context.
@@ -57,7 +49,7 @@ ParseResult ModuleOp::parse(OpAsmParser *parser, OperationState *result) {
     return failure();
 
   // Ensure that this module has a valid terminator.
-  ensureModuleTerminator(*body, parser->getBuilder(), result->location);
+  ensureTerminator(*body, parser->getBuilder(), result->location);
   return success();
 }
 
@@ -88,14 +80,6 @@ LogicalResult ModuleOp::verify() {
   if (body->getNumArguments() != 0)
     return emitOpError("expected body to have no arguments");
 
-  if (body->empty() || !isa<ModuleTerminatorOp>(body->back())) {
-    return emitOpError("expects region to end with '" +
-                       ModuleTerminatorOp::getOperationName() + "'")
-               .attachNote()
-           << "in custom textual format, the absence of terminator implies '"
-           << ModuleTerminatorOp::getOperationName() << "'";
-  }
-
   return success();
 }
 
diff --git a/mlir/test/IR/invalid-module-op.mlir b/mlir/test/IR/invalid-module-op.mlir
index c7e7a32..02ebe92 100644
--- a/mlir/test/IR/invalid-module-op.mlir
+++ b/mlir/test/IR/invalid-module-op.mlir
@@ -3,7 +3,7 @@
 // -----
 
 func @module_op() {
-  // expected-error@+1 {{expected body region to have a single block}}
+  // expected-error@+1 {{expects region #0 to have 0 or 1 blocks}}
   module {
   ^bb1:
     "module_terminator"() : () -> ()
@@ -27,7 +27,7 @@ func @module_op() {
 // -----
 
 func @module_op() {
-  // expected-error@+2 {{expects region to end with 'module_terminator'}}
+  // expected-error@+2 {{expects regions to end with 'module_terminator'}}
   // expected-note@+1 {{the absence of terminator implies 'module_terminator'}}
   module {
     return
-- 
2.7.4