diff --git a/stdlib/public/runtime/LibPrespecialized.cpp b/stdlib/public/runtime/LibPrespecialized.cpp index c48873d417d7c..2d9544403c1a5 100644 --- a/stdlib/public/runtime/LibPrespecialized.cpp +++ b/stdlib/public/runtime/LibPrespecialized.cpp @@ -11,6 +11,7 @@ //===----------------------------------------------------------------------===// #include "swift/Runtime/LibPrespecialized.h" +#include "MetadataCache.h" #include "Private.h" #include "swift/Basic/Lazy.h" #include "swift/Runtime/EnvironmentVariables.h" @@ -71,17 +72,56 @@ static const LibPrespecializedData *findLibPrespecialized() { return data; } +struct LibPrespecializedState { + struct AddressRange { + uintptr_t start, end; + + bool contains(const void *ptr) { + return start <= (uintptr_t)ptr && (uintptr_t)ptr < end; + } + }; + + bool loggingEnabled; + const LibPrespecializedData *data; + AddressRange sharedCacheRange{0, 0}; + AddressRange metadataAllocatorInitialPoolRange{0, 0}; + + LibPrespecializedState() { + loggingEnabled = + runtime::environment::SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED_LOGGING(); + data = findLibPrespecialized(); + +#if DYLD_GET_SWIFT_PRESPECIALIZED_DATA_DEFINED + size_t sharedCacheLength; + sharedCacheRange.start = + (uintptr_t)_dyld_get_shared_cache_range(&sharedCacheLength); + sharedCacheRange.end = sharedCacheRange.start + sharedCacheLength; + + auto [initialPoolStart, initialPoolLength] = + MetadataAllocator::InitialPoolLocation(); + metadataAllocatorInitialPoolRange.start = (uintptr_t)initialPoolStart; + metadataAllocatorInitialPoolRange.end = + metadataAllocatorInitialPoolRange.start + initialPoolLength; +#endif + } +}; + +static Lazy LibPrespecialized; + +#define LOG(fmt, ...) \ + do { \ + if (SWIFT_UNLIKELY(prespecialized.loggingEnabled)) \ + fprintf(stderr, "Prespecializations library: " fmt "\n", __VA_ARGS__); \ + } while (0) + const LibPrespecializedData *swift::getLibPrespecializedData() { return SWIFT_LAZY_CONSTANT(findLibPrespecialized()); } // Returns true if the type has any arguments that aren't plain types (packs or // unknown kinds). -static bool hasNonTypeGenericArguments(const TypeContextDescriptor *description) { - auto generics = description->getGenericContext(); - if (!generics) - return false; - +static bool +hasNonTypeGenericArguments(const TargetGenericContext *generics) { for (auto param : generics->getGenericParams()) if (param.getKind() != GenericParamKind::Type) return true; @@ -89,6 +129,32 @@ static bool hasNonTypeGenericArguments(const TypeContextDescriptor *description) return false; } +static bool +isPotentialPrespecializedPointer(LibPrespecializedState &prespecialized, + const void *pointer) { + // Prespecialized metadata descriptors and arguments are always in the shared + // cache. They're either statically emitted metadata, or they're + // prespecialized metadata. Anything that's dynamically allocated, or + // statically allocated outside the shared cache, is not a possible candidate. + + // If we're loading a debug libprespecialized, we can't do these checks, so + // just say everything is a potential argument. Performance is not so + // important in that case. + if (!prespecialized.sharedCacheRange.contains(prespecialized.data)) + return true; + + // Anything outside the shared cache isn't a potential argument. + if (!prespecialized.sharedCacheRange.contains(pointer)) + return false; + + // Dynamically allocated metadata could be within the shared cache, in the + // initial metadata allocation pool. Reject anything in that region. + if (prespecialized.metadataAllocatorInitialPoolRange.contains(pointer)) + return false; + + return true; +} + static bool disableForValidation = false; Metadata * @@ -97,24 +163,40 @@ swift::getLibPrespecializedMetadata(const TypeContextDescriptor *description, if (disableForValidation) return nullptr; - auto *data = getLibPrespecializedData(); + auto &prespecialized = LibPrespecialized.get(); + + auto *data = prespecialized.data; if (!data) return nullptr; + auto *generics = description->getGenericContext(); + if (!generics) + return nullptr; + // We don't support types with pack parameters yet (and especially not types // with unknown parameter kinds) so don't even try to look those up. - if (hasNonTypeGenericArguments(description)) + if (hasNonTypeGenericArguments(generics)) return nullptr; + if (!isPotentialPrespecializedPointer(prespecialized, description)) { + LOG("Rejecting descriptor %p, not in the shared cache", + (const void *)description); + return nullptr; + } + + auto numKeyArguments = generics->getGenericContextHeader().NumKeyArguments; + for (unsigned i = 0; i < numKeyArguments; i++) { + if (!isPotentialPrespecializedPointer(prespecialized, arguments[i])) { + LOG("Rejecting argument %u %p to descriptor %p, not in the shared cache", + i, arguments[i], (const void *)description); + return nullptr; + } + } + Demangler dem; auto mangleNode = _buildDemanglingForGenericType(description, arguments, dem); if (!mangleNode) { - if (SWIFT_UNLIKELY(runtime::environment:: - SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED_LOGGING())) - fprintf(stderr, - "Prespecializations library: failed to build demangling with " - "descriptor %p.\n", - description); + LOG("failed to build demangling with descriptor %p.", description); return nullptr; } @@ -142,10 +224,7 @@ swift::getLibPrespecializedMetadata(const TypeContextDescriptor *description, auto *metadataMap = data->getMetadataMap(); auto *element = metadataMap->find(key.data(), key.size()); auto *result = element ? element->value : nullptr; - if (SWIFT_UNLIKELY(runtime::environment:: - SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED_LOGGING())) - fprintf(stderr, "Prespecializations library: found %p for key '%.*s'.\n", - result, (int)key.size(), key.data()); + LOG("found %p for key '%.*s'.", result, (int)key.size(), key.data()); return result; } diff --git a/stdlib/public/runtime/Metadata.cpp b/stdlib/public/runtime/Metadata.cpp index 746388096aba5..373e6f7f10298 100644 --- a/stdlib/public/runtime/Metadata.cpp +++ b/stdlib/public/runtime/Metadata.cpp @@ -7420,6 +7420,10 @@ static swift::atomic AllocationPool{PoolRange{InitialAllocationPool.Pool, sizeof(InitialAllocationPool.Pool)}}; +std::tuple MetadataAllocator::InitialPoolLocation() { + return {InitialAllocationPool.Pool, sizeof(InitialAllocationPool.Pool)}; +} + bool swift::_swift_debug_metadataAllocationIterationEnabled = false; const void * const swift::_swift_debug_allocationPoolPointer = &AllocationPool; std::atomic swift::_swift_debug_metadataAllocationBacktraceList; diff --git a/stdlib/public/runtime/MetadataCache.h b/stdlib/public/runtime/MetadataCache.h index 752eab6bc50b2..afd944ead0d73 100644 --- a/stdlib/public/runtime/MetadataCache.h +++ b/stdlib/public/runtime/MetadataCache.h @@ -25,6 +25,7 @@ #include #include #include +#include #ifndef SWIFT_DEBUG_RUNTIME #define SWIFT_DEBUG_RUNTIME 0 @@ -46,6 +47,11 @@ class MetadataAllocator : public llvm::AllocatorBase { void Reset() {} + /// Get the location of the allocator's initial statically allocated pool. + /// The return values are start and size. If there is no statically allocated + /// pool, the return values are NULL, 0. + static std::tuple InitialPoolLocation(); + SWIFT_RETURNS_NONNULL SWIFT_NODISCARD void *Allocate(size_t size, size_t alignment); using AllocatorBase::Allocate; @@ -67,6 +73,9 @@ class MetadataAllocator : public llvm::AllocatorBase { class MetadataAllocator { public: MetadataAllocator(uint16_t tag) {} + static std::tuple InitialPoolLocation() { + return {nullptr, 0}; + } SWIFT_RETURNS_NONNULL SWIFT_NODISCARD void *Allocate(size_t size, size_t alignment) { if (alignment < sizeof(void*)) alignment = sizeof(void*);