[MLIR][Parser] Add `parseBase64Bytes`.

[bzcheeseman]

Summary

This patch adds parseBase64Bytes to the parser. It attempts to avoid double-allocating the buffer by re-using the token's spelling directly and eliding the quotes if they exist. It also avoids extra allocations by using std::vector<char> in the API - something we should change when the llvm::decodeBase64 API changes.

[rriddle]

mlir/lib/AsmParser/AsmParserImpl.h

262

Why does this need to trim whitespace? I'd just consume the string characters and let decodeBase64 error out as necessary.

[bzcheeseman]

mlir/lib/AsmParser/AsmParserImpl.h

262

It's because we're using the token spelling directly, so I couldn't actually get *any* example to work without trim whitespace and trimming the quotes.

If you use token.getStringValue() then it works without trimming, so it's a tradeoff between double-allocating the buffer and having to trim.

[bzcheeseman]

mlir/lib/AsmParser/AsmParserImpl.h

262

I'm happy to revisit this and go the other way if you'd prefer :)

mlir/include/mlir/IR/OpImplementation.h

ParseResult parseString(std::string *string) {
     if (parseOptionalString(string))
       return emitError(loc, "expected string");
     return success();
   }
 
   /// Parse a quoted string token if present.
   virtual ParseResult parseOptionalString(std::string *string) = 0;
 
+  /// Parses a Base64 encoded string of bytes.
+  virtual ParseResult parseBase64Bytes(std::vector<char> *bytes) = 0;
+
   /// Parse a `(` token.
   virtual ParseResult parseLParen() = 0;
 
   /// Parse a `(` token if present.
   virtual ParseResult parseOptionalLParen() = 0;
 
   /// Parse a `)` token.
   virtual ParseResult parseRParen() = 0;

mlir/lib/AsmParser/AsmParserImpl.h

 //===- AsmParserImpl.h - MLIR AsmParserImpl Class ---------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_LIB_ASMPARSER_ASMPARSERIMPL_H
 #define MLIR_LIB_ASMPARSER_ASMPARSERIMPL_H
 
 #include "Parser.h"
 #include "mlir/AsmParser/AsmParserState.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/OpImplementation.h"
+#include "llvm/Support/Base64.h"
 
 namespace mlir {
 namespace detail {
 //===----------------------------------------------------------------------===//
 // AsmParserImpl
 //===----------------------------------------------------------------------===//
 
 /// This class provides the implementation of the generic parser methods within
if (!parser.getToken().is(Token::string))
       return failure();
 
     if (string)
       *string = parser.getToken().getStringValue();
     parser.consumeToken();
     return success();
   }
 
+  /// Parses a Base64 encoded string of bytes.
+  ParseResult parseBase64Bytes(std::vector<char> *bytes) override {
+    auto loc = getCurrentLocation();
+    if (!parser.getToken().is(Token::string))
+      return emitError(loc, "expected string");
+
+    if (bytes) {
+      // decodeBase64 doesn't modify its input so we can use the token spelling
+      // and just slice off the quotes/whitespaces if there are any. Whitespace
+      // and quotes cannot appear as part of a (standard) base64 encoded string,
+      // so this is safe to do.
+      StringRef b64QuotedString = parser.getTokenSpelling();
+      StringRef b64String =
+          b64QuotedString.ltrim("\"  \t\n\v\f\r").rtrim("\" \t\n\v\f\r");
+      if (auto err = llvm::decodeBase64(b64String, *bytes))
+        return emitError(loc, toString(std::move(err)));
+    }
+
+    parser.consumeToken();
+    return success();
+  }
+
   /// Parse a floating point value from the stream.
   ParseResult parseFloat(double &result) override {
     bool isNegative = parser.consumeIf(Token::minus);
     Token curTok = parser.getToken();
     SMLoc loc = curTok.getLoc();
 
     // Check for a floating point value.
     if (curTok.is(Token::floatliteral)) {

mlir/test/IR/parser.mlir

 
 // CHECK-LABEL: func @parse_wrapped_keyword_test
 func.func @parse_wrapped_keyword_test() {
   // CHECK: test.parse_wrapped_keyword foo.keyword
   test.parse_wrapped_keyword foo.keyword
   return
 }
 
+// CHECK-LABEL: func @parse_base64_test
+func.func @parse_base64_test() {
+  // CHECK: test.parse_b64 "hello world"
+  test.parse_b64 "aGVsbG8gd29ybGQ="
+  return
+}
+
 // CHECK-LABEL: func @"\22_string_symbol_reference\22"
 func.func @"\"_string_symbol_reference\""() {
   // CHECK: ref = @"\22_string_symbol_reference\22"
   "foo.symbol_reference"() {ref = @"\"_string_symbol_reference\""} : () -> ()
   return
 }
 
 // CHECK-LABEL: func private @parse_opaque_attr_escape

mlir/test/lib/Dialect/Test/TestDialect.cpp

ParseResult ParseWrappedKeywordOp::parse(OpAsmParser &parser,
   if (parser.parseKeyword(&keyword))
     return failure();
   result.addAttribute("keyword", parser.getBuilder().getStringAttr(keyword));
   return success();
 }
 
 void ParseWrappedKeywordOp::print(OpAsmPrinter &p) { p << " " << getKeyword(); }
 
+ParseResult ParseB64BytesOp::parse(OpAsmParser &parser,
+                                   OperationState &result) {
+  std::vector<char> bytes;
+  if (parser.parseBase64Bytes(&bytes))
+    return failure();
+  result.addAttribute("b64", parser.getBuilder().getStringAttr(
+                                 StringRef(&bytes.front(), bytes.size())));
+  return success();
+}
+
+void ParseB64BytesOp::print(OpAsmPrinter &p) {
+  // Don't print the base64 version to check that we decoded it correctly.
+  p << " \"" << getB64() << "\"";
+}
+
 //===----------------------------------------------------------------------===//
 // Test WrapRegionOp - wrapping op exercising `parseGenericOperation()`.
 
 ParseResult WrappingRegionOp::parse(OpAsmParser &parser,
                                     OperationState &result) {
   if (parser.parseKeyword("wraps"))
     return failure();
 

mlir/test/lib/Dialect/Test/TestOps.td

def ParseIntegerLiteralOp : TEST_Op<"parse_integer_literal"> {
   let hasCustomAssemblyFormat = 1;
 }
 
 def ParseWrappedKeywordOp : TEST_Op<"parse_wrapped_keyword"> {
   let arguments = (ins StrAttr:$keyword);
   let hasCustomAssemblyFormat = 1;
 }
 
+def ParseB64BytesOp : TEST_Op<"parse_b64"> {
+  let arguments = (ins StrAttr:$b64);
+  let hasCustomAssemblyFormat = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Test region argument list parsing.
 
 def IsolatedRegionOp : TEST_Op<"isolated_region", [IsolatedFromAbove]> {
   let summary =  "isolated region operation";
   let description = [{
     Test op with an isolated region, to test passthrough region arguments. Each
     argument is of index type.