[WIP][mlir] Add struct context to parsing.

[ergawy]

Summary

This patch provides a way for dialect parsers to have access to struct
parsing context. This allows, for example, dialects that support
recursive structs (like LLVM, SPIR-V) to properly handle those structs.

[ergawy]

This is an attempt to address a comment on this review: https://reviews.llvm.org/D87206 to provide a common mechanism for maintaining the struct parsing stack.

I hope this is more or less addresses the issue. Since it touches quite a bit of the infrastructure I am not sure if this is a good direction or not. If not, apologies for the extra noise.

Another approach would be to maintain a stack of Type objects instead of StringRefs. This might provide a more general mechanism and enable other uses in addition to recursive structs (don't have a specific use-case in mind though). However, doing this will need different dialects to construct containing/enclosing types before their sub-types and then mutate the enclosing types after parsing is completed. Not sure that will be a good direction.

mlir/include/mlir/IR/DialectImplementation.h

346

TODO: if this is agreed upon as a good approach, add docs here and elsewhere to describe the new bits.

[ergawy]

mlir/lib/Dialect/LLVMIR/IR/LLVMTypeSyntax.cpp

393

Shouldn't we swap this and the above if conditions? At it is right now, we might not properly clean-up the struct stack if parsing the element type fails.

[ergawy]

• Rebase.

mlir/include/mlir/IR/DialectImplementation.h

 // types.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_DIALECTIMPLEMENTATION_H
 #define MLIR_IR_DIALECTIMPLEMENTATION_H
 
 #include "mlir/IR/OpImplementation.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/SMLoc.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace mlir {
 
 class Builder;
 
public:
   ///   dimension-list ::= (dimension `x`)*
   ///   dimension ::= `?` | integer
   ///
   /// When `allowDynamic` is not set, this is used to parse:
   ///
   ///   static-dimension-list ::= (integer `x`)*
   virtual ParseResult parseDimensionList(SmallVectorImpl<int64_t> &dimensions,
                                          bool allowDynamic = true) = 0;
+
+  virtual llvm::SetVector<StringRef> &getStructContext() = 0;
 };
 
 } // end namespace mlir
 
 #endif

mlir/lib/Dialect/LLVMIR/IR/LLVMTypeSyntax.cpp

void mlir::LLVM::detail::printType(LLVMType type, DialectAsmPrinter &printer) {
   llvm::SetVector<StringRef> stack;
   return printTypeImpl(printer.getStream(), type, stack);
 }
 
 //===----------------------------------------------------------------------===//
 // Parsing.
 //===----------------------------------------------------------------------===//
 
-static LLVMType parseTypeImpl(DialectAsmParser &parser,
-                              llvm::SetVector<StringRef> &stack);
+static LLVMType parseTypeImpl(DialectAsmParser &parser);
 
 /// Helper to be chained with other parsing functions.
