Pack expression IDs as single integers or integer arrays (rare) instead of strings

Author: grynspan

Sources/Testing/Expectations/ExpectationContext.swift

@@ -99,6 +99,7 @@ public struct __ExpectationContext: ~Copyable {
     // we've captured during evaluation.
     var expressionGraph = Graph<UInt32, __Expression?>()
     for (id, sourceCode) in sourceCode {
+      print(String(reflecting: id))
       let keyPath = id.keyPath
       expressionGraph.insertValue(__Expression(sourceCode), at: keyPath)
     }

Sources/Testing/SourceAttribution/ExpressionID.swift

@@ -11,68 +11,137 @@
 /// A type providing unique identifiers for expressions captured during
 /// expansion of the `#expect()` and `#require()` macros.
 ///
-/// In the future, this type may use [`StaticBigInt`](https://developer.apple.com/documentation/swift/staticbigint)
-/// as its source representation rather than a string literal.
+/// This type tries to optimize for expressions in shallow syntax trees whose
+/// unique identifiers require 64 bits or fewer. Wider unique identifiers are
+/// stored as arrays of 64-bit words. In the future, this type may use
+/// [`StaticBigInt`](https://developer.apple.com/documentation/swift/staticbigint)
+/// to represent expression identifiers instead.
 ///
 /// - Warning: This type is used to implement the `#expect()` and `#require()`
 ///   macros. Do not use it directly.
 public struct __ExpressionID: Sendable {
   /// The ID of the root node in an expression graph.
   static var root: Self {
-    ""
+    Self(_elements: .none)
   }
 
-  /// The string produced at compile time that encodes the unique identifier of
-  /// the represented expression.
-  var stringValue: String
+  /// An enumeration that attempts to efficiently store the key path elements
+  /// corresponding to an expression ID.
+  fileprivate enum Elements: Sendable {
+    /// This ID does not use any words.
+    ///
+    /// This case represents the root node in a syntax tree. An instance of
+    /// `__ExpressionID` storing this case is implicitly equal to `.root`.
+    case none
 
-  /// The number of bits in a nybble.
-  private static var _bitsPerNybble: Int { 4 }
+    /// This ID packs its corresponding key path value into a single word whose
+    /// value is not `0`.
+    case packed(_ word: UInt64)
+
+    /// This ID contains key path elements that do not fit in a 64-bit integer,
+    /// so they are not packed and map directly to the represented key path.
+    indirect case keyPath(_ keyPath: [UInt32])
+  }
+
+  /// The elements of this identifier.
+  private var _elements: Elements
 
   /// A representation of this instance suitable for use as a key path in an
   /// instance of `Graph` where the key type is `UInt32`.
   ///
   /// The values in this collection, being swift-syntax node IDs, are never more
   /// than 32 bits wide.
   var keyPath: some RandomAccessCollection<UInt32> {
-    let nybbles = stringValue
-      .reversed().lazy
-      .compactMap { UInt8(String($0), radix: 16) }
-
-    return nybbles
-      .enumerated()
-      .flatMap { i, nybble in
-        let nybbleOffset = i * Self._bitsPerNybble
-        return (0 ..< Self._bitsPerNybble).lazy
-          .filter { (nybble & (1 << $0)) != 0 }
-          .map { UInt32(nybbleOffset + $0) }
+    // Helper function to unpack a sequence of words into bit indices for use as
+    // a Graph's key path.
+    func makeKeyPath(from words: some RandomAccessCollection<UInt64>) -> [UInt32] {
+      // Assume approximately 1/4 of the bits are populated. We can always tweak
+      // this guesstimate after gathering more real-world data.
+      var result = [UInt32]()
+      result.reserveCapacity((words.count * UInt64.bitWidth) / 4)
+
+      for (bitOffset, word) in words.enumerated() {
+        var word = word
+        while word != 0 {
+          let bit = word.trailingZeroBitCount
+          result.append(UInt32(bit + bitOffset))
+          word = word & (word &- 1) // Mask off the bit we just counted.
+        }
       }
+
+      return result
+    }
+
+    switch _elements {
+    case .none:
+      return []
+    case let .packed(word):
+      // Assume approximately 1/4 of the bits are populated. We can always tweak
+      // this guesstimate after gathering more real-world data.
+      var result = [UInt32]()
+      result.reserveCapacity(UInt64.bitWidth / 4)
+
+      var word = word
+      while word != 0 {
+        let bit = word.trailingZeroBitCount
+        result.append(UInt32(bit))
+        word = word & (word &- 1) // Mask off the bit we just counted.
+      }
+
+      return result
+    case let .keyPath(keyPath):
+      return keyPath
+    }
   }
 }
 
 // MARK: - Equatable, Hashable
 
 extension __ExpressionID: Equatable, Hashable {}
+extension __ExpressionID.Elements: Equatable, Hashable {}
 
 #if DEBUG
 // MARK: - CustomStringConvertible, CustomDebugStringConvertible
 
 extension __ExpressionID: CustomStringConvertible, CustomDebugStringConvertible {
+  /// The number of bits in a nybble.
+  private static var _bitsPerNybble: Int { 4 }
+
+  /// The number of nybbles in a word.
+  private static var _nybblesPerWord: Int { UInt64.bitWidth / _bitsPerNybble }
+
   public var description: String {
-    stringValue
+    switch _elements {
+    case .none:
+      return "0"
+    case let .packed(word):
+      return String(word, radix: 16)
+    case let .keyPath(keyPath):
+      return keyPath.lazy
+        .map { String($0, radix: 16) }
+        .joined(separator: ",")
+    }
   }
 
   public var debugDescription: String {
-    #""\#(stringValue)" → \#(Array(keyPath))"#
+    #""\#(description)" → \#(Array(keyPath))"#
   }
 }
 #endif
 
-// MARK: - ExpressibleByStringLiteral
+// MARK: - ExpressibleByIntegerLiteral
+
+extension __ExpressionID: ExpressibleByIntegerLiteral {
+  public init(integerLiteral: UInt64) {
+    if integerLiteral == 0 {
+      self.init(_elements: .none)
+    } else {
+      self.init(_elements: .packed(integerLiteral))
+    }
+  }
 
-extension __ExpressionID: ExpressibleByStringLiteral {
-  public init(stringLiteral: String) {
-    stringValue = stringLiteral
+  public init(_ keyPath: UInt32...) {
+    self.init(_elements: .keyPath(keyPath))
   }
 }
 

Sources/TestingMacros/ConditionMacro.swift

@@ -201,12 +201,12 @@ extension ConditionMacro {
 
           // Sort the rewritten nodes. This isn't strictly necessary for
           // correctness but it does make the produced code more consistent.
-          let sortedRewrittenNodes = rewrittenNodes.sorted { $0.id < $1.id }
-          let sourceCodeNodeIDs = sortedRewrittenNodes.compactMap { $0.expressionID(rootedAt: originalArgumentExpr) }
-          let sourceCodeExprs = sortedRewrittenNodes.map { StringLiteralExprSyntax(content: $0.trimmedDescription) }
           let sourceCodeExpr = DictionaryExprSyntax {
-            for (nodeID, sourceCodeExpr) in zip(sourceCodeNodeIDs, sourceCodeExprs) {
-              DictionaryElementSyntax(key: nodeID, value: sourceCodeExpr)
+            for node in (rewrittenNodes.sorted { $0.id < $1.id }) {
+              DictionaryElementSyntax(
+                key: node.expressionID(rootedAt: originalArgumentExpr),
+                value: StringLiteralExprSyntax(content: node.trimmedDescription)
+              )
             }
           }
           checkArguments.append(Argument(label: "sourceCode", expression: sourceCodeExpr))

Sources/TestingMacros/Support/Additions/SyntaxProtocolAdditions.swift

@@ -9,6 +9,7 @@
 //
 
 import SwiftSyntax
+import SwiftSyntaxBuilder
 
 extension SyntaxProtocol {
   /// Get an expression representing the unique ID of this syntax node as well
@@ -26,21 +27,21 @@ extension SyntaxProtocol {
     // rewritten.
     var nodeIDChain = sequence(first: Syntax(self), next: \.parent)
       .map { $0.id.indexInTree.toOpaque() }
-
+    
 #if DEBUG
     assert(nodeIDChain.sorted() == nodeIDChain.reversed(), "Child node had lower ID than parent node in sequence \(nodeIDChain). Please file a bug report at https://github.com/swiftlang/swift-testing/issues/new")
     for id in nodeIDChain {
       assert(id <= UInt32.max, "Node ID \(id) was not a 32-bit integer. Please file a bug report at https://github.com/swiftlang/swift-testing/issues/new")
     }
 #endif
-
+    
+#if DEBUG
     // The highest ID in the chain determines the number of bits needed, and the
     // ID of this node will always be the highest (per the assertion above.)
-    let maxID = id.indexInTree.toOpaque()
-#if DEBUG
-    assert(nodeIDChain.contains(maxID), "ID \(maxID) of syntax node '\(self.trimmed)' was not found in its node ID chain \(nodeIDChain). Please file a bug report at https://github.com/swiftlang/swift-testing/issues/new")
+    let expectedMaxID = id.indexInTree.toOpaque()
+    assert(nodeIDChain.contains(expectedMaxID), "ID \(expectedMaxID) of syntax node '\(self.trimmed)' was not found in its node ID chain \(nodeIDChain). Please file a bug report at https://github.com/swiftlang/swift-testing/issues/new")
 #endif
-
+    
     // Adjust all node IDs downards by the effective root node's ID, then remove
     // the effective root node and its ancestors. This allows us to use lower
     // bit ranges than we would if we always included those nodes.
@@ -51,28 +52,38 @@ extension SyntaxProtocol {
       }
     }
 
-    // Convert the node IDs in the chain to bits in a bit mask.
-    let bitsPerWord = UInt64(UInt64.bitWidth)
-    var words = [UInt64](
-      repeating: 0,
-      count: Int(((maxID + 1) + (bitsPerWord - 1)) / bitsPerWord)
-    )
-    for id in nodeIDChain {
-      let (word, bit) = id.quotientAndRemainder(dividingBy: bitsPerWord)
-      words[Int(word)] |= (1 << bit)
-    }
-
-    // Convert the bits to a hexadecimal string.
-    let bitsPerNybble = 4
-    let nybblesPerWord = UInt64.bitWidth / bitsPerNybble
-    var id: String = words.map { word in
-      let result = String(word, radix: 16)
-      return String(repeating: "0", count: nybblesPerWord - result.count) + result
-    }.joined()
-
-    // Drop any redundant leading zeroes from the string literal.
-    id = String(id.drop { $0 == "0" })
+    let maxID = nodeIDChain.max() ?? 0
+    if maxID < UInt64.bitWidth {
+      // Pack all the node IDs into a single integer value.
+      var word = UInt64(0)
+      for id in nodeIDChain {
+        word |= (1 << id)
+      }
+      let hexWord = "0x\(String(word, radix: 16))"
+      return ExprSyntax(IntegerLiteralExprSyntax(literal: .integerLiteral(hexWord)))
 
-    return ExprSyntax(StringLiteralExprSyntax(content: id))
+    } else {
+      // Some ID exceeds what we can fit in a single literal, so just produce an
+      // array of node IDs instead.
+      let idExprs = nodeIDChain.map { id in
+        IntegerLiteralExprSyntax(literal: .integerLiteral("\(id)"))
+      }
+      return ExprSyntax(
+        FunctionCallExprSyntax(
+          calledExpression: TypeExprSyntax(
+            type: MemberTypeSyntax(
+              baseType: IdentifierTypeSyntax(name: .identifier("Testing")),
+              name: .identifier("__ExpressionID")
+            )
+          ),
+          leftParen: .leftParenToken(),
+          rightParen: .rightParenToken()
+        ) {
+          for idExpr in idExprs {
+            LabeledExprSyntax(expression: idExpr)
+          }
+        }
+      )
+    }
   }
 }