feat: reimplementing effect scope

Author: johnsoncodehk

src/index.ts

@@ -3,6 +3,7 @@ export * from './system.js';
 import { createReactiveSystem, Dependency, Subscriber, SubscriberFlags } from './system.js';
 
 interface Effect extends Subscriber, Dependency {
+	isScope?: true;
 	fn(): void;
 }
 
@@ -49,10 +50,13 @@ const {
 	notifyEffect(e: Effect) {
 		const flags = e.flags;
 		if (
-			flags & SubscriberFlags.Dirty
-			|| (flags & SubscriberFlags.PendingComputed && updateDirtyFlag(e, flags))
+			!e.isScope
+			&& (
+				flags & SubscriberFlags.Dirty
+				|| (flags & SubscriberFlags.PendingComputed && updateDirtyFlag(e, flags))
+			)
 		) {
-			executeEffect(e);
+			runEffect(e);
 		} else {
 			processPendingInnerEffects(e, e.flags);
 		}
@@ -63,6 +67,7 @@ const pauseStack: (Subscriber | undefined)[] = [];
 
 let batchDepth = 0;
 let activeSub: Subscriber | undefined;
+let activeScope: Subscriber | undefined;
 
 //#region Public functions
 export function startBatch() {
@@ -117,14 +122,33 @@ export function effect<T>(fn: () => T): () => void {
 	};
 	if (activeSub !== undefined) {
 		link(e, activeSub);
+	} else if (activeScope !== undefined) {
+		link(e, activeScope);
 	}
-	executeEffect(e);
+	runEffect(e);
+	return effectStop.bind(e);
+}
+
+export function effectScope<T>(fn: () => T): () => void {
+	const e: Effect = {
+		fn,
+		subs: undefined,
+		subsTail: undefined,
+		deps: undefined,
+		depsTail: undefined,
+		flags: SubscriberFlags.Effect,
+		isScope: true,
+	};
+	if (activeSub !== undefined) {
+		link(e, activeSub);
+	}
+	runEffectScope(e);
 	return effectStop.bind(e);
 }
 //#endregion
 
 //#region Internal functions
-function executeEffect(e: Effect): void {
+function runEffect(e: Effect): void {
 	const prevSub = activeSub;
 	activeSub = e;
 	startTracking(e);
@@ -135,6 +159,18 @@ function executeEffect(e: Effect): void {
 		endTracking(e);
 	}
 }
+
+function runEffectScope(e: Effect): void {
+	const prevSub = activeScope;
+	activeScope = e;
+	startTracking(e);
+	try {
+		e.fn();
+	} finally {
+		activeScope = prevSub;
+		endTracking(e);
+	}
+}
 //#endregion
 
 //#region Bound functions
@@ -145,6 +181,8 @@ function computedGetter<T>(this: Computed<T>): T {
 	}
 	if (activeSub !== undefined) {
 		link(this, activeSub);
+	} else if (activeScope !== undefined) {
+		link(this, activeScope);
 	}
 	return this.currentValue!;
 }

tests/effect.spec.ts

@@ -1,5 +1,5 @@
 import { expect, test } from 'vitest';
-import { computed, effect, endBatch, signal, startBatch } from '../src';
+import { computed, effect, effectScope, endBatch, signal, startBatch } from '../src';
 
 test('should clear subscriptions when untracked by all subscribers', () => {
 	let bRunTimes = 0;
@@ -115,3 +115,71 @@ test('should trigger inner effects in sequence', () => {
 
 	expect(order).toEqual(['first inner', 'last inner']);
 });
+
+test('should trigger inner effects in sequence in effect scope', () => {
+	const a = signal(0);
+	const b = signal(0);
+	const order: string[] = [];
+
+	effectScope(() => {
+
+		effect(() => {
+			order.push('first inner');
+			a();
+		});
+
+		effect(() => {
+			order.push('last inner');
+			a();
+			b();
+		});
+	});
+
+	order.length = 0;
+
+	startBatch();
+	b(1);
+	a(1);
+	endBatch();
+
+	expect(order).toEqual(['first inner', 'last inner']);
+});
+
+test('should custom effect support batch', () => {
+	function batchEffect(fn: () => void) {
+		return effect(() => {
+			startBatch();
+			try {
+				return fn();
+			} finally {
+				endBatch();
+			}
+		});
+	}
+
+	const logs: string[] = [];
+	const a = signal(0);
+	const b = signal(0);
+
+	const aa = computed(() => {
+		logs.push('aa-0');
+		if (a() === 0) {
+			b(1);
+		}
+		logs.push('aa-1');
+	});
+
+	const bb = computed(() => {
+		logs.push('bb');
+		return b();
+	});
+
+	batchEffect(() => {
+		bb();
+	});
+	batchEffect(() => {
+		aa();
+	});
+
+	expect(logs).toEqual(['bb', 'aa-0', 'aa-1', 'bb']);
+});

tests/effectScope.spec.ts

@@ -0,0 +1,23 @@
+import { expect, test } from 'vitest';
+import { effect, effectScope, signal } from '../src';
+
+test('should not trigger after stop', () => {
+	const count = signal(1);
+
+	let triggers = 0;
+	let effect1;
+
+	const stopScope = effectScope(() => {
+		effect1 = effect(() => {
+			triggers++;
+			count();
+		});
+	});
+
+	expect(triggers).toBe(1);
+	count(2);
+	expect(triggers).toBe(2);
+	stopScope();
+	count(3);
+	expect(triggers).toBe(2);
+});