[WIP] Use custom hash table for node lookup table

Author: rintaro

Sources/SwiftSyntax/CacheLookupTable.swift

@@ -0,0 +1,200 @@
+//===----------- CacheLookupTable.swift - Swift Syntax Library -----------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2018 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+// This file provides lookup table for cached object.
+//===----------------------------------------------------------------------===//
+
+/// Protocol for self-identifiable object
+protocol Identifiable {
+  associatedtype Identifier : Hashable
+  var id: Identifier { get }
+}
+
+private struct WeakReference<T: AnyObject>: ExpressibleByNilLiteral {
+  weak var value: T?
+  init(_ value: T?) {
+    self.value = value
+  }
+  init(nilLiteral: ()) {
+    self.value = nil
+  }
+}
+
+/// 'Set' like container providing lookup table for objects with identifier.
+/// This doesn't take an ownership of the objects. Instead the objects are
+/// hold by weak reference.
+/// 
+/// References are stored in a hash table with simple open adressing. Because
+/// of weak reference, unlike normal open addressing, erased bucket are simply
+/// turned into 'nil'.
+struct CacheLookupTable<T: Identifiable & AnyObject> {
+
+  private typealias Buffer = Array<WeakReference<T>>
+
+  /// Storage for the hash table.
+  private var buckets: Buffer
+
+  /// Estimated count of inserted values. This is greater than or equal to
+  /// the number of the acually occupied buckets.
+  /// i.e. estimatedCount >= _countOccupiedBuckets()
+  private var estimatedCount: Int
+
+  init(capacity: Int = 0) {
+    buckets = .init(repeating: nil,
+                    count: CacheLookupTable<T>._bucketCount(for: capacity))
+    estimatedCount = 0
+  }
+
+  /// Constant max load factor for hash table.
+  private static var maxLoadFactor: Double {
+    @inline(__always) get {
+      return 0.75
+    }
+  }
+
+  /// Minimal number of bucket count which enough to hold specified capacity
+  /// of values with max load factor taken into account.
+  private static func _minimalBucketCount(for capacity: Int) -> Int {
+    return Int((Double(capacity) / maxLoadFactor).rounded(.up))
+  }
+
+  /// Number of bucket count to allocate for holding specified number of
+  /// values. This is the next power of 2 greater than or equal to
+  /// '_minimalBucketCount(for: capacity)'
+  private static func _bucketCount(for capacity: Int,
+                                   from current: Int = 2) -> Int {
+    // Make sure it's representable.
+    precondition(capacity <= (Int.max >> 1) + 1)
+    // Bucket count must always be power of 2.
+    precondition((current & (current - 1)) == 0)
+    // Minimum is 2 to guarantee at least 1 hole.
+    precondition(current >= 2)
+
+    let minimalBucketCount = _minimalBucketCount(for: capacity)
+    var bucketCount = current
+    while bucketCount < minimalBucketCount {
+      bucketCount &*= 2
+    }
+    return bucketCount
+  }
+
+  private var _bucketMask: Int {
+    @inline(__always) get {
+      return buckets.count &- 1
+    }
+  }
+
+  @inline(__always)
+  private func _idealBucket(for id: T.Identifier) -> Int {
+    return id.hashValue & _bucketMask
+  }
+
+  /// Find the slot where the value with the specified id should be stored to.
+  private func _findHole(_ id: T.Identifier) -> (pos: Int, found: Bool) {
+    var bucket = _idealBucket(for: id)
+
+    // Starting from the ideal bucket for the id, search an available bucket,
+    // or the bucket holding the id.
+    while true {
+      guard let value = buckets[bucket].value else {
+        return (bucket, false)
+      }
+      if value.id == id {
+        return (bucket, true)
+      }
+      bucket = (bucket &+ 1) & _bucketMask
+    }
+  }
+
+  /// Reserves enough space to store the specified number of elements. Returns
+  /// true if resizing happened.
+  mutating func reserveCapacity(_ requiredCapacity: Int) -> Bool {
+    let bucketCount = CacheLookupTable<T>._bucketCount(for: requiredCapacity,
+                                                       from: buckets.count)
+    if (buckets.count >= bucketCount) {
+      return false
+    }
+
+    // Slow path. Resizing.
+    var oldBuckets = buckets
+    buckets = .init(repeating: nil, count: bucketCount)
+
+    // Move all nodes from the old buffer.
+    // TODO: move(), when available.
+    for i in 0..<oldBuckets.count {
+      if let oldValue = oldBuckets[i].value {
+        let pos = _findHole(oldValue.id).pos
+        Swift.swap(&buckets[pos], &oldBuckets[i])
+      }
+    }
+    return true
+  }
+
+  /// Count the actual number of occupied buckets.
+  @inline(__always)
+  private func _countOccupiedBuckets() -> Int {
+    var count = 0
+    for i in 0 ..< buckets.count where buckets[i].value != nil {
+      count &+= 1
+    }
+    return count
+  }
+
+  /// Reserves enough space to store a single new value. Returns true if
+  /// resizing happened.
+  mutating private func _ensurePlusOneCapacity() -> Bool {
+    if buckets.count >= CacheLookupTable<T>
+                          ._minimalBucketCount(for: estimatedCount &+ 1) {
+      return false
+    }
+
+    // Slow path. Re-count the actual number of occupied buckets, then reserve
+    // the capacity.
+    estimatedCount = _countOccupiedBuckets()
+    return reserveCapacity(estimatedCount &+ 1)
+  }
+
+  /// Inserts the given object into the table.
+  @discardableResult
+  mutating func insert(_ obj: T) -> Bool {
+    var (pos, found) = _findHole(obj.id)
+    if found {
+      return false
+    }
+
+    if /*resized=*/_ensurePlusOneCapacity() {
+      pos = _findHole(obj.id).pos
+    }
+    buckets[pos].value = obj
+    estimatedCount &+= 1
+    return true
+  }
+
+  /// Get a value with specified id. Returns 'nil' if the value hasn't been
+  /// insert()-ed or it's already been freed.
+  public subscript(id: T.Identifier) -> T? {
+    // Since we don't fill the bucket when the value is freed (because we don't
+    // know), we can't stop iteration a hole. So in the worst case (i.e. if the
+    // object is not contained in this table), full linear search is required.
+    // However, since we assume the value exists and hasn't been freed yet,
+    // we expect it's stored near the 'idealBucket' anyway.
+    let idealBucket = _idealBucket(for: id)
+    var bucket = idealBucket
+    repeat {
+      if buckets[bucket].value?.id == id {
+        return buckets[bucket].value
+      }
+      bucket = (bucket &+ 1) & _bucketMask
+    } while bucket != idealBucket
+
+    return nil
+  }
+}

