WIP on Collection protocol

Sources/ARTreeModule/ARTree/ARTree+Collection.swift

@@ -0,0 +1,23 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift Collections open source project
+//
+// Copyright (c) 2023 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
+//
+//===----------------------------------------------------------------------===//
+
+@available(macOS 13.3, iOS 16.4, watchOS 9.4, tvOS 16.4, *)
+extension ARTreeImpl {
+  var startIndex: Index  {
+    var idx = Index(forTree: self)
+    idx.descentToLeftMostChild()
+    return idx
+  }
+
+  var endIndex: Index  {
+    return Index(forTree: self)
+  }
+}

Sources/ARTreeModule/ARTree/ARTree+Sequence.swift

@@ -17,7 +17,7 @@ extension ARTreeImpl: Sequence {
     typealias _ChildIndex = InternalNode<Spec>.Index
 
     private let tree: ARTreeImpl<Spec>
-    private var path: [(any InternalNode<Spec>, _ChildIndex?)]
+    private var path: [(any InternalNode<Spec>, _ChildIndex)]
 
     init(tree: ARTreeImpl<Spec>) {
       self.tree = tree
@@ -27,7 +27,7 @@ extension ARTreeImpl: Sequence {
       assert(node.type != .leaf, "root can't be leaf")
       let n: any InternalNode<Spec> = node.toInternalNode()
       if n.count > 0 {
-        self.path = [(n, n.index())]
+        self.path = [(n, n.startIndex)]
       }
     }
   }
@@ -53,13 +53,13 @@ extension ARTreeImpl._Iterator: IteratorProtocol {
       return
     }
 
-    path.append((node, node.next(index: index!)))
+    path.append((node, node.index(after: index)))
   }
 
   mutating func next() -> Element? {
     while !path.isEmpty {
-      while let (node, _index) = path.last {
-        guard let index = _index else {
+      while let (node, index) = path.last {
+        if index == node.endIndex {
           advanceToSibling()
           break
         }
@@ -73,7 +73,7 @@ extension ARTreeImpl._Iterator: IteratorProtocol {
         }
 
         let nextNode: any InternalNode<Spec> = next.toInternalNode()
-        path.append((nextNode, nextNode.index()))
+        path.append((nextNode, nextNode.startIndex))
       }
     }
 

Sources/ARTreeModule/ARTree/ARTree.Index.swift

@@ -0,0 +1,100 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift Collections open source project
+//
+// Copyright (c) 2023 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
+//
+//===----------------------------------------------------------------------===//
+
+@available(macOS 13.3, iOS 16.4, watchOS 9.4, tvOS 16.4, *)
+extension ARTreeImpl {
+  public struct Index {
+    internal typealias _ChildIndex = InternalNode<Spec>.Index
+
+    internal weak var root: RawNodeBuffer? = nil
+    internal var current: (any ArtNode<Spec>)? = nil
+    internal var path: [(any InternalNode<Spec>, _ChildIndex)] = []
+    internal let version: Int
+
+    internal init(forTree tree: ARTreeImpl<Spec>) {
+      self.version = tree.version
+
+      if let root = tree._root {
+        assert(root.type != .leaf, "root can't be leaf")
+        self.root = tree._root?.buf
+        self.current = tree._root?.toArtNode()
+      }
+    }
+  }
+}
+
+@available(macOS 13.3, iOS 16.4, watchOS 9.4, tvOS 16.4, *)
+extension ARTreeImpl.Index {
+  internal var isOnLeaf: Bool {
+    if let current = self.current {
+      return current.type == .leaf
+    }
+
+    return false
+  }
+
+  internal mutating func descentToLeftMostChild() {
+    while !isOnLeaf {
+      descend { $0.startIndex }
+    }
+  }
+
+  internal mutating func descend(_ to: (any InternalNode<Spec>)
+                                   -> (any InternalNode<Spec>).Index) {
+    assert(!isOnLeaf, "can't descent on a leaf node")
+    assert(current != nil, "current node can't be nil")
+
+    let currentNode: any InternalNode<Spec> = current!.rawNode.toInternalNode()
+    let index = to(currentNode)
+    self.path.append((currentNode, index))
+    self.current = currentNode.child(at: index)?.toArtNode()
+  }
+
+  mutating private func advanceToSibling() {
+    let _ = path.popLast()
+    advanceToNextChild()
+  }
+
+  mutating private func advanceToNextChild() {
+    guard let (node, index) = path.popLast() else {
+      return
+    }
+
+    path.append((node, node.index(after: index)))
+  }
+}
+
+
+@available(macOS 13.3, iOS 16.4, watchOS 9.4, tvOS 16.4, *)
+extension ARTreeImpl.Index: Equatable {
+  static func == (lhs: Self, rhs: Self) -> Bool {
+    if case (let lhs?, let rhs?) = (lhs.current, rhs.current) {
+      return lhs.equals(rhs)
+    }
+
+    return false
+  }
+}
+
+@available(macOS 13.3, iOS 16.4, watchOS 9.4, tvOS 16.4, *)
+extension ARTreeImpl.Index: Comparable {
+  static func < (lhs: Self, rhs: Self) -> Bool {
+    for ((_, idxL), (_, idxR)) in zip(lhs.path, rhs.path) {
+      if idxL < idxR {
+        return true
+      } else if idxL > idxR {
+        return false
+      }
+    }
+
+    return false
+  }
+}

Sources/ARTreeModule/ARTree/ARTree.swift

@@ -28,9 +28,11 @@ internal struct ARTreeImpl<Spec: ARTreeSpec> {
 
   @usableFromInline
   internal var _root: RawNode?
+  internal var version: Int
 
   @inlinable
   public init() {
     self._root = nil
+    self.version = -1
   }
 }

Sources/ARTreeModule/ARTree/ArtNode.swift

@@ -40,3 +40,9 @@ extension ArtNode {
   var rawNode: RawNode { RawNode(buf: self.storage.ref.takeUnretainedValue()) }
   var type: NodeType { Self.type }
 }
+
+extension ArtNode {
+  func equals(_ other: any ArtNode<Spec>) -> Bool {
+    return self.rawNode == other.rawNode
+  }
+}

Sources/ARTreeModule/ARTree/InternalNode.swift

@@ -17,23 +17,25 @@ protocol InternalNode<Spec>: ArtNode {
 
   static var size: Int { get }
 
-  var count: Int { get set }
   var partialLength: Int { get }
   var partialBytes: PartialBytes { get set }
 
-  func index(forKey k: KeyPart) -> Index?
-  func index() -> Index?
-  func next(index: Index) -> Index?
+  var count: Int { get set }
+  var startIndex: Index { get }
+  var endIndex: Index { get }
+
+  func index(forKey: KeyPart) -> Index?
+  func index(after: Index) -> Index
 
-  func child(forKey k: KeyPart) -> RawNode?  // TODO: Remove
+  func child(forKey: KeyPart) -> RawNode?  // TODO: Remove
   func child(at: Index) -> RawNode?  // TODO: Remove
 
-  mutating func addChild(forKey k: KeyPart, node: RawNode) -> UpdateResult<RawNode?>
-  mutating func addChild(forKey k: KeyPart, node: some ArtNode<Spec>) -> UpdateResult<RawNode?>
+  mutating func addChild(forKey: KeyPart, node: RawNode) -> UpdateResult<RawNode?>
+  mutating func addChild(forKey: KeyPart, node: some ArtNode<Spec>) -> UpdateResult<RawNode?>
 
-  mutating func removeChild(at index: Index) -> UpdateResult<RawNode?>
+  mutating func removeChild(at: Index) -> UpdateResult<RawNode?>
 
-  mutating func withChildRef<R>(at index: Index, _ body: (RawNode.SlotRef) -> R) -> R
+  mutating func withChildRef<R>(at: Index, _ body: (RawNode.SlotRef) -> R) -> R
 }
 
 struct NodeReference {

Sources/ARTreeModule/ARTree/Node16.swift

@@ -97,6 +97,9 @@ extension Node16: InternalNode {
       * (MemoryLayout<KeyPart>.stride + MemoryLayout<RawNode?>.stride)
   }
 
+  var startIndex: Index { 0 }
+  var endIndex: Index { count }
+
   func index(forKey k: KeyPart) -> Index? {
     for (index, key) in keys[..<count].enumerated() {
       if key == k {
@@ -107,13 +110,13 @@ extension Node16: InternalNode {
     return nil
   }
 
-  func index() -> Index? {
-    return 0
-  }
-
-  func next(index: Index) -> Index? {
+  func index(after index: Index) -> Index {
     let next = index + 1
-    return next < count ? next : nil
+    if next >= count {
+      return count
+    } else {
+      return next
+    }
   }
 
   func _insertSlot(forKey k: KeyPart) -> Int? {
@@ -131,9 +134,10 @@ extension Node16: InternalNode {
     return count
   }
 
-  func child(at: Index) -> RawNode? {
-    assert(at < Self.numKeys, "maximum \(Self.numKeys) childs allowed")
-    return childs[at]
+  func child(at index: Index) -> RawNode? {
+    assert(index < Self.numKeys, "maximum \(Self.numKeys) childs allowed, given index = \(index)")
+    assert(index < count, "not enough childs in node")
+    return childs[index]
   }
 
   mutating func addChild(forKey k: KeyPart, node: RawNode) -> UpdateResult<RawNode?> {

Sources/ARTreeModule/ARTree/Node256.swift

@@ -65,27 +65,33 @@ extension Node256: InternalNode {
     MemoryLayout<InternalNodeHeader>.stride + 256 * MemoryLayout<RawNode?>.stride
   }
 
-  func index(forKey k: KeyPart) -> Index? {
-    return childs[Int(k)] != nil ? Int(k) : nil
+  var startIndex: Index {
+    if count == 0 {
+      return endIndex
+    } else {
+      return index(after: -1)
+    }
   }
 
-  func index() -> Index? {
-    return next(index: -1)
+  var endIndex: Index { 256 }
+
+  func index(forKey k: KeyPart) -> Index? {
+    return childs[Int(k)] != nil ? Int(k) : nil
   }
 
-  func next(index: Index) -> Index? {
-    for idx in index + 1..<256 {
+  func index(after idx: Index) -> Index {
+    for idx in idx + 1..<256 {
       if childs[idx] != nil {
         return idx
       }
     }
 
-    return nil
+    return 256
   }
 
-  func child(at: Index) -> RawNode? {
-    assert(at < 256, "maximum 256 childs allowed")
-    return childs[at]
+  func child(at index: Index) -> RawNode? {
+    assert(index < 256, "maximum 256 childs allowed")
+    return childs[index]
   }
 
   mutating func addChild(forKey k: KeyPart, node: RawNode) -> UpdateResult<RawNode?> {
@@ -96,6 +102,7 @@ extension Node256: InternalNode {
   }
 
   mutating func removeChild(at index: Index) -> UpdateResult<RawNode?> {
+    assert(index < 256, "invalid index")
     childs[index] = nil
     count -= 1
 
@@ -108,6 +115,7 @@ extension Node256: InternalNode {
   }
 
   mutating func withChildRef<R>(at index: Index, _ body: (RawNode.SlotRef) -> R) -> R {
+    assert(index < 256, "invalid index")
     let ref = childs.baseAddress! + index
     return body(ref)
   }

Sources/ARTreeModule/ARTree/Node4.swift

@@ -76,6 +76,9 @@ extension Node4: InternalNode {
       * (MemoryLayout<KeyPart>.stride + MemoryLayout<RawNode?>.stride)
   }
 
+  var startIndex: Index { 0 }
+  var endIndex: Index { count }
+
   func index(forKey k: KeyPart) -> Index? {
     for (index, key) in keys[..<count].enumerated() {
       if key == k {
@@ -86,13 +89,13 @@ extension Node4: InternalNode {
     return nil
   }
 
-  func index() -> Index? {
-    return 0
-  }
-
-  func next(index: Index) -> Index? {
+  func index(after index: Index) -> Index {
     let next = index + 1
-    return next < count ? next : nil
+    if next >= count {
+      return count
+    } else {
+      return next
+    }
   }
 
   func _insertSlot(forKey k: KeyPart) -> Int? {
@@ -109,9 +112,10 @@ extension Node4: InternalNode {
     return count
   }
 
-  func child(at: Index) -> RawNode? {
-    assert(at < Self.numKeys, "maximum \(Self.numKeys) childs allowed, given index = \(at)")
-    return childs[at]
+  func child(at index: Index) -> RawNode? {
+    assert(index < Self.numKeys, "maximum \(Self.numKeys) childs allowed, given index = \(index)")
+    assert(index < count, "not enough childs in node")
+    return childs[index]
   }
 
   mutating func addChild(forKey k: KeyPart, node: RawNode) -> UpdateResult<RawNode?> {