[perf] faster signia capture (again) (#3487)

Describe what your pull request does. If appropriate, add GIFs or images
showing the before and after.

### Change Type

<!-- ❗ Please select a 'Scope' label ❗️ -->

- [x] `sdk` — Changes the tldraw SDK
- [ ] `dotcom` — Changes the tldraw.com web app
- [ ] `docs` — Changes to the documentation, examples, or templates.
- [ ] `vs code` — Changes to the vscode plugin
- [ ] `internal` — Does not affect user-facing stuff

<!-- ❗ Please select a 'Type' label ❗️ -->

- [ ] `bugfix` — Bug fix
- [ ] `feature` — New feature
- [x] `improvement` — Improving existing features
- [ ] `chore` — Updating dependencies, other boring stuff
- [ ] `galaxy brain` — Architectural changes
- [ ] `tests` — Changes to any test code
- [ ] `tools` — Changes to infrastructure, CI, internal scripts,
debugging tools, etc.
- [ ] `dunno` — I don't know


### Test Plan

1. Add a step-by-step description of how to test your PR here.
2.

- [ ] Unit Tests
- [ ] End to end tests

### Release Notes

- Add a brief release note for your PR here.

packages/state/api-report.md

@@ -33,6 +33,8 @@ export interface Computed<Value, Diff = unknown> extends Signal<Value, Diff> {
     readonly parentEpochs: number[];
     // @internal (undocumented)
     readonly parents: Signal<any, any>[];
+    // @internal (undocumented)
+    readonly parentSet: ArraySet<Signal<any, any>>;
 }
 
 // @public

packages/state/src/lib/core/ArraySet.ts

@@ -144,4 +144,29 @@ export class ArraySet<T> {
 
 		throw new Error('no set or array')
 	}
+
+	has(elem: T) {
+		if (this.array) {
+			return this.array.indexOf(elem) !== -1
+		} else {
+			return this.set!.has(elem)
+		}
+	}
+
+	clear() {
+		if (this.set) {
+			this.set.clear()
+		} else {
+			this.arraySize = 0
+			this.array = []
+		}
+	}
+
+	size() {
+		if (this.set) {
+			return this.set.size
+		} else {
+			return this.arraySize
+		}
+	}
 }

packages/state/src/lib/core/Computed.ts

