webstudio/packages/expression/src/index.ts

import {
  type Expression,
  type Identifier,
  type MemberExpression,
  type Pattern,
  parse,
  parseExpressionAt,
  tokenizer,
} from "acorn";
import { simple } from "acorn-walk";

const parseCompleteExpression = (
  expression: string,
  options: { preserveParens?: boolean } = {}
) => {
  const completeNode = parseExpressionAt(expression, 0, {
    ecmaVersion: "latest",
    preserveParens: true,
  });
  const trailingToken = tokenizer(expression.slice(completeNode.end), {
    ecmaVersion: "latest",
  }).getToken();
  if (trailingToken.type.label !== "eof") {
    throw new SyntaxError("Unexpected content after expression");
  }
  if (
    options.preserveParens ||
    completeNode.type !== "ParenthesizedExpression"
  ) {
    return completeNode;
  }
  return parseExpressionAt(expression, 0, {
    ecmaVersion: "latest",
  });
};

export type Diagnostic = {
  from: number;
  to: number;
  severity: "error" | "hint" | "info" | "warning";
  message: string;
};

type ExpressionVisitor = {
  [K in Expression["type"]]: (node: Extract<Expression, { type: K }>) => void;
};

type AssignmentTargetKind = "binding" | "memberObject";

const walkAssignmentTarget = (
  node: Pattern,
  visitor: {
    Identifier?: (node: Identifier, kind: AssignmentTargetKind) => void;
    MemberExpression?: (node: MemberExpression) => void;
    UnsupportedPattern?: (node: Pattern) => void;
  }
) => {
  if (node.type === "Identifier") {
    visitor.Identifier?.(node, "binding");
    return;
  }

  if (node.type === "MemberExpression") {
    visitor.MemberExpression?.(node);
    const { object } = node;
    if (object.type === "Identifier") {
      visitor.Identifier?.(object, "memberObject");
    } else if (object.type === "MemberExpression") {
      walkAssignmentTarget(object, visitor);
    }
    return;
  }

  visitor.UnsupportedPattern?.(node);
};

export type VariableValues =
  | ReadonlyMap<Identifier["name"], unknown>
  | Readonly<Record<Identifier["name"], unknown>>;

export type ExpressionValueKind =
  | "array"
  | "bigint"
  | "boolean"
  | "nullish"
  | "number"
  | "object"
  | "string"
  | "unknown";

const stringMethodReturnKindByName = new Map<string, ExpressionValueKind>([
  ["toLowerCase", "string"],
  ["replace", "string"],
  ["split", "array"],
  ["slice", "string"],
  ["at", "unknown"],
  ["endsWith", "boolean"],
  ["includes", "boolean"],
  ["startsWith", "boolean"],
  ["toString", "string"],
  ["toUpperCase", "string"],
  ["toLocaleLowerCase", "string"],
  ["toLocaleUpperCase", "string"],
]);

const arrayMethodReturnKindByName = new Map<string, ExpressionValueKind>([
  ["at", "unknown"],
  ["includes", "boolean"],
  ["join", "string"],
  ["slice", "array"],
  ["toString", "string"],
]);

export const allowedStringMethods = new Set(
  stringMethodReturnKindByName.keys()
);

export const allowedArrayMethods = new Set(arrayMethodReturnKindByName.keys());

const getVariableValue = (
  variableValues: undefined | VariableValues,
  name: Identifier["name"]
) => {
  if (variableValues === undefined) {
    return;
  }
  const maybeMap = variableValues as Partial<ReadonlyMap<string, unknown>>;
  if (
    typeof maybeMap.has === "function" &&
    typeof maybeMap.get === "function"
  ) {
    if (maybeMap.has(name)) {
      return { value: maybeMap.get(name) };
    }
    return;
  }
  const record = variableValues as Readonly<Record<string, unknown>>;
  if (Object.hasOwn(record, name)) {
    return { value: record[name] };
  }
};

const getValueKind = (value: unknown): ExpressionValueKind => {
  if (Array.isArray(value)) {
    return "array";
  }
  if (value === undefined || value === null) {
    return "nullish";
  }
  switch (typeof value) {
    case "bigint":
      return "bigint";
    case "boolean":
      return "boolean";
    case "number":
      return "number";
    case "string":
      return "string";
    case "object":
      return "object";
    default:
      return "unknown";
  }
};

