caret / packages / math / src / nodes / equation.ts

import { CaretNode, h, t, TokensTag, VNode } from "@caret-js/core";
import { FunctionCallNode } from "./functionCall";
import { FunctionDefinitionNode } from "./functionDefinition";
import { VariableNode } from "./variable";

export class EquationNode extends CaretNode {
  constructor(public expressions: CaretNode[]) {
    super();
  }

  static from(expressions: CaretNode[]): EquationNode | FunctionDefinitionNode {
    const flattenedExpressions: CaretNode[] = expressions.flatMap((node) =>
      node instanceof EquationNode
        ? node.expressions
        : node instanceof FunctionDefinitionNode
        ? [new FunctionCallNode(node.name, node.args), node.body]
        : [node]
    );

    if (flattenedExpressions.length === 2) {
      if (flattenedExpressions[0] instanceof FunctionCallNode) {
        const funcCall = flattenedExpressions[0] as FunctionCallNode;
        const name = funcCall.name;
        const args = funcCall.args;
        if (args.every((arg) => arg instanceof VariableNode)) {
          const body = flattenedExpressions[1];
          return FunctionDefinitionNode.from(
            name,
            args as VariableNode[],
            body
          ).addTag(
            new TokensTag([...(funcCall.getTag(TokensTag)?.ownTokens ?? [])])
          );
        }
      }
    }

    return new EquationNode(flattenedExpressions).mergeTagsFromNodes(
      expressions.filter(
        (node) =>
          node instanceof EquationNode || node instanceof FunctionDefinitionNode
      )
    );
  }

  toDebugMathML(): VNode {
    return h(
      "mrow",
      {},
      ...this.expressions.flatMap((expr, index) => [
        expr.toDebugMathML(),
        index < this.expressions.length - 1 ? h("mo", {}, t("=")) : null,
      ])
    );
  }
}