@@ -123,6 +123,8 @@ export interface Computed<Value, Diff = unknown> extends Signal<Value, Diff> {
 	 */
 	readonly isActivelyListening: boolean
 
+	/** @internal */
+	readonly parentSet: ArraySet<Signal<any, any>>
 	/** @internal */
 	readonly parents: Signal<any, any>[]
 	/** @internal */
@@ -141,6 +143,7 @@ class __UNSAFE__Computed<Value, Diff = unknown> implements Computed<Value, Diff>
 	 */
 	private lastCheckedEpoch = GLOBAL_START_EPOCH
 
+	parentSet = new ArraySet<Signal<any, any>>()
 	parents: Signal<any, any>[] = []
 	parentEpochs: number[] = []
 

packages/state/src/lib/core/EffectScheduler.ts

@@ -1,3 +1,4 @@
+import { ArraySet } from './ArraySet'
 import { startCapturingParents, stopCapturingParents } from './capture'
 import { GLOBAL_START_EPOCH } from './constants'
 import { attach, detach, haveParentsChanged, singleton } from './helpers'
@@ -63,9 +64,11 @@ class __EffectScheduler__<Result> {
 	}
 
 	/** @internal */
-	parentEpochs: number[] = []
+	readonly parentSet = new ArraySet<Signal<any, any>>()
 	/** @internal */
-	parents: Signal<any, any>[] = []
+	readonly parentEpochs: number[] = []
+	/** @internal */
+	readonly parents: Signal<any, any>[] = []
 	private readonly _scheduleEffect?: (execute: () => void) => void
 	constructor(
 		public readonly name: string,

packages/state/src/lib/core/__tests__/arraySet.test.ts

@@ -94,12 +94,16 @@ function runTest(seed: number) {
 	for (let i = 0; i < 1000; i++) {
 		const num = nums[Math.floor(r() * nums.length)]
 
-		if (r() > 0.5) {
+		const choice = r()
+		if (choice < 0.45) {
 			as.add(num)
 			s.add(num)
-		} else {
+		} else if (choice < 0.9) {
 			as.remove(num)
 			s.delete(num)
+		} else {
+			as.clear()
+			s.clear()
 		}
 
 		try {

packages/state/src/lib/core/__tests__/capture.test.ts

@@ -1,3 +1,4 @@
+import { ArraySet } from '../ArraySet'
 import { atom } from '../Atom'
 import { computed } from '../Computed'
 import { react } from '../EffectScheduler'
@@ -14,6 +15,7 @@ const emptyChild = (props: Partial<Child> = {}) =>
 	({
 		parentEpochs: [],
 		parents: [],
+		parentSet: new ArraySet(),
 		isActivelyListening: false,
 		lastTraversedEpoch: 0,
 		...props,

packages/state/src/lib/core/__tests__/fuzz.tlstate.test.ts

@@ -349,7 +349,7 @@ class Test {
 	}
 }
 
-const NUM_TESTS = 100
+const NUM_TESTS = 20
 const NUM_OPS_PER_TEST = 1000
 
 function runTest(seed: number) {
@@ -367,3 +367,7 @@ for (let i = 0; i < NUM_TESTS; i++) {
 		runTest(seed)
 	})
 }
+
+test('regression 728608', () => {
+	runTest(728608)
+})

packages/state/src/lib/core/capture.ts

@@ -3,7 +3,6 @@ import type { Child, Signal } from './types'
 
 class CaptureStackFrame {
 	offset = 0
-	numNewParents = 0
 
 	maybeRemoved?: Signal<any>[]
 
@@ -50,33 +49,29 @@ export function unsafe__withoutCapture<T>(fn: () => T): T {
 
 export function startCapturingParents(child: Child) {
 	inst.stack = new CaptureStackFrame(inst.stack, child)
+	child.parentSet.clear()
 }
 
 export function stopCapturingParents() {
 	const frame = inst.stack!
 	inst.stack = frame.below
 
-	const didParentsChange = frame.numNewParents > 0 || frame.offset !== frame.child.parents.length
-
-	if (!didParentsChange) {
-		return
-	}
-
-	for (let i = frame.offset; i < frame.child.parents.length; i++) {
-		const p = frame.child.parents[i]
-		const parentWasRemoved = frame.child.parents.indexOf(p) >= frame.offset
-		if (parentWasRemoved) {
-			detach(p, frame.child)
+	if (frame.offset < frame.child.parents.length) {
+		for (let i = frame.offset; i < frame.child.parents.length; i++) {
+			const maybeRemovedParent = frame.child.parents[i]
+			if (!frame.child.parentSet.has(maybeRemovedParent)) {
+				detach(maybeRemovedParent, frame.child)
+			}
 		}
+
+		frame.child.parents.length = frame.offset
+		frame.child.parentEpochs.length = frame.offset
 	}
 
-	frame.child.parents.length = frame.offset
-	frame.child.parentEpochs.length = frame.offset
-
-	if (inst.stack?.maybeRemoved) {
-		for (let i = 0; i < inst.stack.maybeRemoved.length; i++) {
-			const maybeRemovedParent = inst.stack.maybeRemoved[i]
-			if (frame.child.parents.indexOf(maybeRemovedParent) === -1) {
+	if (frame.maybeRemoved) {
+		for (let i = 0; i < frame.maybeRemoved.length; i++) {
+			const maybeRemovedParent = frame.maybeRemoved[i]
+			if (!frame.child.parentSet.has(maybeRemovedParent)) {
 				detach(maybeRemovedParent, frame.child)
 			}
 		}
@@ -86,34 +81,35 @@ export function stopCapturingParents() {
 // this must be called after the parent is up to date
 export function maybeCaptureParent(p: Signal<any, any>) {
 	if (inst.stack) {
-		const idx = inst.stack.child.parents.indexOf(p)
+		const wasCapturedAlready = inst.stack.child.parentSet.has(p)
 		// if the child didn't deref this parent last time it executed, then idx will be -1
 		// if the child did deref this parent last time but in a different order relative to other parents, then idx will be greater than stack.offset
 		// if the child did deref this parent last time in the same order, then idx will be the same as stack.offset
 		// if the child did deref this parent already during this capture session then 0 <= idx < stack.offset
 
-		if (idx < 0) {
-			inst.stack.numNewParents++
-			if (inst.stack.child.isActivelyListening) {
-				attach(p, inst.stack.child)
-			}
+		if (wasCapturedAlready) {
+			return
 		}
 
-		if (idx < 0 || idx >= inst.stack.offset) {
-			if (idx !== inst.stack.offset && idx > 0) {
-				const maybeRemovedParent = inst.stack.child.parents[inst.stack.offset]
+		inst.stack.child.parentSet.add(p)
+		if (inst.stack.child.isActivelyListening) {
+			attach(p, inst.stack.child)
+		}
 
+		if (inst.stack.offset < inst.stack.child.parents.length) {
+			const maybeRemovedParent = inst.stack.child.parents[inst.stack.offset]
+			if (maybeRemovedParent !== p) {
 				if (!inst.stack.maybeRemoved) {
 					inst.stack.maybeRemoved = [maybeRemovedParent]
-				} else if (inst.stack.maybeRemoved.indexOf(maybeRemovedParent) === -1) {
+				} else {
 					inst.stack.maybeRemoved.push(maybeRemovedParent)
 				}
 			}
-
-			inst.stack.child.parents[inst.stack.offset] = p
-			inst.stack.child.parentEpochs[inst.stack.offset] = p.lastChangedEpoch
-			inst.stack.offset++
 		}
+
+		inst.stack.child.parents[inst.stack.offset] = p
+		inst.stack.child.parentEpochs[inst.stack.offset] = p.lastChangedEpoch
+		inst.stack.offset++
 	}
 }
 

packages/state/src/lib/core/types.ts

@@ -51,8 +51,9 @@ export interface Signal<Value, Diff = unknown> {
 /** @internal */
 export type Child = {
 	lastTraversedEpoch: number
-	parents: Signal<any, any>[]
-	parentEpochs: number[]
+	readonly parentSet: ArraySet<Signal<any, any>>
+	readonly parents: Signal<any, any>[]
+	readonly parentEpochs: number[]
 	isActivelyListening: boolean
 }