Misc cleanup

src/Environment.ts

@@ -19,7 +19,7 @@ import { ProbingContext } from './ApiTypes';
 export type DisposeOp = () => void;
 export interface Environment {
     _onDispose: (op: DisposeOp) => void;
-    _getProbingContext: () => ProbingContext | undefined;
+    _getProbingContext: () => ProbingContext;
 }
 
 let _currentEnv: Environment | null = null;
@@ -50,10 +50,5 @@ export function useProbingContext(): ProbingContext {
         throw new Error('Environment underflow');
     }
 
-    const result = _currentEnv!._getProbingContext();
-    if (process.env.NODE_ENV !== 'production' && !result) {
-        throw new Error('There is no active probing context');
-    }
-
-    return result!;
+    return _currentEnv!._getProbingContext();
 }

src/Node.ts

@@ -25,8 +25,8 @@ export interface NodeBuildData {
     _cb: (...arg: unknown[]) => unknown;
     _args: unknown[];
 
-    _next?: IPNode;
-    _resolveAs?: IPNode;
+    _next?: BaseNode;
+    _resolveAs: BaseNode;
     _prober: IProber;
     _context: ProbingContext;
 }
@@ -43,17 +43,24 @@ export abstract class BaseNode implements IPNode {
     }
 
     dispose(): void {
-        if (this._onDispose) {
-            this._onDispose.forEach((c) => c());
-            this._onDispose = undefined;
-        }
+        // Nodes should only ever be disposed once
+        this._onDispose.forEach((c) => c());
+    }
+
+    _addToDispose(ops: DisposeOp[]): void {
+        this._onDispose.push(...ops);
+    }
+
+    _assign(rhs: BaseNode): void {
+        this._result = rhs._result;
+        this._onDispose.push(...rhs._onDispose);
     }
 
     abstract get result(): unknown;
 
     _result?: unknown;
     _probed_pnodetype?: number;
-    _onDispose?: DisposeOp[];
+    _onDispose: DisposeOp[] = [];
 
     _buildData?: NodeBuildData;
     _uniqueNodeId?: number;

src/Prober.ts

@@ -20,14 +20,6 @@ const isIntrinsic = <I>(cb: IKeys<I> | ((...args: unknown[]) => unknown)): cb is
     return typeof cb === 'string';
 };
 
