From 7aadb446489838aab10ac13b80ac1849870047ec Mon Sep 17 00:00:00 2001
From: Liam Griffiths <liamgriffiths@gmail.com>
Date: Mon, 24 Jun 2024 13:34:25 -0400
Subject: [PATCH] Add `depends` option to set edges without future inputs

---
 src/Node.ts             | 19 ++++++++++++++++--
 src/Substrate.ts        | 13 +++++++++++-
 tests/Node.test.ts      | 11 ++++++-----
 tests/Substrate.test.ts | 44 +++++++++++++++++++++++++++++++++++++++++
 4 files changed, 79 insertions(+), 8 deletions(-)

diff --git a/src/Node.ts b/src/Node.ts
index f6edb0c..66809a8 100644
--- a/src/Node.ts
+++ b/src/Node.ts
@@ -11,6 +11,8 @@ export type Options = {
   id?: Node["id"];
   /** When true the server will omit this node's output. Default: false */
   hide?: boolean;
+  /** Specify nodes that this node depends on. */
+  depends?: Node[];
 };
 
 export abstract class Node {
@@ -22,6 +24,8 @@ export abstract class Node {
   args: Object;
   /** When true the server will omit this node's output. Default: false */
   hide: boolean;
+  /** Specify nodes that this node depends on. */
+  depends: Node[];
 
   /** TODO this field stores the last response, but it's just temporary until the internals are refactored */
   protected _response: SubstrateResponse | undefined;
@@ -29,8 +33,9 @@ export abstract class Node {
   constructor(args: Object = {}, opts?: Options) {
     this.node = this.constructor.name;
     this.args = args;
-    this.id = opts?.id || generator(this.node);
-    this.hide = opts?.hide || false;
+    this.id = opts?.id ?? generator(this.node);
+    this.hide = opts?.hide ?? false;
+    this.depends = opts?.depends ?? [];
   }
 
   /**
@@ -116,6 +121,16 @@ export abstract class Node {
 
     nodes.add(this);
 
+    for (let node of this.depends) {
+      const references = node.references();
+      for (let node of references.nodes) {
+        nodes.add(node);
+      }
+      for (let future of references.futures) {
+        futures.add(future);
+      }
+    }
+    
     const collectFutures = (obj: any) => {
       if (Array.isArray(obj)) {
         for (let item of obj) {
diff --git a/src/Substrate.ts b/src/Substrate.ts
index a103ffb..ab70438 100644
--- a/src/Substrate.ts
+++ b/src/Substrate.ts
@@ -198,10 +198,21 @@ export class Substrate {
       }
     }
 
+    const allEdges: Record<string, Set<string>> = {};
+    for (let n of allNodes) {
+      allEdges[n.id] = new Set<string>();
+      for (let d of n.depends) {
+        allEdges[n.id]!.add(d.id);
+      }
+    }
+
     return {
       nodes: Array.from(allNodes).map((node) => node.toJSON()),
       futures: Array.from(allFutures).map((future) => future.toJSON()),
-      edges: [], // @deprecated
+      edges: Object.keys(allEdges).flatMap((toId: string) => {
+        let fromIds: string[] = Array.from(allEdges[toId] as Set<string>);
+        return fromIds.map((fromId: string) => [fromId, toId]);
+      }),
       initial_args: {}, // @deprecated
     };
   }
diff --git a/tests/Node.test.ts b/tests/Node.test.ts
index 6e912a8..ae2bf8c 100644
--- a/tests/Node.test.ts
+++ b/tests/Node.test.ts
@@ -31,17 +31,18 @@ describe("Node", () => {
   });
 
   test(".references", () => {
-    const a = new FooNode({ x: "x" });
+    const a = new FooNode({ x: "x" }, { id: "a" });
     const f1 = a.future.get("x");
     const f2 = a.future.get("y");
-    const b = new FooNode({ x: f1, z: f2 });
+    const b = new FooNode({ x: f1, z: f2 }, { id: "b" });
     const f3 = b.future.get("x");
-    const c = new FooNode({ x: f3 });
+    const c = new FooNode({ x: f3 }, { id: "c" });
+    const d = new FooNode({}, { id: "d", depends: [c] });
 
     // @ts-ignore (protected)
-    const { nodes, futures } = c.references();
+    const { nodes, futures } = d.references();
 
-    expect(nodes).toEqual(new Set([a, b, c]));
+    expect(nodes).toEqual(new Set([a, b, c, d]));
     expect(futures).toEqual(new Set([f1, f2, f3]));
   });
 });
diff --git a/tests/Substrate.test.ts b/tests/Substrate.test.ts
index b99bc5d..a21c6b9 100644
--- a/tests/Substrate.test.ts
+++ b/tests/Substrate.test.ts
@@ -144,5 +144,49 @@ describe("Substrate", () => {
         ],
       });
     });
+
+    test("when there are nodes and we use the `depends` key", () => {
+      const a = new FooNode({ a: 123 }, { id: "a" });
+      const b = new FooNode({ b: 456 }, { id: "b", depends: [a] });
+      const c = new FooNode({ c: 789 }, { id: "c", depends: [a, b] });
+
+      const result = Substrate.serialize(a, b, c);
+
+      expect(result).toEqual({
+        edges: [
+          ["a", "b"],
+          ["a", "c"],
+          ["b", "c"],
+        ],
+        initial_args: {},
+        nodes: [
+          {
+            node: "FooNode",
+            id: a.id,
+            args: {
+              a: 123,
+            },
+            _should_output_globally: true,
+          },
+          {
+            node: "FooNode",
+            id: b.id,
+            args: {
+              b: 456,
+            },
+            _should_output_globally: true,
+          },
+          {
+            node: "FooNode",
+            id: c.id,
+            args: {
+              c: 789,
+            },
+            _should_output_globally: true,
+          },
+        ],
+        futures: [],
+      });
+    });
   });
 });