mlir/lib/IR/DialectResourceBlobManager.cpp

   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
  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
  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
  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 }