[mlir:PassOptions] Fix parsing of nested option values/better handle escaping

[rriddle]

Summary

The option parser currently does not properly handle nested options, meaning that
in some cases we can print pass pipelines that we can't actually parse back in. For example,
from #52885 we currently can't parse in inliner pipelines that have nested pipeline strings.

This commit adds handling for string values (e.g. "...") and nested options
(e.g. foo{baz{bar=10 pizza=11}}).

Fixes #52885

[mehdi_amini]

mlir/lib/Pass/PassRegistry.cpp

192

Could write this as a while(true) since it isn't supposed to finish?

199

This wouldn't support a string with an escaped double quote right?

"foo\"bar" would hit the \" at the middle as the end of the quote I think

213

Similar question about parsing pipeline = { option = "}" }

[rriddle]

mlir/lib/Pass/PassRegistry.cpp

192

I used a for because it was simpler to define the variables + update them. Do you have a preference here? Maybe add a true to the condition to make it more obvious?

199

Correct. I'm not sure what our policy here should be given that right now we don't support character escaping of any kind (also gets kind of weird from a command line perspective). I guess the thing interpreting the option would know that it's escaped?

213

Conceptually yes, but there is another of the parser that needs to be updated to support that.

[jpienaar]

This is going to help a lot :)

mlir/lib/Pass/PassRegistry.cpp

160

This is different than most of the stringref manipulating ones - in this case you might get a std::pair<std::pair<StringRef, StringRef>, StringRef> and have the string updated that way. I think the majority returns the remainder again rather than ref param.

192

SGTM (keeps the variable scope small and easier to spot)

199

Would the reproducer dump produce such a string?

[mehdi_amini]

mlir/lib/Pass/PassRegistry.cpp

192

In general I reserve for to regular iteration space ; I feel that using for for other purposes is abusing the language construct.

199

what about a string option containing a non-escaped " character though?

213

OK, I guess this patch is already one step forward...

[rriddle]

mlir/lib/Pass/PassRegistry.cpp

199

Such a string option (at this point) wouldn't work given that we've only really considered use cases for simple things (largely given the terminal
command line origins/llvm CL library usage). String options are printed raw (i.e. no escaping), so I'm not sure that would even parse properly in most shells.

[mehdi_amini]

mlir/lib/Pass/PassRegistry.cpp

199

But we could escape things couldn't we?

My main target isn't "shells" convenience, but the reproducer that opt known how to load natively :)

[rriddle]

mlir/lib/Pass/PassRegistry.cpp

199

We could yes, I was just answering on if it is currently possible. I didn't mean to imply that we shouldn't just because it's convenient from a shell point of view, just that adding support requires reworking other parts of the infra (namely printing) to make it work. Was originally intending to explore this in a followup (to allow for the current bug to be fixed first), unless you have a preference for rolling it into this one?

[mehdi_amini]

mlir/lib/Pass/PassRegistry.cpp

199

Same as in the other thread: this patch seems like a step forward, still LGTM

mlir/lib/Pass/PassRegistry.cpp

 void detail::PassOptions::copyOptionValuesFrom(const PassOptions &other) {
   assert(options.size() == other.options.size());
   if (options.empty())
     return;
   for (auto optionsIt : llvm::zip(options, other.options))
     std::get<0>(optionsIt)->copyValueFrom(*std::get<1>(optionsIt));
 }
 