const getMethodReturnKind = (
  receiverKind: ExpressionValueKind,
  methodName: string
): ExpressionValueKind => {
  if (receiverKind === "array") {
    return arrayMethodReturnKindByName.get(methodName) ?? "unknown";
  }
  if (receiverKind === "string") {
    return stringMethodReturnKindByName.get(methodName) ?? "unknown";
  }
  if (receiverKind === "unknown") {
    const stringReturnKind = stringMethodReturnKindByName.get(methodName);
    const arrayReturnKind = arrayMethodReturnKindByName.get(methodName);
    if (stringReturnKind && arrayReturnKind === undefined) {
      return stringReturnKind;
    }
    if (arrayReturnKind && stringReturnKind === undefined) {
      return arrayReturnKind;
    }
    if (stringReturnKind === arrayReturnKind) {
      return stringReturnKind ?? "unknown";
    }
    return "unknown";
  }
  return methodName === "toString" ? "string" : "unknown";
};

const isMethodSupported = (
  receiverKind: ExpressionValueKind,
  methodName: string
) => {
  if (receiverKind === "unknown") {
    return (
      allowedStringMethods.has(methodName) ||
      allowedArrayMethods.has(methodName)
    );
  }
  if (receiverKind === "string") {
    return stringMethodReturnKindByName.has(methodName);
  }
  if (receiverKind === "array") {
    return arrayMethodReturnKindByName.has(methodName);
  }
  if (receiverKind === "nullish") {
    return false;
  }
  return methodName === "toString";
};

const getExpressionNodeValueKind = (
  node: Expression,
  variableValues: undefined | VariableValues
): ExpressionValueKind => {
  if (node.type === "Identifier") {
    if (node.name === "undefined") {
      return "nullish";
    }
    const variable = getVariableValue(variableValues, node.name);
    return variable ? getValueKind(variable.value) : "unknown";
  }
  if (node.type === "Literal") {
    return getValueKind(node.value);
  }
  if (node.type === "ArrayExpression") {
    return "array";
  }
  if (node.type === "ObjectExpression") {
    return "object";
  }
  if (node.type === "TemplateLiteral") {
    return "string";
  }
  if (node.type === "BinaryExpression" && node.operator === "+") {
    const leftKind =
      node.left.type === "PrivateIdentifier"
        ? "unknown"
        : getExpressionNodeValueKind(node.left, variableValues);
    const rightKind = getExpressionNodeValueKind(node.right, variableValues);
    if (leftKind === "string" || rightKind === "string") {
      return "string";
    }
  }
  if (
    node.type === "ChainExpression" ||
    node.type === "ParenthesizedExpression"
  ) {
    return getExpressionNodeValueKind(node.expression, variableValues);
  }
  if (
    node.type === "CallExpression" &&
    node.callee.type === "MemberExpression" &&
    node.callee.object.type !== "Super" &&
    node.callee.property.type === "Identifier"
  ) {
    const receiverKind = getExpressionNodeValueKind(
      node.callee.object,
      variableValues
    );
    const methodName = node.callee.property.name;
    if (isMethodSupported(receiverKind, methodName)) {
      return getMethodReturnKind(receiverKind, methodName);
    }
  }
  return "unknown";
};

export const getExpressionValueKind = ({
  expression,
  variableValues,
}: {
  expression: string;
  variableValues?: VariableValues;
}): ExpressionValueKind => {
  try {
    const node = parseCompleteExpression(expression);
    return getExpressionNodeValueKind(node, variableValues);
  } catch {
    return "unknown";
  }
};

