Adding constant count support.

Author: mirek

balance-inlined.ts

@@ -1,44 +1,44 @@
-/**
- * Inlined version, small improvement, 1mln random integer inserts:
- *
- * 1.83s
- * rss 248.05 MB
- * heapTotal 219.61 MB
- * heapUsed 117.88 MB
- * external 0.25 MB
- * arrayBuffers 0.01 MB
- *
- * vs original:
- *
- * 1.85s
- * rss 248.34 MB
- * heapTotal 219.36 MB
- * heapUsed 126.33 MB
- * external 0.25 MB
- * arrayBuffers 0.01 MB
- *
- */
+// /**
+//  * Inlined version, small improvement, 1mln random integer inserts:
+//  *
+//  * 1.83s
+//  * rss 248.05 MB
+//  * heapTotal 219.61 MB
+//  * heapUsed 117.88 MB
+//  * external 0.25 MB
+//  * arrayBuffers 0.01 MB
+//  *
+//  * vs original:
+//  *
+//  * 1.85s
+//  * rss 248.34 MB
+//  * heapTotal 219.36 MB
+//  * heapUsed 126.33 MB
+//  * external 0.25 MB
+//  * arrayBuffers 0.01 MB
+//  *
+//  */
 
-import { R, B, BB, M } from './prelude.js'
+// import { R, B, BB, M } from './prelude.js'
 
-const balance =
-  <T>(_: M<T>): M<T> => {
-    switch (true) {
-      case _.c === B && _.l?.c === R && _.l.l?.c === R:
-        return { c: R, l: { c: B, l: _.l!.l!.l, v: _.l!.l!.v, r: _.l!.l!.r }, v: _.l!.v, r: { c: B, l: _.l!.r, v: _.v, r: _.r } }
-      case _.c === B && _.l?.c === R && _.l.r?.c === R:
-        return { c: R, l: { c: B, l: _.l!.l, v: _.l!.v, r: _.l!.r!.l }, v: _.l!.r!.v, r: { c: B, l: _.l!.r!.r, v: _.v, r: _.r } }
-      case _.c === B && _.r?.c === R && _.r.l?.c === R:
-        return { c: R, l: { c: B, l: _.l, v: _.v, r: _.r!.l!.l }, v: _.r!.l!.v, r: { c: B, l: _.r!.l!.r, v: _.r!.v, r: _.r!.r } }
-      case _.c === B && _.r?.c === R && _.r.r?.c === R:
-        return { c: R, l: { c: B, l: _.l, v: _.v, r: _.r!.l }, v: _.r!.v, r: { c: B, l: _.r!.r!.l, v: _.r!.r!.v, r: _.r!.r!.r } }
-      case _.c === BB && _.r?.c === R && _.r.l?.c === R:
-        return { c: B, l: { c: B, l: _.l, v: _.v, r: _.r!.l!.l }, v: _.r!.l!.v, r: { c: B, l: _.r!.l!.r, v: _.r!.v, r: _.r!.r } }
-      case _.c === BB && _.l?.c === R && _.l.r?.c === R:
-        return { c: B, l: { c: B, l: _.l!.l, v: _.l!.v, r: _.l!.r!.l }, v: _.l!.r!.v, r: { c: B, l: _.l!.r!.r, v: _.v, r: _.r } }
-      default:
-        return _
-    }
-  }
+// const balance =
+//   <T>(_: M<T>): M<T> => {
+//     switch (true) {
+//       case _.c === B && _.l?.c === R && _.l.l?.c === R:
+//         return { c: R, l: { c: B, l: _.l!.l!.l, v: _.l!.l!.v, r: _.l!.l!.r }, v: _.l!.v, r: { c: B, l: _.l!.r, v: _.v, r: _.r } }
+//       case _.c === B && _.l?.c === R && _.l.r?.c === R:
+//         return { c: R, l: { c: B, l: _.l!.l, v: _.l!.v, r: _.l!.r!.l }, v: _.l!.r!.v, r: { c: B, l: _.l!.r!.r, v: _.v, r: _.r } }
+//       case _.c === B && _.r?.c === R && _.r.l?.c === R:
+//         return { c: R, l: { c: B, l: _.l, v: _.v, r: _.r!.l!.l }, v: _.r!.l!.v, r: { c: B, l: _.r!.l!.r, v: _.r!.v, r: _.r!.r } }
+//       case _.c === B && _.r?.c === R && _.r.r?.c === R:
+//         return { c: R, l: { c: B, l: _.l, v: _.v, r: _.r!.l }, v: _.r!.v, r: { c: B, l: _.r!.r!.l, v: _.r!.r!.v, r: _.r!.r!.r } }
+//       case _.c === BB && _.r?.c === R && _.r.l?.c === R:
+//         return { c: B, l: { c: B, l: _.l, v: _.v, r: _.r!.l!.l }, v: _.r!.l!.v, r: { c: B, l: _.r!.l!.r, v: _.r!.v, r: _.r!.r } }
+//       case _.c === BB && _.l?.c === R && _.l.r?.c === R:
+//         return { c: B, l: { c: B, l: _.l!.l, v: _.l!.v, r: _.l!.r!.l }, v: _.l!.r!.v, r: { c: B, l: _.l!.r!.r, v: _.v, r: _.r } }
+//       default:
+//         return _
+//     }
+//   }
 
