refactor: implement effect scope as subscriber

lib/computed.ts

@@ -22,7 +22,7 @@ export class Computed<T = any> implements Dependency, Subscriber {
 	) { }
 
 	get(): T {
-		Dependency.link(this);
+		Dependency.linkSubscriber(this);
 		const dirtyLevel = this.versionOrDirtyLevel;
 		if (dirtyLevel === DirtyLevels.MaybeDirty) {
 			Subscriber.resolveMaybeDirty(this);

lib/effect.ts

@@ -1,16 +1,19 @@
-import { currentEffectScope } from './effectScope';
-import { DirtyLevels, IEffect, Subscriber } from './system';
+import { Dependency, DirtyLevels, IEffect, Subscriber } from './system';
 
 export function effect(fn: () => void) {
 	const e = new Effect(fn);
 	e.run();
 	return e;
 }
 
-export class Effect implements IEffect, Subscriber {
-	scope = currentEffectScope;
+export class Effect implements IEffect, Dependency, Subscriber {
 	nextNotify = undefined;
 
+	// Dependency
+	subs = undefined;
+	subsTail = undefined;
+	subVersion = -1;
+
 	// Subscriber
 	deps = undefined;
 	depsTail = undefined;
@@ -19,7 +22,7 @@ export class Effect implements IEffect, Subscriber {
 	constructor(
 		private fn: () => void
 	) {
-		currentEffectScope?.subs.push(this);
+		Dependency.linkSubscriberScope(this);
 	}
 
 	notify() {
@@ -33,12 +36,7 @@ export class Effect implements IEffect, Subscriber {
 
 	run() {
 		const lastActiveSub = Subscriber.startTrack(this);
-		if (this.scope !== undefined) {
-			this.scope.run(this.fn);
-		}
-		else {
-			this.fn();
-		}
+		this.fn();
 		Subscriber.endTrack(this, lastActiveSub);
 	}
 }

lib/effectScope.ts

@@ -1,27 +1,23 @@
-import { Subscriber } from './system';
-
-export let currentEffectScope: EffectScope | undefined = undefined;
+import { DirtyLevels, Subscriber } from './system';
 
 export function effectScope() {
 	return new EffectScope();
 }
 
-export class EffectScope {
-	subs: Subscriber[] = [];
+export class EffectScope implements Subscriber {
+	// Subscriber
+	deps = undefined;
+	depsTail = undefined;
+	versionOrDirtyLevel = DirtyLevels.NotDirty;
 
 	run<T>(fn: () => T) {
-		const lastEffectScope = currentEffectScope;
-		try {
-			currentEffectScope = this;
-			return fn();
-		} finally {
-			currentEffectScope = lastEffectScope;
-		}
+		const prevActiveSub = Subscriber.startScopeTrack(this);
+		const res = fn();
+		Subscriber.endScopeTrack(this, prevActiveSub);
+		return res;
 	}
 
 	stop() {
-		for (const sub of this.subs) {
-			Subscriber.clearTrack(sub);
-		}
+		Subscriber.clearTrack(this);
 	}
 }

lib/signal.ts

@@ -17,7 +17,7 @@ export class Signal<T = any> implements Dependency {
 	) { }
 
 	get() {
-		Dependency.link(this);
+		Dependency.linkSubscriber(this);
 		return this.currentValue!;
 	}
 

lib/system.ts

@@ -7,6 +7,7 @@ export interface Dependency {
 	subs: Link | undefined;
 	subsTail: Link | undefined;
 	subVersion: number;
+	run?(): void;
 }
 
 export interface Subscriber {
@@ -19,7 +20,6 @@ export interface Subscriber {
 	versionOrDirtyLevel: number | DirtyLevels;
 	deps: Link | undefined;
 	depsTail: Link | undefined;
-	run(): void;
 }
 
 export interface Link {
@@ -40,6 +40,7 @@ export const enum DirtyLevels {
 export namespace System {
 
 	export let activeSub: Subscriber | undefined = undefined;
+	export let activeSubScope: Subscriber | undefined = undefined;
 	export let activeSubsDepth = 0;
 	export let batchDepth = 0;
 	export let subVersion = DirtyLevels.Dirty + 1;
@@ -142,7 +143,7 @@ export namespace Dependency {
 
 	const system = System;
 
-	export function link(dep: Dependency) {
+	export function linkSubscriber(dep: Dependency) {
 		if (system.activeSubsDepth === 0) {
 			return;
 		}
@@ -182,6 +183,49 @@ export namespace Dependency {
 		}
 	}
 
+	/**
+	 * @deprecated TODO: Reuse linkSubscriber without performance regression
+	 */
+	export function linkSubscriberScope(dep: Dependency) {
+		if (system.activeSubScope === undefined) {
+			return;
+		}
+		const sub = system.activeSubScope!;
+		const subVersion = sub.versionOrDirtyLevel;
+		if (dep.subVersion === subVersion) {
+			return;
+		}
+		dep.subVersion = subVersion;
+
+		const depsTail = sub.depsTail;
+		const old = depsTail !== undefined
+			? depsTail.nextDep
+			: sub.deps;
+
+		if (old === undefined || old.dep !== dep) {
+			const newLink = Link.get(dep, sub);
+			if (old !== undefined) {
+				newLink.nextDep = old;
+			}
+			if (depsTail === undefined) {
+				sub.depsTail = sub.deps = newLink;
+			} else {
+				sub.depsTail = depsTail.nextDep = newLink;
+			}
+			if (dep.subs === undefined) {
+				dep.subs = newLink;
+				dep.subsTail = newLink;
+			} else {
+				const oldTail = dep.subsTail!;
+				newLink.prevSubOrUpdate = oldTail;
+				oldTail.nextSub = newLink;
+				dep.subsTail = newLink;
+			}
+		} else {
+			sub.depsTail = old;
+		}
+	}
+
 	export function propagate(dep: Dependency) {
 		let dirtyLevel = DirtyLevels.Dirty;
 		let currentSubs = dep.subs;
@@ -240,7 +284,7 @@ export namespace Subscriber {
 
 	const system = System;
 
-	export function resolveMaybeDirty(sub: Subscriber) {
+	export function resolveMaybeDirty(sub: Dependency & Subscriber) {
 		let link = sub.deps;
 
 		top: while (true) {
@@ -258,7 +302,7 @@ export namespace Subscriber {
 
 						continue top;
 					} else if (depDirtyLevel === DirtyLevels.Dirty) {
-						dep.run();
+						dep.run!();
 
 						if ((sub.versionOrDirtyLevel as DirtyLevels) === DirtyLevels.Dirty) {
 							break;
@@ -282,11 +326,11 @@ export namespace Subscriber {
 
 				if (prevLink !== undefined) {
 					if (dirtyLevel === DirtyLevels.Dirty) {
-						sub.run();
+						sub.run!();
 					}
 
 					subSubs.prevSubOrUpdate = undefined;
-					sub = prevLink.sub;
+					sub = prevLink.sub as Dependency & Subscriber;
 					link = prevLink.nextDep;
 
 					continue;
@@ -311,6 +355,18 @@ export namespace Subscriber {
 		system.activeSub = lastActiveSub;
 	}
 
+	export function startScopeTrack(sub: Subscriber) {
+		const lastActiveSub = system.activeSubScope;
+		system.activeSubScope = sub;
+		preTrack(sub);
+		return lastActiveSub;
+	}
+
+	export function endScopeTrack(sub: Subscriber, lastActiveSub: Subscriber | undefined) {
+		postTrack(sub);
+		system.activeSubScope = lastActiveSub;
+	}
+
 	export function clearTrack(sub: Subscriber) {
 		preTrack(sub);
 		postTrack(sub);

unstable/vue.ts

@@ -1,6 +1,5 @@
 import {
 	Computed,
-	currentEffectScope,
 	Dependency,
 	DirtyLevels,
 	Effect,
@@ -16,9 +15,19 @@ export function effect(fn: () => void) {
 	return new ReactiveEffect(fn);
 }
 
+let currentEffectScope: VueEffectScope | undefined = undefined;
+
 class VueEffectScope extends EffectScope {
 	onDispose: (() => void)[] = [];
 
+	run<T>(fn: () => T): T {
+		const prevScope = currentEffectScope;
+		currentEffectScope = this;
+		const res = super.run(fn);
+		currentEffectScope = prevScope;
+		return res;
+	}
+
 	stop() {
 		super.stop();
 		this.onDispose.forEach(cb => cb());
@@ -93,7 +102,5 @@ export class ReactiveEffect extends Effect {
 }
 
 export function onScopeDispose(cb: () => void) {
-	if (currentEffectScope instanceof VueEffectScope) {
-		currentEffectScope.onDispose.push(cb);
-	}
+	currentEffectScope?.onDispose.push(cb);
 }