xref: /llvm-project/mlir/lib/IR/DialectResourceBlobManager.cpp (revision 0a81ace0047a2de93e71c82cdf0977fc989660df)
1 //===- DialectResourceBlobManager.cpp - Dialect Blob Management -----------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/IR/DialectResourceBlobManager.h"
10 #include "llvm/ADT/SmallString.h"
11 #include <optional>
12 
13 using namespace mlir;
14 
15 //===----------------------------------------------------------------------===//
16 // DialectResourceBlobManager
17 //===---------------------------------------------------------------------===//
18 
lookup(StringRef name)19 auto DialectResourceBlobManager::lookup(StringRef name) -> BlobEntry * {
20   llvm::sys::SmartScopedReader<true> reader(blobMapLock);
21 
22   auto it = blobMap.find(name);
23   return it != blobMap.end() ? &it->second : nullptr;
24 }
25 
update(StringRef name,AsmResourceBlob && newBlob)26 void DialectResourceBlobManager::update(StringRef name,
27                                         AsmResourceBlob &&newBlob) {
28   BlobEntry *entry = lookup(name);
29   assert(entry && "`update` expects an existing entry for the provided name");
30   entry->setBlob(std::move(newBlob));
31 }
32 
insert(StringRef name,std::optional<AsmResourceBlob> blob)33 auto DialectResourceBlobManager::insert(StringRef name,
34                                         std::optional<AsmResourceBlob> blob)
35     -> BlobEntry & {
36   llvm::sys::SmartScopedWriter<true> writer(blobMapLock);
37 
38   // Functor used to attempt insertion with a given name.
39   auto tryInsertion = [&](StringRef name) -> BlobEntry * {
40     auto it = blobMap.try_emplace(name, BlobEntry());
41     if (it.second) {
42       it.first->second.initialize(it.first->getKey(), std::move(blob));
43       return &it.first->second;
44     }
45     return nullptr;
46   };
47 
48   // Try inserting with the name provided by the user.
49   if (BlobEntry *entry = tryInsertion(name))
50     return *entry;
51 
52   // If an entry already exists for the user provided name, tweak the name and
53   // re-attempt insertion until we find one that is unique.
54   llvm::SmallString<32> nameStorage(name);
55   nameStorage.push_back('_');
56   size_t nameCounter = 1;
57   do {
58     Twine(nameCounter++).toVector(nameStorage);
59 
60     // Try inserting with the new name.
61     if (BlobEntry *entry = tryInsertion(nameStorage))
62       return *entry;
63     nameStorage.resize(name.size() + 1);
64   } while (true);
65 }
66