From f62fa6912202a3b4278d2be5ada20af7cde8e08f Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Sat, 29 Apr 2023 11:42:16 -0700
Subject: [PATCH] Autogenerate layout from graph

---
 src/lib/components/ComfyUIPane.svelte     |  51 ++----
 src/lib/components/WidgetContainer.svelte | 168 ++++++++++++--------
 src/lib/stores/layoutState.ts             | 179 ++++++++++++++++++++--
 src/lib/stores/uiState.ts                 |   2 +-
 src/lib/widgets/ComboWidget.svelte        |   2 +-
 5 files changed, 290 insertions(+), 112 deletions(-)

diff --git a/src/lib/components/ComfyUIPane.svelte b/src/lib/components/ComfyUIPane.svelte
index 29df71d..c06972d 100644
--- a/src/lib/components/ComfyUIPane.svelte
+++ b/src/lib/components/ComfyUIPane.svelte
@@ -10,9 +10,6 @@
 
  export let app: ComfyApp;
  let dragConfigured: boolean = false;
- export let dragItems: DragItem[][] = []
-
- export let totalId = 0;
  //
  //  function addUIForNewNode(node: LGraphNode, paneIndex?: number) {
  //      if (!paneIndex)
@@ -32,48 +29,32 @@
  }
 
  export function restore(panels: SerializedPanes) {
-     const a: ContainerLayout = {
-         type: "container",
-         id: "0",
-         attrs: {
-             title: "root!",
-             direction: "horizontal"
-         }
+     const id = 0;
+     $layoutState.root = layoutState.addContainer(null, { direction: "horizontal", showTitle: false });
+     const left = layoutState.addContainer($layoutState.root.id, { direction: "vertical", showTitle: false });
+     const right = layoutState.addContainer($layoutState.root.id, { direction: "vertical", showTitle: false });
+
+     for (const node of app.lGraph.computeExecutionOrder(false, null)) {
+         layoutState.nodeAdded(node)
      }
-     $layoutState.root = a;
-     $layoutState.children[0] = [
-         {
-             type: "widget",
-             id: "1",
-             nodeId: 7,
-             widgetName: "text",
-             attrs: {
-
-             }
-         },
-         {
-             type: "widget",
-             id: "2",
-             nodeId: 6,
-             widgetName: "text",
-             attrs: {
-
-             }
-         },
-     ]
-     $layoutState.children[1] = []
-     $layoutState.children[2] = []
+     console.warn($layoutState)
  }
 </script>
 
 <div id="comfy-ui-panes" >
-    <WidgetContainer bind:dragItem={$layoutState.root} />
+    <WidgetContainer bind:dragItem={$layoutState.root} classes={["root-container"]} />
 </div>
 
-<style>
+<style lang="scss">
  #comfy-ui-panes {
      width: 100%;
      height: 100%;
      overflow: scroll;
+
  }
+
+ #comfy-ui-panes > :global(.root-container > .block) {
+     padding: 0px;
+ }
+
 </style>
diff --git a/src/lib/components/WidgetContainer.svelte b/src/lib/components/WidgetContainer.svelte
index 6c70a56..ce554a1 100644
--- a/src/lib/components/WidgetContainer.svelte
+++ b/src/lib/components/WidgetContainer.svelte
@@ -1,8 +1,5 @@
 <script lang="ts">
- import { onDestroy } from "svelte";
- import { get } from "svelte/store"
  import { Block, BlockTitle } from "@gradio/atoms";
- import { Move } from 'radix-icons-svelte';
  import queueState from "$lib/stores/queueState";
  import nodeState, { type WidgetUIState } from "$lib/stores/nodeState";
  import uiState from "$lib/stores/uiState";
@@ -13,14 +10,15 @@
  // notice - fade in works fine but don't add svelte's fade-out (known issue)
  import {cubicIn} from 'svelte/easing';
  import { flip } from 'svelte/animate';
- import ComfyApp from "./ComfyApp";
- import type { LGraphNode } from "@litegraph-ts/core";
  import layoutState, { type ContainerLayout, type WidgetLayout, type IDragItem } from "$lib/stores/layoutState";
  import { getComponentForWidgetState } from "$lib/utils"
 
