Refactor Attribute uniquing to use StorageUniquer instead of being hard coded in the MLIRContext. This allows for attributes to be uniqued similarly to types. This is the second step towards allowing dialects to define attributes.

--

PiperOrigin-RevId: 245974705

1 files changed, 26 insertions, 0 deletions

diff --git a/mlir/lib/Support/StorageUniquer.cpp b/mlir/lib/Support/StorageUniquer.cpp
index 14d8f393e08..6fd55e77777 100644
--- a/mlir/lib/Support/StorageUniquer.cpp
+++ b/mlir/lib/Support/StorageUniquer.cpp
@@ -105,6 +105,23 @@ struct StorageUniquerImpl {
     return result = initializeStorage(kind, ctorFn);
   }
 
+  /// Erase an instance of a complex derived type.
+  void erase(unsigned kind, unsigned hashValue,
+             llvm::function_ref<bool(const BaseStorage *)> isEqual,
+             llvm::function_ref<void(BaseStorage *)> cleanupFn) {
+    LookupKey lookupKey{kind, hashValue, isEqual};
+
+    // Acquire a writer-lock so that we can safely erase the type instance.
+    llvm::sys::SmartScopedWriter<true> typeLock(mutex);
+    auto existing = storageTypes.find_as(lookupKey);
+    if (existing == storageTypes.end())
+      return;
+
+    // Cleanup the storage and remove it from the map.
+    cleanupFn(existing->storage);
+    storageTypes.erase(existing);
+  }
+
   //===--------------------------------------------------------------------===//
   // Instance Storage
   //===--------------------------------------------------------------------===//
@@ -179,3 +196,12 @@ auto StorageUniquer::getImpl(
     -> BaseStorage * {
   return impl->getOrCreate(kind, ctorFn);
 }
+
+/// Implementation for erasing an instance of a derived type with complex
+/// storage.
+void StorageUniquer::eraseImpl(
+    unsigned kind, unsigned hashValue,
+    llvm::function_ref<bool(const BaseStorage *)> isEqual,
+    std::function<void(BaseStorage *)> cleanupFn) {
+  impl->erase(kind, hashValue, isEqual, cleanupFn);
+}