Auto arrange and auto centering support

packages/editor/src/editor.ts

@@ -3,6 +3,7 @@ import ConnectionPlugin from 'rete-connection-plugin'
 import { Plugin } from 'rete/types/core/plugin'
 import gridimg from './grid.png'
 import CommentPlugin from './plugins/commentPlugin'
+import CommentManager from './plugins/commentPlugin/manager'
 import ContextMenuPlugin from './plugins/contextMenu'
 import {
   CachePlugin,
@@ -38,6 +39,7 @@ import {
 } from '@magickml/core'
 
 import AreaPlugin from './plugins/areaPlugin'
+import AutoArrangePlugin from './plugins/autoArrangePlugin'
 import { initSharedEngine, MagickEngine } from '@magickml/core'
 
 /**
@@ -147,9 +149,21 @@ export const initEditor = function ({
     snap: true,
   })
 
+  // Set up the CommentManager
+  const commentManager = new CommentManager(editor)
+
   // Use CommentPlugin
   editor.use(CommentPlugin, {
     margin: 30,
+    commentManager,
+  })
+
+  editor.use(AutoArrangePlugin, {
+    margin: { x: 50, y: 50 },
+    depth: 0,
+    arrangeHotkey: { key: '/', ctrl: true },
+    centerHotkey: { key: '.', ctrl: true },
+    commentManager,
   })
 
   editor.use(KeyCodePlugin)

packages/editor/src/plugins/autoArrangePlugin/auto-arrange.js

@@ -0,0 +1,177 @@
+import { Board } from './board'
+import { Cache } from './cache'
+
+export class AutoArrange {
+  constructor(editor, margin, depth, vertical, offset, commentManager) {
+    this.editor = editor
+    this.margin = margin
+    this.depth = depth
+    this.vertical = vertical
+    this.offset = offset
+    this.commentManager = commentManager
+  }
+
+  getNodes(node, type = 'output') {
+    const nodes = []
+    const key = `${type}s`
+
+    for (let io of node[key].values())
+      for (let connection of io.connections.values())
+        nodes.push(connection[type === 'input' ? 'output' : 'input'].node)
+
+    return nodes
+  }
+
+  getNodesBoard(
+    node,
+    options,
+    cache = new Cache(),
+    board = new Board(),
+    depth = 0
+  ) {
+    if (options.depth && depth > options.depth) return board
+    if (options.skip && options.skip(node)) return board
+    if (cache.track(node)) return board
+
+    board.add(depth, node)
+
+    const outputNodes =
+      (options.substitution && options.substitution.output(node)) ||
+      this.getNodes(node, 'output')
+    const inputNodes =
+      (options.substitution && options.substitution.input(node)) ||
+      this.getNodes(node, 'input')
+
+    outputNodes.map(n =>
+      this.getNodesBoard(n, options, cache, board, depth + 1)
+    )
+    inputNodes.map(n => this.getNodesBoard(n, options, cache, board, depth - 1))
+
+    return board
+  }
+
+  getNodeSize(node, vertical = this.vertical) {
+    const el = this.editor.view.nodes.get(node).el
+
+    return vertical
+      ? {
+          height: el.clientWidth,
+          width: el.clientHeight,
+        }
+      : {
+          width: el.clientWidth,
+          height: el.clientHeight,
+        }
+  }
+
+  translateNode(node, { x, y, vertical = this.vertical }) {
+    const position = vertical ? [y, x] : [x, y]
+
+    this.editor.view.nodes.get(node).translate(...position)
+    this.editor.view.updateConnections({ node })
+  }
+
+  arrange(node = this.editor.nodes[0], options = {}) {
+    const {
+      margin = this.margin,
+      vertical = this.vertical,
+      depth = this.depth,
+      offset = this.offset,
+      skip,
+      substitution,
+    } = options
+    const board = this.getNodesBoard(node, {
+      depth,
+      skip,
+      substitution,
+    }).toArray()
+    const currentMargin = vertical ? { x: margin.y, y: margin.x } : margin
+    const currentOffset = vertical ? { x: offset.y, y: offset.x } : offset
+
+    let x = currentOffset.x
+
+    // Access the commentManager
+    const commentManager = this.editor.plugins.get('comment').commentManager
+
+    for (let column of board) {
+      const sizes = column.map(node => this.getNodeSize(node, vertical))
+      const columnWidth = Math.max(...sizes.map(size => size.width))
+      const fullHeight = sizes.reduce(
+        (sum, node) => sum + node.height + currentMargin.y,
+        0
+      )
+
+      let y = currentOffset.y
+
+      for (let node of column) {
+        const position = { x, y: y - fullHeight / 2, vertical }
+        const { height } = this.getNodeSize(node, vertical)
+
+        this.translateNode(node, position, vertical)
+
+        // Update the position of the comment related to the current node
+        const relatedComments = commentManager.comments.filter(comment =>
+          comment.linkedTo(node)
+        )
+        relatedComments.forEach(comment => {
+          const newPosition = vertical
+            ? { x: y - fullHeight / 2, y: x }
+            : { x, y: y - fullHeight / 2 }
+          comment.x = newPosition.x
+          comment.y = newPosition.y
+          comment.update()
+        })
+
+        y += height + currentMargin.y
+      }
+
+      x += columnWidth + currentMargin.x
+    }
+  }
+
+  markVisitedNodes(node, visited) {
+    if (visited.has(node)) return
+
+    visited.add(node)
+    const connectedNodes = [
+      ...this.getNodes(node, 'output'),
+      ...this.getNodes(node, 'input'),
+    ]
+
+    connectedNodes.forEach(connectedNode =>
+      this.markVisitedNodes(connectedNode, visited)
+    )
+  }
+
+  centerProject() {
+    const maxX = Math.max(...this.editor.nodes.map(node => node.position[0]))
+    const maxY = Math.max(...this.editor.nodes.map(node => node.position[1]))
+    const minX = Math.min(...this.editor.nodes.map(node => node.position[0]))
+    const minY = Math.min(...this.editor.nodes.map(node => node.position[1]))
+
+    const offsetX = (maxX + minX) / 2
+    const offsetY = (maxY + minY) / 2
+
+    const adjustmentX = -offsetX
+    const adjustmentY = -offsetY
+
+    // Translate nodes
+    for (const node of this.editor.nodes) {
+      const [x, y] = node.position
+      const newPosition = [x + adjustmentX, y + adjustmentY]
+      this.translateNode(node, { x: newPosition[0], y: newPosition[1] })
+    }
+
+    // Translate comments
+    const commentManager = this.editor.plugins.get('comment').commentManager
+    if (commentManager) {
+      for (const comment of commentManager.comments) {
+        const [x, y] = [comment.x, comment.y]
+        const newPosition = [x + adjustmentX, y + adjustmentY]
+        comment.x = newPosition[0]
+        comment.y = newPosition[1]
+        comment.update()
+      }
+    }
+  }
+}

packages/editor/src/plugins/autoArrangePlugin/board.js

@@ -0,0 +1,19 @@
+export class Board {
+    constructor() {
+        this._cols = [];
+    }
+
+    add(columnIndex, value) {
+        if (!this._cols[columnIndex]) this._cols[columnIndex] = [];
+
+        this._cols[columnIndex].push(value);
+    }
+
+    toArray() {
+        const normalized = Object.keys(this._cols)
+            .sort((i1, i2) => +i1 - +i2)
+            .map(key => this._cols[key]);
+
+        return normalized;
+    }
+}

packages/editor/src/plugins/autoArrangePlugin/cache.js

@@ -0,0 +1,10 @@
+export class Cache {
+    constructor() {
+        this._map = new WeakMap();
+    }
+
+    track(value) {
+        if (this._map.has(value)) return true;
+        this._map.set(value, true);
+    }
+}

packages/editor/src/plugins/autoArrangePlugin/index.js

@@ -0,0 +1,74 @@
+import { AutoArrange } from './auto-arrange'
+
+function install(
+  editor,
+  {
+    margin = { x: 50, y: 50 },
+    depth = null,
+    vertical = false,
+    offset = { x: 0, y: 0 },
+    arrangeHotkey = { key: '/', ctrl: true },
+    centerHotkey = { key: '.', ctrl: true },
+  }
+) {
+  editor.bind('arrange')
+
+  const ar = new AutoArrange(editor, margin, depth, vertical, offset)
+
+  editor.on('arrange', ({ node, ...options }) => ar.arrange(node, options))
+
+  editor.arrange = (node, options) => {
+    console.log(`Deprecated: use editor.trigger('arrange', { node }) instead`)
+    ar.arrange(node, options)
+  }
+
+  if (arrangeHotkey) {
+    const { key, ctrl } = arrangeHotkey
+
+    window.addEventListener('keydown', event => {
+      if (ctrl && event.ctrlKey && event.key === key) {
+        // Find all unvisited nodes and arrange them separately
+        const visited = new Set()
+        let currentOffset = { ...offset } // Initialize currentOffset with the provided offset
+
+        for (const node of editor.nodes) {
+          if (!visited.has(node)) {
+            const options = {
+              depth,
+              margin,
+              vertical,
+              skip: undefined,
+              substitution: undefined,
+              offset: currentOffset,
+            }
+            ar.arrange(node, options) // Pass the currentOffset
+            ar.markVisitedNodes(node, visited)
+
+            // Update currentOffset with the maxY value from the last arranged group
+            const maxY = Math.max(
+              ...Array.from(visited).map(
+                n => n.position[1] + ar.getNodeSize(n).height + 400
+              )
+            )
+            currentOffset = { x: offset.x, y: maxY + margin.y }
+          }
+        }
+      }
+    })
+  }
+
+  if (centerHotkey) {
+    const { key, ctrl } = centerHotkey
+
+    window.addEventListener('keydown', event => {
+      if (ctrl && event.ctrlKey && event.key === key) {
+        ar.centerProject()
+      }
+    })
+  }
+}
+
+export default {
+  name: 'auto-arrange',
+  install,
+}

packages/editor/src/plugins/commentPlugin/index.js

@@ -7,6 +7,7 @@ import { nodesBBox, listenWindow } from './utils'
 function install(
   editor,
   {
+    commentManager = null,
     margin = 30,
     disableBuiltInEdit = false,
     frameCommentKeys = {
@@ -37,7 +38,8 @@ function install(
   editor.bind('removecomment')
   editor.bind('editcomment')
 
-  const manager = new CommentManager(editor)
+  // const manager = new CommentManager(editor)
+  const manager = commentManager || new CommentManager(editor)
 
   if (!disableBuiltInEdit) {
     editor.on('editcomment', comment => {