[PersistentCollections] Reduce popcount usages where possible

# Conflicts: # Sources/PersistentCollections/Node/_Node+Primitive Removals.swift
lorentey · Oct 12, 2022 · e153ebf · e153ebf
1 parent 785caa4
commit e153ebf
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 13 deletions.
diff --git a/Sources/PersistentCollections/Node/_Bitmap.swift b/Sources/PersistentCollections/Node/_Bitmap.swift
@@ -86,6 +86,11 @@ extension _Bitmap {
   @inlinable @inline(__always)
   internal var isEmpty: Bool { _value == 0 }
 
+  @inlinable @inline(__always)
+  internal var hasExactlyOneMember: Bool {
+    _value != 0 && _value & (_value &- 1) == 0
+  }
+
   @inlinable @inline(__always)
   internal var first: _Bucket? {
     guard !isEmpty else { return nil }

diff --git a/Sources/PersistentCollections/Node/_Node+Primitive Removals.swift b/Sources/PersistentCollections/Node/_Node+Primitive Removals.swift
@@ -127,7 +127,7 @@ extension _Node {
     assert(count == 1)
     count = 0
     return update {
-      assert($0.itemCount == 1 && $0.childCount == 0)
+      assert($0.hasSingletonItem)
       let old = $0._removeItem(at: .zero) { $0.move() }
       $0.clear()
       return old
@@ -138,7 +138,7 @@ extension _Node {
   internal mutating func removeSingletonChild() -> _Node {
     defer { _invariantCheck() }
     let child: _Node = update {
-      assert($0.itemCount == 0 && $0.childCount == 1)
+      assert($0.hasSingletonChild)
       let child = $0._removeChild(at: .zero)
       $0.childMap = .empty
       return child

diff --git a/Sources/PersistentCollections/Node/_Node+Subtree Removals.swift b/Sources/PersistentCollections/Node/_Node+Subtree Removals.swift
@@ -165,8 +165,9 @@ extension _Node {
       assert(slot == $0.itemMap.slot(of: bucket))
 
       let willAtrophy = (
-        $0.itemCount == 1
-        && $0.childCount == 1
+        !$0.isCollisionNode
+        && $0.itemMap.hasExactlyOneMember
+        && $0.childMap.hasExactlyOneMember
         && $0[child: .zero].isCollisionNode)
       if willAtrophy {
         // Compression
@@ -175,8 +176,8 @@ extension _Node {
         return (old, .collisionNode(level, child))
       }
 
-      let willEvaporate = ($0.itemCount == 2 && $0.childCount == 0)
-      if willEvaporate {
+      if $0.itemMap.count == 2 && $0.childMap.isEmpty {
+        // Evaporating node
         let remainder = _Slot(1 &- slot.value)
 
         var map = $0.itemMap
@@ -203,7 +204,7 @@ extension _Node {
     read {
       assert(!$0.isCollisionNode && $0.childMap.contains(bucket))
       let willAtrophy = (
-        $0.itemCount == 0
+        $0.itemMap.isEmpty
         && $0.childCount == 2
         && $0[child: _Slot(1 &- slot.value)].isCollisionNode
       )
@@ -212,12 +213,11 @@ extension _Node {
         let child = $0[child: _Slot(1 &- slot.value)]
         return .collisionNode(level, child)
       }
-      let willTurnIntoItem = ($0.itemCount == 1 && $0.childCount == 1)
-      if willTurnIntoItem {
+      if $0.itemMap.hasExactlyOneMember && $0.childMap.hasExactlyOneMember {
         return .item(level, $0[item: .zero], at: $0.itemMap.first!)
       }
-      let willEvaporate = ($0.itemCount == 0 && $0.childCount == 1)
-      if willEvaporate {
+      if $0.hasSingletonChild {
+        // Evaporate node
         return .empty(level)
       }
 

diff --git a/Sources/PersistentCollections/Node/_Node+UnsafeHandle.swift b/Sources/PersistentCollections/Node/_Node+UnsafeHandle.swift
@@ -280,12 +280,12 @@ extension _Node.UnsafeHandle {
 extension _Node.UnsafeHandle {
   @inlinable
   internal var hasSingletonItem: Bool {
-    itemCount == 1 && childCount == 0
+    _header.pointee.hasSingletonItem
   }
 
   @inlinable
   internal var hasSingletonChild: Bool {
-    itemMap.isEmpty && childCount == 1
+    _header.pointee.hasSingletonChild
   }
 
   @inlinable

diff --git a/Sources/PersistentCollections/Node/_StorageHeader.swift b/Sources/PersistentCollections/Node/_StorageHeader.swift
@@ -87,6 +87,19 @@ extension _StorageHeader {
      : itemMap.count)
   }
 
+  @inlinable
+  internal var hasSingletonChild: Bool {
+    itemMap.isEmpty && childMap.hasExactlyOneMember
+  }
+
+  @inlinable
+  internal var hasSingletonItem: Bool {
+    if itemMap == childMap {
+      return itemMap._value == 1
+    }
+    return childMap.isEmpty && itemMap.hasExactlyOneMember
+  }
+
   @inlinable @inline(__always)
   internal var childrenEndSlot: _Slot {
     _Slot(childCount)