Sources/SwiftSyntax/SwiftSyntax.swift

@@ -46,13 +46,8 @@ public enum SerializationFormat {
   case byteTree
 }
 
-fileprivate struct WeakReference<T: AnyObject> {
-  weak private(set) var value: T?
-
-  init(_ value: T) {
-    self.value = value
-  }
-}
+// For CacheLookupTable.
+extension RawSyntax : Identifiable {}
 
 /// Deserializes the syntax tree from its serialized form to an object tree in
 /// Swift. To deserialize incrementally transferred syntax trees, the same
@@ -64,7 +59,7 @@ public final class SyntaxTreeDeserializer {
 
   /// Syntax nodes that have already been parsed and are able to be reused if
   /// they were omitted in an incremental syntax tree transfer
-  private var nodeLookupTable: [SyntaxNodeId: WeakReference<RawSyntax>] = [:]
+  private var nodeLookupTable: CacheLookupTable<RawSyntax> = .init()
 
   /// Keep a strong reference to the syntax tree that contains the nodes in the
   /// `nodeLookupTable`. Because `nodeLookupTable` only holds a weak reference 
@@ -131,17 +126,11 @@ public final class SyntaxTreeDeserializer {
 
   private func lookupNode(id: SyntaxNodeId) -> RawSyntax? {
     reusedNodeIds.insert(id)
-    guard let weakRef = nodeLookupTable[id] else {
-      return nil
-    }
-    guard let value = weakRef.value else {
-      fatalError("Trying to retrieve a node that has since been deallocated")
-    }
-    return value
+    return nodeLookupTable[id]
   }
 
   private func addToLookupTable(_ node: RawSyntax) {
-    nodeLookupTable[node.id] = WeakReference(node)
+    nodeLookupTable.insert(node)
   }
 }