- export let dragItem: DragItem | null = null;
+ export let dragItem: IDragItem | null = null;
+ export let zIndex: number = 100;
+ export let classes: string[] = [];
  let container: ContainerLayout | null = null;
- let widget: WidgetUIState | null = null;
+ let widget: WidgetLayout | null = null;
+ let widgetState: WidgetUIState | null = null;
  let children: IDragItem[] | null = null;
  let dragDisabled = true;
  const flipDurationMs = 200;
@@ -28,12 +26,12 @@
  $: if (dragItem) {
      if (dragItem.type === "container") {
          container = dragItem as ContainerLayout;
-         children = $layoutState.children[dragItem.id];
+         children = $layoutState.allItems[dragItem.id].children;
          widget = null;
      }
      else if (dragItem.type === "widget") {
-         const widgetLayout = dragItem as WidgetLayout;
-         widget = nodeState.findWidgetByName(widgetLayout.nodeId, widgetLayout.widgetName)
+         widget = dragItem as WidgetLayout;
+         widgetState = nodeState.findWidgetByName(widget.nodeId, widget.widgetName)
          children = null;
          container = null;
      }
@@ -41,14 +39,13 @@
 
  $: dragDisabled = !$uiState.unlocked;
 
- const handleConsider = evt => {
-     $layoutState.children[dragItem.id] = evt.detail.items;
-     children = $layoutState.children[dragItem.id];
+ function handleConsider(evt: any) {
+     children = layoutState.updateChildren(dragItem, evt.detail.items)
      // console.log(dragItems);
  };
- const handleFinalize = evt => {
-     $layoutState.children[dragItem.id] = evt.detail.items;
-     children = $layoutState.children[dragItem.id];
+
+ function handleFinalize(evt: any) {
+     children = layoutState.updateChildren(dragItem, evt.detail.items)
      // Ensure dragging is stopped on drag finish
      // dragDisabled = true;
  };
@@ -64,66 +61,106 @@
      // dragDisabled = true;
  };
 
- const unsubscribe = nodeState.subscribe(state => {
-     if (container) {
-         $layoutState.children[container.id] = $layoutState.children[container.id].filter(item => item.node.id in state);
-         children = $layoutState.children[container.id];
-     }
- });
-
- onDestroy(unsubscribe);
-
  $: if ($queueState && widget) {
      widget.isNodeExecuting = $queueState.runningNodeId === widget.nodeId;
-     children = $layoutState.children[widget.nodeId];
- }
-
- function updateNodeName(node: LGraphNode, value: string) {
-     nodeState.nodeStateChanged(node);
+     children = $layoutState.allItems[widget.id].children;
  }
 </script>
 
 
-{#if container}
-    {@const node = container.node}
+{#if container && children}
     {@const id = container.id}
-    <Block>
-        <label for={String(id)} class={$uiState.unlocked ? "edit-title-label" : ""}>
-            <BlockTitle>
-                {#if $uiState.unlocked}
-                    <input class="edit-title" bind:value={container.attrs.title} type="text" minlength="1" on:input="{(v) => { updateNodeName(node, v) }}"/>
-                {:else}
-                    {container.attrs.title}
-                {/if}
-            </BlockTitle>
-        </label>
-        <div class="v-pane"
-             use:dndzone="{{ items: children, dragDisabled, flipDurationMs }}"
-             on:consider="{handleConsider}"
-             on:finalize="{handleFinalize}"
-        >
-			{#each children.filter(item => item.id !== SHADOW_PLACEHOLDER_ITEM_ID) as item(item.id)}
-                <div class="animation-wrapper" class:is-executing={item.isNodeExecuting} animate:flip={{duration:flipDurationMs}}>
-                    <svelte:self dragItem={item}/>
-                    {#if item[SHADOW_ITEM_MARKER_PROPERTY_NAME]}
-                        <div in:fade={{duration:200, easing: cubicIn}} class='drag-item-shadow'/>
-                    {/if}
-                </div>
-            {/each}
-        </div>
-    </Block>
+    <div class="container {container.attrs.direction} {container.attrs.classes} {classes.join(' ')}">
+        <Block>
+            {#if container.attrs.showTitle}
+                <label for={String(id)} class={$uiState.unlocked ? "edit-title-label" : ""}>
+                    <BlockTitle>
+                        {#if $uiState.unlocked}
+                            <input class="edit-title" bind:value={container.attrs.title} type="text" minlength="1" />
+                        {:else}
+                            {container.attrs.title}
+                        {/if}
+                    </BlockTitle>
+                </label>
+            {/if}
+            <div class="v-pane"
+                 class:empty={children.length === 0}
+                 use:dndzone="{{ items: children, dragDisabled, flipDurationMs }}"
+                 on:consider="{handleConsider}"
+                 on:finalize="{handleFinalize}"
+            >
+                {#each children.filter(item => item.id !== SHADOW_PLACEHOLDER_ITEM_ID) as item(item.id)}
+                    <div class="animation-wrapper" class:is-executing={item.isNodeExecuting} animate:flip={{duration:flipDurationMs}}>
+                        <svelte:self dragItem={item} zIndex={zIndex+1} />
+                        {#if item[SHADOW_ITEM_MARKER_PROPERTY_NAME]}
+                            <div in:fade={{duration:200, easing: cubicIn}} class='drag-item-shadow'/>
+                        {/if}
+                    </div>
+                {/each}
+            </div>
+            {#if $uiState.unlocked}
+                <div class="handle handle-container" style="z-index: {zIndex}" on:mousedown={startDrag} on:touchstart={startDrag} on:mouseup={stopDrag} on:touchend={stopDrag}/>
+            {/if}
+        </Block>
+    </div>
 {:else if widget}
-    <svelte:component this={getComponentForWidgetState(widget)} item={widget} />
+    <svelte:component this={getComponentForWidgetState(widgetState)} item={widgetState} />
     {#if $uiState.unlocked}
-        <div class="handle" on:mousedown={startDrag} on:touchstart={startDrag} on:mouseup={stopDrag} on:touchend={stopDrag}/>
+        <div class="handle handle-widget" style="z-index: {zIndex}" on:mousedown={startDrag} on:touchstart={startDrag} on:mouseup={stopDrag} on:touchend={stopDrag}/>
     {/if}
 {/if}
 
-<style>
+<style lang="scss">
  .v-pane {
      height: 100%;
      width: 100%;
-     overflow-y: auto ;
+     overflow: visible;
+     display: flex;
+
+     &.empty {
+         border-width: 3px;
+         border-color: var(--color-grey-400);
+         border-radius: var(--block-radius);
+         background: var(--color-grey-300);
+         min-height: 50px;
+         border-style: dashed;
+     }
+ }
+
+ .container {
+     display: flex;
+
+     :global(.block) {
+         height: fit-content;
+     }
+
+     &.horizontal {
+         flex-wrap: wrap;
+         gap: var(--layout-gap);
+         width: var(--size-full);
+
+         .v-pane {
+             flex-direction: row;
+         }
+
+         > :global(*), > :global(.form > *) {
+             flex: 1 1 0%;
+             flex-wrap: wrap;
+             min-width: min(160px, 100%);
+         }
+     }
+
+     &.vertical {
+         position: relative;
+
+         .v-pane {
+             flex-direction: column;
+         }
+
+         > :global(*), > :global(.form > *), .v-pane {
+             width: var(--size-full);
+         }
+     }
  }
 
  .is-executing :global(.block) {
@@ -132,8 +169,7 @@
 
  .animation-wrapper {
      position: relative;
-     width: 100%;
-     height: 100%;
+     flex-grow: 1;
  }
 
  .handle {
@@ -146,10 +182,14 @@
      height: 100%;
  }
 
- .handle:hover {
+ .handle-widget:hover {
      background-color: #add8e680;
  }
 
+ .handle-container:hover {
+     background-color: #d8ade680;
+ }
+
  .drag-item-shadow {
      position: absolute;
      top: 0; left:0; right: 0; bottom: 0;
diff --git a/src/lib/stores/layoutState.ts b/src/lib/stores/layoutState.ts
index a169c46..86f8ead 100644
--- a/src/lib/stores/layoutState.ts
+++ b/src/lib/stores/layoutState.ts
@@ -2,21 +2,34 @@ import { get, writable } from 'svelte/store';
 import type { Readable, Writable } from 'svelte/store';
 import type ComfyApp from "$lib/components/ComfyApp"
 import type { LGraphNode, IWidget } from "@litegraph-ts/core"
+import nodeState from "$lib/state/nodeState";
+import type { NodeStateStore } from './nodeState';
+
+type DragItemEntry = {
+    dragItem: IDragItem,
+    children: IDragItem[] | null,
+    parent: IDragItem | null
+}
 
 export type LayoutState = {
     root: IDragItem | null,
-    children: Record<number, IDragItem[]>,
+    allItems: Record<DragItemID, DragItemEntry>,
+    currentId: number,
 }
 
-export type Properties = {
+export type Attributes = {
     direction: string,
+    title: string,
+    showTitle: boolean,
+    classes: string,
+    associatedNode: number | null
 }
 
 export interface IDragItem {
     type: string,
-    id: number,
+    id: DragItemID,
     isNodeExecuting?: boolean,
-    properties: Properties
+    attrs: Attributes
 }
 
 export interface ContainerLayout extends IDragItem {
@@ -29,37 +42,181 @@ export interface WidgetLayout extends IDragItem {
     widgetName: string
 }
 
+type DragItemID = string;
+
 type LayoutStateOps = {
+    addContainer: (parentId: DragItemID, attrs: Partial<Attributes>) => ContainerLayout,
     findDefaultContainerForInsertion: () => ContainerLayout | null,
-    reset: () => void,
+    addWidget: (parentId: DragItemID, node: LGraphNode, widget: IWidget<any, any>, attrs: Partial<Attributes>) => WidgetLayout,
+    updateChildren: (parent: IDragItem, children: IDragItem[]) => IDragItem[],
+    nodeAdded: (node: LGraphNode) => void,
+    nodeRemoved: (node: LGraphNode) => void,
+    clear: () => void,
+    resetLayout: () => void,
 }
 
 export type WritableLayoutStateStore = Writable<LayoutState> & LayoutStateOps;
 const store: Writable<LayoutState> = writable({
     root: null,
-    children: []
+    allItems: [],
+    currentId: 0
 })
 
 function findDefaultContainerForInsertion(): ContainerLayout | null {
     const state = get(store);
-    if ("children" in state.root) {
+
+    if (state.root === null) {
+        // Should never happen
+        throw "Root container was null!";
+    }
+
+    if (state.root.type === "container") {
         const container = state.root as ContainerLayout;
-        const found = state.children[container.id].find((di) => {"children" in di})
-        if (found && "children" in found)
+        const children: IDragItem[] = state.allItems[container.id]?.children || []
+        const found = children.find((di) =>  di.type === "container" )
+        if (found && found.type === "container")
             return found as ContainerLayout;
         return container;
     }
+
     return null
 }
 
-function reset() {
+function addContainer(parentId: DragItemID | null, attrs: Partial<Attributes> = {}): ContainerLayout {
+    const state = get(store);
+    const dragItem: ContainerLayout = {
+        type: "container",
+        id: `${state.currentId++}`,
+        attrs: {
+            title: "Container",
+            showTitle: true,
+            direction: "vertical",
+            classes: "",
+            associatedNode: null,
+            ...attrs
+        }
+    }
+    const parent = parentId ? state.allItems[parentId] : null;
+    const entry: DragItemEntry = { dragItem, children: [], parent: parent?.dragItem };
+    state.allItems[dragItem.id] = entry;
+    if (parent) {
+        parent.children ||= []
+        parent.children.push(dragItem)
+    }
+    store.set(state)
+    return dragItem;
+}
+
+function addWidget(parentId: DragItemID, node: LGraphNode, widget: IWidget<any, any>, attrs: Partial<Attributes> = {}): WidgetLayout {
+    const state = get(store);
+    const dragItem: WidgetLayout = {
+        type: "widget",
+        id: `${state.currentId++}`,
+        nodeId: node.id,
+        widgetName: widget.name, // TODO name and displayName
+        attrs: {
+            title: widget.name,
+            showTitle: true,
+            direction: "horizontal",
+            classes: "",
+            associatedNode: null,
+            ...attrs
+        }
+    }
+    const parent = state.allItems[parentId]
+    const entry: DragItemEntry = { dragItem, children: [], parent: parent.dragItem };
+    state.allItems[dragItem.id] = entry;
+    parent.children ||= []
+    parent.children.push(dragItem)
+    store.set(state)
+    return dragItem;
+}
+
+function updateChildren(parent: IDragItem, children: IDragItem[]): IDragItem[] {
+    const state = get(store);
+    state.allItems[parent.id].children = children;
+    for (const child of children) {
+        state.allItems[child.id].parent = parent;
+    }
+    store.set(state)
+    return children
+}
+
+function nodeAdded(node: LGraphNode) {
+    const parent = findDefaultContainerForInsertion();
+    // Add default node panel containing all widgets
+    if (node.widgets && node.widgets.length > 0) {
+        const container = addContainer(parent.id, { title: node.title, direction: "vertical", associatedNode: node.id });
+        for (const widget of node.widgets) {
+            addWidget(container.id, node, widget, { associatedNode: node.id });
+        }
+    }
+}
+
+function removeEntry(state: LayoutState, id: DragItemID) {
+    const entry = state.allItems[id]
+    if (entry.children && entry.children.length > 0) {
+        console.error(entry)
+        throw `Tried removing entry ${id} but it still had children!`
+    }
+    const parent = entry.parent;
+    if (parent) {
+        const parentEntry = state.allItems[parent.id];
+        parentEntry.children = parentEntry.children.filter(item => item.id !== id)
+    }
+    delete state.allItems[id]
+}
+
+function nodeRemoved(node: LGraphNode) {
+    const state = get(store)
+
+    // Remove widgets bound to the node
+    let del = Object.entries(state.allItems).filter(pair =>
+        pair[1].dragItem.type === "widget"
+        && pair[1].dragItem.attrs.associatedNode === node.id)
+    for (const id in del) {
+        console.debug("[layoutState] Remove widget", id, state.allItems[id])
+        removeEntry(state, id)
+    }
+
+    const isAssociatedContainer = (dragItem: IDragItem) =>
+        dragItem.type === "container"
+        && dragItem.attrs.associatedNode === node.id;
+
+    let delContainers = []
+
+    // Remove widget from all children lists
+    for (const entry of Object.values(state.allItems)) {
+        if (entry.children?.length === 0 && isAssociatedContainer(entry.dragItem))
+            delContainers.push(entry.dragItem.id)
+    }
+
+    // Remove empty containers bound to the node
+    for (const id in delContainers) {
+        console.debug("[layoutState] Remove container", id, state.allItems[id])
+        removeEntry(state, id)
+    }
+
+    store.set(state)
+}
+
+function clear() {
+}
+
+function resetLayout() {
     // TODO
 }
 
 const uiStateStore: WritableLayoutStateStore =
 {
     ...store,
+    addContainer,
+    addWidget,
     findDefaultContainerForInsertion,
-    reset
+    updateChildren,
+    nodeAdded,
+    nodeRemoved,
+    clear,
+    resetLayout
 }
 export default uiStateStore;
diff --git a/src/lib/stores/uiState.ts b/src/lib/stores/uiState.ts
index 507232f..43cd4d8 100644
--- a/src/lib/stores/uiState.ts
+++ b/src/lib/stores/uiState.ts
@@ -14,7 +14,7 @@ const store: WritableUIStateStore = writable(
     {
         graphLocked: true,
         nodesLocked: false,
-        unlocked: true,
+        unlocked: false,
     })
 
 const uiStateStore: WritableUIStateStore =
diff --git a/src/lib/widgets/ComboWidget.svelte b/src/lib/widgets/ComboWidget.svelte
index e49693a..b8a36bb 100644
--- a/src/lib/widgets/ComboWidget.svelte
+++ b/src/lib/widgets/ComboWidget.svelte
@@ -45,6 +45,6 @@
  }
 
  :global(.svelte-select-list) {
-     z-index: var(--layer-5) !important;
+     z-index: var(--layer-top) !important;
  }
 </style>