Blame for frontends/react/src/StepFlow.tsx

import { useMemo, useSyncExternalStore } from "react";

import { formatNumber, nodeArgs } from "./formula";
import { edgePath, expressionOf, OUTPUT, SLOT_OFFSET, type GraphLayout, type PlacedEdge, type PlacedShape } from "./layout";
import {
  APPLY_MS, FADE_AT, IN_AT, IN_MS, LAND_AT, OUT_MS, PRODUCT_AT, RESULT_AT,
  SEED_HOLD_MS, SEED_TRAVEL_MS, SETTLE_AT, SETTLE_MS, SUM_AT, TRAVEL_MS,
} from "./flowTiming";
import type { Graph, GraphNode, Step } from "./types";
import type { Flow, Timeline } from "./Timeline";

const POP_MS = 150;
const CHAR_WIDTH = 7.2;
const TOKEN_HEIGHT = 20;
// A token's text baseline below its centre; a slot's text uses it too, so a settling token lands exactly on it
export const TOKEN_BASELINE = 4;
const SVG_NS = "http://www.w3.org/2000/svg";

function easeInOut(t: number): number {
  return t < 0.5 ? 2 * t * t : 1 - (1 - t) ** 2 * 2;
}

function progress(elapsed: number, from: number, duration: number): number {
  return Math.min(Math.max((elapsed - from) / duration, 0), 1);
}

function shapeOf(layout: GraphLayout, id: string): PlacedShape {
  return layout.shapes.find((s) => s.id === id)!;
}

const valueSlot = (shape: PlacedShape) => ({ x: shape.x, y: shape.y - SLOT_OFFSET });
const derivativeSlot = (shape: PlacedShape) => ({ x: shape.x, y: shape.y + SLOT_OFFSET });

// A path for a token: along an edge, joining it from one point and leaving it for another
function travelPath(edge: PlacedEdge, start: { x: number; y: number }, end: { x: number; y: number }) {
  const path = document.createElementNS(SVG_NS, "path");
  path.setAttribute("d", `${edgePath(edge).replace("M", `M ${start.x} ${start.y} L`)} L ${end.x} ${end.y}`);
  return { path, length: path.getTotalLength() };
}

// Milliseconds since the animation started, from the timeline's playhead
function useElapsed({ timeline, flow }: FlowProps): number {
  const time = useSyncExternalStore(timeline.subscribeTime, timeline.getTime);
  return Math.max(time - flow.start, 0);
}

interface TokenProps {
  x: number;
  y: number;
  text: string;
  scale?: number;
  strong?: boolean;
  // How much of the pill shows; at 0, only the text is left
  outline?: number;
  // Hide what's behind the pill, such as a number that's being replaced
  backdrop?: boolean;
}

// A number or calculation in a pill, centred on a point
function Token({ x, y, text, scale = 1, strong = false, outline = 1, backdrop = false }: TokenProps) {
  const width = text.length * CHAR_WIDTH + 12;
  const pill = { x: -width / 2, y: -TOKEN_HEIGHT / 2, width, height: TOKEN_HEIGHT, rx: TOKEN_HEIGHT / 2 };
  return (
    <g transform={`translate(${x} ${y}) scale(${scale})`}>
      {backdrop && <rect {...pill} className="fill-white" />}
      <rect {...pill} strokeWidth={1.5} opacity={outline} className={`stroke-ink ${strong ? "fill-marker" : "fill-marker-soft"}`} />
      {/* As the pill fades, the text gets the same outline as a number resting in its slot */}
      <text y={TOKEN_BASELINE} textAnchor="middle" strokeWidth={3} strokeOpacity={1 - outline} paintOrder="stroke" className="fill-ink stroke-white font-mono text-xs">
        {text}
      </text>
    </g>
  );
}