export const lintExpression = ({
  expression,
  availableVariables = new Set(),
  allowAssignment = false,
  variableValues,
}: {
  expression: string;
  availableVariables?: Set<Identifier["name"]>;
  allowAssignment?: boolean;
  variableValues?: VariableValues;
}): Diagnostic[] => {
  const diagnostics: Diagnostic[] = [];
  const addMessage = (
    message: string,
    severity: "error" | "warning" = "error"
  ) => {
    return (node: { start: number; end: number }) => {
      diagnostics.push({
        // tune error position after wrapping expression with parentheses
        from: node.start - 1,
        to: node.end - 1,
        severity,
        message: message,
      });
    };
  };
  // allow empty expression
  if (expression.trim().length === 0) {
    diagnostics.push({
      from: 0,
      to: 0,
      severity: "error",
      message: "Expression cannot be empty",
    });
    return diagnostics;
  }
  try {
    // wrap expression with parentheses to force acorn parse whole expression
    // instead of just first valid part
    const root = parse(`(${expression})`, {
      ecmaVersion: "latest",
      // support parsing import to forbid explicitly
      sourceType: "module",
    });

    simple(root, {
      Identifier(node) {
        if (availableVariables.has(node.name) === false) {
          addMessage(
            `"${node.name}" is not defined in the scope`,
            "warning"
          )(node);
        }
      },
      Literal() {},
      ArrayExpression() {},
      ObjectExpression() {},
      UnaryExpression() {},
      BinaryExpression() {},
      LogicalExpression() {},
      MemberExpression() {},
      ConditionalExpression() {},
      TemplateLiteral() {},
      ChainExpression() {},
      ParenthesizedExpression() {},
      AssignmentExpression(node) {
        if (allowAssignment === false) {
          addMessage("Assignment is not supported in this context")(node);
          return;
        }
        walkAssignmentTarget(node.left, {
          Identifier(node, kind) {
            if (kind !== "binding") {
              return;
            }
            if (availableVariables.has(node.name) === false) {
              addMessage(
                `"${node.name}" is not defined in the scope`,
                "warning"
              )(node);
            }
          },
          UnsupportedPattern: addMessage(
            "Destructuring assignment is not supported"
          ),
        });
      },
      // parser forbids to yield inside module
      YieldExpression() {},
      ThisExpression: addMessage(`"this" keyword is not supported`),
      FunctionExpression: addMessage("Functions are not supported"),
      UpdateExpression: addMessage("Increment and decrement are not supported"),
      CallExpression(node) {
        let calleeName;
        if (node.callee.type === "MemberExpression") {
          if (node.callee.property.type === "Identifier") {
            const methodName = node.callee.property.name;
            const receiverKind =
              node.callee.object.type === "Super"
                ? "unknown"
                : getExpressionNodeValueKind(
                    node.callee.object,
                    variableValues
                  );
            if (isMethodSupported(receiverKind, methodName)) {
              return;
            }
            calleeName = methodName;
          }
        } else if (node.callee.type === "Identifier") {
          calleeName = node.callee.name;
        }
        if (calleeName) {
          addMessage(`"${calleeName}" function is not supported`)(node);
        } else {
          addMessage("Functions are not supported")(node);
        }
      },
      NewExpression: addMessage("Classes are not supported"),
      SequenceExpression: addMessage(`Only single expression is supported`),
      ArrowFunctionExpression: addMessage("Functions are not supported"),
      TaggedTemplateExpression: addMessage("Tagged template is not supported"),
      ClassExpression: addMessage("Classes are not supported"),
      MetaProperty: addMessage("Imports are not supported"),
      AwaitExpression: addMessage(`"await" keyword is not supported`),
      ImportExpression: addMessage("Imports are not supported"),
    } satisfies ExpressionVisitor);
  } catch (error) {
    const castedError = error as { message: string; pos: number };
    diagnostics.push({
      // tune error position after wrapping expression with parentheses
      from: castedError.pos - 1,
      to: castedError.pos - 1,
      severity: "error",
      // trim auto generated error location
      // to not conflict with tuned position
      message: castedError.message.replaceAll(/\s+\(\d+:\d+\)$/g, ""),
    });
  }
  return diagnostics;
};

const isLiteralNode = (node: Expression): boolean => {
  if (node.type === "Identifier" && node.name === "undefined") {
    return true;
  }
  if (node.type === "Literal") {
    return true;
  }
  if (node.type === "ArrayExpression") {
    return node.elements.every((node) => {
      if (node === null || node.type === "SpreadElement") {
        return false;
      }
      return isLiteralNode(node);
    });
  }
  if (node.type === "ObjectExpression") {
    return node.properties.every((property) => {
      if (property.type === "SpreadElement") {
        return false;
      }
      const key = property.key;
      const isIdentifierKey =
        key.type === "Identifier" && property.computed === false;
      const isLiteralKey = key.type === "Literal";
      return (isLiteralKey || isIdentifierKey) && isLiteralNode(property.value);
    });
  }
  return false;
};

