[llvm-branch-commits] [mlir] 2a27a98 - [mlir][AsmPrinter] Properly escape strings when printing locations

[River Riddle via llvm-branch-commits]

Author: River Riddle
Date: 2021-01-15T17:14:57-08:00
New Revision: 2a27a9819a1e3371e3e08d7fbad6235cf50d336d

URL: https://github.com/llvm/llvm-project/commit/2a27a9819a1e3371e3e08d7fbad6235cf50d336d
DIFF: https://github.com/llvm/llvm-project/commit/2a27a9819a1e3371e3e08d7fbad6235cf50d336d.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/IR/AsmPrinter.cpp
    mlir/test/IR/locations.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index c56fdff2cde5..3182e906ce93 100644
--- 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 950b67827a0f..5ad854eedc10 100644
--- 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)