Refactor the parsing/printing of the top-level module.

This changes the top-level module parser to handle the case where the top-level module is defined with the module operation syntax, i.e:
  module ... {
  }

The printer is also updated to always print the top-level module in this form. This allows for cleanly round-tripping the location and attributes of the top-level module.

PiperOrigin-RevId: 257492069

mlir/lib/IR/AsmPrinter.cpp

@@ -1616,14 +1616,9 @@ void ModulePrinter::print(ModuleOp module) {
   state.printAttributeAliases(os);
   state.printTypeAliases(os);
 
-  // Print the module body without the terminator. The terminator is not printed
-  // as part of the custom syntax for modules.
-  OperationPrinter opPrinter(module, *this);
-  auto *moduleBody = module.getBody();
-  for (auto &op : llvm::make_range(moduleBody->begin(), --moduleBody->end())) {
-    opPrinter.print(&op);
-    os << "\n\n";
-  }
+  // Print the module.
+  OperationPrinter(module, *this).print(module);
+  os << '\n';
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Parser/Parser.cpp

@@ -3951,8 +3951,30 @@ ParseResult ModuleParser::parseModule(ModuleOp module) {
       break;
 
     // If we got to the end of the file, then we're done.
-    case Token::eof:
-      return opParser.finalize();
+    case Token::eof: {
+      if (opParser.finalize())
+        return failure();
+
+      // Handle the case where the top level module was explicitly defined.
+      auto &bodyBlocks = module.getBodyRegion().getBlocks();
+      auto &operations = bodyBlocks.front().getOperations();
+      assert(!operations.empty() && "expected a valid module terminator");
+
+      // Check that the first operation is a module, and it is the only
+      // non-terminator operation.
+      ModuleOp nested = dyn_cast<ModuleOp>(operations.front());
+      if (nested && std::next(operations.begin(), 2) == operations.end()) {
+        // Merge the data of the nested module operation into 'module'.
+        module.setLoc(nested.getLoc());
+        module.setAttrs(nested.getOperation()->getAttrList());
+        bodyBlocks.splice(bodyBlocks.end(), nested.getBodyRegion().getBlocks());
+
+        // Erase the original module body.
+        bodyBlocks.pop_front();
+      }
+
+      return success();
+    }
 
     // If we got an error token, then the lexer already emitted an error, just
     // stop.  Someday we could introduce error recovery if there was demand

mlir/test/IR/module-op.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -mlir-print-debuginfo | FileCheck %s
 
 // CHECK: module {
 module {
@@ -19,3 +19,17 @@ module {
   // CHECK-NEXT: "foo.result_op"() : () -> i32
   %result = "foo.result_op"() : () -> i32
 }
+
+// -----
+
+// Check that a top-level module is always created, with location info.
+// CHECK: module {
+// CHECK-NEXT: } loc({{.*}}module-op.mlir{{.*}})
+
+// -----
+
+// Check that the top-level module can be defined via a single module operation.
+// CHECK: module {
+// CHECK-NOT: module {
+module {
+}