ART: Adaptive Radix Tree implementation #307

vishesh · 2023-08-03T00:25:39Z

This PR implements Adaptive Radix Tree (ART)[1], a compact key-value search tree. It maps byte-sequences to arbitrary values while maintaining data in sorted order. ART can compress multiple nodes into single node and adaptively expand or shrink its internal nodes to avoid worse-case space consumption.

Within the broader swift-collection project, it can be considered a compact variant of SortedSet/SortedDictionary but at the cost of keys being convertible to byte-sequences.

[1] https://ieeexplore.ieee.org/abstract/document/6544812

Checklist

I've read the Contribution Guidelines
My contributions are licensed under the Swift license.
I've followed the coding style of the rest of the project.
I've added tests covering all new code paths my change adds to the project (if appropriate).
I've added benchmarks covering new functionality (if appropriate).
I've verified that my change does not break any existing tests or introduce unexplained benchmark regressions.
I've updated the documentation if necessary.

This PR implements Adaptive Radix Tree (ART)[1], a compact key-value search tree. It maps byte-sequences to arbitrary values while maintaining data in sorted order. ART can compress multiple nodes into single node and adaptively expand or shrink its internal nodes to avoid worse-case space consumption. Within the broader swift-collection project, it can be considered a compact variant of SortedSet/SortedDictionary but at the cost of keys being convertible to byte-sequences. [1] https://ieeexplore.ieee.org/abstract/document/6544812

Instead of using UnsafeMutablePointer<Node?> to locate where parent stores the child pointer everywhere, localize the unsafe behaviour using the `updateChild` method.

ktoso

I'll be honest that this is a bit outside of the area of my expertise and I wasn't able to provide a lot of valuable input -- I know Karoy has some specific tricks to pull for such collections. However, all those hints are about to become invalidated by the arrival of moveonly types. At the same time, the allocation semantics this tree has I understand are something we specifically want here -- so it is somewhat novel and not much precedent for it.

One thought is to maybe consider using the benchmark package and measure the ARC traffic before we start doing these Unmanaged dances -- it really seems a bit iffy -- also because the manual retain/releases done on the class references are on a strong class reference (buf) to a ManagedBuffer class so that would incur ARC trafi anyway -- probably best to kick off the benchmark package mentioned inline and measure ARC traffic before optimizing.

We discussed with @lorentey and the collection should probably incubate for now outside of swift-collections but we'd like to try to help out polish it up -- even if for a specific use case; which may be easier than bringing it up to the long-term compatibility standards swift-collections would require...

I think that perhaps once we put the collection through some usage that'd be a good driver for brushing it up?

Thanks for the awesome work here so far @vishesh 🙏

ktoso · 2023-10-17T11:05:24Z

Sources/ARTreeModule/ARTree/Consts and Specs.swift

