Merge pull request #81186 from jabrylg/add-json-summary-to-swift-inspect-dump-raw-metadata

[swift-inspect] Add json and summary options to DumpRawMetadata

Author: mikeash

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

@@ -88,6 +88,18 @@ internal struct Output: TextOutputStream {
   }
 }
 
+internal func dumpJson(of: (any Encodable), outputFile: String?) throws {
+  let encoder = JSONEncoder()
+  encoder.outputFormatting = [.prettyPrinted, .sortedKeys]
+  let data = try encoder.encode(of)
+  let jsonOutput = String(data: data, encoding: .utf8)!
+  if let outputFile = outputFile {
+    try jsonOutput.write(toFile: outputFile, atomically: true, encoding: .utf8)
+  } else {
+    print(jsonOutput)
+  }
+}
+
 internal struct DumpGenericMetadata: ParsableCommand {
   static let configuration = CommandConfiguration(
     abstract: "Print the target's generic metadata allocations.")
@@ -99,7 +111,10 @@ internal struct DumpGenericMetadata: ParsableCommand {
   var backtraceOptions: BacktraceOptions
 
   @OptionGroup()
-  var genericMetadataOptions: GenericMetadataOptions
+  var metadataOptions: MetadataOptions
+  
+  @Flag(help: "Show allocations in mangled form")
+  var mangled: Bool = false
 
   func run() throws {
     disableStdErrBuffer()
@@ -126,7 +141,7 @@ internal struct DumpGenericMetadata: ParsableCommand {
 
         return Metadata(ptr: pointer,
                         allocation: allocation,
-                        name: process.context.name(type: pointer, mangled: genericMetadataOptions.mangled) ?? "<unknown>",
+                        name: process.context.name(type: pointer, mangled: mangled) ?? "<unknown>",
                         isArrayOfClass: process.context.isArrayOfClass(pointer),
                         garbage: garbage,
                         backtrace: currentBacktrace)
@@ -146,30 +161,30 @@ internal struct DumpGenericMetadata: ParsableCommand {
         }
       }
 
-      if genericMetadataOptions.json {
+      if metadataOptions.json {
         let processMetadata = ProcessMetadata(name: process.processName,
                                               pid: process.processIdentifier as! ProcessIdentifier,
                                               metadata: generics)
         allProcesses.append(processMetadata)
-      } else if !genericMetadataOptions.summary {
+      } else if !metadataOptions.summary {
         try dumpText(process: process, generics: generics)
         }
     } // inspect
 
-    if genericMetadataOptions.json {
-      if genericMetadataOptions.summary {
-        try dumpJson(of: metadataSummary)
+    if metadataOptions.json {
+      if metadataOptions.summary {
+        try dumpJson(of: metadataSummary, outputFile: metadataOptions.outputFile)
       } else {
-        try dumpJson(of: allProcesses)
+        try dumpJson(of: allProcesses, outputFile: metadataOptions.outputFile)
         }
-    } else if genericMetadataOptions.summary {
+    } else if metadataOptions.summary {
       try dumpTextSummary(of: metadataSummary)
     }
   }
 
   private func dumpText(process: any RemoteProcess, generics: [Metadata]) throws {
     var erroneousMetadata: [(ptr: swift_reflection_ptr_t, name: String)] = []
-    var output = try Output(genericMetadataOptions.outputFile)
+    var output = try Output(metadataOptions.outputFile)
     print("\(process.processName)(\(process.processIdentifier)):\n", to: &output)
     print("Address", "Allocation", "Size", "Offset", "isArrayOfClass", "Name", separator: "\t", to: &output)
     generics.forEach {
@@ -198,20 +213,8 @@ internal struct DumpGenericMetadata: ParsableCommand {
     print("", to: &output)
   }
 
-  private func dumpJson(of: (any Encodable)) throws {
-    let encoder = JSONEncoder()
-    encoder.outputFormatting = [.prettyPrinted, .sortedKeys]
-    let data = try encoder.encode(of)
-    let jsonOutput = String(data: data, encoding: .utf8)!
-    if let outputFile = genericMetadataOptions.outputFile {
-      try jsonOutput.write(toFile: outputFile, atomically: true, encoding: .utf8)
-    } else {
-      print(jsonOutput)
-    }
-  }
-
   private func dumpTextSummary(of: [String: MetadataSummary]) throws {
-    var output = try Output(genericMetadataOptions.outputFile)
+    var output = try Output(metadataOptions.outputFile)
     print("Size", "Owners", "Name", separator: "\t", to: &output)
     var totalSize = 0
     var unknownSize = 0

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

@@ -12,6 +12,23 @@
 
 import ArgumentParser
 import SwiftRemoteMirror
+import Foundation
+
+private struct AllocatorTagTotal: Encodable {
+  let name: String
+  let tag: Int
+  var totalBytes: Int
+}
+
+private struct Summary: Encodable {
+  let totalBytesAllocated: Int
+  let allocatorTags: [AllocatorTagTotal]
+}
+
+private struct RawMetadataOutput: Encodable {
+  let allocationList: [swift_metadata_allocation_t]?
+  let summary: Summary?
+}
 
 internal struct DumpRawMetadata: ParsableCommand {
   static let configuration = CommandConfiguration(
@@ -22,8 +39,15 @@ internal struct DumpRawMetadata: ParsableCommand {
 
   @OptionGroup()
   var backtraceOptions: BacktraceOptions
+  
+ @OptionGroup()
+  var metadataOptions: MetadataOptions
 
   func run() throws {
+    var allocatorTagTotals = [Int: AllocatorTagTotal]()
+    var total: Int = 0
+    var allocationList: [swift_metadata_allocation_t] = []
+
     try inspect(options: options) { process in
       let stacks: [swift_reflection_ptr_t:[swift_reflection_ptr_t]]? =
           backtraceOptions.style == nil
@@ -33,6 +57,18 @@ internal struct DumpRawMetadata: ParsableCommand {
       try process.context.allocations.forEach { allocation in
         let name: String = process.context.name(allocation: allocation.tag) ??  "<unknown>"
         print("Metadata allocation at: \(hex: allocation.ptr) size: \(allocation.size) tag: \(allocation.tag) (\(name))")
+        
+        if metadataOptions.summary {
+          if var allocatorTagTotal = allocatorTagTotals[Int(allocation.tag)] {
+            allocatorTagTotal.totalBytes += allocation.size
+            allocatorTagTotals[Int(allocation.tag)] = allocatorTagTotal
+          } else {
+            allocatorTagTotals[Int(allocation.tag)] = AllocatorTagTotal(name: name, tag: Int(allocation.tag), totalBytes: allocation.size)
+          }
+          
+          total += allocation.size
+        }
+        allocationList.append(allocation)
         if let style = backtraceOptions.style {
           if let stack = stacks?[allocation.ptr] {
             print(backtrace(stack, style: style, process.symbolicate))
@@ -42,5 +78,27 @@ internal struct DumpRawMetadata: ParsableCommand {
         }
       }
     }
+    
+    if metadataOptions.json {
+      let jsonStruct: RawMetadataOutput
+      let allocatorTagArray = Array(allocatorTagTotals.values).sorted(by: {$0.totalBytes > $1.totalBytes})
+      
+      if metadataOptions.summary {
+        let summaryStruct = Summary(totalBytesAllocated: total, allocatorTags: allocatorTagArray)
+        jsonStruct = RawMetadataOutput(allocationList: allocationList, summary: summaryStruct)
+      } else {
+        jsonStruct = RawMetadataOutput(allocationList: allocationList, summary: nil)
+      }
+      try dumpJson(of: jsonStruct, outputFile: metadataOptions.outputFile)
+    } else if metadataOptions.summary {
+      let allocatorTagArray = Array(allocatorTagTotals.values).sorted(by: {$0.totalBytes > $1.totalBytes})
+      
+      print("Metadata allocation summary:")
+      for tag in allocatorTagArray {
+        print("Tag: \(tag.tag) (\(tag.name)) Size: \(tag.totalBytes) bytes")
+      }
+      
+      print("\nTotal bytes allocated: \(total)")
+    }
   }
 }

tools/swift-inspect/Sources/swift-inspect/main.swift

@@ -62,10 +62,7 @@ internal struct BacktraceOptions: ParsableArguments {
   }
 }
 
-internal struct GenericMetadataOptions: ParsableArguments {
-  @Flag(help: "Show allocations in mangled form")
-  var mangled: Bool = false
-
+internal struct MetadataOptions: ParsableArguments {
   @Flag(help: "Output JSON")
   var json: Bool = false