// Copies of a step's arguments flow along their edges and combine at the node
function ForwardFlow(props: FlowProps) {
  const { graph, layout, flow } = props;
  const { step, before } = flow;
  const node = graph.nodes.find((n) => n.id === step.node)!;
  const shape = shapeOf(layout, step.node);

  const forwardMode = step.derivative !== undefined;
  const { tokens, toOutput, tangents, tangentToOutput } = useMemo(() => {
    const target = shapeOf(layout, step.node);
    const outputEdge = step.node === graph.output ? layout.edges.find((edge) => edge.to === OUTPUT) : undefined;
    return {
      tokens: layout.edges
        .filter((edge) => edge.to === step.node)
        .map((edge) => ({ id: edge.id, text: formatNumber(before.values[edge.from]), ...travelPath(edge, valueSlot(shapeOf(layout, edge.from)), target) })),
      // The last result also carries on to f
      toOutput: outputEdge && travelPath(outputEdge, valueSlot(target), valueSlot(shapeOf(layout, OUTPUT))),
      // Forward mode's tangents follow the same way
      tangents: step.derivative === undefined ? undefined : derivativeMoves(graph, layout, step, before.derivatives),
      tangentToOutput: outputEdge && travelPath(outputEdge, derivativeSlot(target), derivativeSlot(shapeOf(layout, OUTPUT))),
    };
  }, [graph, layout, step, before]);

  const elapsed = useElapsed(props);
  const pop = (from: number) => 0.7 + 0.3 * progress(elapsed, from, POP_MS);
  const calculation = expressionOf(node, (arg) => (arg in before.values ? formatNumber(before.values[arg]) : arg));
  const pulse = (elapsed - TRAVEL_MS) / APPLY_MS;
  const settled = easeInOut(progress(elapsed, SETTLE_AT, SETTLE_MS));
  const result = formatNumber(step.value);
  // How far the tangents have got, once they've set off
  const t = elapsed - SETTLE_AT;
  const valuePhase = elapsed < TRAVEL_MS ? "travel" : elapsed < RESULT_AT ? "apply" : elapsed < SETTLE_AT ? "result" : "settle";

  return (
    <g pointerEvents="none" data-flow={forwardMode && t >= 0 ? `tangent-${derivativePhase(t)}` : valuePhase}>
      {elapsed < TRAVEL_MS &&
        tokens.map(({ id, path, length, text }) => {
          const { x, y } = path.getPointAtLength(length * easeInOut(elapsed / TRAVEL_MS));
          return <Token key={id} x={x} y={y} text={text} />;
        })}
      {pulse >= 0 && pulse < 1 && (
        <ellipse
          cx={shape.x}
          cy={shape.y}
          rx={(shape.width / 2) * (1 + pulse * 0.4)}
          ry={(shape.height / 2) * (1 + pulse * 0.4)}
          opacity={1 - pulse}
          fill="none"
          strokeWidth={2}
          className="stroke-ink"
        />
      )}
      {elapsed >= TRAVEL_MS && elapsed < RESULT_AT && <Token x={shape.x} y={shape.y} text={calculation} scale={pop(TRAVEL_MS)} />}
      {elapsed >= RESULT_AT && (
        <Token x={shape.x} y={shape.y - SLOT_OFFSET * settled} text={result} scale={pop(RESULT_AT)} strong outline={1 - settled} />
      )}
      {toOutput && elapsed >= SETTLE_AT && (() => {
        const { x, y } = toOutput.path.getPointAtLength(toOutput.length * settled);
        return <Token x={x} y={y} text={result} strong outline={1 - settled ** 4} />;
      })()}
      {tangents && t >= 0 && <DerivativeTokens {...tangents} t={t} />}
      {tangents && tangentToOutput && t >= FADE_AT && (() => {
        const moved = easeInOut(progress(t, FADE_AT, SETTLE_MS));
        const { x, y } = tangentToOutput.path.getPointAtLength(tangentToOutput.length * moved);
        return <Token x={x} y={y} text={formatNumber(step.derivative!)} strong outline={1 - moved ** 4} />;
      })()}
    </g>
  );
}

// The seed leaves f, keeping a copy there, and settles below the output node
function SeedFlow(props: FlowProps) {
  const { layout, flow } = props;
  const output = shapeOf(layout, OUTPUT);
  const { path, length } = useMemo(() => {
    const edge = layout.edges.find((e) => e.to === OUTPUT)!;
    return travelPath(edge, derivativeSlot(shapeOf(layout, edge.from)), derivativeSlot(shapeOf(layout, OUTPUT)));
  }, [layout]);
  const elapsed = useElapsed(props);
  const moved = easeInOut(progress(elapsed, SEED_HOLD_MS, SEED_TRAVEL_MS));
  const { x, y } = path.getPointAtLength(length * (1 - moved));
  const text = formatNumber(flow.step.derivative!);
  return (
    <g pointerEvents="none" data-flow={elapsed < SEED_HOLD_MS ? "seed" : elapsed < SEED_HOLD_MS + SEED_TRAVEL_MS ? "travel" : "settle"}>
      <text x={output.x} y={output.y + SLOT_OFFSET + TOKEN_BASELINE} textAnchor="middle" strokeWidth={3} paintOrder="stroke" className="fill-ink stroke-white font-mono text-xs">
        {text}
      </text>
      <Token x={x} y={y} text={text} strong outline={1 - progress(elapsed, SEED_HOLD_MS + SEED_TRAVEL_MS, SETTLE_MS)} />
    </g>
  );
}