/**
 * check whether provided expression is a literal value
 * like "", 0 or { param: "value" }
 * which does not depends on any variable
 */
export const isLiteralExpression = (expression: string) => {
  try {
    const node = parseCompleteExpression(expression);
    return isLiteralNode(node);
  } catch {
    // treat invalid expression as non-literal
    return false;
  }
};

/** Check whether the complete source is one valid JavaScript expression. */
export const isValidExpression = (expression: string) => {
  try {
    parseCompleteExpression(expression);
    return true;
  } catch {
    return false;
  }
};

export const parseStringLiteralExpression = (expression: string) => {
  try {
    const value = JSON.parse(expression);
    return typeof value === "string" && JSON.stringify(value) === expression
      ? value
      : undefined;
  } catch {
    return;
  }
};

const getStaticMemberPath = (node: Expression): string[] | undefined => {
  if (node.type === "Identifier") {
    return [node.name];
  }
  if (
    node.type === "ChainExpression" ||
    node.type === "ParenthesizedExpression"
  ) {
    return getStaticMemberPath(node.expression);
  }
  if (node.type !== "MemberExpression" || node.object.type === "Super") {
    return;
  }
  const objectPath = getStaticMemberPath(node.object);
  if (objectPath === undefined) {
    return;
  }
  const property = node.property;
  const propertyName = node.computed
    ? property.type === "Literal" &&
      (typeof property.value === "string" || typeof property.value === "number")
      ? String(property.value)
      : undefined
    : property.type === "Identifier"
      ? property.name
      : undefined;
  return propertyName === undefined ? undefined : [...objectPath, propertyName];
};

/**
 * An expression containing only an identifier and statically known member
 * access. Unlike an arbitrary expression, it identifies one data location.
 * The owner of that location must still explicitly support write-back before
 * callers may treat it as writable.
 */
export type DirectPathExpression = Readonly<{
  type: "direct-path";
  expression: string;
  path: readonly [root: string, ...segments: string[]];
}>;

export const parseDirectPathExpression = (
  expression: string
): DirectPathExpression | undefined => {
  try {
    const node = parseCompleteExpression(expression);
    const path = getStaticMemberPath(node);
    if (path === undefined || path.length === 0) {
      return;
    }
    return {
      type: "direct-path",
      expression,
      path: path as [string, ...string[]],
    };
  } catch {
    return;
  }
};

/** Parse an identifier/member expression whose complete property path is static. */
export const parseStaticMemberPath = (expression: string) => {
  try {
    const path = getStaticMemberPath(parseCompleteExpression(expression));
    return path === undefined || path.length === 0 ? undefined : path;
  } catch {
    return;
  }
};

export const getExpressionIdentifiers = (expression: string) => {
  const identifiers = new Set<string>();
  try {
    const root = parseCompleteExpression(expression);
    simple(root, {
      Identifier: (node) => identifiers.add(node.name),
      AssignmentExpression(node) {
        walkAssignmentTarget(node.left, {
          Identifier: (node) => identifiers.add(node.name),
        });
      },
    });
  } catch {
    // empty block
  }
  return identifiers;
};

/**
 * transpile expression into executable one
 *
 * add optional chaining operator to every member expression
 * to access any field without runtime errors
 *
 * replace variable names if necessary
 */
