[mlir] Support for mutable types

[ftynse]

Summary

Introduce support for mutable storage in the StorageUniquer infrastructure.
This makes MLIR have key-value storage instead of just uniqued key storage. A
storage instance now contains a unique immutable key and a mutable value, both
stored in the arena allocator that belongs to the context. This is a
preconditio for supporting recursive types that require delayed initialization,
in particular LLVM structure types. The functionality is exercised in the test
pass with trivial self-recursive type. So far, recursive types can only be
printed in parsed in a closed type system. Removing this restriction is left
for future work.

[rriddle]

Thanks Alex, I think this is an elegant solution.

Could you update https://mlir.llvm.org/docs/Tutorials/DefiningAttributesAndTypes/? It's gotten quite crusty and needs a good cleanup, so I can add docs when I clean it up if you don't want to do it now. This also might be something that we could include in the Toy dialect.

Just one conceptual question, what are the intended use cases for a failed mutation? In the current use cases we have, the type would realistically assert that the body was either not set or identical to the newly provided body.
For the cases where you envision this being used, what would a type do with the result of mutate?

Adding approval here, because the result of the above discussion doesn't change much of anything user facing.

mlir/include/mlir/IR/TypeSupport.h

136

Can you add the same to the AttributeUniquer for uniformity?

mlir/test/lib/IR/TestTypes.cpp

28

nit: Can you move the body of this function out to a new function?

[ftynse]

Address review

[ftynse]

In D84171#2165365, @rriddle wrote:

Thanks Alex, I think this is an elegant solution.

Could you update https://mlir.llvm.org/docs/Tutorials/DefiningAttributesAndTypes/? It's gotten quite crusty and needs a good cleanup, so I can add docs when I clean it up if you don't want to do it now. This also might be something that we could include in the Toy dialect.

Sure, I was waiting to make sure we agree on the approach before updating that. I added a description of types with mutable components, but haven't touched the rest. I want to make sure LLVM dialect type rework is on tracks first, and then can look into updating the rest if you don't beat me to it.

Just one conceptual question, what are the intended use cases for a failed mutation? In the current use cases we have, the type would realistically assert that the body was either not set or identical to the newly provided body.
For the cases where you envision this being used, what would a type do with the result of mutate?

One thing is to have a graceful failure mechanism similar to getChecked. A bit stretched example: imagine creating two types with the same name by parsing them from strings (we intend to expose this through C API), it's better to be able to return null to the caller instead of asserting inside the library. One could argue that it's a wrong use of the library, but it can be triggered by user input so having an option not to assert sounds better.

Another thing is partly provisioning for the type auto-renaming scheme along the lines of