+//
+//===----------------------------------------------------------------------===//
+
+internal struct Const {


filename: a bit weird to have filename with spaces... it should "work fine" but just in case I'd suggest avoiding that.

ktoso · 2023-10-17T11:05:59Z

Sources/ARTreeModule/ARTree/Consts and Specs.swift

+  static var testPrintAddr = false
+}
+
+public protocol ARTreeSpec {


What is this for? 🤔

ktoso · 2023-10-17T11:06:04Z

Sources/ARTreeModule/ARTree/Consts and Specs.swift

@@ -0,0 +1,25 @@
+//===----------------------------------------------------------------------===//


filename: a bit weird to have filename with spaces... it should "work fine" but just in case I'd suggest avoiding that.

ktoso · 2023-10-17T11:09:28Z

Sources/ARTreeModule/CMakeLists.txt

+  "ARTree/Node48.swift"
+  "ARTree/NodeHeader.swift"
+  "ARTree/NodeLeaf.swift"
+  "ARTree/NodeStorage.swif"


Suggested change

"ARTree/NodeStorage.swif"

"ARTree/NodeStorage.swift"

ktoso · 2023-10-17T11:10:23Z

Sources/ARTreeModule/ARTree/ArtNode.swift

+public typealias KeyPart = UInt8
+public typealias Key = [KeyPart]
+
+protocol ArtNode<Spec> {


naming consistency: ARTNode

ktoso · 2023-10-17T11:20:28Z

Sources/ARTreeModule/ARTree/ARTree+Collection.swift

+  }
+
+  var endIndex: Index  {
+    return Index(forTree: self)


Is this right?

ktoso · 2023-10-17T11:26:39Z

Sources/ARTreeModule/ARTree/FixedArray+impl.swift

+  }
+
+  internal var endIndex: Index {
+    return capacity


Hmmm... is it? In other places we bound by the count, so the active element number -- is this ok here?

ktoso · 2023-10-17T11:34:56Z

Sources/ARTreeModule/ARTree/FixedArray+Storage.swift

+  static var capacity: Int {
+    @inline(__always) get { 256 }
+  }
+}


Yeah that's the way without going unsafe I think.

ktoso · 2023-10-17T11:47:02Z

Sources/ARTreeModule/RadixTree.swift

+///
+/// `RadixTree` has the same functionality as a standard `Dictionary`, and it largely implements the
+/// same APIs. However, `RadixTree` is optimized specifically for use cases where underlying keys
+/// share common prefixes. The underlying data-structure is a _persistent variant of _Adaptive Radix


Suggested change

/// share common prefixes. The underlying data-structure is a _persistent variant of _Adaptive Radix

/// share common prefixes. The underlying data-structure is a _persistent_ variant of _Adaptive Radix

ktoso · 2023-10-17T12:39:32Z

Sources/ARTreeModule/ARTree/InternalNode.swift

+  static func retainChildren(_ children: Children, count: Int) {
+    for idx in 0..<count {
+      if let c = children[idx] {
+        _ = Unmanaged.passRetained(c.buf)


I'm not really comfortable with the somewhat manual retain dances we're doing here.
On the other hand, I understand the desire -- perhaps the right next step would be to write a benchmark using https://github.com/ordo-one/package-benchmark and observe the refcounts; this way you can actually measure and optimize them.

// It'll report slightly unbalanced retain/release counts, but this is okey: swiftlang/swift#64636 as at least you'll be able to view the numers changing as you tweak

ktoso · 2023-10-19T07:40:53Z

I discussed the refernce counting a bit with @lorentey and here's some notes:

The preferred setup is to either use strong references and automatic retain/release, or to keep child pointers represented with Unmanaged and use _withUnsafeGuaranteedRef and explicit retain/release operations.

Both options have significant drawbacks, but either option is better than trying to do an ad-hoc mix of the two.

FWIW, in the prefix tree implementation of HashTreeCollections, the tree representation uses regular strong references to store the children of each node. There are no explicit retain/release operations anywhere in the code base — everything relies on regular strong reference semantics. (It tries to carefully avoid making any unnecessary copies, to avoid retain/release traffic. For example, child references aren’t loaded into local variables; they are accessed directly via storage pointers.)

(In the CHAMP data structure, child references and leaf items are mixed into the same storage buffer, and that heavily obscures the actual layout.)

That module does also define a separate internal type that wraps an Unmanaged node reference. However, that has very limited functionality — it is primarily designed for use within the Index representation, and it requires extraordinary care to avoid accidentally dereferencing a stale value.

So it really feels like we should try to use a class to do the refcounting and use the package-benchmark to measure and optimize 🤔

vishesh and others added 13 commits August 2, 2023 17:03

Add license headers

b6c1de1

Use ManagedBuffer for underlying storage

66a6727

Fix test formatting

5346d29

Add FixedArray type using tuples

dd98a04

Fixed some test case formatting

3fd67ac

Make NodeLeaf follow ManagedBuffer semantics

21ebcb6

Remove Header pointer

579aabe

Initialize and bindMemory properly

8d3a75d

Update size test

f52d685

Some refactor, cleanup and regression fixes

0ed8b2b

Iterate empty tree

69ca513

Remove some TODO's

439c6a1

vishesh force-pushed the develop-artree branch from b77d93c to 654b5e4 Compare August 19, 2023 01:35

Move around and some cleanup

91e2c7c

vishesh force-pushed the develop-artree branch from 654b5e4 to 91e2c7c Compare August 19, 2023 01:53

vishesh added 2 commits August 18, 2023 19:08

Fix deletion when node4/16 is full

f9a5452

Cleanup some NodeLeaf generic and some warnings

4174087

vishesh force-pushed the develop-artree branch from 2e56de5 to 4174087 Compare August 19, 2023 03:20

vishesh added 11 commits August 18, 2023 23:50

Replace unsafe parent updating in delete functions

accbd63

Instead of using UnsafeMutablePointer<Node?> to locate where parent stores the child pointer everywhere, localize the unsafe behaviour using the `updateChild` method.

Replace unsafe update of parent in insertion

8ade1e4

Fix warnings

8634916

Make insert work, but still using unsafe ChildSlotPtr

79e1966

Removed unused child(at: .., ref: ...) methods

9d313d1

Don't use ChildSlotPtr out as InternalNode method params directly

5d76b25

Break ARTree.insert into two parts

1ce7b8f

Use Node.rawNode instead of constructor

4d5a9c9

Add NodeReference type to avoid exposing pointer directly

dd52ccb

Generalize Tree a bit more

784b788

Refactor ARTree into ARTreeImpl

e29d1da

Add few test and fixed a bug when setting partial length

17604fb

vishesh force-pushed the develop-artree branch 2 times, most recently from 58b664a to 977d7dd Compare August 25, 2023 02:33

Refactor to use getters

5fa86d5

vishesh force-pushed the develop-artree branch from 977d7dd to 5fa86d5 Compare August 25, 2023 02:42

vishesh and others added 9 commits August 24, 2023 23:57

Implement replace when inserting

7641a5a

Add some refcounting tests

e6fc5b7

More testing

d0c680e

Move ARTree into separate submodule

9f74116

Add support for unsigned integers and RadixTree type

1490717

Support signed integer

f7d27b1

Reorganize some tests

bd143f9

Reorganize some files

a67393a

Remove old FixedArray implementation

c9578a2

vishesh force-pushed the develop-artree branch from 70fcdb1 to c9578a2 Compare September 4, 2023 14:22

Add README

27bfed1

vishesh force-pushed the develop-artree branch from 0054aa1 to 27bfed1 Compare September 4, 2023 16:13

vishesh and others added 10 commits September 4, 2023 09:15

Rename FixedStorage to FixedArrayStorage

6ecc661

Supress some warnings

1b74fbc

Use collection test library

6ef8f08

change testRefCountBasic

5b18390

Fix IntMapTest unique number generation

79abfd9

Add String as BinaryComparableBytes

46f2054

Add subscript and dictionary literal

cf652bc

Some documentation

562a42c

More documentation

2613cb2

WIP on Collection protocol

1cf3b88

vishesh force-pushed the develop-artree branch from 085d821 to 1cf3b88 Compare September 7, 2023 04:15

ktoso reviewed Oct 17, 2023

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

ART: Adaptive Radix Tree implementation #307

ART: Adaptive Radix Tree implementation #307

vishesh commented Aug 3, 2023 •

edited

Loading

ktoso left a comment •

edited

Loading

ktoso Oct 17, 2023

ktoso Oct 17, 2023

ktoso Oct 17, 2023

ktoso Oct 17, 2023

ktoso Oct 17, 2023

ktoso Oct 17, 2023

ktoso Oct 17, 2023

ktoso Oct 17, 2023

ktoso Oct 17, 2023

ktoso Oct 17, 2023

ktoso commented Oct 19, 2023

		@@ -0,0 +1,25 @@
		//===----------------------------------------------------------------------===//

	/// share common prefixes. The underlying data-structure is a _persistent variant of _Adaptive Radix
	/// share common prefixes. The underlying data-structure is a _persistent_ variant of _Adaptive Radix

ART: Adaptive Radix Tree implementation #307

Are you sure you want to change the base?

ART: Adaptive Radix Tree implementation #307

Conversation

vishesh commented Aug 3, 2023 • edited Loading

Checklist

ktoso left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ktoso commented Oct 19, 2023

vishesh commented Aug 3, 2023 •

edited

Loading

ktoso left a comment •

edited

Loading