smartbotic/webui/src/utils/loopMembership.ts

import { Node, Edge } from 'reactflow'

/**
 * Maps every node inside a loop body to the id of the loop node that owns
 * it. A node not reachable from any loop node's 'loop' output is absent
 * from the map, including loop nodes themselves - a loop node is not a
 * member of its own body.
 *
 * Shared by ExecutionResultsPanel (grouping loop-body results by iteration)
 * and WorkflowEditorPage (deciding which cached upstream output belongs to
 * the same iteration as a node being re-run). Both need the identical
 * answer to "is this node inside this loop's body", so this is the one
 * place that walks the graph to find out.
 */
export function computeNodeToLoopMap(nodes: Node[], edges: Edge[]): Map<string, string> {
  const nodeToLoopMap = new Map<string, string>()
  const loopNodes = nodes.filter(n => n.data?.type === 'loop')

  const outgoingEdges = new Map<string, Array<{ target: string; handle: string }>>()
  nodes.forEach(n => outgoingEdges.set(n.id, []))
  edges.forEach(e => {
    outgoingEdges.get(e.source)?.push({ target: e.target, handle: e.sourceHandle || 'main' })
  })

  const findLoopBodyNodes = (loopId: string, nodeId: string, visited: Set<string>) => {
    if (visited.has(nodeId)) return
    visited.add(nodeId)

    if (nodeId !== loopId) {
      nodeToLoopMap.set(nodeId, loopId)
    }

    const outEdges = outgoingEdges.get(nodeId) || []
    for (const edge of outEdges) {
      if (nodeId === loopId && edge.handle === 'done') continue
      const targetNode = nodes.find(n => n.id === edge.target)
      if (targetNode?.data?.type === 'loop' && edge.target !== loopId) continue
      findLoopBodyNodes(loopId, edge.target, visited)
    }
  }

  loopNodes.forEach(loopNode => {
    const outEdges = outgoingEdges.get(loopNode.id) || []
    const loopEdges = outEdges.filter(e => e.handle === 'loop')
    const visited = new Set<string>()
    loopEdges.forEach(edge => findLoopBodyNodes(loopNode.id, edge.target, visited))
  })

  return nodeToLoopMap
}