Persistent collections updates (part 5)

Sources/PersistentCollections/Node/_AncestorOffsets.swift

@@ -0,0 +1,104 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift Collections open source project
+//
+// Copyright (c) 2022 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
+//
+//===----------------------------------------------------------------------===//
+
+/// A collection of slot values logically addressing a particular node in a
+/// hash tree. The collection is (logically) extended with zero slots up to
+/// the maximum depth of the tree -- to unambiguously address a single node,
+/// this therefore needs to be augmented with a `_Level` value.
+///
+/// This construct can only be used to identify a particular node in the tree;
+/// it does not necessarily have room to include an item offset in the addressed
+/// node. (See `_Path` if you need to address a particular item.)
+@usableFromInline
+@frozen
+struct _AncestorSlots {
+  @usableFromInline
+  internal var path: UInt
+
+  @inlinable @inline(__always)
+  internal init(_ path: UInt) {
+    self.path = path
+  }
+}
+
+extension _AncestorSlots: Equatable {
+  @inlinable @inline(__always)
+  internal static func ==(left: Self, right: Self) -> Bool {
+    left.path == right.path
+  }
+}
+
+extension _AncestorSlots {
+  @inlinable @inline(__always)
+  internal static var empty: Self { Self(0) }
+}
+
+extension _AncestorSlots {
+  /// Return or set the slot value at the specified level.
+  /// If this is used to mutate the collection, then the original value
+  /// on the given level must be zero.
+  @inlinable @inline(__always)
+  internal subscript(_ level: _Level) -> _Slot {
+    get {
+      assert(level.shift < UInt.bitWidth)
+      return _Slot((path &>> level.shift) & _Bucket.bitMask)
+    }
+    set {
+      assert(newValue._value < UInt.bitWidth)
+      assert(self[level] == .zero)
+      path |= (UInt(truncatingIfNeeded: newValue._value) &<< level.shift)
+    }
+  }
+
+  @inlinable @inline(__always)
+  internal func appending(_ slot: _Slot, at level: _Level) -> Self {
+    var result = self
+    result[level] = slot
+    return result
+  }
+
+  /// Clear the slot at the specified level, by setting it to zero.
+  @inlinable
+  internal mutating func clear(_ level: _Level) {
+    guard level.shift < UInt.bitWidth else { return }
+    path &= ~(_Bucket.bitMask &<< level.shift)
+  }
+
+  /// Truncate this path to the specified level.
+  /// Slots at or beyond the specified level are cleared.
+  @inlinable
+  internal func truncating(to level: _Level) -> _AncestorSlots {
+    assert(level.shift <= UInt.bitWidth)
+    guard level.shift < UInt.bitWidth else { return self }
+    return _AncestorSlots(path & ((1 &<< level.shift) &- 1))
+  }
+
+  /// Returns true if this path contains non-zero slots at or beyond the
+  /// specified level, otherwise returns false.
+  @inlinable
+  internal func hasDataBelow(_ level: _Level) -> Bool {
+    guard level.shift < UInt.bitWidth else { return false }
+    return (path &>> level.shift) != 0
+  }
+
+  /// Compares this path to `other` up to but not including the specified level.
+  /// Returns true if the path prefixes compare equal, otherwise returns false.
+  @inlinable
+  internal func isEqual(to other: Self, upTo level: _Level) -> Bool {
+    if level.isAtRoot { return true }
+    if level.isAtBottom { return self == other }
+    let s = UInt(UInt.bitWidth) - level.shift
+    let v1 = self.path &<< s
+    let v2 = other.path &<< s
+    return v1 == v2
+  }
+}
+

Sources/PersistentCollections/Node/_Bitmap.swift

@@ -105,13 +105,13 @@ extension _Bitmap {
   }
 
   @inlinable @inline(__always)
