Flow graph node caching (#14453)

* v1 * v2 * common cache class * change how cache is saved so it doesn't grow data infinitely * Update packages/dev/core/src/FlowGraph/flowGraphContext.ts Co-authored-by: Popov72 <github@evpopov.com> * Update packages/dev/core/src/FlowGraph/flowGraphContext.ts Co-authored-by: Popov72 <github@evpopov.com> * review changes --------- Co-authored-by: Popov72 <github@evpopov.com>
BabylonJS · Oct 26, 2023 · 5a14b51 · 5a14b51
1 parent f87bd90
commit 5a14b51
Show file tree

Hide file tree

Showing 8 changed files with 126 additions and 26 deletions.
diff --git a/packages/dev/core/src/FlowGraph/Blocks/Data/flowGraphBinaryOperationBlock.ts b/packages/dev/core/src/FlowGraph/Blocks/Data/flowGraphBinaryOperationBlock.ts
@@ -1,17 +1,16 @@
 import type { FlowGraphDataConnection } from "../../flowGraphDataConnection";
-import { FlowGraphBlock } from "../../flowGraphBlock";
 import type { RichType } from "../../flowGraphRichTypes";
 import type { FlowGraphContext } from "../../flowGraphContext";
 import type { IFlowGraphBlockConfiguration } from "../../flowGraphBlock";
+import { FlowGraphCachedOperationBlock } from "./flowGraphCachedOperationBlock";
 /**
  * @experimental
  * The base block for all binary operation blocks. Receives an input of type
  * LeftT, one of type RightT, and outputs a value of type ResultT.
  */
-export class FlowGraphBinaryOperationBlock<LeftT, RightT, ResultT> extends FlowGraphBlock {
+export class FlowGraphBinaryOperationBlock<LeftT, RightT, ResultT> extends FlowGraphCachedOperationBlock<ResultT> {
     leftInput: FlowGraphDataConnection<LeftT>;
     rightInput: FlowGraphDataConnection<RightT>;
-    output: FlowGraphDataConnection<ResultT>;
 
     constructor(
         leftRichType: RichType<LeftT>,
@@ -21,14 +20,13 @@ export class FlowGraphBinaryOperationBlock<LeftT, RightT, ResultT> extends FlowG
         private _className: string,
         config?: IFlowGraphBlockConfiguration
     ) {
-        super(config);
+        super(resultRichType, config);
         this.leftInput = this._registerDataInput("leftInput", leftRichType);
         this.rightInput = this._registerDataInput("rightInput", rightRichType);
-        this.output = this._registerDataOutput("Output", resultRichType);
     }
 
-    public _updateOutputs(_context: FlowGraphContext): void {
-        this.output.setValue(this._operation(this.leftInput.getValue(_context), this.rightInput.getValue(_context)), _context);
+    public override _doOperation(context: FlowGraphContext): ResultT {
+        return this._operation(this.leftInput.getValue(context), this.rightInput.getValue(context));
     }
 
     public getClassName(): string {

diff --git a/packages/dev/core/src/FlowGraph/Blocks/Data/flowGraphCachedOperationBlock.ts b/packages/dev/core/src/FlowGraph/Blocks/Data/flowGraphCachedOperationBlock.ts
@@ -0,0 +1,44 @@
+import type { IFlowGraphBlockConfiguration } from "../../flowGraphBlock";
+import { FlowGraphBlock } from "../../flowGraphBlock";
+import type { FlowGraphContext } from "../../flowGraphContext";
+import type { FlowGraphDataConnection } from "../../flowGraphDataConnection";
+import type { RichType } from "../../flowGraphRichTypes";
+
+const CACHE_NAME = "cachedOperationValue";
+const CACHE_EXEC_ID_NAME = "cachedExecutionId";
+
+/**
+ * @experimental
+ */
+export abstract class FlowGraphCachedOperationBlock<OutputT> extends FlowGraphBlock {
+    /**
+     * The output of the operation
+     */
+    public readonly output: FlowGraphDataConnection<OutputT>;
+
+    constructor(outputRichType: RichType<OutputT>, config?: IFlowGraphBlockConfiguration) {
+        super(config);
+
+        this.output = this._registerDataOutput("output", outputRichType);
+    }
+
+    /**
+     * @internal
+     * Operation to realize
+     * @param context the graph context
+     */
+    public abstract _doOperation(context: FlowGraphContext): OutputT;
+
+    public _updateOutputs(context: FlowGraphContext) {
+        const cachedExecutionId = context._getExecutionVariable(this, CACHE_EXEC_ID_NAME);
+        const cachedValue = context._getExecutionVariable(this, CACHE_NAME);
+        if (cachedValue !== undefined && cachedExecutionId === context.executionId) {
+            this.output.setValue(cachedValue, context);
+        } else {
+            const calculatedValue = this._doOperation(context);
+            context._setExecutionVariable(this, CACHE_NAME, calculatedValue);
+            context._setExecutionVariable(this, CACHE_EXEC_ID_NAME, context.executionId);
+            this.output.setValue(calculatedValue, context);
+        }
+    }
+}
diff --git a/packages/dev/core/src/FlowGraph/Blocks/Data/flowGraphConstantOperationBlock.ts b/packages/dev/core/src/FlowGraph/Blocks/Data/flowGraphConstantOperationBlock.ts
@@ -1,22 +1,18 @@
 import type { FlowGraphContext } from "../../flowGraphContext";
-import { FlowGraphBlock } from "../../flowGraphBlock";
-import type { FlowGraphDataConnection } from "../../flowGraphDataConnection";
 import type { RichType } from "../../flowGraphRichTypes";
 import type { IFlowGraphBlockConfiguration } from "../../flowGraphBlock";
+import { FlowGraphCachedOperationBlock } from "./flowGraphCachedOperationBlock";
 /**
  * @experimental
  * Block that outputs a value of type ResultT, resulting of an operation with no inputs.
  */
-export class FlowGraphConstantOperationBlock<ResultT> extends FlowGraphBlock {
-    public output: FlowGraphDataConnection<ResultT>;
-
+export class FlowGraphConstantOperationBlock<ResultT> extends FlowGraphCachedOperationBlock<ResultT> {
     constructor(richType: RichType<ResultT>, private _operation: () => ResultT, private _className: string, config?: IFlowGraphBlockConfiguration) {
-        super(config);
-        this.output = this._registerDataOutput("output", richType);
+        super(richType, config);
     }
 
-    public _updateOutputs(context: FlowGraphContext): void {
-        this.output.setValue(this._operation(), context);
+    public override _doOperation(_context: FlowGraphContext): ResultT {
+        return this._operation();
     }
 
     public getClassName(): string {

diff --git a/packages/dev/core/src/FlowGraph/Blocks/Data/flowGraphUnaryOperationBlock.ts b/packages/dev/core/src/FlowGraph/Blocks/Data/flowGraphUnaryOperationBlock.ts
@@ -1,16 +1,15 @@
 import type { FlowGraphDataConnection } from "../../flowGraphDataConnection";
 import type { IFlowGraphBlockConfiguration } from "../../flowGraphBlock";
-import { FlowGraphBlock } from "../../flowGraphBlock";
 import type { RichType } from "../../flowGraphRichTypes";
 import type { FlowGraphContext } from "../../flowGraphContext";
+import { FlowGraphCachedOperationBlock } from "./flowGraphCachedOperationBlock";
 
 /**
  * @experimental
  * The base block for all unary operation blocks. Receives an input of type InputT, and outputs a value of type ResultT.
  */
-export class FlowGraphUnaryOperationBlock<InputT, ResultT> extends FlowGraphBlock {
+export class FlowGraphUnaryOperationBlock<InputT, ResultT> extends FlowGraphCachedOperationBlock<ResultT> {
     input: FlowGraphDataConnection<InputT>;
-    output: FlowGraphDataConnection<ResultT>;
 
     constructor(
         inputRichType: RichType<InputT>,
@@ -19,13 +18,11 @@ export class FlowGraphUnaryOperationBlock<InputT, ResultT> extends FlowGraphBloc
         private _className: string,
         config?: IFlowGraphBlockConfiguration
     ) {
-        super(config);
+        super(resultRichType, config);
         this.input = this._registerDataInput("input", inputRichType);
-        this.output = this._registerDataOutput("resultOutput", resultRichType);
     }
-
-    public _updateOutputs(_context: FlowGraphContext): void {
-        this.output.setValue(this._operation(this.input.getValue(_context)), _context);
+    public override _doOperation(context: FlowGraphContext): ResultT {
+        return this._operation(this.input.getValue(context));
     }
 
     public getClassName(): string {

diff --git a/packages/dev/core/src/FlowGraph/flowGraphContext.ts b/packages/dev/core/src/FlowGraph/flowGraphContext.ts
@@ -90,6 +90,11 @@ export class FlowGraphContext {
      * These are blocks that have currently pending tasks/listeners that need to be cleaned up.
      */
     private _pendingBlocks: FlowGraphAsyncExecutionBlock[] = [];
+    /**
+     * A monotonically increasing ID for each execution.
+     * Incremented for every block executed.
+     */
+    private _executionId = 0;
 
     constructor(params: IFlowGraphContextConfiguration) {
         this._configuration = params;
@@ -243,6 +248,17 @@ export class FlowGraphContext {
         this._pendingBlocks.length = 0;
     }
 
+    /**
+     * @internal
+     */
+    public _increaseExecutionId() {
+        this._executionId++;
+    }
+
+    public get executionId() {
+        return this._executionId;
+    }
+
     /**
      * Serializes a context
      * @param serializationObject the object to write the values in

diff --git a/packages/dev/core/src/FlowGraph/flowGraphSignalConnection.ts b/packages/dev/core/src/FlowGraph/flowGraphSignalConnection.ts
@@ -24,6 +24,7 @@ export class FlowGraphSignalConnection extends FlowGraphConnection<FlowGraphExec
     public _activateSignal(context: FlowGraphContext): void {
         if (this.connectionType === FlowGraphConnectionType.Input) {
             this._ownerBlock._execute(context, this);
+            context._increaseExecutionId();
         } else {
             this._connectedPoint[0]?._activateSignal(context);
         }

diff --git a/packages/dev/core/test/unit/FlowGraph/flowGraphDataNodes.test.ts b/packages/dev/core/test/unit/FlowGraph/flowGraphDataNodes.test.ts
@@ -1,7 +1,14 @@
 import type { Engine } from "core/Engines";
 import { NullEngine } from "core/Engines";
 import type { FlowGraph, FlowGraphContext } from "core/FlowGraph";
-import { FlowGraphCoordinator, FlowGraphGetVariableBlock, FlowGraphSceneReadyEventBlock, FlowGraphLogBlock } from "core/FlowGraph";
+import {
+    FlowGraphCoordinator,
+    FlowGraphGetVariableBlock,
+    FlowGraphSceneReadyEventBlock,
+    FlowGraphLogBlock,
+    FlowGraphAddNumberBlock,
+    FlowGraphRandomNumberBlock,
+} from "core/FlowGraph";
 import { Scene } from "core/scene";
 
 describe("Flow Graph Data Nodes", () => {
@@ -51,4 +58,46 @@ describe("Flow Graph Data Nodes", () => {
         expect(console.log).toHaveBeenCalledWith(42);
         expect(console.log).toHaveBeenCalledWith(43);
     });
+
+    it("Values are cached for the same execution id", () => {
+        const sceneReady = new FlowGraphSceneReadyEventBlock({ name: "SceneReady" });
+        flowGraph.addEventBlock(sceneReady);
+
+        const add = new FlowGraphAddNumberBlock();
+
+        const rnd = new FlowGraphRandomNumberBlock();
+        rnd.leftInput.setValue(0, flowGraphContext);
+        rnd.rightInput.setValue(1, flowGraphContext);
+
+        // add a number to itself, which should only trigger the random number block once and cache the result
+        add.leftInput.connectTo(rnd.output);
+        add.rightInput.connectTo(rnd.output);
+
+        // log ther result
+        const log = new FlowGraphLogBlock();
+        log.message.connectTo(add.output);
+        sceneReady.onDone.connectTo(log.onStart);
+
+        flowGraph.start();
+
+        let mockRandomIndex = 1;
+        const mockedRandom = (): number => {
+            return mockRandomIndex++;
+        };
+
+        // clear the random mock before calling
+        const random = jest.spyOn(global.Math, "random").mockImplementation(mockedRandom);
+
+        scene.onReadyObservable.notifyObservers(scene);
+
+        expect(random).toHaveBeenCalledTimes(1);
+        expect(console.log).toHaveBeenCalledWith(2); // 1 + 1
+
+        random.mockRestore();
+    });
+
+    afterEach(() => {
+        scene.dispose();
+        engine.dispose();
+    });
 });
diff --git a/packages/dev/core/test/unit/FlowGraph/flowGraphExecutionNodes.test.ts b/packages/dev/core/test/unit/FlowGraph/flowGraphExecutionNodes.test.ts
@@ -155,7 +155,6 @@ describe("Flow Graph Execution Nodes", () => {
         const sceneReady = new FlowGraphSceneReadyEventBlock();
         flowGraph.addEventBlock(sceneReady);
 
-        debugger;
         const switchBlock = new FlowGraphSwitchBlock({ cases: [1, 2, 3] });
         sceneReady.onDone.connectTo(switchBlock.onStart);
         switchBlock.selection.setValue(2, flowGraphContext);