-export default balance
+// export default balance

insert.ts

@@ -1,6 +1,6 @@
 import { R, E, EE, N, mk } from './prelude.js'
 import * as Cmp from '@prelude/cmp'
-import balance from './balance-inlined.js'
+import balance from './balance.js'
 
 const insert =
   <T, K>(_: N<T>, key: (value: T) => K, cmp: Cmp.t<K>, x: T, merge: (a: T, b: T) => T): NonNullable<N<T>> => {

prelude.ts

@@ -6,10 +6,17 @@ export type P = typeof R | typeof B | typeof BB
 
 export const E = undefined
 export const EE = null
-export type N<T> = typeof E | typeof EE | { c: 1 | 2 | 3, l: N<T>, v: T, r: N<T> }
+export type N<T> = typeof E | typeof EE | {
+  c: 1 | 2 | 3,
+  l: N<T>,
+  v: T,
+  r: N<T>,
+  n: number
+}
+
 export type M<T> = NonNullable<N<T>>
 
 export const mk =
   <T>(c: P, l: N<T>, v: T, r: N<T>): M<T> => ({
-    c, l, v, r
+    c, l, v, r, n: 1 + (l?.n ?? 0) + (r?.n ?? 0)
   })

rb-tree.ts

@@ -108,6 +108,10 @@ export const each =
   <T, K>(tree: RbTree<T, K>): Generator<T> =>
     each_(tree.root)
 
+export const count =
+  <T, K>(tree: RbTree<T, K>): number =>
+    tree.root?.n ?? 0
+
 const linearCount_ =
   <T>(_: N<T>): number =>
     _ ? 1 + linearCount_(_.l) + linearCount_(_.r) : 0

rotate.ts

@@ -1,5 +1,5 @@
 import { R, B, BB, E, EE, M, mk } from './prelude.js'
-import balance from './balance-inlined.js'
+import balance from './balance.js'
 
 const rotate =
   <T>(_: M<T>): M<T> => {

test/rbtree.test.ts

@@ -38,13 +38,16 @@ test('random numbers', () => {
 
 describe('pop', () => {
 
-  const n = 1_000_000
+  const n = 1_000
   const rb = RbTree.of(RbTree.Cmp.numbers, (_: number) => _)
   const xs: number[] = []
 
   test(`insert ${n}`, () => {
     for (let i = 0; i < n; i++) {
       RbTree.insert(rb, Math.random())
+      if (rb.root?.n !== i + 1) {
+        throw new Error(`${rb.root?.n} != ${i + 1}`)
+      }
     }
   })
 
@@ -83,13 +86,16 @@ describe('deletes', () => {
 
   test('deletions', () => {
     expect(RbTree.linearCount(rb)).toBe(n)
+    expect(RbTree.count(rb)).toBe(n)
     for (let i = 0; i < n; i++) {
       const value = Arrays.deleteSwapRandom(xs)
       expect(RbTree.has(rb, value)).toBe(true)
       RbTree.delete(rb, value)
       expect(RbTree.has(rb, value)).toBe(false)
+      expect(RbTree.count(rb)).toBe(n - (i + 1))
     }
     expect(RbTree.linearCount(rb)).toBe(0)
+    expect(RbTree.count(rb)).toBe(0)
   })
 
 })