[mlir][AsmPrinter] Properly escape strings when printing locations

This fixes errors when location strings contains newlines, or other non-ascii characters.

Differential Revision: https://reviews.llvm.org/D94847

diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index c56fdff..3182e90 100644 (file)
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -1257,12 +1257,19 @@ void ModulePrinter::printLocationInternal(LocationAttr loc, bool pretty) {
           os << "unknown";
       })
       .Case<FileLineColLoc>([&](FileLineColLoc loc) {
-        StringRef mayQuote = pretty ? "" : "\"";
-        os << mayQuote << loc.getFilename() << mayQuote << ':' << loc.getLine()
-           << ':' << loc.getColumn();
+        if (pretty) {
+          os << loc.getFilename();
+        } else {
+          os << "\"";
+          printEscapedString(loc.getFilename(), os);
+          os << "\"";
+        }
+        os << ':' << loc.getLine() << ':' << loc.getColumn();
       })
       .Case<NameLoc>([&](NameLoc loc) {
-        os << '\"' << loc.getName() << '\"';
+        os << '\"';
+        printEscapedString(loc.getName(), os);
+        os << '\"';
 
         // Print the child if it isn't unknown.
         auto childLoc = loc.getChildLoc();

diff --git a/mlir/test/IR/locations.mlir b/mlir/test/IR/locations.mlir
index 950b678..5ad854e 100644 (file)
--- a/mlir/test/IR/locations.mlir
+++ b/mlir/test/IR/locations.mlir
@@ -27,6 +27,20 @@ func @inline_notation() -> i32 {
 // CHECK-LABEL: func private @loc_attr(i1 {foo.loc_attr = loc(callsite("foo" at "mysource.cc":10:8))})
 func private @loc_attr(i1 {foo.loc_attr = loc(callsite("foo" at "mysource.cc":10:8))})
 
+  // Check that locations get properly escaped.
+// CHECK-LABEL: func @escape_strings()
+func @escape_strings() {
+  // CHECK: loc("escaped\0A")
+  "foo"() : () -> () loc("escaped\n")
+
+  // CHECK: loc("escaped\0A")
+  "foo"() : () -> () loc("escaped\0A")
+
+  // CHECK: loc("escaped\0A":0:0)
+  "foo"() : () -> () loc("escaped\n":0:0)
+  return
+}
+
 // CHECK-ALIAS: "foo.op"() : () -> () loc(#[[LOC:.*]])
 "foo.op"() : () -> () loc(#loc)