codepod-io · lihebi · Apr 27, 2023 · Apr 27, 2023
diff --git a/ui/src/components/nodes/Scope.tsx b/ui/src/components/nodes/Scope.tsx
@@ -33,6 +33,7 @@ import Grid from "@mui/material/Grid";
 import DeleteIcon from "@mui/icons-material/Delete";
 import ContentCutIcon from "@mui/icons-material/ContentCut";
 import ContentCopyIcon from "@mui/icons-material/ContentCopy";
+import PlayCircleOutlineIcon from "@mui/icons-material/PlayCircleOutline";
 import Moveable from "react-moveable";
 
 import { useStore } from "zustand";
@@ -51,6 +52,7 @@ function MyFloatingToolbar({ id }: { id: string }) {
   const reactFlowInstance = useReactFlow();
   // const selected = useStore(store, (state) => state.pods[id]?.selected);
   const isGuest = useStore(store, (state) => state.role === "GUEST");
+  const wsRunScope = useStore(store, (state) => state.wsRunScope);
   const clonePod = useStore(store, (state) => state.clonePod);
 
   const onCopy = useCallback(
@@ -81,6 +83,18 @@ function MyFloatingToolbar({ id }: { id: string }) {
   );
   return (
     <Box>
+      {!isGuest && (
+        <Tooltip title="Run (shift-enter)">
+          <IconButton
+            size="small"
+            onClick={() => {
+              wsRunScope(id);
+            }}
+          >
+            <PlayCircleOutlineIcon fontSize="inherit" />
+          </IconButton>
+        </Tooltip>
+      )}
       <CopyToClipboard
         text="dummy"
         options={{ debug: true, format: "text/plain", onCopy } as any}

diff --git a/ui/src/lib/store/canvasSlice.tsx b/ui/src/lib/store/canvasSlice.tsx
@@ -275,6 +275,8 @@ export interface CanvasSlice {
   onEdgesChange: (client: ApolloClient<any>) => OnEdgesChange;
   onConnect: (client: ApolloClient<any>) => OnConnect;
 
+  node2children: Map<string, string[]>;
+  buildNode2Children: () => void;
   autoLayout: () => void;
 }
 
@@ -808,9 +810,16 @@ export const createCanvasSlice: StateCreator<MyState, [], [], CanvasSlice> = (
   },
   setPaneFocus: () => set({ isPaneFocused: true }),
   setPaneBlur: () => set({ isPaneFocused: false }),
-  autoLayout: () => {
-    // Auto layout the nodes.
-    // Find all the scope nodes, and change its width and height to fit its children nodes.
+  /**
+   * This node2children is maintained with the canvas reactflow states, not with
+   * the pods. This mapping may be used by other components, e.g. the runtime.
+   *
+   * TODO we should optimize the performance of this function, maybe only update
+   * the mapping when the structure is changed.
+   */
+  node2children: new Map<string, string[]>(),
+  buildNode2Children: () => {
+    // build a map from node to its children
     let nodesMap = get().ydoc.getMap<Node>("pods");
     let nodes: Node[] = Array.from(nodesMap.values());
     let node2children = new Map<string, string[]>();
@@ -825,12 +834,20 @@ export const createCanvasSlice: StateCreator<MyState, [], [], CanvasSlice> = (
         node2children.get(node.parentNode)?.push(node.id);
       }
     });
+    set({ node2children });
+  },
+  autoLayout: () => {
+    // Auto layout the nodes.
+    // Find all the scope nodes, and change its width and height to fit its children nodes.
+    let nodesMap = get().ydoc.getMap<Node>("pods");
+    let nodes: Node[] = Array.from(nodesMap.values());
+    get().buildNode2Children();
     // fit the children.
     nodes
       // sort the children so that the inner scope gets processed first.
       .sort((a: Node, b: Node) => b.data.level - a.data.level)
       .forEach((node) => {
-        let newSize = fitChildren(node, node2children, nodesMap);
+        let newSize = fitChildren(node, get().node2children, nodesMap);
         if (newSize === null) return;
         let { x, y, width, height } = newSize;
         let newNode = {
@@ -849,7 +866,7 @@ export const createCanvasSlice: StateCreator<MyState, [], [], CanvasSlice> = (
         };
         nodesMap.set(node.id, newNode);
         // I actually need to set the children's position as well, because they are relative to the parent.
-        let children = node2children.get(node.id);
+        let children = get().node2children.get(node.id);
         children?.forEach((child) => {
           let n = nodesMap.get(child)!;
           let newChild = {

diff --git a/ui/src/lib/store/runtimeSlice.tsx b/ui/src/lib/store/runtimeSlice.tsx
@@ -2,7 +2,7 @@ import produce from "immer";
 import { ApolloClient, gql } from "@apollo/client";
 import { createStore, StateCreator, StoreApi } from "zustand";
 
-import { Edge } from "reactflow";
+import { Edge, Node } from "reactflow";
 
 // FIXME cyclic import
 import { MyState } from ".";
@@ -167,6 +167,7 @@ export interface RuntimeSlice {
   resolveAllPods: () => void;
   runningId: string | null;
   wsRun: (id: string) => void;
+  wsRunScope: (id: string) => void;
   wsSendRun: (id: string) => void;
   wsRunNext: () => void;
   wsRunNoRewrite: (id: string) => void;
@@ -314,13 +315,44 @@ export const createRuntimeSlice: StateCreator<MyState, [], [], RuntimeSlice> = (
    * Add a pod to the chain and run it.
    */
   wsRun: async (id) => {
-    // Add to the chain
-    get().clearResults(id);
-    get().setRunning(id);
-    set({ chain: [...get().chain, id] });
-    // if there's nothing running, run it
+    // If this pod is a code pod, add it.
+    if (get().pods[id].type === "CODE") {
+      // Add to the chain
+      get().clearResults(id);
+      get().setRunning(id);
+      set({ chain: [...get().chain, id] });
+    } else if (get().pods[id].type === "DECK") {
+      // If this pod is a scope, run all pods inside a scope by geographical order.
+      // get the pods in the scope
+      let children = get().node2children.get(id);
+      if (!children) return;
+      // The reactflow nodesMap stored in Yjs
+      let nodesMap = get().ydoc.getMap<Node>("pods");
+      // Sort by x and y positions, with the leftmost and topmost first.
+      children = [...children].sort((a, b) => {
+        let nodeA = nodesMap.get(a);
+        let nodeB = nodesMap.get(b);
+        if (nodeA && nodeB) {
+          if (nodeA.position.y === nodeB.position.y) {
+            return nodeA.position.x - nodeB.position.x;
+          } else {
+            return nodeA.position.y - nodeB.position.y;
+          }
+        } else {
+          return 0;
+        }
+      });
+      // add to the chain
+      // set({ chain: [...get().chain, ...children.map(({ id }) => id)] });
+      children.forEach((id) => get().wsRun(id));
+    }
     get().wsRunNext();
   },
+  wsRunScope: async (id) => {
+    // This is a separate function only because we need to build the node2children map first.
+    get().buildNode2Children();
+    get().wsRun(id);
+  },
   /**
    * Add the pod and all its downstream pods (defined by edges) to the chain and run the chain.
    * @param id the id of the pod to start the chain