function sumText(terms: number[]): string {
  return terms.map(formatNumber).join(" + ").replaceAll("+ -", "− ");
}

// A derivative's moves along one step's edges: one token per partial, and where they land
function derivativeMoves(graph: Graph, layout: GraphLayout, step: Step, derivatives: Record<string, number>) {
  const node = graph.nodes.find((n) => n.id === step.node) as GraphNode & { kind: "op" };
  const forward = step.phase === "forward";
  const args = nodeArgs(graph, node);
  const tokens = step.partials.map((p, n) => {
    const edge = layout.edges.find((e) => e.id === `${p.arg}-${step.node}-${args[n].index}`)!;
    // A tangent comes from the argument; an adjoint comes from the node
    const start = forward ? derivatives[p.arg] : step.derivative!;
    return {
      id: edge.id,
      start: formatNumber(start),
      multiply: `${formatNumber(start)} · ${formatNumber(p.partial)}`,
      product: formatNumber(p.amount),
      reversed: !forward,
      ...travelPath(edge, derivativeSlot(shapeOf(layout, p.arg)), derivativeSlot(shapeOf(layout, step.node))),
    };
  });
  // Forward: everything adds up at the node. Backward: at each argument, on top of what it has.
  const amounts = (arg: string) => step.partials.filter((p) => p.arg === arg).map((p) => p.amount);
  const landing = (id: string, terms: number[]) => ({
    id,
    ...derivativeSlot(shapeOf(layout, id)),
    sum: sumText(terms),
    total: formatNumber(terms.reduce((a, b) => a + b, 0)),
  });
  const landings = forward
    ? [landing(step.node, step.partials.map((p) => p.amount))]
    : [...new Set(step.partials.map((p) => p.arg))].map((arg) =>
        landing(arg, [...(arg in derivatives ? [derivatives[arg]] : []), ...amounts(arg)]),
      );
  return { tokens, landings };
}

function derivativePhase(t: number): string {
  return t < OUT_MS ? "out" : t < PRODUCT_AT ? "multiply" : t < IN_AT ? "product" : t < LAND_AT ? "in" : t < SUM_AT ? "land" : "sum";
}

// t is the time since the derivative set off
function DerivativeTokens({ tokens, landings, t }: ReturnType<typeof derivativeMoves> & { t: number }) {
  const pop = (from: number) => 0.7 + 0.3 * progress(t, from, POP_MS);
  const phase = derivativePhase(t);
  // How far along the edge a token has come, from where it set off (0) to where it lands (1)
  const along = 0.5 * easeInOut(progress(t, 0, OUT_MS)) + 0.5 * easeInOut(progress(t, IN_AT, IN_MS));
  return (
    <>
      {t < LAND_AT &&
        tokens.map(({ id, path, length, start, multiply, product, reversed }) => {
          const { x, y } = path.getPointAtLength(length * (reversed ? 1 - along : along));
          if (phase === "multiply") {
            return <Token key={id} x={x} y={y} text={multiply} scale={pop(OUT_MS)} />;
          }
          return phase === "out" ? <Token key={id} x={x} y={y} text={start} /> : <Token key={id} x={x} y={y} text={product} strong />;
        })}
      {t >= LAND_AT &&
        landings.map(({ id, x, y, sum, total }) =>
          t < SUM_AT ? (
            <Token key={id} x={x} y={y} text={sum} scale={pop(LAND_AT)} backdrop />
          ) : (
            <Token key={id} x={x} y={y} text={total} strong outline={1 - progress(t, FADE_AT, SETTLE_MS)} backdrop />
          ),
        )}
    </>
  );
}

// The node's adjoint leaves along each edge, picks up the edge's partial on the way, and adds to the argument's adjoint
function BackwardFlow(props: FlowProps) {
  const { graph, layout, flow } = props;
  const { step, before } = flow;
  const moves = useMemo(() => derivativeMoves(graph, layout, step, before.derivatives), [graph, layout, step, before]);
  const t = useElapsed(props);
  return (
    <g pointerEvents="none" data-flow={derivativePhase(t)}>
      <DerivativeTokens {...moves} t={t} />
    </g>
  );
}

interface FlowProps {
  graph: Graph;
  layout: GraphLayout;
  flow: Flow;
  timeline: Timeline;
}

export default function StepFlow(props: FlowProps) {
  const { step } = props.flow;
  if (step.phase === "forward") {
    return <ForwardFlow {...props} />;
  }
  return step.partials.length === 0 ? <SeedFlow {...props} /> : <BackwardFlow {...props} />;
}