[BitCollections] Move shared utilities to _CollectionsUtilities

Sources/BitCollections/BitArray/BitArray+Copy.swift

@@ -9,6 +9,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+import _CollectionsUtilities
+
 extension BitArray {
   internal mutating func _copy(
     from range: Range<Int>,

Sources/BitCollections/BitArray/BitArray+CustomStringConvertible.swift

@@ -23,7 +23,7 @@ extension BitArray: CustomStringConvertible {
         let b1: UInt8 = 49 // ASCII 1
         var i = 0
         for v in self {
-          target._initialize(at: i, to: v ? b1 : b0)
+          target.initializeElement(at: i, to: v ? b1 : b0)
           i &+= 1
         }
         return i

Sources/BitCollections/BitSet/BitSet+Initializers.swift

@@ -172,16 +172,16 @@ extension BitSet {
     _storage = Array(unsafeUninitializedCapacity: capacity) { buffer, count in
       let sharedCount = Swift.min(w1.count, w2.count)
       for w in 0 ..< sharedCount {
-        buffer._initialize(at: w, to: function(w1[w], w2[w]))
+        buffer.initializeElement(at: w, to: function(w1[w], w2[w]))
       }
       if includingTail {
         if w1.count < w2.count {
           for w in w1.count ..< w2.count {
-            buffer._initialize(at: w, to: function(_Word.empty, w2[w]))
+            buffer.initializeElement(at: w, to: function(_Word.empty, w2[w]))
           }
         } else {
           for w in w2.count ..< w1.count {
-            buffer._initialize(at: w, to: function(w1[w], _Word.empty))
+            buffer.initializeElement(at: w, to: function(w1[w], _Word.empty))
           }
         }
       }

Sources/BitCollections/BitSet/BitSet._UnsafeHandle.swift

@@ -9,6 +9,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+import _CollectionsUtilities
+
 extension BitSet {
   /// An unsafe-unowned bitset view over `UInt` storage, providing bit set
   /// primitives.
@@ -273,7 +275,7 @@ extension BitSet._UnsafeHandle: BidirectionalCollection {
   internal var startIndex: Index {
     let word = _words.firstIndex { !$0.isEmpty }
     guard let word = word else { return endIndex }
-    return Index(word: word, bit: _words[word].firstMember)
+    return Index(word: word, bit: _words[word].firstMember!)
   }
 
   @inlinable
@@ -299,7 +301,7 @@ extension BitSet._UnsafeHandle: BidirectionalCollection {
       }
       w = _words[word]
     }
-    return Index(word: word, bit: w.firstMember)
+    return Index(word: word, bit: w.firstMember!)
   }
 
   @usableFromInline
@@ -318,7 +320,7 @@ extension BitSet._UnsafeHandle: BidirectionalCollection {
       precondition(word >= 0, "Can't advance below startIndex")
       w = _words[word]
     }
-    return Index(word: word, bit: w.lastMember)
+    return Index(word: word, bit: w.lastMember!)
   }
 
   @usableFromInline
@@ -501,9 +503,7 @@ extension BitSet._UnsafeHandle {
     _mutableWords[l.word].formIntersection(_Word(upTo: l.bit).complement())
     if u.word < wordCount {
       _mutableWords[u.word].formIntersection(_Word(upTo: u.bit))
-      _mutableWords[(u.word + 1)...]
-        ._rebased()
-        .assign(repeating: .empty)
+      _mutableWords[(u.word + 1)...].update(repeating: .empty)
     }
   }
 
@@ -542,9 +542,7 @@ extension BitSet._UnsafeHandle {
     }
 
     _mutableWords[l.word].subtract(_Word(upTo: l.bit).complement())
-    _mutableWords[(l.word + 1) ..< u.word]
-      ._rebased()
-      .assign(repeating: .empty)
+    _mutableWords[(l.word + 1) ..< u.word].update(repeating: .empty)
     if u.word < wordCount {
       _mutableWords[u.word].subtract(_Word(upTo: u.bit))
     }

Sources/BitCollections/Shared/UInt+Tricks.swift

@@ -9,36 +9,9 @@
 //
 //===----------------------------------------------------------------------===//
 
-extension UInt {
-  @inlinable @inline(__always)
-  internal var _nonzeroBitCount: UInt {
-    Self(truncatingIfNeeded: nonzeroBitCount)
-  }
-
-  internal var _reversed: UInt {
-    // https://graphics.stanford.edu/~seander/bithacks.html#ReverseParallel
-    var shift: UInt = UInt(UInt.bitWidth)
-    var mask: UInt = ~0;
-    var result = self
-    while true {
-      shift &>>= 1
-      guard shift > 0 else { break }
-      mask ^= mask &<< shift
-      result = ((result &>> shift) & mask) | ((result &<< shift) & ~mask)
-    }
-    return result
-  }
-
-  @inlinable @inline(__always)
-  internal var _lastSetBit: UInt {
-    UInt(truncatingIfNeeded: self.trailingZeroBitCount)
-  }
-
-  @inlinable @inline(__always)
-  internal var _firstSetBit: UInt {
-    UInt(truncatingIfNeeded: UInt.bitWidth &- 1 &- self.leadingZeroBitCount)
-  }
+import _CollectionsUtilities
 
+extension UInt {
   /// Returns the position of the `n`th set bit in `self`.
   ///
   /// - Parameter n: The  to retrieve. This value is
@@ -48,52 +21,14 @@ extension UInt {
   /// - Returns: If this integer contains enough set bits to satisfy the
   ///    request, then this function returns the position of the bit found.
   ///    Otherwise it returns nil.
-  internal func _nthSetBit(_ n: inout UInt) -> UInt? {
-    // FIXME: Use bit deposit instruction when available (PDEP on Intel).
-    assert(UInt.bitWidth == 64 || UInt.bitWidth == 32, "Unsupported UInt bitWidth")
-
-    var shift: UInt = 0
-
-    let c = self._nonzeroBitCount
+  internal func _rank(ofBit n: inout UInt) -> UInt? {
+    let c = self.nonzeroBitCount
     guard n < c else {
-      n &-= c
+      n &-= UInt(bitPattern: c)
       return nil
     }
-
-    if UInt.bitWidth == 64 {
-      let c32 = (self & 0xFFFFFFFF)._nonzeroBitCount
-      if n >= c32 {
-        shift &+= 32
-        n &-= c32
-      }
-    }
-    let c16 = ((self &>> shift) & 0xFFFF)._nonzeroBitCount
-    if n >= c16 {
-      shift &+= 16
-      n &-= c16
-    }
-    let c8 = ((self &>> shift) & 0xFF)._nonzeroBitCount
-    if n >= c8 {
-      shift &+= 8
-      n &-= c8
-    }
-    let c4 = ((self &>> shift) & 0xF)._nonzeroBitCount
-    if n >= c4 {
-      shift &+= 4
-      n &-= c4
-    }
-    let c2 = ((self &>> shift) & 0x3)._nonzeroBitCount
-    if n >= c2 {
-      shift &+= 2
-      n &-= c2
-    }
-    let c1 = (self &>> shift) & 0x1
-    if n >= c1 {
-      shift &+= 1
-      n &-= c1
-    }
-    precondition(n == 0 && (self &>> shift) & 0x1 == 1)
-    return shift
+    let m = Int(bitPattern: n)
+    n = 0
+    return _bit(ranked: m)!
   }
 }
-

Sources/BitCollections/Shared/_Word.swift

@@ -9,6 +9,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+import _CollectionsUtilities
+
 @usableFromInline
 @frozen
 internal struct _Word {
@@ -52,12 +54,12 @@ extension _Word {
   ///    then this function returns the member found. Otherwise it returns nil.
   @inline(never)
   internal func nthElement(_ n: inout UInt) -> UInt? {
-    value._nthSetBit(&n)
+    value._rank(ofBit: &n)
   }
 
   @inline(never)
   internal func nthElementFromEnd(_ n: inout UInt) -> UInt? {
-    guard let i = value._reversed._nthSetBit(&n) else { return nil }
+    guard let i = value._reversed._rank(ofBit: &n) else { return nil }
     return UInt(UInt.bitWidth) &- 1 &- i
   }
 }
@@ -102,13 +104,13 @@ extension _Word {
 
   @inlinable
   @inline(__always)
-  internal var firstMember: UInt {
+  internal var firstMember: UInt? {
     value._lastSetBit
   }
 
   @inlinable
   @inline(__always)
-  internal var lastMember: UInt {
+  internal var lastMember: UInt? {
     value._firstSetBit
   }
 

Sources/_CollectionsUtilities/Integer tricks/FixedWidthInteger+roundUpToPowerOfTwo.swift

@@ -0,0 +1,21 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+extension FixedWidthInteger {
+  /// Round up `self` to the nearest power of two, assuming it's representable.
+  /// Returns 0 if `self` isn't positive.
+  @inlinable
+  public func _roundUpToPowerOfTwo() -> Self {
+    guard self > 0 else { return 0 }
+    let l = Self.bitWidth - (self &- 1).leadingZeroBitCount
+    return 1 << l
+  }
+}