Merge pull request #84446 from mikeash/swift-inspect-scan-search

[swift-inspect] Add a mode to search for and dump generic metadata without metadata iteration enabled.

Author: mikeash

include/swift/Remote/MetadataReader.h

@@ -3215,6 +3215,23 @@ class MetadataReader {
     if (!descriptor)
       return BuiltType();
 
+    // Make sure the kinds match, to catch bad data from not reading an actual
+    // metadata.
+    auto metadataKind = metadata.getLocalBuffer()->getKind();
+    auto descriptorKind = descriptor.getLocalBuffer()->getKind();
+    if (metadataKind == MetadataKind::Class &&
+        descriptorKind != ContextDescriptorKind::Class)
+      return BuiltType();
+    if (metadataKind == MetadataKind::Struct &&
+        descriptorKind != ContextDescriptorKind::Struct)
+      return BuiltType();
+    if (metadataKind == MetadataKind::Enum &&
+        descriptorKind != ContextDescriptorKind::Enum)
+      return BuiltType();
+    if (metadataKind == MetadataKind::Optional &&
+        descriptorKind != ContextDescriptorKind::Enum)
+      return BuiltType();
+
     // From that, attempt to resolve a nominal type.
     BuiltTypeDecl typeDecl = buildNominalTypeDecl(descriptor);
     if (!typeDecl)

include/swift/Runtime/Debug.h

@@ -266,6 +266,14 @@ bool _swift_debug_metadataAllocationIterationEnabled;
 SWIFT_RUNTIME_STDLIB_SPI
 const void * const _swift_debug_allocationPoolPointer;
 
+SWIFT_RUNTIME_STDLIB_SPI
+const size_t _swift_debug_allocationPoolSize;
+
+// The size of the pages the metadata allocator allocates on the heap. May be
+// used to filter out possible metadata pages when examining the heap.
+SWIFT_RUNTIME_STDLIB_SPI
+const size_t _swift_debug_metadataAllocatorPageSize;
+
 SWIFT_RUNTIME_STDLIB_SPI
 std::atomic<const void *> _swift_debug_metadataAllocationBacktraceList;
 

stdlib/public/RemoteInspection/TypeLowering.cpp

@@ -2324,8 +2324,9 @@ class LowerType
     if (auto N = dyn_cast<NominalTypeRef>(TR)) {
       Demangler Dem;
       auto Node = N->getDemangling(Dem);
-      if (Node->getKind() == Node::Kind::Type && Node->getNumChildren() == 1) {
-	auto Alias = Node->getChild(0);
+      if (Node && Node->getKind() == Node::Kind::Type &&
+          Node->getNumChildren() == 1) {
+        auto Alias = Node->getChild(0);
 	if (Alias->getKind() == Node::Kind::TypeAlias && Alias->getNumChildren() == 2) {
 	  auto Module = Alias->getChild(0);
 	  auto Name = Alias->getChild(1);

stdlib/public/runtime/Metadata.cpp

@@ -7903,6 +7903,8 @@ std::tuple<const void *, size_t> MetadataAllocator::InitialPoolLocation() {
 
 bool swift::_swift_debug_metadataAllocationIterationEnabled = false;
 const void * const swift::_swift_debug_allocationPoolPointer = &AllocationPool;
+const size_t swift::_swift_debug_allocationPoolSize = InitialPoolSize;
+const size_t swift::_swift_debug_metadataAllocatorPageSize = PoolRange::PageSize;
 std::atomic<const void *> swift::_swift_debug_metadataAllocationBacktraceList;
 
 static void recordBacktrace(void *allocation) {

test/abi/macOS/arm64/stdlib.swift

@@ -1139,3 +1139,7 @@ Added: _$ss7UnicodeO27_RandomAccessWordRecognizerVN
 // Obsolete/broken SPIs removed in 6.3
 Removed: _$sSS17_nearestWordIndex9atOrBelowSS0C0VAD_tF
 Removed: _$sSS10_wordIndex6beforeSS0B0VAD_tF
+
+// Internal info exposed for swift-inspect.
+Added: __swift_debug_allocationPoolSize
+Added: __swift_debug_metadataAllocatorPageSize

test/abi/macOS/x86_64/stdlib.swift

@@ -1139,3 +1139,7 @@ Added: _$ss7UnicodeO27_RandomAccessWordRecognizerVN
 // Obsolete/broken SPIs removed in 6.3
 Removed: _$sSS17_nearestWordIndex9atOrBelowSS0C0VAD_tF
 Removed: _$sSS10_wordIndex6beforeSS0B0VAD_tF
+
+// Internal info exposed for swift-inspect.
+Added: __swift_debug_allocationPoolSize
+Added: __swift_debug_metadataAllocatorPageSize

tools/swift-inspect/Sources/swift-inspect/DarwinRemoteProcess.swift

@@ -87,17 +87,30 @@ internal final class DarwinRemoteProcess: RemoteProcess {
     return task_peek(task.value, address, mach_vm_size_t(size))
   }
 
-  func getAddr(symbolName: String) -> swift_addr_t {
+  func getAddr(symbolName: String) -> swift_addr_t? {
     // FIXME: use `__USER_LABEL_PREFIX__` instead of the hardcoded `_`.
     let fullName = "_\(symbolName)"
     var symbol = CSSymbolOwnerGetSymbolWithMangledName(swiftCore, fullName)
     if CSIsNull(symbol) {
       symbol = CSSymbolOwnerGetSymbolWithMangledName(swiftConcurrency, fullName)
     }
+    if CSIsNull(symbol) {
+      return nil
+    }
     let range = CSSymbolGetRange(symbol)
     return swift_addr_t(range.location)
   }
 
+  private func readGlobalVariable<T>(named name: String) -> T? {
+    guard let globalPointer = getAddr(symbolName: name) else {
+      return nil
+    }
+    guard let readPointer = read(address: globalPointer, size: MemoryLayout<T>.size) else {
+      return nil
+    }
+    return readPointer.load(as: T.self)
+  }
+
   static var Free: FreeFunction? { return nil }
 
   static var ReadBytes: ReadBytesFunction {
@@ -125,7 +138,7 @@ internal final class DarwinRemoteProcess: RemoteProcess {
         let buffer = UnsafeBufferPointer(start: $0, count: Int(length))
         return String(decoding: buffer, as: UTF8.self)
       }
-      return process.getAddr(symbolName: name)
+      return process.getAddr(symbolName: name) ?? 0
     }
   }
 
@@ -224,6 +237,21 @@ internal final class DarwinRemoteProcess: RemoteProcess {
       }
     }
   }
+
+  internal func iteratePotentialMetadataPages(_ body: (swift_addr_t, UInt64) -> Void) {
+    if let initialPoolPointer: UInt = readGlobalVariable(named: "_swift_debug_allocationPoolPointer"),
+       let initialPoolSize: UInt = readGlobalVariable(named: "_swift_debug_allocationPoolSize") {
+      body(swift_reflection_ptr_t(initialPoolPointer), UInt64(initialPoolSize));
+    }
+
+    if let pageSize: UInt = readGlobalVariable(named: "_swift_debug_metadataAllocatorPageSize") {
+      iterateHeap { address, size in
+        if size == pageSize {
+          body(address, size)
+        }
+      }
+    }
+  }
 }
 
 extension DarwinRemoteProcess {

tools/swift-inspect/Sources/swift-inspect/LinuxRemoteProcess.swift

@@ -154,5 +154,9 @@
     internal func iterateHeap(_ body: (swift_addr_t, UInt64) -> Void) {
       fatalError("heap iteration is not supported on Linux")
     }
+
+    internal func iteratePotentialMetadataPages(_ body: (swift_addr_t, UInt64) -> Void) {
+      fatalError("metadata page iteration is not supported on Linux")
+    }
   }
 #endif  // os(Linux)

tools/swift-inspect/Sources/swift-inspect/Operations/DumpConcurrency.swift

@@ -92,8 +92,8 @@ fileprivate class ConcurrencyDumper {
     self.process = process
 
     func getMetadata(symbolName: String) -> swift_reflection_ptr_t? {
-      let addr = process.getAddr(symbolName: symbolName)
-      if let ptr = process.read(address: addr, size: MemoryLayout<UInt>.size) {
+      if let addr = process.getAddr(symbolName: symbolName),
+         let ptr = process.read(address: addr, size: MemoryLayout<UInt>.size) {
         return swift_reflection_ptr_t(ptr.load(as: UInt.self))
       }
       return nil

tools/swift-inspect/Sources/swift-inspect/Operations/DumpGenericMetadata.swift

@@ -116,36 +116,17 @@ internal struct DumpGenericMetadata: ParsableCommand {
   @Flag(help: "Show allocations in mangled form")
   var mangled: Bool = false
 
+  @Flag(help: "Scan for data that looks like Swift type metadata.")
+  var scanSearch: Bool = false
+
   func run() throws {
     disableStdErrBuffer()
     var metadataSummary = [String: MetadataSummary]()
     var allProcesses = [ProcessMetadata]()
     try inspect(options: options) { process in
-      let allocations: [swift_metadata_allocation_t] =
-          try process.context.allocations.sorted()
-
-      let stacks: [swift_reflection_ptr_t:[swift_reflection_ptr_t]] =
-          backtraceOptions.style == nil
-          ? [swift_reflection_ptr_t:[swift_reflection_ptr_t]]()
-          : try process.context.allocationStacks
-
-      let generics: [Metadata] = allocations.compactMap { allocation -> Metadata? in
-        let pointer = swift_reflection_allocationMetadataPointer(process.context, allocation)
-        if pointer == 0 { return nil }
-        let allocation = allocations.last(where: { pointer >= $0.ptr && pointer < $0.ptr + swift_reflection_ptr_t($0.size) })
-        let garbage = (allocation == nil && swift_reflection_ownsAddressStrict(process.context, UInt(pointer)) == 0)
-        var currentBacktrace: String?
-        if let style = backtraceOptions.style, let allocation, let stack = stacks[allocation.ptr] {
-          currentBacktrace = backtrace(stack, style: style, process.symbolicate)
-        }
-
-        return Metadata(ptr: pointer,
-                        allocation: allocation,
-                        name: process.context.name(type: pointer, mangled: mangled) ?? "<unknown>",
-                        isArrayOfClass: process.context.isArrayOfClass(pointer),
-                        garbage: garbage,
-                        backtrace: currentBacktrace)
-      } // generics
+      let generics = scanSearch
+          ? try metadataFromScanning(process: process)
+          : try metadataFromAllocations(process: process)
 
       // Update summary
       generics.forEach { metadata in
@@ -182,6 +163,57 @@ internal struct DumpGenericMetadata: ParsableCommand {
     }
   }
 
+  private func metadataFromAllocations(process: any RemoteProcess) throws -> [Metadata] {
+    let allocations: [swift_metadata_allocation_t] =
+        try process.context.allocations.sorted()
+
+    let stacks: [swift_reflection_ptr_t:[swift_reflection_ptr_t]] =
+        backtraceOptions.style == nil
+        ? [swift_reflection_ptr_t:[swift_reflection_ptr_t]]()
+        : try process.context.allocationStacks
+
+    return allocations.compactMap { allocation -> Metadata? in
+      let pointer = swift_reflection_allocationMetadataPointer(process.context, allocation)
+      if pointer == 0 { return nil }
+      let allocation = allocations.last(where: { pointer >= $0.ptr && pointer < $0.ptr + swift_reflection_ptr_t($0.size) })
+      let garbage = (allocation == nil && swift_reflection_ownsAddressStrict(process.context, UInt(pointer)) == 0)
+      var currentBacktrace: String?
+      if let style = backtraceOptions.style, let allocation, let stack = stacks[allocation.ptr] {
+        currentBacktrace = backtrace(stack, style: style, process.symbolicate)
+      }
+
+      return Metadata(ptr: pointer,
+                      allocation: allocation,
+                      name: process.context.name(type: pointer, mangled: mangled) ?? "<unknown>",
+                      isArrayOfClass: process.context.isArrayOfClass(pointer),
+                      garbage: garbage,
+                      backtrace: currentBacktrace)
+    }
+  }
+
+  private func metadataFromScanning(process: any RemoteProcess) throws -> [Metadata] {
+    var metadata: [Metadata] = []
+
+    func scanMemory(address: swift_reflection_ptr_t, size: UInt64) {
+      for candidate in stride(from: address, to: address + swift_reflection_ptr_t(size), by: MemoryLayout<UInt>.size) {
+        guard let name = process.context.name(type: candidate, mangled: mangled) else {
+          continue
+        }
+        let m = Metadata(ptr: candidate,
+                         allocation: nil,
+                         name: name,
+                         isArrayOfClass: process.context.isArrayOfClass(candidate),
+                         garbage: false,
+                         backtrace: nil)
+        metadata.append(m)
+      }
+    }
+
+    process.iteratePotentialMetadataPages(scanMemory)
+
+    return metadata
+  }
+
   private func dumpText(process: any RemoteProcess, generics: [Metadata]) throws {
     var erroneousMetadata: [(ptr: swift_reflection_ptr_t, name: String)] = []
     var output = try Output(metadataOptions.outputFile)