-static ParseResult parseTypeImpl(DialectAsmParser &parser,
-                                 llvm::SetVector<StringRef> &stack,
-                                 LLVMType &result) {
-  result = parseTypeImpl(parser, stack);
+static ParseResult parseTypeImpl(DialectAsmParser &parser, LLVMType &result) {
+  result = parseTypeImpl(parser);
   return success(result != nullptr);
 }
 
 /// Parses an LLVM dialect function type.
 ///   llvm-type :: = `func<` llvm-type `(` llvm-type-list `...`? `)>`
-static LLVMFunctionType parseFunctionType(DialectAsmParser &parser,
-                                          llvm::SetVector<StringRef> &stack) {
+static LLVMFunctionType parseFunctionType(DialectAsmParser &parser) {
   Location loc = parser.getEncodedSourceLoc(parser.getCurrentLocation());
   LLVMType returnType;
-  if (parser.parseLess() || parseTypeImpl(parser, stack, returnType) ||
+  if (parser.parseLess() || parseTypeImpl(parser, returnType) ||
       parser.parseLParen())
     return LLVMFunctionType();
 
   // Function type without arguments.
   if (succeeded(parser.parseOptionalRParen())) {
     if (succeeded(parser.parseGreater()))
       return LLVMFunctionType::getChecked(loc, returnType, {},
                                           /*isVarArg=*/false);
     return LLVMFunctionType();
   }
 
   // Parse arguments.
   SmallVector<LLVMType, 8> argTypes;
   do {
     if (succeeded(parser.parseOptionalEllipsis())) {
       if (parser.parseOptionalRParen() || parser.parseOptionalGreater())
         return LLVMFunctionType();
       return LLVMFunctionType::getChecked(loc, returnType, argTypes,
                                           /*isVarArg=*/true);
     }
 
-    argTypes.push_back(parseTypeImpl(parser, stack));
+    argTypes.push_back(parseTypeImpl(parser));
     if (!argTypes.back())
       return LLVMFunctionType();
   } while (succeeded(parser.parseOptionalComma()));
 
   if (parser.parseOptionalRParen() || parser.parseOptionalGreater())
     return LLVMFunctionType();
   return LLVMFunctionType::getChecked(loc, returnType, argTypes,
                                       /*isVarArg=*/false);
 }
 
 /// Parses an LLVM dialect pointer type.
 ///   llvm-type ::= `ptr<` llvm-type (`,` integer)? `>`
-static LLVMPointerType parsePointerType(DialectAsmParser &parser,
-                                        llvm::SetVector<StringRef> &stack) {
+static LLVMPointerType parsePointerType(DialectAsmParser &parser) {
   Location loc = parser.getEncodedSourceLoc(parser.getCurrentLocation());
   LLVMType elementType;
-  if (parser.parseLess() || parseTypeImpl(parser, stack, elementType))
+  if (parser.parseLess() || parseTypeImpl(parser, elementType))
     return LLVMPointerType();
 
   unsigned addressSpace = 0;
   if (succeeded(parser.parseOptionalComma()) &&
       failed(parser.parseInteger(addressSpace)))
     return LLVMPointerType();
   if (failed(parser.parseGreater()))
     return LLVMPointerType();
   return LLVMPointerType::getChecked(loc, elementType, addressSpace);
 }
 
 /// Parses an LLVM dialect vector type.
 ///   llvm-type ::= `vec<` `? x`? integer `x` llvm-type `>`
 /// Supports both fixed and scalable vectors.
-static LLVMVectorType parseVectorType(DialectAsmParser &parser,
-                                      llvm::SetVector<StringRef> &stack) {
+static LLVMVectorType parseVectorType(DialectAsmParser &parser) {
   SmallVector<int64_t, 2> dims;
   llvm::SMLoc dimPos;
   LLVMType elementType;
   Location loc = parser.getEncodedSourceLoc(parser.getCurrentLocation());
   if (parser.parseLess() || parser.getCurrentLocation(&dimPos) ||
       parser.parseDimensionList(dims, /*allowDynamic=*/true) ||
-      parseTypeImpl(parser, stack, elementType) || parser.parseGreater())
+      parseTypeImpl(parser, elementType) || parser.parseGreater())
     return LLVMVectorType();
 
   // We parsed a generic dimension list, but vectors only support two forms:
   //  - single non-dynamic entry in the list (fixed vector);
   //  - two elements, the first dynamic (indicated by -1) and the second
   //    non-dynamic (scalable vector).
   if (dims.empty() || dims.size() > 2 ||
       ((dims.size() == 2) ^ (dims[0] == -1)) ||
       (dims.size() == 2 && dims[1] == -1)) {
     parser.emitError(dimPos)
         << "expected '? x <integer> x <type>' or '<integer> x <type>'";
     return LLVMVectorType();
   }
 
   bool isScalable = dims.size() == 2;
   if (isScalable)
     return LLVMScalableVectorType::getChecked(loc, elementType, dims[1]);
   return LLVMFixedVectorType::getChecked(loc, elementType, dims[0]);
 }
 
 /// Parses an LLVM dialect array type.
 ///   llvm-type ::= `array<` integer `x` llvm-type `>`
-static LLVMArrayType parseArrayType(DialectAsmParser &parser,
-                                    llvm::SetVector<StringRef> &stack) {
+static LLVMArrayType parseArrayType(DialectAsmParser &parser) {
   SmallVector<int64_t, 1> dims;
   llvm::SMLoc sizePos;
   LLVMType elementType;
   Location loc = parser.getEncodedSourceLoc(parser.getCurrentLocation());
   if (parser.parseLess() || parser.getCurrentLocation(&sizePos) ||
       parser.parseDimensionList(dims, /*allowDynamic=*/false) ||
-      parseTypeImpl(parser, stack, elementType) || parser.parseGreater())
+      parseTypeImpl(parser, elementType) || parser.parseGreater())
     return LLVMArrayType();
 
   if (dims.size() != 1) {
     parser.emitError(sizePos) << "expected ? x <type>";
     return LLVMArrayType();
   }
 
   return LLVMArrayType::getChecked(loc, elementType, dims[0]);
 }
 
 /// Attempts to set the body of an identified structure type. Reports a parsing
 /// error at `subtypesLoc` in case of failure, uses `stack` to make sure the
 /// types printed in the error message look like they did when parsed.
 static LLVMStructType trySetStructBody(LLVMStructType type,
                                        ArrayRef<LLVMType> subtypes,
                                        bool isPacked, DialectAsmParser &parser,
-                                       llvm::SMLoc subtypesLoc,
-                                       llvm::SetVector<StringRef> &stack) {
+                                       llvm::SMLoc subtypesLoc) {
   for (LLVMType t : subtypes) {
     if (!LLVMStructType::isValidElementType(t)) {
       parser.emitError(subtypesLoc)
           << "invalid LLVM structure element type: " << t;
       return LLVMStructType();
     }
   }
 
   if (succeeded(type.setBody(subtypes, isPacked)))
     return type;
 
   std::string currentBody;
   llvm::raw_string_ostream currentBodyStream(currentBody);
+  llvm::SetVector<StringRef> stack;
   printStructTypeBody(currentBodyStream, type, stack);
   auto diag = parser.emitError(subtypesLoc)
               << "identified type already used with a different body";
   diag.attachNote() << "existing body: " << currentBodyStream.str();
   return LLVMStructType();
 }
 
 /// Parses an LLVM dialect structure type.
 ///   llvm-type ::= `struct<` (string-literal `,`)? `packed`?
 ///                 `(` llvm-type-list `)` `>`
 ///               | `struct<` string-literal `>`
 ///               | `struct<` string-literal `, opaque>`
-static LLVMStructType parseStructType(DialectAsmParser &parser,
-                                      llvm::SetVector<StringRef> &stack) {
+static LLVMStructType parseStructType(DialectAsmParser &parser) {
   Location loc = parser.getEncodedSourceLoc(parser.getCurrentLocation());
 
   if (failed(parser.parseLess()))
     return LLVMStructType();
 
   // If we are parsing a self-reference to a recursive struct, i.e. the parsing
   // stack already contains a struct with the same identifier, bail out after
   // the name.
   StringRef name;
   bool isIdentified = succeeded(parser.parseOptionalString(&name));
   if (isIdentified) {
-    if (stack.count(name)) {
+    if (parser.getStructContext().count(name)) {
       if (failed(parser.parseGreater()))
         return LLVMStructType();
       return LLVMStructType::getIdentifiedChecked(loc, name);
     }
     if (failed(parser.parseComma()))
       return LLVMStructType();
   }
 
static LLVMStructType parseStructType(DialectAsmParser &parser) {
 
   // Fast pass for structs with zero subtypes.
   if (succeeded(parser.parseOptionalRParen())) {
     if (failed(parser.parseGreater()))
       return LLVMStructType();
     if (!isIdentified)
       return LLVMStructType::getLiteralChecked(loc, {}, isPacked);
     auto type = LLVMStructType::getIdentifiedChecked(loc, name);
-    return trySetStructBody(type, {}, isPacked, parser, kwLoc, stack);
+    return trySetStructBody(type, {}, isPacked, parser, kwLoc);
   }
 
   // Parse subtypes. For identified structs, put the identifier of the struct on
   // the stack to support self-references in the recursive calls.
   SmallVector<LLVMType, 4> subtypes;
   llvm::SMLoc subtypesLoc = parser.getCurrentLocation();
   do {
     if (isIdentified)
-      stack.insert(name);
-    LLVMType type = parseTypeImpl(parser, stack);
+      parser.getStructContext().insert(name);
+    LLVMType type = parseTypeImpl(parser);
     if (!type)
       return LLVMStructType();
     subtypes.push_back(type);
     if (isIdentified)
-      stack.pop_back();
+      parser.getStructContext().pop_back();
   } while (succeeded(parser.parseOptionalComma()));
 
   if (parser.parseRParen() || parser.parseGreater())
     return LLVMStructType();
 
   // Construct the struct with body.
   if (!isIdentified)
     return LLVMStructType::getLiteralChecked(loc, subtypes, isPacked);
   auto type = LLVMStructType::getIdentifiedChecked(loc, name);
-  return trySetStructBody(type, subtypes, isPacked, parser, subtypesLoc, stack);
+  return trySetStructBody(type, subtypes, isPacked, parser, subtypesLoc);
 }
 
 /// Parses one of the LLVM dialect types.
-static LLVMType parseTypeImpl(DialectAsmParser &parser,
-                              llvm::SetVector<StringRef> &stack) {
+static LLVMType parseTypeImpl(DialectAsmParser &parser) {
   // Special case for integers (i[1-9][0-9]*) that are literals rather than
   // keywords for the parser, so they are not caught by the main dispatch below.
   // Try parsing it a built-in integer type instead.
   Type maybeIntegerType;
   MLIRContext *ctx = parser.getBuilder().getContext();
   llvm::SMLoc keyLoc = parser.getCurrentLocation();
   Location loc = parser.getEncodedSourceLoc(keyLoc);
   OptionalParseResult result = parser.parseOptionalType(maybeIntegerType);
return llvm::StringSwitch<function_ref<LLVMType()>>(key)
       .Case("double", [&] { return LLVMDoubleType::get(ctx); })
       .Case("fp128", [&] { return LLVMFP128Type::get(ctx); })
       .Case("x86_fp80", [&] { return LLVMX86FP80Type::get(ctx); })
       .Case("ppc_fp128", [&] { return LLVMPPCFP128Type::get(ctx); })
       .Case("x86_mmx", [&] { return LLVMX86MMXType::get(ctx); })
       .Case("token", [&] { return LLVMTokenType::get(ctx); })
       .Case("label", [&] { return LLVMLabelType::get(ctx); })
       .Case("metadata", [&] { return LLVMMetadataType::get(ctx); })
-      .Case("func", [&] { return parseFunctionType(parser, stack); })
-      .Case("ptr", [&] { return parsePointerType(parser, stack); })
-      .Case("vec", [&] { return parseVectorType(parser, stack); })
-      .Case("array", [&] { return parseArrayType(parser, stack); })
-      .Case("struct", [&] { return parseStructType(parser, stack); })
+      .Case("func", [&] { return parseFunctionType(parser); })
+      .Case("ptr", [&] { return parsePointerType(parser); })
+      .Case("vec", [&] { return parseVectorType(parser); })
+      .Case("array", [&] { return parseArrayType(parser); })
+      .Case("struct", [&] { return parseStructType(parser); })
       .Default([&] {
         parser.emitError(keyLoc) << "unknown LLVM type: " << key;
         return LLVMType();
       })();
 }
 
 LLVMType mlir::LLVM::detail::parseType(DialectAsmParser &parser) {
-  llvm::SetVector<StringRef> stack;
-  return parseTypeImpl(parser, stack);
+  return parseTypeImpl(parser);
 }

mlir/lib/Parser/DialectSymbolParser.cpp

ParseResult parseType(Type &result) override {
     return success(static_cast<bool>(result));
   }
 
   ParseResult parseDimensionList(SmallVectorImpl<int64_t> &dimensions,
                                  bool allowDynamic) override {
     return parser.parseDimensionListRanked(dimensions, allowDynamic);
   }
 
+  llvm::SetVector<StringRef> &getStructContext() override {
+    return parser.getStructContext();
+  }
+
   OptionalParseResult parseOptionalType(Type &result) override {
     return parser.parseOptionalType(result);
   }
 
 private:
   /// The full symbol specification.
   StringRef fullSpec;
 

mlir/lib/Parser/Parser.h

public:
   /// Parse strided layout specification.
   ParseResult parseStridedLayout(int64_t &offset,
                                  SmallVectorImpl<int64_t> &strides);
 
   // Parse a brace-delimiter list of comma-separated integers with `?` as an
   // unknown marker.
   ParseResult parseStrideList(SmallVectorImpl<int64_t> &dimensions);
 
+  llvm::SetVector<StringRef> &getStructContext();
+
   //===--------------------------------------------------------------------===//
   // Attribute Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse an arbitrary attribute with an optional type.
   Attribute parseAttribute(Type type = {});
 
   /// Parse an optional attribute with the provided type.

mlir/lib/Parser/ParserState.h

 //===- ParserState.h - MLIR ParserState -------------------------*- 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_PARSER_PARSERSTATE_H
 #define MLIR_LIB_PARSER_PARSERSTATE_H
 
 #include "Lexer.h"
 #include "mlir/IR/Attributes.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringMap.h"
 
 namespace mlir {
 namespace detail {
 
 //===----------------------------------------------------------------------===//
 // SymbolState
 //===----------------------------------------------------------------------===//
struct ParserState {
   /// This is the next token that hasn't been consumed yet.
   Token curToken;
 
   /// The current state for symbol parsing.
   SymbolState &symbols;
 
   /// The depth of this parser in the nested parsing stack.
   size_t parserDepth;
+
+  llvm::SetVector<StringRef> structContext;
 };
 
 } // end namespace detail
 } // end namespace mlir
 
 #endif // MLIR_LIB_PARSER_PARSERSTATE_H

mlir/lib/Parser/TypeParser.cpp

while (true) {
     }
     if (!consumeIf(Token::comma))
       break;
   }
   if (!consumeIf(Token::r_square))
     return failure();
   return success();
 }
+
+llvm::SetVector<StringRef> &Parser::getStructContext() {
+  return state.structContext;
+}