ART: Adaptive Radix Tree implementation

Package.swift

@@ -278,6 +278,16 @@ let targets: [CustomTarget] = [
     name: "SortedCollectionsTests",
     dependencies: ["SortedCollections", "_CollectionsTestSupport"]),
 
+  .target(
+    kind: .exported,
+    name: "ARTreeModule",
+    dependencies: ["_CollectionsUtilities"],
+    exclude: ["CMakeLists.txt"]),
+  .target(
+    kind: .test,
+    name: "ARTreeModuleTests",
+    dependencies: ["ARTreeModule", "_CollectionsTestSupport"]),
+
   .target(
     kind: .exported,
     name: "Collections",

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

@@ -0,0 +1,82 @@
+//===----------------------------------------------------------------------===//
+//
+// 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: Sequence {
+  public typealias Iterator = _Iterator
+
+  public struct _Iterator {
+    typealias _ChildIndex = InternalNode<Spec>.Index
+
+    private let tree: ARTreeImpl<Spec>
+    private var path: [(any InternalNode<Spec>, _ChildIndex)]
+
+    init(tree: ARTreeImpl<Spec>) {
+      self.tree = tree
+      self.path = []
+      guard let node = tree._root else { return }
+
+      assert(node.type != .leaf, "root can't be leaf")
+      let n: any InternalNode<Spec> = node.toInternalNode()
+      if n.count > 0 {
+        self.path = [(n, n.startIndex)]
+      }
+    }
+  }
+
+  public func makeIterator() -> Iterator {
+    return Iterator(tree: self)
+  }
+}
+
+// TODO: Instead of index, use node iterators, to advance to next child.
+@available(macOS 13.3, iOS 16.4, watchOS 9.4, tvOS 16.4, *)
+extension ARTreeImpl._Iterator: IteratorProtocol {
+  public typealias Element = (Key, Spec.Value)  // TODO: Why just Value fails?
+
+  // Exhausted childs on the tip of path. Forward to sibling.
+  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)))
+  }
+
+  mutating func next() -> Element? {
+    while !path.isEmpty {
+      while let (node, index) = path.last {
+        if index == node.endIndex {
+          advanceToSibling()
+          break
+        }
+
+        let next = node.child(at: index)!
+        if next.type == .leaf {
+          let leaf: NodeLeaf<Spec> = next.toLeafNode()
+          let result = (leaf.key, leaf.value)
+          advanceToNextChild()
+          return result
+        }
+
+        let nextNode: any InternalNode<Spec> = next.toInternalNode()
+        path.append((nextNode, nextNode.startIndex))
+      }
+    }
+
+    return nil
+  }
+}

Sources/ARTreeModule/ARTree/ARTree+delete.swift

@@ -0,0 +1,64 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 mutating func delete(key: Key) {
+    if _root == nil {
+      return
+    }
+
+    let isUnique = _root!.isUnique
+    var child = _root
+    switch _delete(child: &child, key: key, depth: 0, isUniquePath: isUnique) {
+    case .noop:
+      return
+    case .replaceWith(let newValue):
+      _root = newValue
+    }
+  }
+
+  public mutating func deleteRange(start: Key, end: Key) {
+    // TODO
+    fatalError("not implemented")
+  }
+
+  private mutating func _delete(
+    child: inout RawNode?,
+    key: Key,
+    depth: Int,
+    isUniquePath: Bool
+  ) -> UpdateResult<RawNode?> {
+    if child?.type == .leaf {
+      let leaf: NodeLeaf<Spec> = child!.toLeafNode()
+      if !leaf.keyEquals(with: key, depth: depth) {
+        return .noop
+      }
+
+      return .replaceWith(nil)
+    }
+
+    assert(!Const.testCheckUnique || isUniquePath, "unique path is expected in this test")
+    var node: any InternalNode<Spec> = child!.toInternalNode()
+    var newDepth = depth
+
+    if node.partialLength > 0 {
+      let matchedBytes = node.prefixMismatch(withKey: key, fromIndex: depth)
+      assert(matchedBytes <= node.partialLength)
+      newDepth += matchedBytes
+    }
+
+    return node.updateChild(forKey: key[newDepth], isUniquePath: isUniquePath) {
+      var child = $0
+      return _delete(child: &child, key: key, depth: newDepth + 1, isUniquePath: $1)
+    }
+  }
+}

Sources/ARTreeModule/ARTree/ARTree+get.swift

@@ -0,0 +1,50 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 func getValue(key: Key) -> Value? {
+    var current = _root
+    var depth = 0
+    while depth <= key.count {
+      guard let _rawNode = current else {
+        return nil
+      }
+
+      if _rawNode.type == .leaf {
+        let leaf: NodeLeaf<Spec> = _rawNode.toLeafNode()
+        return leaf.keyEquals(with: key)
+          ? leaf.value
+          : nil
+      }
+
+      let node: any InternalNode<Spec> = _rawNode.toInternalNode()
+      if node.partialLength > 0 {
+        let prefixLen = node.prefixMismatch(withKey: key, fromIndex: depth)
+        assert(prefixLen <= Const.maxPartialLength, "partial length is always bounded")
+        if prefixLen != node.partialLength {
+          return nil
+        }
+        depth = depth + node.partialLength
+      }
+
+      current = node.child(forKey: key[depth])
+      depth += 1
+    }
+
+    return nil
+  }
+
+  public mutating func getRange(start: Key, end: Key) {
+    // TODO
+    fatalError("not implemented")
+  }
+}