74 lines · 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-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/IR/DialectResourceBlobManager.h"10#include "llvm/ADT/SmallString.h"11#include <optional>12 13using namespace mlir;14 15//===----------------------------------------------------------------------===//16// DialectResourceBlobManager17//===---------------------------------------------------------------------===//18 19auto 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 26void 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 33auto 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 and53 // 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 67void DialectResourceBlobManager::getBlobMap(68 llvm::function_ref<void(const llvm::StringMap<BlobEntry> &)> accessor)69 const {70 llvm::sys::SmartScopedReader<true> reader(blobMapLock);71 72 accessor(blobMap);73}74