+/// Parse in the next argument from the given options string. Returns a tuple
+/// containing [the key of the option, the value of the option, updated
+/// `options` string pointing after the parsed option].
+static std::tuple<StringRef, StringRef, StringRef>
+parseNextArg(StringRef options) {
+  // Functor used to extract an argument from 'options' and update it to point
+  // after the arg.
+  auto extractArgAndUpdateOptions = [&](size_t argSize) {
+    StringRef str = options.take_front(argSize).trim();
+    options = options.drop_front(argSize).ltrim();
+    return str;
+  };
+  // Try to process the given punctuation, properly escaping any contained
+  // characters.
+  auto tryProcessPunct = [&](size_t &currentPos, char punct) {
+    if (options[currentPos] != punct)
+      return false;
+    size_t nextIt = options.find_first_of(punct, currentPos + 1);
+    if (nextIt != StringRef::npos)
+      currentPos = nextIt;
+    return true;
+  };
+
+  // Parse the argument name of the option.
+  StringRef argName;
+  for (size_t argEndIt = 0, optionsE = options.size();; ++argEndIt) {
+    // Check for the end of the full option.
+    if (argEndIt == optionsE || options[argEndIt] == ' ') {
+      argName = extractArgAndUpdateOptions(argEndIt);
+      return std::make_tuple(argName, StringRef(), options);
+    }
+
+    // Check for the end of the name and the start of the value.
+    if (options[argEndIt] == '=') {
+      argName = extractArgAndUpdateOptions(argEndIt);
+      options = options.drop_front();
+      break;
+    }
+  }
+
+  // Parse the value of the option.
+  for (size_t argEndIt = 0, optionsE = options.size();; ++argEndIt) {
+    // Handle the end of the options string.
+    if (argEndIt == optionsE || options[argEndIt] == ' ') {
+      StringRef value = extractArgAndUpdateOptions(argEndIt);
+      return std::make_tuple(argName, value, options);
+    }
+
+    // Skip over escaped sequences.
+    char c = options[argEndIt];
+    if (tryProcessPunct(argEndIt, '\'') || tryProcessPunct(argEndIt, '"'))
+      continue;
+    // '{...}' is used to specify options to passes, properly escape it so
+    // that we don't accidentally split any nested options.
+    if (c == '{') {
+      size_t braceCount = 1;
+      for (++argEndIt; argEndIt != optionsE; ++argEndIt) {
+        // Allow nested punctuation.
+        if (tryProcessPunct(argEndIt, '\'') || tryProcessPunct(argEndIt, '"'))
+          continue;
+        if (options[argEndIt] == '{')
+          ++braceCount;
+        else if (options[argEndIt] == '}' && --braceCount == 0)
+          break;
+      }
+      // Account for the increment at the top of the loop.
+      --argEndIt;
+    }
+  }
+  llvm_unreachable("unexpected control flow in pass option parsing");
+}
+
 LogicalResult detail::PassOptions::parseFromString(StringRef options) {
-  // TODO: Handle escaping strings.
   // NOTE: `options` is modified in place to always refer to the unprocessed
   // part of the string.
   while (!options.empty()) {
-    size_t spacePos = options.find(' ');
-    StringRef arg = options;
-    if (spacePos != StringRef::npos) {
-      arg = options.substr(0, spacePos);
-      options = options.substr(spacePos + 1);
-    } else {
-      options = StringRef();
-    }
-    if (arg.empty())
+    StringRef key, value;
+    std::tie(key, value, options) = parseNextArg(options);
+    if (key.empty())
       continue;
 
-    // At this point, arg refers to everything that is non-space in options
-    // upto the next space, and options refers to the rest of the string after
-    // that point.
-
-    // Split the individual option on '=' to form key and value. If there is no
-    // '=', then value is `StringRef()`.
-    size_t equalPos = arg.find('=');
-    StringRef key = arg;
-    StringRef value;
-    if (equalPos != StringRef::npos) {
-      key = arg.substr(0, equalPos);
-      value = arg.substr(equalPos + 1);
-    }
     auto it = OptionsMap.find(key);
     if (it == OptionsMap.end()) {
       llvm::errs() << "<Pass-Options-Parser>: no such option " << key << "\n";
       return failure();
     }
     if (llvm::cl::ProvidePositionalOption(it->second, value, 0))
       return failure();
   }

mlir/test/Pass/pipeline-options-parsing.mlir

 // RUN: not mlir-opt %s -pass-pipeline='builtin.module(test-module-pass{)' 2>&1 | FileCheck --check-prefix=CHECK_ERROR_1 %s
 // RUN: not mlir-opt %s -pass-pipeline='builtin.module(test-module-pass{test-option=3})' 2>&1 | FileCheck --check-prefix=CHECK_ERROR_2 %s
 // RUN: not mlir-opt %s -pass-pipeline='builtin.module(builtin.func(test-options-pass{list=3}), test-module-pass{invalid-option=3})' 2>&1 | FileCheck --check-prefix=CHECK_ERROR_3 %s
 // RUN: not mlir-opt %s -pass-pipeline='test-options-pass{list=3 list=notaninteger}' 2>&1 | FileCheck --check-prefix=CHECK_ERROR_4 %s
 // RUN: not mlir-opt %s -pass-pipeline='builtin.func(test-options-pass{list=1,2,3,4 list=5 string=value1 string=value2})' 2>&1 | FileCheck --check-prefix=CHECK_ERROR_5 %s
-// RUN: mlir-opt %s -verify-each=false -pass-pipeline='builtin.func(test-options-pass{string-list=a list=1,2,3,4 string-list=b,c list=5 string-list=d string=some_value})' -test-dump-pipeline 2>&1 | FileCheck --check-prefix=CHECK_1 %s
+// RUN: mlir-opt %s -verify-each=false -pass-pipeline='builtin.func(test-options-pass{string-list=a list=1,2,3,4 string-list=b,c list=5 string-list=d string=nested_pipeline{arg1=10 arg2=" {} " arg3=true}})' -test-dump-pipeline 2>&1 | FileCheck --check-prefix=CHECK_1 %s
 // RUN: mlir-opt %s -verify-each=false -test-options-pass-pipeline='list=1 string-list=a,b' -test-dump-pipeline 2>&1 | FileCheck --check-prefix=CHECK_2 %s
 // RUN: mlir-opt %s -verify-each=false -pass-pipeline='builtin.module(builtin.func(test-options-pass{list=3}), builtin.func(test-options-pass{list=1,2,3,4}))' -test-dump-pipeline 2>&1 | FileCheck --check-prefix=CHECK_3 %s
 
 // CHECK_ERROR_1: missing closing '}' while processing pass options
 // CHECK_ERROR_2: no such option test-option
 // CHECK_ERROR_3: no such option invalid-option
 // CHECK_ERROR_4: 'notaninteger' value invalid for integer argument
 // CHECK_ERROR_5: string option: may only occur zero or one times
 
-// CHECK_1: test-options-pass{list=1,2,3,4,5 string=some_value string-list=a,b,c,d}
+// CHECK_1: test-options-pass{list=1,2,3,4,5 string=nested_pipeline{arg1=10 arg2=" {} " arg3=true} string-list=a,b,c,d}
 // CHECK_2: test-options-pass{list=1 string= string-list=a,b}
 // CHECK_3: builtin.module(builtin.func(test-options-pass{list=3 string= }), builtin.func(test-options-pass{list=1,2,3,4 string= }))