Implement a zero-allocation collection abstraction for expression IDs

Author: grynspan

Sources/Testing/SourceAttribution/ExpressionID.swift

@@ -22,7 +22,7 @@
 public struct __ExpressionID: Sendable {
   /// The ID of the root node in an expression graph.
   static var root: Self {
-    Self(_elements: .none)
+    Self(elements: .none)
   }
 
   /// An enumeration that attempts to efficiently store the key path elements
@@ -44,74 +44,102 @@ public struct __ExpressionID: Sendable {
   }
 
   /// The elements of this identifier.
-  private var _elements: Elements
+  fileprivate 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`.
+// MARK: - Equatable, Hashable
+
+extension __ExpressionID: Equatable, Hashable {}
+extension __ExpressionID.Elements: Equatable, Hashable {}
+
+// MARK: - Collection
+
+extension __ExpressionID {
+  /// A type representing the elements in a key path produced from the unique
+  /// identifier of an expression.
   ///
-  /// The values in this collection, being swift-syntax node IDs, are never more
-  /// than 32 bits wide.
-  var keyPath: some RandomAccessCollection<UInt32> {
-    // 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.
-        }
+  /// Instances of this type can be used to produce keys and key paths for an
+  /// instance of `Graph` whose key type is `UInt32`.
+  fileprivate struct KeyPath: Collection {
+    /// Underlying storage for the collection.
+    var elements: __ExpressionID.Elements
+
+    var underestimatedCount: Int {
+      count
+    }
+
+    var count: Int {
+      switch elements {
+      case .none:
+        0
+      case let .packed(word):
+        word.nonzeroBitCount
+      case let .keyPath(keyPath):
+        keyPath.count
       }
+    }
 
-      return result
+    var startIndex: Int {
+      switch elements {
+      case .none, .keyPath:
+        0
+      case let .packed(word):
+        word.trailingZeroBitCount
+      }
     }
 
-    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.
+    var endIndex: Int {
+      switch elements {
+      case .none:
+        0
+      case .packed:
+        UInt64.bitWidth
+      case let .keyPath(keyPath):
+        keyPath.count
       }
+    }
 
-      return result
-    case let .keyPath(keyPath):
-      return keyPath
+    func index(after i: Int) -> Int {
+      switch elements {
+      case .none, .keyPath:
+        return i + 1
+      case let .packed(word):
+        // Mask off the low bits including the one at `i`. The trailing zero
+        // count of the resulting value equals the next actual bit index.
+        let nextIndex = i + 1
+        let maskedWord = word & (~0 << nextIndex)
+        return maskedWord.trailingZeroBitCount
+      }
     }
-  }
-}
 
-// MARK: - Equatable, Hashable
+    subscript(position: Int) -> UInt32 {
+      switch elements {
+      case .none:
+        fatalError("Unreachable")
+      case .packed:
+        UInt32(position)
+      case let .keyPath(keyPath):
+        keyPath[position]
+      }
+    }
+  }
 
-extension __ExpressionID: Equatable, Hashable {}
-extension __ExpressionID.Elements: Equatable, Hashable {}
+  /// 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 Collection<UInt32> {
+    KeyPath(elements: elements)
+  }
+}
 
 #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 {
-    switch _elements {
+    switch elements {
     case .none:
       return "0"
     case let .packed(word):
@@ -134,14 +162,14 @@ extension __ExpressionID: CustomStringConvertible, CustomDebugStringConvertible
 extension __ExpressionID: ExpressibleByIntegerLiteral {
   public init(integerLiteral: UInt64) {
     if integerLiteral == 0 {
-      self.init(_elements: .none)
+      self.init(elements: .none)
     } else {
-      self.init(_elements: .packed(integerLiteral))
+      self.init(elements: .packed(integerLiteral))
     }
   }
 
   public init(_ keyPath: UInt32...) {
-    self.init(_elements: .keyPath(keyPath))
+    self.init(elements: .keyPath(keyPath))
   }
 }