Merge pull request #6 from nielsenko/fix-add-bug

fix add bug

lib/src/node.dart

@@ -2,27 +2,23 @@
 // SPDX-License-Identifier: BSD-3-Clause
 import 'split.dart';
 
-class Node<T> {
+sealed class Node<T> {
   final T item;
   final int priority;
-  final int _size;
-  final Node<T>? left, right;
+  int get _size; // size is an extension method on Node<T>?
+  Node<T>? get left;
+  Node<T>? get right;
 
-  Node(this.item, this.priority, {this.left, this.right})
-      : _size = 1 + left.size + right.size {
+  Node(this.item, this.priority) {
     checkInvariant();
   }
 
-  /// Replace the left child with [left].
-  Node<T> withLeft(Node<T>? left) =>
-      Node(item, priority, left: left, right: right);
+  Node<T> withLeft(Node<T>? left);
 
-  /// Replace the right child with [right].
-  Node<T> withRight(Node<T>? right) =>
-      Node(item, priority, left: left, right: right);
+  Node<T> withRight(Node<T>? right);
 
-  /// Create a copy of this node without children.
-  Node<T> withoutChildren() => Node(item, priority);
+  /// x
+  Node<T> withoutChildren();
 
   /// The minimum item in the treap.
   T get min => left == null ? item : left!.min;
@@ -42,6 +38,31 @@ class Node<T> {
   }
 }
 
+class PersistentNode<T> extends Node<T> {
+  @override
+  final PersistentNode<T>? left, right;
+  @override
+  final int _size;
+
+  PersistentNode(super.item, super.priority, {this.left, this.right})
+      : _size = 1 + left.size + right.size;
+
+  /// Create a copy with the left child set to [left].
+  @override
+  PersistentNode<T> withLeft(covariant PersistentNode<T>? left) =>
+      PersistentNode(item, priority, left: left, right: right);
+
+  /// Create a copy with the right child set to [right].
+  @override
+  PersistentNode<T> withRight(covariant PersistentNode<T>? right) =>
+      PersistentNode(item, priority, left: left, right: right);
+
+  /// Create a copy of this node without children.
+  /// @super
+  @override
+  PersistentNode<T> withoutChildren() => PersistentNode(item, priority);
+}
+
 extension NodeEx<T> on Node<T>? {
   int get size => this?._size ?? 0;
 
@@ -57,7 +78,8 @@ extension NodeEx<T> on Node<T>? {
       final s = self.right.split(pivot, compare);
       return s.withLow(self.withRight(s.low));
     }
-    return Split((low: self.left, middle: self, high: self.right)) // order == 0
+    // order == 0
+    return Split((low: self.left, middle: self, high: self.right))
       ..checkInvariant(compare);
   }
 
@@ -73,11 +95,16 @@ extension NodeEx<T> on Node<T>? {
     return self ?? other; // zero or one
   }
 
-  Node<T> add(Node<T> child, Comparator<T> compare) =>
-      split(child.item, compare).withMiddle(child).join()!;
+  ({Node<T> root, T? old}) upsert(Node<T> child, Comparator<T> compare) {
+    final item = child.item;
+    final s = split(item, compare);
+    return (root: s.withMiddle(child).join()!, old: s.middle?.item);
+  }
 
-  Node<T>? erase(T dead, Comparator<T> compare) =>
-      split(dead, compare).withMiddle(null).join();
+  ({Node<T>? root, T? old}) erase(T dead, Comparator<T> compare) {
+    final s = split(dead, compare);
+    return (root: s.withMiddle(null).join(), old: s.middle?.item);
+  }
 
   Node<T>? union(Node<T>? other, Comparator<T> compare) {
     final self = this; // for type promotion

lib/src/treap_base.dart

@@ -41,6 +41,8 @@ class Treap<T> {
   const Treap([Comparator<T>? compare])
       : this._(null, compare ?? Comparable.compare as Comparator<T>);
 
+  Treap<T> _new(Node<T>? root) => root == _root ? this : Treap._(root, compare);
+
   /// Build a treap containing the [items].
   ///
   /// Constructs the treap by folding [add] over the [items], adding each one
@@ -49,28 +51,42 @@ class Treap<T> {
   /// This method is `O(N log(N))` in complexity. An `O(N)` algorithm exists if the
   /// items are sorted. However, this works in all cases.
   factory Treap.of(Iterable<T> items, [Comparator<T>? comparator]) =>
-      items.fold(Treap(comparator), (acc, i) => acc.add(i));
+      Treap(comparator).addAll(items);
 
   /// Adds an [item].
   ///
-  /// Creates a new node with the [item] and a random priority. Returns a new treap
-  /// with the added node.
-  Treap<T> add(T item) => Treap<T>._(
-      _root.add(
-        Node<T>(item, _rnd.nextInt(1 << 32)),
-        compare,
-      ),
-      compare);
+  /// If the [item] is already present in the treap, the original treap is returned.
+  /// Otherwise, a new treap is returned with the item added.
+  Treap<T> add(T item) {
+    final (:root, :old) = _root.upsert(_createNode(item), compare);
+    return old == null ? _new(root) : this;
+  }
+
+  /// Adds or updates an [item].
+  ///
+  /// Returns a new treap, with [item] either added or updated.
+  Treap<T> addOrUpdate(T item) {
+    final (:root, :old) = _root.upsert(_createNode(item), compare);
+    return _new(root);
+  }
+
+  Node<T> _createNode(T item) => PersistentNode<T>(item, _rnd.nextInt(1 << 32));
 
   /// Adds a range of [items].
   ///
-  /// Returns a new treap with the added [items].
-  Treap<T> addRange(Iterable<T> items) => union(Treap.of(items, compare));
+  /// Returns a new treap with the added [items]. If all the [items] are already
+  /// present, the original treap is returned.
+  Treap<T> addAll(Iterable<T> items) =>
+      items.fold(this, (acc, i) => acc.add(i));
 
   /// Erases an [item] from the treap, if it exists.
   ///
-  /// Returns a new treap without the erased [item].
-  Treap<T> erase(T item) => Treap._(_root.erase(item, compare), compare);
+  /// Returns a new treap without the erased [item]. If the [item] was not present,
+  /// the original treap is returned.
+  Treap<T> erase(T item) {
+    final (:root, :old) = _root.erase(item, compare);
+    return old != null ? _new(root) : this;
+  }
 
   /// Whether this treap is empty.
   bool get isEmpty => _root == null;
@@ -85,8 +101,6 @@ class Treap<T> {
   T? find(T item) => _root.find(item, compare)?.item;
 
   /// Whether an[item] exists in this treap.
-  ///
-  /// Returns `true` if `find` re
   bool has(T item) => find(item) != null;
 
   /// The rank of an [item].
@@ -132,9 +146,11 @@ class Treap<T> {
   /// Returns a new treap with the first [n] items.
   ///
   /// Returns the original treap, if [n] is greater than the [size] of this treap.
-  Treap<T> take(int n) => n < size
-      ? Treap._(_root.split(_root.select(n).item, compare).low, compare)
-      : this;
+  Treap<T> take(int n) {
+    if (n == 0) return Treap(compare); // empty;
+    if (n >= size) return this;
+    return _new(_root.split(_root.select(n).item, compare).low);
+  }
 
   /// Skips the first [n] items and returns a new treap with the remaining items.
   ///
@@ -143,23 +159,19 @@ class Treap<T> {
   Treap<T> skip(int n) {
     if (n == 0) return this;
     if (n >= size) return Treap(compare); // empty
-    return Treap._(
-      _root.split(_root.select(n - 1).item, compare).high,
-      compare,
-    );
+    return _new(_root.split(_root.select(n - 1).item, compare).high);
   }
 
   /// Returns a new treap that is the union of this treap and the [other] treap.
-  Treap<T> union(Treap<T> other) =>
-      Treap._(_root.union(other._root, compare), compare);
+  Treap<T> union(Treap<T> other) => _new(_root.union(other._root, compare));
 
   /// Returns a new treap that is the intersection of this treap and the [other] treap.
   Treap<T> intersection(Treap<T> other) =>
-      Treap._(_root.intersection(other._root, compare), compare);
+      _new(_root.intersection(other._root, compare));
 
   /// Returns a new treap that is the difference of this treap minus the [other] treap.
   Treap<T> difference(Treap<T> other) =>
-      Treap._(_root.difference(other._root, compare), compare);
+      _new(_root.difference(other._root, compare));
 
   /// Operator overload for [add]ing an [item] to the treap.
   Treap<T> operator +(T item) => add(item);
@@ -176,3 +188,5 @@ class Treap<T> {
   /// Operator overload for [select]ing an item in the treap by its [index].
   T operator [](int index) => select(index);
 }
+
+class MutableTreap<T> {}

lib/src/treap_map.dart

@@ -16,7 +16,7 @@ class TreapMap<K, V> extends MapBase<K, V> {
 
   @override
   void operator []=(K key, V value) =>
-      _root = _root.add((key: key, value: value));
+      _root = _root.addOrUpdate((key: key, value: value));
 
   @override
   void clear() => _root = Treap(_root.compare);

lib/src/treap_set.dart

@@ -22,7 +22,7 @@ class TreapSet<T> extends SetBase<T> {
   }
 
   @override
-  void addAll(Iterable<T> elements) => _root = _root.addRange(elements);
+  void addAll(Iterable<T> elements) => _root = _root.addAll(elements);
 
   @override
   bool contains(covariant T element) => lookup(element) != null;

test/set_test.dart

@@ -150,5 +150,11 @@ void main() {
       expect(it.current, 5);
       expect(it.moveNext(), isFalse);
     });
+
+    test("add don't update", () {
+      final s = TreapSet.of([(1, 2)], (a, b) => a.$1 - b.$1);
+      expect(s.add((1, 3)), isFalse);
+      expect(s.lookup((1, 0)), (1, 2));
+    });
   });
 }

test/treap_test.dart

@@ -12,7 +12,9 @@ void main() {
       test('add, erase, build', () {
         final x = Treap<num>() + 1;
         final y = x.add(1);
-        expect(x, isNot(y)); // add gives new Treap, even for existing items
+        final z = x.addOrUpdate(1);
+        expect(x, y);
+        expect(x, isNot(z));
 
         // Be aware, that chaining with .. operator probably don't do what you want
         x
@@ -34,6 +36,11 @@ void main() {
         expect(t.isEmpty, isTrue);
         expect(t.values, []);
       });
+
+      test('duplicates', () {
+        final t = Treap<num>.of([1, 1, 1, 1, 1]);
+        expect(t.values, [1]);
+      });
     });
 
     group('retrieval', () {
@@ -165,15 +172,17 @@ void main() {
 
   group('Node', () {
     final rnd = Random(42 ^ 42);
-    Node<num> node(num value) => Node<num>(value, rnd.nextInt(1 << 32));
+    Node<num> node(num value) =>
+        PersistentNode<num>(value, rnd.nextInt(1 << 32));
 
     test('add, find, erase, inOrder', () {
       final first = node(1);
-      final again = first.add(node(1), Comparable.compare);
-      final second = first.add(node(2), Comparable.compare);
-      final third = second.add(node(3), Comparable.compare);
-      final another = third.add(node(3), Comparable.compare);
-      final forth = third.add(node(0), Comparable.compare);
+      final (root: again, old: _) = first.upsert(node(1), Comparable.compare);
+      final (root: second, old: _) = first.upsert(node(2), Comparable.compare);
+      final (root: third, old: _) = second.upsert(node(3), Comparable.compare);
+      final (root: another, old: _) =
+          second.upsert(node(3), Comparable.compare);
+      final (root: forth, old: _) = third.upsert(node(0), Comparable.compare);
 
       expect(first.inOrder().map((n) => n.item), [1]);
       expect(again.inOrder().map((n) => n.item), [1]);
@@ -190,18 +199,19 @@ void main() {
       expect(second.find(1, Comparable.compare), isNotNull);
       expect(second.find(2, Comparable.compare), isNotNull);
 
-      final fifth = forth.erase(0, Comparable.compare);
+      final (root: fifth, old: _) = forth.erase(0, Comparable.compare);
       expect(fifth!.inOrder().map((n) => n.item), [1, 2, 3]);
       expect(identical(third, fifth), false);
 
-      expect(forth.erase(2, Comparable.compare)!.inOrder().map((n) => n.item),
-          [0, 1, 3]);
+      expect(
+        forth.erase(2, Comparable.compare).root!.inOrder().map((n) => n.item),
+        [0, 1, 3],
+      );
     });
 
     test('rank, select', () {
-      final top = [1, 2, 3, 4, 5, 6, 7, 8, 9]
-          .reversed
-          .fold(node(0), (acc, i) => acc.add(node(i), Comparable.compare));
+      final top = [1, 2, 3, 4, 5, 6, 7, 8, 9].reversed.fold(
+          node(0), (acc, i) => acc.upsert(node(i), Comparable.compare).root);
       expect(top.inOrder().map((n) => n.item), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
       expect(top.inOrder().map((n) => top.rank(n.item, Comparable.compare)),
           [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
@@ -217,9 +227,10 @@ void main() {
       // Deterministic shaped treap despite shuffle
       final items = [5, 6, 3, 9, 1, 8, 2, 4, 7]; // evil order
       print(items);
-      final top = items.fold(
-        Node<num>(0, 0), // will have same priority (5,0)
-        (acc, i) => acc.add(Node(i, 5 - i), Comparable.compare),
+      final top = items.fold<Node<num>>(
+        PersistentNode<num>(0, 0), // will have same priority (5,0)
+        (acc, i) =>
+            acc.upsert(PersistentNode<num>(i, 5 - i), Comparable.compare).root,
       );
       expect(
         top.inOrder().map((n) => n.item),