-  internal func offset(of bucket: _Bucket) -> Int {
-    _value._rank(ofBit: bucket.value)
+  internal func slot(of bucket: _Bucket) -> _Slot {
+    _Slot(_value._rank(ofBit: bucket.value))
   }
 
   @inlinable @inline(__always)
-  internal func bucket(at offset: Int) -> _Bucket {
-    _Bucket(_value._bit(ranked: offset)!)
+  internal func bucket(at slot: _Slot) -> _Bucket {
+    _Bucket(_value._bit(ranked: slot.value)!)
   }
 }
 

Sources/PersistentCollections/Node/_Bucket.swift

@@ -15,12 +15,23 @@
 @frozen
 internal struct _Bucket {
   @usableFromInline
-  internal var value: UInt
+  internal var _value: UInt8
+
+  @inlinable @inline(__always)
+  internal init(_value: UInt8) {
+    assert(_value < _Bitmap.capacity || _value == .max)
+    self._value = _value
+  }
+}
+
+extension _Bucket {
+  @inlinable @inline(__always)
+  internal var value: UInt { UInt(truncatingIfNeeded: _value) }
 
   @inlinable @inline(__always)
   internal init(_ value: UInt) {
     assert(value < _Bitmap.capacity || value == .max)
-    self.value = value
+    self._value = UInt8(truncatingIfNeeded: value)
   }
 }
 
@@ -32,29 +43,29 @@ extension _Bucket {
   static var bitMask: UInt { UInt(bitPattern: _Bitmap.capacity) &- 1 }
 
   @inlinable @inline(__always)
-  static var invalid: _Bucket { _Bucket(.max) }
+  static var invalid: _Bucket { _Bucket(_value: .max) }
 
   @inlinable @inline(__always)
-  var isInvalid: Bool { value == .max }
+  var isInvalid: Bool { _value == .max }
 }
 
 extension _Bucket: Equatable {
   @inlinable @inline(__always)
   internal static func ==(left: Self, right: Self) -> Bool {
-    left.value == right.value
+    left._value == right._value
   }
 }
 
 extension _Bucket: Comparable {
   @inlinable @inline(__always)
   internal static func <(left: Self, right: Self) -> Bool {
-    left.value < right.value
+    left._value < right._value
   }
 }
 
 extension _Bucket: CustomStringConvertible {
   @usableFromInline
   internal var description: String {
-    String(value, radix: _Bitmap.capacity)
+    String(_value, radix: _Bitmap.capacity)
   }
 }

Sources/PersistentCollections/Node/_Level.swift

@@ -9,15 +9,28 @@
 //
 //===----------------------------------------------------------------------===//
 
+/// Identifies a particular level within the hash tree.
+///
+/// Hash trees have a maximum depth of ⎡`UInt.bitWidth / _Bucket.bitWidth`⎤, so
+/// the level always fits in an `UInt8` value.
 @usableFromInline
 @frozen
 internal struct _Level {
+  /// The bit position within a hash value that begins the hash slice that is
+  /// associated with this level. For collision nodes, this can be larger than
+  /// `UInt.bitWidth`.
   @usableFromInline
-  internal var shift: UInt
+  internal var _shift: UInt8
+
+  @inlinable @inline(__always)
+  init(_shift: UInt8) {
+    self._shift = _shift
+  }
 
   @inlinable @inline(__always)
   init(shift: UInt) {
-    self.shift = shift
+    assert(shift <= UInt8.max)
+    self._shift = UInt8(truncatingIfNeeded: shift)
   }
 }
 
@@ -28,37 +41,50 @@ extension _Level {
   }
 
   @inlinable @inline(__always)
-  internal static var _step: UInt {
-    UInt(bitPattern: _Bitmap.bitWidth)
+  internal static var _step: UInt8 {
+    UInt8(truncatingIfNeeded: _Bitmap.bitWidth)
   }
 
   @inlinable @inline(__always)
   internal static var top: _Level {
     _Level(shift: 0)
   }
 
+  /// The bit position within a hash value that begins the hash slice that is
+  /// associated with this level. For collision nodes, this can be larger than
+  /// `UInt.bitWidth`.
+  @inlinable @inline(__always)
+  internal var shift: UInt { UInt(truncatingIfNeeded: _shift) }
+
   @inlinable @inline(__always)
-  internal var isAtRoot: Bool { shift == 0 }
+  internal var isAtRoot: Bool { _shift == 0 }
 
   @inlinable @inline(__always)
-  internal var isAtBottom: Bool { shift >= UInt.bitWidth }
+  internal var isAtBottom: Bool { _shift >= UInt.bitWidth }
 
   @inlinable @inline(__always)
   internal func descend() -> _Level {
     // FIXME: Consider returning nil when we run out of bits
-    _Level(shift: shift &+ Self._step)
+    _Level(_shift: _shift &+ Self._step)
   }
 
   @inlinable @inline(__always)
   internal func ascend() -> _Level {
     assert(!isAtRoot)
-    return _Level(shift: shift &+ Self._step)
+    return _Level(_shift: _shift &- Self._step)
   }
 }
 
 extension _Level: Equatable {
   @inlinable @inline(__always)
   internal static func ==(left: Self, right: Self) -> Bool {
-    left.shift == right.shift
+    left._shift == right._shift
+  }
+}
+
+extension _Level: Comparable {
+  @inlinable @inline(__always)
+  internal static func <(left: Self, right: Self) -> Bool {
+    left._shift < right._shift
   }
 }

Sources/PersistentCollections/Node/_Node+CustomStringConvertible.swift

@@ -19,7 +19,7 @@ extension _Node: CustomStringConvertible {
     var result = "["
     var first = true
     read {
-      for (key, value) in $0._items {
+      for (key, value) in $0.reverseItems.reversed() {
         if first {
           first = false
         } else {