-const addToDisposeQueue = (node: BaseNode, ops: DisposeOp[]) => {
-    if (!node._onDispose) {
-        node._onDispose = ops;
-    } else {
-        node._onDispose = node._onDispose.concat(ops);
-    }
-};
-
 let _NextUniqueNodeId = 0;
 
 interface NodeQueue {
@@ -45,14 +37,14 @@ class Prober<I extends FuncMap> implements IProber {
     private _intrinsics: Partial<I>;
     private _fallback?: IntrinsicFallback<I>;
     private _stack: ProberStackFrame[] = [];
-    private _current: ProberStackFrame = { _disposeOps: [] };
+    private _currentFrame: ProberStackFrame = { _disposeOps: [] };
 
     _onDispose(op: DisposeOp): void {
-        this._current._disposeOps.push(op);
+        this._currentFrame._disposeOps.push(op);
     }
 
-    _getProbingContext(): ProbingContext | undefined {
-        return this._current!._node!._buildData!._context;
+    _getProbingContext(): ProbingContext {
+        return this._currentFrame!._node!._buildData!._context;
     }
 
     constructor(intrinsics: Partial<I>, fallback?: IntrinsicFallback<I>) {
@@ -79,14 +71,15 @@ class Prober<I extends FuncMap> implements IProber {
             newNode._uniqueNodeId = _NextUniqueNodeId++;
         }
 
-        if (!this._current._announced) {
-            this._current._announced = { _head: newNode, _tail: newNode };
+        if (!this._currentFrame._announced) {
+            this._currentFrame._announced = { _head: newNode, _tail: newNode };
         } else {
-            this._current._announced._tail._buildData!._next = newNode;
-            this._current._announced._tail = newNode;
+            this._currentFrame._announced._tail._buildData!._next = newNode;
+            this._currentFrame._announced._tail = newNode;
         }
 
         newNode._buildData = {
+            _resolveAs: newNode,
             _cb,
             _prober: this,
             _args,
@@ -100,41 +93,33 @@ class Prober<I extends FuncMap> implements IProber {
 
     _finalizeNode(node: BaseNode) {
         // If a component returns a Node (as opposed to a value), then we short-circuit to the parent.
-        let destinationNode = node;
-        while (destinationNode._buildData && destinationNode._buildData._resolveAs) {
-            destinationNode = destinationNode._buildData!._resolveAs;
-        }
+        const bd = node._buildData!;
+        const destinationNode = bd._resolveAs;
 
-        this._current._node = node;
-        const { _cb, _args } = node._buildData!;
-        const cbResult = _cb(..._args, node._buildData!._context);
+        this._currentFrame._node = node;
+        const cbResult = bd._cb(...bd._args, bd._context);
 
         if (isPNode(cbResult)) {
             if (cbResult.finalized) {
-                // Post-ex-facto proxying.
-                destinationNode._result = cbResult._result;
-                if (cbResult._onDispose) {
-                    addToDisposeQueue(destinationNode, cbResult._onDispose);
-                    cbResult._onDispose = [];
-                }
+                destinationNode._assign(cbResult);
             } else {
                 cbResult._buildData!._resolveAs = destinationNode;
             }
         } else {
             destinationNode._result = cbResult;
         }
 
-        if (this._current._disposeOps.length > 0) {
-            addToDisposeQueue(destinationNode, this._current._disposeOps);
-            this._current._disposeOps = [];
+        if (this._currentFrame._disposeOps.length > 0) {
+            destinationNode._addToDispose(this._currentFrame._disposeOps);
+            this._currentFrame._disposeOps = [];
         }
     }
 
     _finalize(target: IPNode): void {
         if (process.env.NODE_ENV === 'check') {
             let lookup: BaseNode | undefined;
-            if (this._current._announced) {
-                lookup = this._current._announced._head;
+            if (this._currentFrame._announced) {
+                lookup = this._currentFrame._announced._head;
             }
             while (lookup && lookup !== target) {
                 lookup = lookup._buildData!._next;
@@ -144,58 +129,51 @@ class Prober<I extends FuncMap> implements IProber {
             }
         }
 
-        let node: BaseNode = this._current._announced!._head!;
-        let end: BaseNode = target as BaseNode;
+        let node = this._currentFrame._announced!._head!;
+        let end = target as BaseNode;
 
         if (end._buildData!._next) {
-            this._current._announced!._head = end._buildData!._next;
+            this._currentFrame._announced!._head = end._buildData!._next;
         } else {
-            this._current._announced = undefined;
+            this._currentFrame._announced = undefined;
         }
         end._buildData!._next = undefined;
 
-        /*
-        //These two steps are, I suspect, Technically unnecessary
-
-        if (!this._current._announced._head) {
-            this._current._announced = {};
-        }
-        */
         pushEnv(this);
-        this._stack.push(this._current);
-        this._current = { _node: node, _disposeOps: [] };
+        this._stack.push(this._currentFrame);
+        this._currentFrame = { _node: node, _disposeOps: [] };
 
         let done = false;
         while (!done) {
             this._finalizeNode(node);
 
-            // Queue up any work that was discovered in the process.
-            if (this._current._announced) {
-                end = this._current._announced._tail;
-                node._buildData!._next = this._current._announced._head;
-                this._current._announced = undefined;
+            // Queue up any work that was discovered in the process,
+            // and update our end target so that we complete it before
+            // returning.
+            if (this._currentFrame._announced) {
+                end = this._currentFrame._announced._tail;
+                node._buildData!._next = this._currentFrame._announced._head;
+                this._currentFrame._announced = undefined;
             }
 
             done = node === end;
-            const nextNode = node._buildData!._next as BaseNode;
+            const nextNode = node._buildData!._next;
             node._buildData = undefined;
-            node = nextNode;
+            node = nextNode!;
         }
 
-        this._current = this._stack.pop()!;
+        this._currentFrame = this._stack.pop()!;
         popEnv();
     }
 
     private _getCb<T extends IKeys<I> | ComponentCb>(what: T): { _cb: ComponentCb; _name: string } {
         if (isIntrinsic<I>(what)) {
             let _cb: ComponentCb | undefined = this._intrinsics[what];
-            const _name = what.toString();
             if (!_cb) {
-                // This is safe, it's caught at the start of _announce()
                 _cb = this._fallback!;
             }
 
-            return { _cb: _cb!, _name };
+            return { _cb: _cb!, _name: what.toString() };
         } else {
             return { _cb: what as ComponentCb, _name: (what as ComponentCb).name };
         }

tests/dynamicList.test.ts

@@ -27,7 +27,7 @@ const cleanup = () => {
 beforeEach(() => {
     pushEnv({
         _onDispose: (op: DisposeOp) => disposeQueue.push(op),
-        _getProbingContext: () => undefined,
+        _getProbingContext: () => ({ componentName: '' }),
     });
 });
 

tests/dynamicOperations.test.ts

@@ -27,7 +27,7 @@ const cleanup = () => {
 beforeEach(() => {
     pushEnv({
         _onDispose: (op: DisposeOp) => disposeQueue.push(op),
-        _getProbingContext: () => undefined,
+        _getProbingContext: () => ({ componentName: '' }),
     });
 });
 

tests/dynamicVal.test.ts

@@ -26,7 +26,7 @@ const cleanup = () => {
 beforeEach(() => {
     pushEnv({
         _onDispose: (op: DisposeOp) => disposeQueue.push(op),
-        _getProbingContext: () => undefined,
+        _getProbingContext: () => ({ componentName: '' }),
     });
 });
 

tests/environment.test.ts

@@ -20,12 +20,12 @@ describe('useOnDispose ', () => {
 
 class TestEnv implements Environment {
     count = 0;
-    probeContext?: ProbingContext;
+    probeContext: ProbingContext = { componentName: '' };
     _onDispose(): void {
         this.count += 1;
     }
 
-    _getProbingContext(): ProbingContext | undefined {
+    _getProbingContext(): ProbingContext {
         return this.probeContext;
     }
 }
@@ -77,13 +77,6 @@ describe('useProbingContext ', () => {
         expectThrowInNotProd(useProbingContext);
     });
 
-    it('Fails if probing context is set', () => {
-        const env = new TestEnv();
-        pushEnv(env);
-        expectThrowInNotProd(useProbingContext);
-        popEnv();
-    });
-
     it('Works if a probing context is set', () => {
         const env = new TestEnv();
         env.probeContext = { componentName: 'aaa' };