From 66ec24d83310dd623d79619840313f98e9a72644 Mon Sep 17 00:00:00 2001
From: Sean Silva <silvasean@google.com>
Date: Mon, 28 Oct 2019 15:11:00 -0700
Subject: [PATCH] Parse locations in parseGenericOperation

For ops that recursively re-enter the parser to parse an operation (such as
ops with a "wraps" pretty form), this ensures that the wrapped op will parse
its location, which can then be used for the locations of the wrapping op
and any other implicit ops.

PiperOrigin-RevId: 277152636
---
 mlir/lib/Parser/Parser.cpp                | 17 ++++++++++-------
 mlir/test/IR/wrapping_op.mlir             | 12 ++++++------
 mlir/test/lib/TestDialect/TestDialect.cpp |  6 +++++-
 3 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/mlir/lib/Parser/Parser.cpp b/mlir/lib/Parser/Parser.cpp
index c3f9081..af7e0b6 100644
--- a/mlir/lib/Parser/Parser.cpp
+++ b/mlir/lib/Parser/Parser.cpp
@@ -271,8 +271,7 @@ public:
   ///
   ///   trailing-location     ::= location?
   ///
-  template <typename Owner>
-  ParseResult parseOptionalTrailingLocation(Owner *owner) {
+  ParseResult parseOptionalTrailingLocation(Location &loc) {
     // If there is a 'loc' we parse a trailing location.
     if (!getToken().is(Token::kw_loc))
       return success();
@@ -281,7 +280,7 @@ public:
     LocationAttr directLoc;
     if (parseLocation(directLoc))
       return failure();
-    owner->setLoc(directLoc);
+    loc = directLoc;
     return success();
   }
 
@@ -3186,10 +3185,6 @@ ParseResult OperationParser::parseOperation() {
     }
   }
 
-  // Try to parse the optional trailing location.
-  if (parseOptionalTrailingLocation(op))
-    return failure();
-
   return success();
 }
 
@@ -3351,6 +3346,10 @@ Operation *OperationParser::parseGenericOperation() {
     result.addSuccessor(successor, operands);
   }
 
+  // Parse a location if one is present.
+  if (parseOptionalTrailingLocation(result.location))
+    return nullptr;
+
   return opBuilder.createOperation(result);
 }
 
@@ -3885,6 +3884,10 @@ Operation *OperationParser::parseCustomOperation() {
   if (opAsmParser.didEmitError())
     return nullptr;
 
+  // Parse a location if one is present.
+  if (parseOptionalTrailingLocation(opState.location))
+    return nullptr;
+
   // Otherwise, we succeeded.  Use the state it parsed as our op information.
   return opBuilder.createOperation(opState);
 }
diff --git a/mlir/test/IR/wrapping_op.mlir b/mlir/test/IR/wrapping_op.mlir
index f7b838c..9212e89 100644
--- a/mlir/test/IR/wrapping_op.mlir
+++ b/mlir/test/IR/wrapping_op.mlir
@@ -1,14 +1,14 @@
 // RUN: mlir-opt %s | FileCheck %s
-// RUN: mlir-opt -mlir-print-op-generic %s | FileCheck %s --check-prefix=CHECK-GENERIC
+// RUN: mlir-opt -mlir-print-op-generic -mlir-print-debuginfo %s | FileCheck %s --check-prefix=CHECK-GENERIC
 
 // CHECK-LABEL: func @wrapping_op
 // CHECK-GENERIC-LABEL: func @wrapping_op
 func @wrapping_op(%arg0 : i32, %arg1 : f32) -> (i3, i2, i1) {
 // CHECK: %0:3 = test.wrapping_region wraps "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3)
 // CHECK-GENERIC: "test.wrapping_region"() ( {
-// CHECK-GENERIC:   %[[NESTED_RES:.*]]:3 = "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3)
-// CHECK-GENERIC:   "test.return"(%[[NESTED_RES]]#0, %[[NESTED_RES]]#1, %[[NESTED_RES]]#2) : (i1, i2, i3) -> ()
-// CHECK-GENERIC: }) : () -> (i1, i2, i3)
-  %res:3 = test.wrapping_region wraps "some.op"(%arg1, %arg0) { test.attr = "attr" } : (f32, i32) -> (i1, i2, i3)
+// CHECK-GENERIC:   %[[NESTED_RES:.*]]:3 = "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3) loc("some_NameLoc")
+// CHECK-GENERIC:   "test.return"(%[[NESTED_RES]]#0, %[[NESTED_RES]]#1, %[[NESTED_RES]]#2) : (i1, i2, i3) -> () loc("some_NameLoc")
+// CHECK-GENERIC: }) : () -> (i1, i2, i3) loc("some_NameLoc")
+  %res:3 = test.wrapping_region wraps "some.op"(%arg1, %arg0) { test.attr = "attr" } : (f32, i32) -> (i1, i2, i3) loc("some_NameLoc")
   return %res#2, %res#1, %res#0 : i3, i2, i1
-}
\ No newline at end of file
+}
diff --git a/mlir/test/lib/TestDialect/TestDialect.cpp b/mlir/test/lib/TestDialect/TestDialect.cpp
index 2e3d97b..496cfc6 100644
--- a/mlir/test/lib/TestDialect/TestDialect.cpp
+++ b/mlir/test/lib/TestDialect/TestDialect.cpp
@@ -197,12 +197,16 @@ static ParseResult parseWrappingRegionOp(OpAsmParser &parser,
   SmallVector<Value *, 8> return_operands(wrapped_op->getResults());
   OpBuilder builder(parser.getBuilder().getContext());
   builder.setInsertionPointToEnd(&block);
-  builder.create<TestReturnOp>(result.location, return_operands);
+  builder.create<TestReturnOp>(wrapped_op->getLoc(), return_operands);
 
   // Get the results type for the wrapping op from the terminator operands.
   Operation &return_op = body.back().back();
   result.types.append(return_op.operand_type_begin(),
                       return_op.operand_type_end());
+
+  // Use the location of the wrapped op for the "test.wrapping_region" op.
+  result.location = wrapped_op->getLoc();
+
   return success();
 }
 
-- 
2.7.4