Type getTypeWithRenaming(StringRef name, ...) {
  Type t;
  do {
    t = RecursiveType::get(name);
    if (succeeded(t.setBody(...)) break;
    name = nextPossibleName(name);
  } while (true);
  return t ;
}

[ftynse]

upd

[ftynse]

rebase

mlir/docs/Tutorials/DefiningAttributesAndTypes.md

 `PRIVATE_EXPERIMENTAL_0` range. These definitions will provide a range in the
 Type::Kind enum to use when defining the derived types.
 
 ```c++
 namespace MyTypes {
 enum Kinds {
   // These kinds will be used in the examples below.
   Simple = Type::Kind::FIRST_PRIVATE_EXPERIMENTAL_0_TYPE,
-  Complex
+  Complex,
+  Recursive
 };
 }
 ```
 
 ### Defining the type class
 
 As described above, `Type` objects in MLIR are value-typed and rely on having an
 implicitly internal storage object that holds the actual data for the type. When
 defining a new `Type` it isn't always necessary to define a new storage class.
 So before defining the derived `Type`, it's important to know which of the two
-classes of `Type` we are defining. Some types are `primitives` meaning they do
+classes of `Type` we are defining. Some types are _primitives_ meaning they do
 not have any parameters and are singletons uniqued by kind, like the
 [`index` type](LangRef.md#index-type). Parametric types on the other hand, have
 additional information that differentiates different instances of the same
 `Type` kind. For example the [`integer` type](LangRef.md#integer-type) has a
 bitwidth, making `i8` and `i16` be different instances of
-[`integer` type](LangRef.md#integer-type).
+[`integer` type](LangRef.md#integer-type). Types can also have a mutable
+component, which can be used, for example, to construct self-referring recursive
+types. The mutable component _cannot_ be used to differentiate types within the
+same kind, so usually such types are also parametric where the parameters serve
+to identify them.
 
 #### Simple non-parametric types
 
 For simple parameterless types, we can jump straight into defining the derived
 type class. Given that these types are uniqued solely on `kind`, we don't need
 to provide our own storage class.
 
 ```c++
public:
   /// Return the integer parameter type.
   IntegerType getParameterType() {
     // 'getImpl' returns a pointer to our internal storage instance.
     return getImpl()->integerType;
   }
 };
 ```
 
+#### Types with a mutable component
+
+Types with a mutable component require defining a type storage class regardless
+of being parametric. The storage contains both the parameters and the mutable
+component and is accessed in a thread-safe way by the type support
+infrastructure.
+
+##### Defining a type storage
+
+In addition to the requirements for the type storage class for parametric types,
+the storage class for types with a mutable component must additionally obey the
+following.
+
+*   The mutable component must not participate in the storage key.
+*   Provide a mutation method that is used to modify an existing instance of the
+    storage. This method modifies the mutable component based on arguments,
+    using `allocator` for any new dynamically-allocated storage, and indicates
+    whether the modification was successful.
+    -   `LogicalResult mutate(StorageAllocator &allocator, Args ...&& args)`
+
+Let's define a simple storage for recursive types, where a type is identified by
+its name and can contain another type including itself.
+
+```c++
+/// Here we define a storage class for a RecursiveType that is identified by its
+/// name and contains another type.
+struct RecursiveTypeStorage : public TypeStorage {
+  /// The type is uniquely identified by its name. Note that the contained type
+  /// is _not_ a part of the key.
+  using KeyTy = StringRef;
+
+  /// Construct the storage from the type name. Explicitly initialize the
+  /// containedType to nullptr, which is used as marker for the mutable
+  /// component being not yet initialized.
+  RecursiveTypeStorage(StringRef name) : name(name), containedType(nullptr) {}
+
+  /// Define the comparison function.
+  bool operator==(const KeyTy &key) const { return key == name; }
+
+  /// Define a construction method for creating a new instance of the storage.
+  static RecursiveTypeStorage *construct(StorageAllocator &allocator,
+                                         const KeyTy &key) {
+    // Note that the key string is copied into the allocator to ensure it
+    // remains live as long as the storage itself.
+    return new (allocator.allocate<RecursiveTypeStorage>())
+        RecursiveTypeStorage(allocator.copyInto(key));
+  }
+
+  /// Define a mutation method for changing the type after it is created. In
+  /// many cases, we only want to set the mutable component once and reject
+  /// any further modification, which can be achieved by returning failure from
+  /// this function.
+  LogicalResult mutate(StorageAllocator &, Type body) {
+    // If the contained type has been initialized already, and the call tries
+    // to change it, reject the change.
+    if (containedType && containedType != body)
+      return failure();
+
+    // Change the body successfully.
+    containedType = body;
+    return success();
+  }
+
+  StringRef name;
+  Type containedType;
+};
+```
+
+##### Type class definition
+
+Having defined the storage class, we can define the type class itself. This is
+similar to parametric types. `Type::TypeBase` provides a `mutate` method that
+forwards its arguments to the `mutate` method of the storage and ensures the
+modification happens under lock.
+
+```c++
+class RecursiveType : public Type::TypeBase<RecursiveType, Type,
+                                            RecursiveTypeStorage> {
+public:
+  /// Inherit parent constructors.
+  using Base::Base;
+
+  /// This static method is used to support type inquiry through isa, cast,
+  /// and dyn_cast.
+  static bool kindof(unsigned kind) { return kind == MyTypes::Recursive; }
+
+  /// Creates an instance of the Recursive type. This only takes the type name
+  /// and returns the type with uninitialized body.
+  static RecursiveType get(MLIRContext *ctx, StringRef name) {
+    // Call into the base to get a uniqued instance of this type. The parameter
+    // (name) is passed after the kind.
+    return Base::get(ctx, MyTypes::Recursive, name);
+  }
+
+  /// Now we can change the mutable component of the type. This is an instance
+  /// method callable on an already existing RecursiveType.
+  void setBody(Type body) {
+    // Call into the base to mutate the type.
+    LogicalResult result = Base::mutate(body);
+    // Most types expect mutation to always succeed, but types can implement
+    // custom logic for handling mutation failures.
+    assert(succeeded(result) &&
+           "attempting to change the body of an already-initialized type");
+    // Avoid unused-variable warning when building without assertions.
+    (void) result;
+  }
+
+  /// Returns the contained type, which may be null if it has not been
+  /// initialized yet.
+  Type getBody() {
+    return getImpl()->containedType;
+  }
+
+  /// Returns the name.
+  StringRef getName() {
+    return getImpl()->name;
+  }
+};
+```
+
 ### Registering types with a Dialect
 
 Once the dialect types have been defined, they must then be registered with a
 `Dialect`. This is done via similar mechanism to
 [operations](LangRef.md#operations), `addTypes`.
 
 ```c++
 struct MyDialect : public Dialect {

mlir/include/mlir/IR/AttributeSupport.h

public:
   static T get(MLIRContext *ctx, unsigned kind, Args &&... args) {
     return ctx->getAttributeUniquer().get<typename T::ImplType>(
         [ctx](AttributeStorage *storage) {
           initializeAttributeStorage(storage, ctx, T::getTypeID());
         },
         kind, std::forward<Args>(args)...);
   }
 
+  template <typename ImplType, typename... Args>
+  static LogicalResult mutate(MLIRContext *ctx, ImplType *impl,
+                              Args &&...args) {
+    assert(impl && "cannot mutate null attribute");
+    return ctx->getAttributeUniquer().mutate(impl, std::forward<Args>(args)...);
+  }
+
 private:
   /// Initialize the given attribute storage instance.
   static void initializeAttributeStorage(AttributeStorage *storage,
                                          MLIRContext *ctx, TypeID attrID);
 };
 } // namespace detail
 
 } // end namespace mlir
 
 #endif

mlir/include/mlir/IR/Attributes.h

 struct DenseIntOrFPElementsAttributeStorage;
 struct DenseStringElementsAttributeStorage;
 struct OpaqueElementsAttributeStorage;
 struct SparseElementsAttributeStorage;
 } // namespace detail
 
 /// Attributes are known-constant values of operations and functions.
 ///
-/// Instances of the Attribute class are references to immutable, uniqued,
-/// and immortal values owned by MLIRContext. As such, an Attribute is a thin
-/// wrapper around an underlying storage pointer. Attributes are usually passed
-/// by value.
+/// Instances of the Attribute class are references to immortal key-value pairs
+/// with immutable, uniqued key owned by MLIRContext. As such, an Attribute is a
+/// thin wrapper around an underlying storage pointer. Attributes are usually
+/// passed by value.
 class Attribute {
 public:
   /// Integer identifier for all the concrete attribute kinds.
   enum Kind {
   // Reserve attribute kinds for dialect specific extensions.
 #define DEFINE_SYM_KIND_RANGE(Dialect)                                         \
   FIRST_##Dialect##_ATTR, LAST_##Dialect##_ATTR = FIRST_##Dialect##_ATTR + 0xff,
 #include "DialectSymbolRegistry.def"

mlir/include/mlir/IR/StorageUniquerSupport.h

protected:
   template <typename LocationT, typename... Args>
   static ConcreteT getChecked(LocationT loc, unsigned kind, Args... args) {
     // If the construction invariants fail then we return a null attribute.
     if (failed(ConcreteT::verifyConstructionInvariants(loc, args...)))
       return ConcreteT();
     return UniquerT::template get<ConcreteT>(loc.getContext(), kind, args...);
   }
 
+  /// Mutate the current storage instance. This will not change the unique key.
+  /// The arguments are forwarded to 'ConcreteT::mutate'.
+  template <typename... Args>
+  LogicalResult mutate(Args &&...args) {
+    return UniquerT::mutate(this->getContext(), getImpl(),
+                            std::forward<Args>(args)...);
+  }
+
   /// Default implementation that just returns success.
   template <typename... Args>
   static LogicalResult verifyConstructionInvariants(Args... args) {
     return success();
   }
 
   /// Utility for easy access to the storage instance.
   ImplType *getImpl() const { return static_cast<ImplType *>(this->impl); }
 };
 } // namespace detail
 } // namespace mlir
 
 #endif

mlir/include/mlir/IR/TypeSupport.h

struct TypeUniquer {
   template <typename T, typename... Args>
   static T get(MLIRContext *ctx, unsigned kind, Args &&... args) {
     return ctx->getTypeUniquer().get<typename T::ImplType>(
         [&](TypeStorage *storage) {
           storage->initialize(AbstractType::lookup(T::getTypeID(), ctx));
         },
         kind, std::forward<Args>(args)...);
   }
+
+  /// Change the mutable component of the given type instance in the provided
+  /// context.
+  template <typename ImplType, typename... Args>
+  static LogicalResult mutate(MLIRContext *ctx, ImplType *impl,
+                              Args &&...args) {
+    assert(impl && "cannot mutate null type");
+    return ctx->getTypeUniquer().mutate(impl, std::forward<Args>(args)...);
+  }
 };
 } // namespace detail
 
 } // end namespace mlir
 
 #endif

mlir/include/mlir/IR/Types.h

 class MLIRContext;
 class TypeStorage;
 
 namespace detail {
 struct FunctionTypeStorage;
 struct OpaqueTypeStorage;
 } // namespace detail
 
-/// Instances of the Type class are immutable and uniqued.  They wrap a pointer
-/// to the storage object owned by MLIRContext.  Therefore, instances of Type
-/// are passed around by value.
+/// Instances of the Type class are uniqued, have an immutable identifier and an
+/// optional mutable component.  They wrap a pointer to the storage object owned
+/// by MLIRContext.  Therefore, instances of Type are passed around by value.
 ///
 /// Some types are "primitives" meaning they do not have any parameters, for
 /// example the Index type.  Parametric types have additional information that
 /// differentiates the types of the same kind between them, for example the
 /// Integer type has bitwidth, making i8 and i16 belong to the same kind by be
-/// different instances of the IntegerType.
+/// different instances of the IntegerType.  Type parameters are part of the
+/// unique immutable key.  The mutable component of the type can be modified
+/// after the type is created, but cannot affect the identity of the type.
 ///
 /// Types are constructed and uniqued via the 'detail::TypeUniquer' class.
 ///
 /// Derived type classes are expected to implement several required
 /// implementation hooks:
 ///  * Required:
 ///    - static bool kindof(unsigned kind);
 ///      * Returns if the provided type kind corresponds to an instance of the

 ///      * 'args' must correspond with the arguments passed into the
 ///        'TypeBase::get' call after the type kind.
 ///
 ///
 /// Type storage objects inherit from TypeStorage and contain the following:
 ///    - The type kind (for LLVM-style RTTI).
 ///    - The dialect that defined the type.
 ///    - Any parameters of the type.
+///    - An optional mutable component.
 /// For non-parametric types, a convenience DefaultTypeStorage is provided.
 /// Parametric storage types must derive TypeStorage and respect the following:
 ///    - Define a type alias, KeyTy, to a type that uniquely identifies the
 ///      instance of the type within its kind.
 ///      * The key type must be constructible from the values passed into the
 ///        detail::TypeUniquer::get call after the type kind.
 ///      * If the KeyTy does not have an llvm::DenseMapInfo specialization, the
 ///        storage class must define a hashing method:
 ///         'static unsigned hashKey(const KeyTy &)'
 ///
 ///    - Provide a method, 'bool operator==(const KeyTy &) const', to
 ///      compare the storage instance against an instance of the key type.
 ///
-///    - Provide a construction method:
+///    - Provide a static construction method:
 ///        'DerivedStorage *construct(TypeStorageAllocator &, const KeyTy &key)'
 ///      that builds a unique instance of the derived storage. The arguments to
 ///      this function are an allocator to store any uniqued data within the
 ///      context and the key type for this storage.
+///
+///    - If they have a mutable component, this component must not be a part of
+//       the key.
 class Type {
 public:
   /// Integer identifier for all the concrete type kinds.
   /// Note: This is not an enum class as each dialect will likely define a
   /// separate enumeration for the specific types that they define. Not being an
   /// enum class also simplifies the handling of type kinds by not requiring
   /// casts for each use.
   enum Kind {

mlir/include/mlir/Support/StorageUniquer.h

 //===- StorageUniquer.h - Common Storage Class Uniquer ----------*- 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_SUPPORT_STORAGEUNIQUER_H
 #define MLIR_SUPPORT_STORAGEUNIQUER_H
 
 #include "mlir/Support/LLVM.h"
+#include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/Allocator.h"
 
 namespace mlir {
 namespace detail {
 struct StorageUniquerImpl;
 
 /// Trait to check if ImplTy provides a 'getKey' method with types 'Args'.

 ///      this function are an allocator to store any uniqued data and the key
 ///      type for this storage.
 ///
 ///    - Provide a cleanup method:
 ///        'void cleanup()'
 ///      that is called when erasing a storage instance. This should cleanup any
 ///      fields of the storage as necessary and not attempt to free the memory
 ///      of the storage itself.
+///
+/// Storage classes may have an optional mutable component, which must not take
+/// part in the unique immutable key. In this case, storage classes may be
+/// mutated with `mutate` and must additionally respect the following:
+///    - Provide a mutation method:
+///        'LogicalResult mutate(StorageAllocator &, <...>)'
+///      that is called when mutating a storage instance. The first argument is
+///      an allocator to store any mutable data, and the remaining arguments are
+///      forwarded from the call site. The storage can be mutated at any time
+///      after creation. Care must be taken to avoid excessive mutation since
+///      the allocated storage can keep containing previous states. The return
+///      value of the function is used to indicate whether the mutation was
+///      successful, e.g., to limit the number of mutations or enable deferred
+///      one-time assignment of the mutable component.
 class StorageUniquer {
 public:
   StorageUniquer();
   ~StorageUniquer();
 
   /// Set the flag specifying if multi-threading is disabled within the uniquer.
   void disableMultithreading(bool disable = true);
 
auto ctorFn = [&](StorageAllocator &allocator) {
       auto *storage = new (allocator.allocate<Storage>()) Storage();
       if (initFn)
         initFn(storage);
       return storage;
     };
     return static_cast<Storage *>(getImpl(kind, ctorFn));
   }
 
+  /// Changes the mutable component of 'storage' by forwarding the trailing
+  /// arguments to the 'mutate' function of the derived class.
+  template <typename Storage, typename... Args>
+  LogicalResult mutate(Storage *storage, Args &&...args) {
+    auto mutationFn = [&](StorageAllocator &allocator) -> LogicalResult {
+      return static_cast<Storage &>(*storage).mutate(
+          allocator, std::forward<Args>(args)...);
+    };
+    return mutateImpl(mutationFn);
+  }
+
   /// Erases a uniqued instance of 'Storage'. This function is used for derived
   /// types that have complex storage or uniquing constraints.
   template <typename Storage, typename Arg, typename... Args>
   void erase(unsigned kind, Arg &&arg, Args &&... args) {
     // Construct a value of the derived key type.
     auto derivedKey =
         getKey<Storage>(std::forward<Arg>(arg), std::forward<Args>(args)...);
 
BaseStorage *getImpl(unsigned kind,
                        function_ref<BaseStorage *(StorageAllocator &)> ctorFn);
 
   /// Implementation for erasing an instance of a derived type with complex
   /// storage.
   void eraseImpl(unsigned kind, unsigned hashValue,
                  function_ref<bool(const BaseStorage *)> isEqual,
                  function_ref<void(BaseStorage *)> cleanupFn);
 
+  /// Implementation for mutating an instance of a derived storage.
+  LogicalResult
+  mutateImpl(function_ref<LogicalResult(StorageAllocator &)> mutationFn);
+
   /// The internal implementation class.
   std::unique_ptr<detail::StorageUniquerImpl> impl;
 
   //===--------------------------------------------------------------------===//
   // Key Construction
   //===--------------------------------------------------------------------===//
 
   /// Used to construct an instance of 'ImplTy::KeyTy' if there is an

mlir/lib/Support/StorageUniquer.cpp

void erase(unsigned kind, unsigned hashValue,
     if (existing == storageTypes.end())
       return;
 
     // Cleanup the storage and remove it from the map.
     cleanupFn(existing->storage);
     storageTypes.erase(existing);
   }
 
+  /// Mutates an instance of a derived storage in a thread-safe way.
+  LogicalResult
+  mutate(function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
+    if (!threadingIsEnabled)
+      return mutationFn(allocator);
+
+    llvm::sys::SmartScopedWriter<true> lock(mutex);
+    return mutationFn(allocator);
+  }
+
   //===--------------------------------------------------------------------===//
   // Instance Storage
   //===--------------------------------------------------------------------===//
 
   /// Utility to create and initialize a storage instance.
   BaseStorage *
   initializeStorage(unsigned kind,
                     function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {

 
 /// Implementation for erasing an instance of a derived type with complex
 /// storage.
 void StorageUniquer::eraseImpl(unsigned kind, unsigned hashValue,
                                function_ref<bool(const BaseStorage *)> isEqual,
                                function_ref<void(BaseStorage *)> cleanupFn) {
   impl->erase(kind, hashValue, isEqual, cleanupFn);
 }
+
+/// Implementation for mutating an instance of a derived storage.
+LogicalResult StorageUniquer::mutateImpl(
+    function_ref<LogicalResult(StorageAllocator &)> mutationFn) {
+  return impl->mutate(mutationFn);
+}

mlir/test/IR/recursive-type.mlir

+// RUN: mlir-opt %s -test-recursive-types | FileCheck %s
+
+// CHECK-LABEL: @roundtrip
+func @roundtrip() {
+  // CHECK: !test.test_rec<a, test_rec<b, test_type>>
+  "test.dummy_op_for_roundtrip"() : () -> !test.test_rec<a, test_rec<b, test_type>>
+  // CHECK: !test.test_rec<c, test_rec<c>>
+  "test.dummy_op_for_roundtrip"() : () -> !test.test_rec<c, test_rec<c>>
+  return
+}
+
+// CHECK-LABEL: @create
+func @create() {
+  // CHECK: !test.test_rec<some_long_and_unique_name, test_rec<some_long_and_unique_name>>
+  return
+}

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

 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringSwitch.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // TestDialect Interfaces
 //===----------------------------------------------------------------------===//
 

 TestDialect::TestDialect(MLIRContext *context)
     : Dialect(getDialectNamespace(), context) {
   addOperations<
 #define GET_OP_LIST
 #include "TestOps.cpp.inc"
       >();
   addInterfaces<TestOpAsmInterface, TestOpFolderDialectInterface,
                 TestInlinerInterface>();
-  addTypes<TestType>();
+  addTypes<TestType, TestRecursiveType>();
   allowUnknownOperations();
 }
 
-Type TestDialect::parseType(DialectAsmParser &parser) const {
-  if (failed(parser.parseKeyword("test_type")))
+static Type parseTestType(DialectAsmParser &parser,
+                          llvm::SetVector<Type> &stack) {
+  StringRef typeTag;
+  if (failed(parser.parseKeyword(&typeTag)))
+    return Type();
+
+  if (typeTag == "test_type")
+    return TestType::get(parser.getBuilder().getContext());
+
+  if (typeTag != "test_rec")
+    return Type();
+
+  StringRef name;
+  if (parser.parseLess() || parser.parseKeyword(&name))
     return Type();
-  return TestType::get(getContext());
+  auto rec = TestRecursiveType::create(parser.getBuilder().getContext(), name);
+
+  // If this type already has been parsed above in the stack, expect just the
+  // name.
+  if (stack.contains(rec)) {
+    if (failed(parser.parseGreater()))
+      return Type();
+    return rec;
   }
 
-void TestDialect::printType(Type type, DialectAsmPrinter &printer) const {
-  assert(type.isa<TestType>() && "unexpected type");
+  // Otherwise, parse the body and update the type.
+  if (failed(parser.parseComma()))
+    return Type();
+  stack.insert(rec);
+  Type subtype = parseTestType(parser, stack);
+  stack.pop_back();
+  if (!subtype || failed(parser.parseGreater()) || failed(rec.setBody(subtype)))
+    return Type();
+
+  return rec;
+}
+
+Type TestDialect::parseType(DialectAsmParser &parser) const {
+  llvm::SetVector<Type> stack;
+  return parseTestType(parser, stack);
+}
+
+static void printTestType(Type type, DialectAsmPrinter &printer,
+                          llvm::SetVector<Type> &stack) {
+  if (type.isa<TestType>()) {
     printer << "test_type";
+    return;
+  }
+
+  auto rec = type.cast<TestRecursiveType>();
+  printer << "test_rec<" << rec.getName();
+  if (!stack.contains(rec)) {
+    printer << ", ";
+    stack.insert(rec);
+    printTestType(rec.getBody(), printer, stack);
+    stack.pop_back();
+  }
+  printer << ">";
+}
+
+void TestDialect::printType(Type type, DialectAsmPrinter &printer) const {
+  llvm::SetVector<Type> stack;
+  printTestType(type, printer, stack);
 }
 
 LogicalResult TestDialect::verifyOperationAttribute(Operation *op,
                                                     NamedAttribute namedAttr) {
   if (namedAttr.first == "test.invalid_attr")
     return op->emitError() << "invalid to use 'test.invalid_attr'";
   return success();
 }

mlir/test/lib/Dialect/Test/TestTypes.h

static TestType get(MLIRContext *context) {
     return Base::get(context, Type::Kind::FIRST_PRIVATE_EXPERIMENTAL_9_TYPE);
   }
 
   /// Provide a definition for the necessary interface methods.
   void printTypeC(Location loc) const {
     emitRemark(loc) << *this << " - TestC";
   }
 };
+
+/// Storage for simple named recursive types, where the type is identified by
+/// its name and can "contain" another type, including itself.
+struct TestRecursiveTypeStorage : public TypeStorage {
+  using KeyTy = StringRef;
+
+  explicit TestRecursiveTypeStorage(StringRef key) : name(key), body(Type()) {}
+
+  bool operator==(const KeyTy &other) const { return name == other; }
+
+  static TestRecursiveTypeStorage *construct(TypeStorageAllocator &allocator,
+                                             const KeyTy &key) {
+    return new (allocator.allocate<TestRecursiveTypeStorage>())
+        TestRecursiveTypeStorage(allocator.copyInto(key));
+  }
+
+  LogicalResult mutate(TypeStorageAllocator &allocator, Type newBody) {
+    // Cannot set a different body than before.
+    if (body && body != newBody)
+      return failure();
+
+    body = newBody;
+    return success();
+  }
+
+  StringRef name;
+  Type body;
+};
+
+/// Simple recursive type identified by its name and pointing to another named
+/// type, potentially itself. This requires the body to be mutated separately
+/// from type creation.
+class TestRecursiveType
+    : public Type::TypeBase<TestRecursiveType, Type, TestRecursiveTypeStorage> {
+public:
+  using Base::Base;
+
+  static bool kindof(unsigned kind) {
+    return kind == Type::Kind::FIRST_PRIVATE_EXPERIMENTAL_9_TYPE + 1;
+  }
+
+  static TestRecursiveType create(MLIRContext *ctx, StringRef name) {
+    return Base::get(ctx, Type::Kind::FIRST_PRIVATE_EXPERIMENTAL_9_TYPE + 1,
+                     name);
+  }
+
+  /// Body getter and setter.
+  LogicalResult setBody(Type body) { return Base::mutate(body); }
+  Type getBody() { return getImpl()->body; }
+
+  /// Name/key getter.
+  StringRef getName() { return getImpl()->name; }
+};
+
 } // end namespace mlir
 
 #endif // MLIR_TESTTYPES_H

mlir/test/lib/IR/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestIR
   TestFunc.cpp
   TestInterfaces.cpp
   TestMatchers.cpp
   TestSideEffects.cpp
   TestSymbolUses.cpp
+  TestTypes.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRPass
   MLIRTestDialect
   )
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../Dialect/Test)
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/../Dialect/Test)

mlir/test/lib/IR/TestTypes.cpp

+//===- TestTypes.cpp - Test passes for MLIR types -------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestTypes.h"
+#include "TestDialect.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+struct TestRecursiveTypesPass
+    : public PassWrapper<TestRecursiveTypesPass, FunctionPass> {
+  LogicalResult createIRWithTypes();
+
+  void runOnFunction() override {
+    FuncOp func = getFunction();
+
+    // Just make sure recurisve types are printed and parsed.
+    if (func.getName() == "roundtrip")
+      return;
+
+    // Create a recursive type and print it as a part of a dummy op.
+    if (func.getName() == "create") {
+      if (failed(createIRWithTypes()))
+        signalPassFailure();
+      return;
+    }
+
+    // Unknown key.
+    func.emitOpError() << "unexpected function name";
+    signalPassFailure();
+  }
+};
+} // end namespace
+
+LogicalResult TestRecursiveTypesPass::createIRWithTypes() {
+  MLIRContext *ctx = &getContext();
+  FuncOp func = getFunction();
+  auto type = TestRecursiveType::create(ctx, "some_long_and_unique_name");
+  if (failed(type.setBody(type)))
+    return func.emitError("expected to be able to set the type body");
+
+  // Setting the same body is fine.
+  if (failed(type.setBody(type)))
+    return func.emitError(
+        "expected to be able to set the type body to the same value");
+
+  // Setting a different body is not.
+  if (succeeded(type.setBody(IndexType::get(ctx))))
+    return func.emitError(
+        "not expected to be able to change function body more than once");
+
+  // Expecting to get the same type for the same name.
+  auto other = TestRecursiveType::create(ctx, "some_long_and_unique_name");
+  if (type != other)
+    return func.emitError("expected type name to be the uniquing key");
+
+  // Create the op to check how the type is printed.
+  OperationState state(func.getLoc(), "test.dummy_type_test_op");
+  state.addTypes(type);
+  func.getBody().front().push_front(Operation::create(state));
+
+  return success();
+}
+
+namespace mlir {
+
+void registerTestRecursiveTypesPass() {
+  PassRegistration<TestRecursiveTypesPass> reg(
+      "test-recursive-types", "Test support for recursive types");
+}
+
+} // end namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestLoopFusion();
 void registerTestLoopMappingPass();
 void registerTestLoopUnrollingPass();
 void registerTestMatchers();
 void registerTestMemRefDependenceCheck();
 void registerTestMemRefStrideCalculation();
 void registerTestOpaqueLoc();
 void registerTestPreparationPassWithAllowedMemrefResults();
+void registerTestRecursiveTypesPass();
 void registerTestReducer();
 void registerTestGpuParallelLoopMappingPass();
 void registerTestSCFUtilsPass();
 void registerTestVectorConversions();
 void registerVectorizerTestPass();
 } // namespace mlir
 
 static cl::opt<std::string>
#endif
   registerTestLoopFusion();
   registerTestLoopMappingPass();
   registerTestLoopUnrollingPass();
   registerTestMatchers();
   registerTestMemRefDependenceCheck();
   registerTestMemRefStrideCalculation();
   registerTestOpaqueLoc();
   registerTestPreparationPassWithAllowedMemrefResults();
+  registerTestRecursiveTypesPass();
   registerTestReducer();
   registerTestGpuParallelLoopMappingPass();
   registerTestSCFUtilsPass();
   registerTestVectorConversions();
   registerVectorizerTestPass();
 }
 #endif