export const transpileExpression = ({
  expression,
  executable = false,
  replaceVariable,
}: {
  expression: string;
  executable?: boolean;
  replaceVariable?: (
    identifier: string,
    assignee: boolean
  ) => string | undefined | void;
}) => {
  let root;
  try {
    root = parseCompleteExpression(expression);
  } catch (error) {
    const message = (error as Error).message;
    // throw new error to trace error in our code instead of acorn
    throw Error(`${message} in ${JSON.stringify(expression)}`);
  }
  const assignmentTargetMemberRanges: [start: number, end: number][] = [];
  if (executable) {
    simple(root, {
      AssignmentExpression(node) {
        walkAssignmentTarget(node.left, {
          MemberExpression(node) {
            assignmentTargetMemberRanges.push([node.start, node.end]);
          },
        });
      },
    });
  }
  const replacements: [start: number, end: number, fragment: string][] = [];
  const replacementIndexByRange = new Map<string, number>();
  const addReplacement = (
    start: number,
    end: number,
    fragment: string,
    { replaceExisting = false }: { replaceExisting?: boolean } = {}
  ) => {
    const range = `${start}:${end}`;
    const existingIndex = replacementIndexByRange.get(range);
    if (existingIndex !== undefined) {
      if (replaceExisting) {
        replacements[existingIndex] = [start, end, fragment];
      }
      return;
    }
    replacementIndexByRange.set(range, replacements.length);
    replacements.push([start, end, fragment]);
  };
  const replaceIdentifier = (node: Identifier, assignee: boolean) => {
    const newName = replaceVariable?.(node.name, assignee);
    if (newName) {
      addReplacement(node.start, node.end, newName, {
        replaceExisting: assignee,
      });
    }
  };
  simple(root, {
    Identifier: (node) => replaceIdentifier(node, false),
    AssignmentExpression(node) {
      walkAssignmentTarget(node.left, {
        Identifier: (node) => replaceIdentifier(node, true),
      });
    },
    MemberExpression(node) {
      if (executable === false || node.optional) {
        return;
      }
      if (
        assignmentTargetMemberRanges.some(
          ([start, end]) => start === node.start && end === node.end
        )
      ) {
        return;
      }
      // a . b -> a ?. b
      if (node.computed === false) {
        const dotIndex = expression.indexOf(".", node.object.end);
        addReplacement(dotIndex, dotIndex, "?");
      }
      // a [b] -> a ?.[b]
      if (node.computed === true) {
        const dotIndex = expression.indexOf("[", node.object.end);
        addReplacement(dotIndex, dotIndex, "?.");
      }
    },
    CallExpression(node) {
      if (executable === false || node.optional) {
        return;
      }
      // Add optional chaining to method calls: obj.method() -> obj?.method?.()
      if (node.callee.type === "MemberExpression") {
        // Find the opening parenthesis after the method name
        const openParenIndex = expression.indexOf("(", node.callee.end);
        if (openParenIndex !== -1) {
          addReplacement(openParenIndex, openParenIndex, "?.");
        }
      }
    },
  });
  // order from the latest to the first insertion to not break other positions
  replacements.sort(([leftStart], [rightStart]) => rightStart - leftStart);
  for (const [start, end, fragment] of replacements) {
    const before = expression.slice(0, start);
    const after = expression.slice(end);
    expression = before + fragment + after;
  }
  return expression;
};

/** Parse an array expression into its individual element expressions. */
export const parseArrayExpression = (expression: string) => {
  const items: string[] = [];
  let root;
  try {
    root = parseCompleteExpression(expression);
  } catch {
    return;
  }
  if (root.type !== "ArrayExpression") {
    return;
  }
  for (const element of root.elements) {
    if (element === null || element.type === "SpreadElement") {
      return;
    }
    items.push(expression.slice(element.start, element.end));
  }
  return items;
};

/** Parse a complete object expression into its property expressions. */
export const parseExpressionObject = (expression: string) => {
  const fields = new Map<string, string>();
  let root;
  try {
    root = parseCompleteExpression(expression, { preserveParens: true });
  } catch {
    return;
  }
  while (root.type === "ParenthesizedExpression") {
    root = root.expression;
  }
  if (root.type !== "ObjectExpression") {
    return;
  }
  for (const property of root.properties) {
    if (property.type === "SpreadElement" || property.computed) {
      return;
    }
    const key =
      property.key.type === "Identifier"
        ? property.key.name
        : property.key.type === "Literal" &&
            typeof property.key.value === "string"
          ? property.key.value
          : undefined;
    if (key === undefined || fields.has(key)) {
      return;
    }
    fields.set(key, expression.slice(property.value.start, property.value.end));
  }
  return fields;
};

