[signia] perf thing again (#3645)

Will explain tomorrow

### 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/src/lib/core/Computed.ts

@@ -4,7 +4,7 @@ import { HistoryBuffer } from './HistoryBuffer'
 import { maybeCaptureParent, startCapturingParents, stopCapturingParents } from './capture'
 import { GLOBAL_START_EPOCH } from './constants'
 import { EMPTY_ARRAY, equals, haveParentsChanged, singleton } from './helpers'
-import { getGlobalEpoch } from './transactions'
+import { getGlobalEpoch, getIsReacting, getReactionEpoch } from './transactions'
 import { Child, ComputeDiff, RESET_VALUE, Signal } from './types'
 import { logComputedGetterWarning } from './warnings'
 
@@ -189,8 +189,17 @@ class __UNSAFE__Computed<Value, Diff = unknown> implements Computed<Value, Diff>
 	__unsafe__getWithoutCapture(ignoreErrors?: boolean): Value {
 		const isNew = this.lastChangedEpoch === GLOBAL_START_EPOCH
 
-		if (!isNew && (this.lastCheckedEpoch === getGlobalEpoch() || !haveParentsChanged(this))) {
-			this.lastCheckedEpoch = getGlobalEpoch()
+		const globalEpoch = getGlobalEpoch()
+
+		if (
+			!isNew &&
+			(this.lastCheckedEpoch === globalEpoch ||
+				(this.isActivelyListening &&
+					getIsReacting() &&
+					this.lastTraversedEpoch < getReactionEpoch()) ||
+				!haveParentsChanged(this))
+		) {
+			this.lastCheckedEpoch = globalEpoch
 			if (this.error) {
 				if (!ignoreErrors) {
 					throw this.error.thrownValue

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

@@ -144,8 +144,12 @@ class __EffectScheduler__<Result> {
 	execute(): Result {
 		try {
 			startCapturingParents(this)
+			// Important! We have to make a note of the current epoch before running the effect.
+			// We allow atoms to be updated during effects, which increments the global epoch,
+			// so if we were to wait until after the effect runs, the this.lastReactedEpoch value might get ahead of itself.
+			const currentEpoch = getGlobalEpoch()
 			const result = this.runEffect(this.lastReactedEpoch)
-			this.lastReactedEpoch = getGlobalEpoch()
+			this.lastReactedEpoch = currentEpoch
 			return result
 		} finally {
 			stopCapturingParents()

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

@@ -17,7 +17,7 @@ describe('reactors', () => {
 		expect(r.scheduler.isActivelyListening).toBe(true)
 	})
 
-	it('can not set atom values directly yet', () => {
+	it('can not set atom values indefinitely', () => {
 		const a = atom('', 1)
 		const r = reactor('', () => {
 			if (a.get() < +Infinity) {
@@ -25,7 +25,7 @@ describe('reactors', () => {
 			}
 		})
 		expect(() => r.start()).toThrowErrorMatchingInlineSnapshot(
-			`"cannot change atoms during reaction cycle"`
+			`"Reaction update depth limit exceeded"`
 		)
 	})
 

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

@@ -200,3 +200,91 @@ describe('transact', () => {
 		expect.assertions(3)
 	})
 })
+
+describe('setting atoms during a reaction', () => {
+	it('should work', () => {
+		const a = atom('', 0)
+		const b = atom('', 0)
+
+		react('', () => {
+			b.set(a.get() + 1)
+		})
+
+		expect(a.get()).toBe(0)
+		expect(b.get()).toBe(1)
+	})
+
+	it('should throw an error if it gets into a loop', () => {
+		expect(() => {
+			const a = atom('', 0)
+
+			react('', () => {
+				a.set(a.get() + 1)
+			})
+		}).toThrowErrorMatchingInlineSnapshot(`"Reaction update depth limit exceeded"`)
+	})
+
+	it('should work with a transaction running', () => {
+		const a = atom('', 0)
+
+		react('', () => {
+			transact(() => {
+				if (a.get() < 10) {
+					a.set(a.get() + 1)
+				}
+			})
+		})
+
+		expect(a.get()).toBe(10)
+	})
+
+	it('[regression 1] should allow computeds to be updated properly', () => {
+		const a = atom('', 0)
+		const b = atom('', 0)
+		const c = computed('', () => b.get() * 2)
+
+		let cValue = 0
+
+		react('', () => {
+			b.set(a.get() + 1)
+			cValue = c.get()
+		})
+
+		expect(a.get()).toBe(0)
+		expect(b.get()).toBe(1)
+		expect(cValue).toBe(2)
+
+		transact(() => {
+			a.set(1)
+		})
+		expect(cValue).toBe(4)
+	})
+
+	it('[regression 2] should allow computeds to be updated properly', () => {
+		const a = atom('', 0)
+		const b = atom('', 1)
+		const c = atom('', 0)
+		const d = computed('', () => a.get() * 2)
+
+		let dValue = 0
+		react('', () => {
+			// update a, causes a and d to be traversed (but not updated)
+			a.set(b.get())
+			// update c
+			c.set(a.get())
+			// make sure that when we get d, it is updated properly
+			dValue = d.get()
+		})
+
+		expect(a.get()).toBe(1)
+		expect(b.get()).toBe(1)
+		expect(c.get()).toBe(1)
+
+		expect(dValue).toBe(2)
+
+		transact(() => {
+			b.set(2)
+		})
+		expect(dValue).toBe(4)
+	})
+})

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

@@ -1,6 +1,6 @@
 import { _Atom } from './Atom'
-import { GLOBAL_START_EPOCH } from './constants'
 import { EffectScheduler } from './EffectScheduler'
+import { GLOBAL_START_EPOCH } from './constants'
 import { singleton } from './helpers'
 import { Child, Signal } from './types'
 
@@ -25,12 +25,10 @@ class Transaction {
 	 */
 	commit() {
 		if (this.isRoot) {
-			// For root transactions, flush changes to each of the atom's initial values.
-			const atoms = this.initialAtomValues
-			this.initialAtomValues = new Map()
-			flushChanges(atoms.keys())
+			// For root transactions, flush changed atoms
+			flushChanges(this.initialAtomValues.keys())
 		} else {
-			// For transaction's with parents, add the transaction's initial values to the parent's.
+			// For transactions with parents, add the transaction's initial values to the parent's.
 			this.initialAtomValues.forEach((value, atom) => {
 				if (!this.parent!.initialAtomValues.has(atom)) {
 					this.parent!.initialAtomValues.set(atom, value)
@@ -64,12 +62,37 @@ const inst = singleton('transactions', () => ({
 	// Whether any transaction is reacting.
 	globalIsReacting: false,
 	currentTransaction: null as Transaction | null,
+
+	cleanupReactors: null as null | Set<EffectScheduler<unknown>>,
+	reactionEpoch: GLOBAL_START_EPOCH + 1,
 }))
 
+export function getReactionEpoch() {
+	return inst.reactionEpoch
+}
+
 export function getGlobalEpoch() {
 	return inst.globalEpoch
 }
 
+export function getIsReacting() {
+	return inst.globalIsReacting
+}
+
+function traverse(reactors: Set<EffectScheduler<unknown>>, child: Child) {
+	if (child.lastTraversedEpoch === inst.globalEpoch) {
+		return
+	}
+
+	child.lastTraversedEpoch = inst.globalEpoch
+
+	if (child instanceof EffectScheduler) {
+		reactors.add(child)
+	} else {
+		;(child as any as Signal<any>).children.visit((c) => traverse(reactors, c))
+	}
+}
+
 /**
  * Collect all of the reactors that need to run for an atom and run them.
  *
@@ -82,34 +105,33 @@ function flushChanges(atoms: Iterable<_Atom>) {
 
 	try {
 		inst.globalIsReacting = true
+		inst.reactionEpoch = inst.globalEpoch
 
 		// Collect all of the visited reactors.
 		const reactors = new Set<EffectScheduler<unknown>>()
 
-		// Visit each descendant of the atom, collecting reactors.
-		const traverse = (node: Child) => {
-			if (node.lastTraversedEpoch === inst.globalEpoch) {
-				return
-			}
-
-			node.lastTraversedEpoch = inst.globalEpoch
-
-			if (node instanceof EffectScheduler) {
-				reactors.add(node)
-			} else {
-				;(node as any as Signal<any>).children.visit(traverse)
-			}
-		}
-
 		for (const atom of atoms) {
-			atom.children.visit(traverse)
+			atom.children.visit((child) => traverse(reactors, child))
 		}
 
 		// Run each reactor.
 		for (const r of reactors) {
 			r.maybeScheduleEffect()
 		}
+
+		let updateDepth = 0
+		while (inst.cleanupReactors?.size) {
+			if (updateDepth++ > 1000) {
+				throw new Error('Reaction update depth limit exceeded')
+			}
+			const reactors = inst.cleanupReactors
+			inst.cleanupReactors = null
+			for (const r of reactors) {
+				r.maybeScheduleEffect()
+			}
+		}
 	} finally {
+		inst.cleanupReactors = null
 		inst.globalIsReacting = false
 	}
 }
@@ -123,9 +145,19 @@ function flushChanges(atoms: Iterable<_Atom>) {
  * @internal
  */
 export function atomDidChange(atom: _Atom, previousValue: any) {
-	if (!inst.currentTransaction) {
+	if (inst.globalIsReacting) {
+		// If the atom changed during the reaction phase of flushChanges
+		// then we are past the point where a transaction can be aborted
+		// so we don't need to note down the previousValue.
+		const rs = (inst.cleanupReactors ??= new Set())
+		atom.children.visit((child) => traverse(rs, child))
+	} else if (!inst.currentTransaction) {
+		// If there is no transaction, flush the changes immediately.
 		flushChanges([atom])
 	} else if (!inst.currentTransaction.initialAtomValues.has(atom)) {
+		// If we are in a transaction, then all we have to do is preserve
+		// the value of the atom at the start of the transaction in case
+		// we need to roll back.
 		inst.currentTransaction.initialAtomValues.set(atom, previousValue)
 	}
 }
@@ -213,24 +245,26 @@ export function transaction<T>(fn: (rollback: () => void) => T) {
 	inst.currentTransaction = txn
 
 	try {
+		let result = undefined as T | undefined
 		let rollback = false
 
-		// Run the function.
-		const result = fn(() => (rollback = true))
+		try {
+			// Run the function.
+			result = fn(() => (rollback = true))
+		} catch (e) {
+			// Abort the transaction if the function throws.
+			txn.abort()
+			throw e
+		}
 
 		if (rollback) {
 			// If the rollback was triggered, abort the transaction.
 			txn.abort()
 		} else {
-			// Otherwise, commit the transaction.
 			txn.commit()
 		}
 
 		return result
-	} catch (e) {
-		// Abort the transaction if the function throws.
-		txn.abort()
-		throw e
 	} finally {
 		// Set the current transaction to the transaction's parent.
 		inst.currentTransaction = inst.currentTransaction.parent