CSE machine: support for StreamViz drawing arrows

src/createContext.ts

@@ -139,6 +139,8 @@
     runtime: createEmptyRuntime(),
     numberOfOuterEnvironments: 1,
     prelude: null,
+    pendingStreamFnStack: [],
+    streamLineage: new Map<string, string[]>(),
     debugger: createEmptyDebugger(),
     nativeStorage: createNativeStorage(),
     executionMethod: 'auto',
@@ -363,7 +365,7 @@
         'call_cc(f)',
         context.variant === Variant.EXPLICIT_CONTROL
           ? call_with_current_continuation
           : (f: any) => {
               throw new Error('call_cc is only available in Explicit-Control variant');
             },
       );

src/cse-machine/__tests__/cse-machine-runtime-context.test.ts

@@ -210,3 +210,43 @@ test('Avoid unnescessary environment instruction 4', () => {
     expect(state.control.getNumEnvDependentItems()).toMatchSnapshot();
   }
 });
+
+test('CSE Machine correctly updates streamLineage with simple streams', async () => {
+  const context = await getContextFrom(
+    stripIndent`
+      const s1 = pair(1, () => pair(2, () => pair(3, () => null)));
+      eval_stream(s1, 3);
+    `,
+    172,
+  );
+
+  expect(context.streamLineage.size).toEqual(2);
+});
+
+test('CSE Machine correctly updates streamLineage with self-referential streams', async () => {
+  const context = await getContextFrom(
+    stripIndent`
+      const s1 = pair(1, () => s1);
+      eval_stream(s1, 3);
+    `,
+    164,
+  );
+
+  expect((context.streamLineage.values().next().value as Array<string>).length).toEqual(2);
+});
+
+test('CSE Machine does not update streamLineage when nullary function does not return a stream', async () => {
+  const context = await getContextFrom(
+    stripIndent`
+      const s1 = pair(1, () => [0, 1, 2, () => 1]);
+      stream_tail(s1);
+      const s2 = pair(1, x => s2);
+      tail(s2)(1);
+      const s3 = pair(1, () => 2);
+      stream_tail(s3);
+    `,
+    115,
+  );
+
+  expect(context.streamLineage.size).toEqual(0);
+});

src/cse-machine/interpreter.ts

@@ -359,6 +359,25 @@ export function* generateCSEMachineStateStream(
       context.runtime.changepointSteps.push(steps + 1);
     }
 
+    const evalResult = context.runtime.stash?.peek();
+    const mostRecentControlHeight =
+      context.pendingStreamFnStack[context.pendingStreamFnStack.length - 1]?.[1];
+
+    if (
+      Array.isArray(evalResult) &&
+      evalResult.length === 2 &&
+      mostRecentControlHeight !== undefined &&
+      context.runtime.control?.size() === mostRecentControlHeight - 1
+    ) {
+      const mostRecentNullaryFnId = context.pendingStreamFnStack.pop()?.[0];
+      if (mostRecentNullaryFnId !== undefined) {
+        if (!context.streamLineage.get(mostRecentNullaryFnId)) {
+          context.streamLineage.set(mostRecentNullaryFnId, []);
+        }
+        context.streamLineage.get(mostRecentNullaryFnId)?.push((evalResult as any).id);
+      }
+    }
+
     control.pop();
     if (isNode(command)) {
       context.runtime.nodes.shift();
@@ -834,6 +853,16 @@ const cmdEvaluators: CommandEvaluators = {
         context.runtime.environments.unshift(func.environment);
       }
 
+      // Streams Visualisation:
+      // If CALL0, push the nullary fn to pendingStreamFnStack
+      // (to map which nullary fn produces which pairs; for nested nullary fn)
+      if (func.node.params.length === 0) {
+        context.pendingStreamFnStack.push([
+          func.id,
+          context.runtime.control !== null ? context.runtime.control.size() : 0,
+        ]);
+      }
+
       // Handle special case if function is simple
       if (isSimpleFunction(func.node)) {
         // Closures convert ArrowExpressionStatements to BlockStatements

src/types.ts

@@ -65,6 +65,11 @@ export interface Context<T = any> {
     changepointSteps: number[];
   };
 
+  // STREAM VISUALISATION
+  streamLineage: Map<string, string[]>;
+
+  pendingStreamFnStack: [string, number][];
+
   numberOfOuterEnvironments: number;
 
   prelude: string | null;