const formatObjectPropertyKey = (key: string) => {
  const source = `{${key}: 0}`;
  try {
    const root = parseExpressionAt(source, 0, { ecmaVersion: "latest" });
    const property =
      root.type === "ObjectExpression" ? root.properties[0] : undefined;
    if (
      root.type === "ObjectExpression" &&
      root.end === source.length &&
      root.properties.length === 1 &&
      property?.type === "Property" &&
      property.computed === false &&
      property.method === false &&
      property.shorthand === false &&
      property.key.type === "Identifier" &&
      property.key.start === 1 &&
      property.key.end === key.length + 1
    ) {
      return key;
    }
  } catch {
    // Property names outside JavaScript's identifier grammar remain quoted.
  }
  return JSON.stringify(key);
};

/** Generate property expressions as a readable JavaScript object expression. */
export const generateObjectExpression = (
  fields: ReadonlyMap<string, string>
) => {
  let expression = "{\n";
  for (const [key, value] of fields) {
    const indentedValue = value.replaceAll("\n", "\n  ");
    expression += `  ${formatObjectPropertyKey(key)}: ${indentedValue},\n`;
  }
  return `${expression}}`;
};

/** Generate a JSON value as readable JavaScript expression source. */
export const generateJsonExpression = (value: unknown): string => {
  if (
    value === null ||
    typeof value === "boolean" ||
    typeof value === "string"
  ) {
    return JSON.stringify(value);
  }
  if (typeof value === "number" && Number.isFinite(value)) {
    return JSON.stringify(value);
  }
  if (Array.isArray(value)) {
    return `[${value.map(generateJsonExpression).join(", ")}]`;
  }
  if (typeof value === "object") {
    const prototype = Object.getPrototypeOf(value);
    if (prototype !== Object.prototype && prototype !== null) {
      throw new Error("Only JSON values can be converted to expression source");
    }
    return generateObjectExpression(
      new Map(
        Object.entries(value).map(([key, item]) => [
          key,
          generateJsonExpression(item),
        ])
      )
    );
  }
  throw new Error("Only JSON values can be converted to expression source");
};

type ParsedJsonExpression =
  | { success: true; value: unknown }
  | { success: false };

const parseJsonNode = (node: Expression): ParsedJsonExpression => {
  if (node.type === "Literal") {
    if (
      node.regex !== undefined ||
      typeof node.value === "bigint" ||
      (typeof node.value === "number" && Number.isFinite(node.value) === false)
    ) {
      return { success: false };
    }
    return { success: true, value: node.value };
  }
  if (
    node.type === "UnaryExpression" &&
    node.operator === "-" &&
    node.argument.type === "Literal" &&
    typeof node.argument.value === "number" &&
    Number.isFinite(node.argument.value)
  ) {
    return { success: true, value: -node.argument.value };
  }
  if (node.type === "ArrayExpression") {
    const value = [];
    for (const element of node.elements) {
      if (element === null || element.type === "SpreadElement") {
        return { success: false };
      }
      const parsed = parseJsonNode(element);
      if (parsed.success === false) {
        return parsed;
      }
      value.push(parsed.value);
    }
    return { success: true, value };
  }
  if (node.type === "ObjectExpression") {
    const entries: [string, unknown][] = [];
    const keys = new Set<string>();
    for (const property of node.properties) {
      if (
        property.type === "SpreadElement" ||
        property.computed ||
        property.method ||
        property.shorthand
      ) {
        return { success: false };
      }
      const key =
        property.key.type === "Identifier"
          ? property.key.name
          : property.key.type === "Literal" &&
              typeof property.key.value === "string"
            ? property.key.value
            : undefined;
      const parsed = parseJsonNode(property.value);
      if (key === undefined || keys.has(key) || parsed.success === false) {
        return { success: false };
      }
      keys.add(key);
      entries.push([key, parsed.value]);
    }
    return { success: true, value: Object.fromEntries(entries) };
  }
  return { success: false };
};

/** Parse JSON-compatible JavaScript expression source without evaluating it. */
export const parseJsonExpression = (expression: string | undefined) => {
  if (expression === undefined) {
    return;
  }
  try {
    const node = parseCompleteExpression(expression);
    const parsed = parseJsonNode(node);
    return parsed.success ? parsed.value : undefined;
  } catch {
    return;
  }
};