a naive and WIP babel transform for inline "use server" closures

babel-plugin-inline-actions.cjs

// @ts-check
/* eslint-disable @typescript-eslint/no-var-requires */
const { declare: declarePlugin } = require("@babel/helper-plugin-utils");
const { addNamed } = require("@babel/helper-module-imports");

const crypto = require("node:crypto");
const { pathToFileURL } = require("node:url");

// TODO: handle inline actions calling each other...? sounds tricky...

// duplicated from packages/core/src/build/build.ts

const getHash = (s) =>
  crypto.createHash("sha1").update(s).digest().toString("hex");

// FIXME: this is can probably result in weird bugs --
// we're only looking at the name of the module,
// so we'll give it the same id even if the contents changed completely!
// this id should probably look at some kind of source-hash...
const getServerActionModuleId = (resource) =>
  getHash(pathToFileURL(resource).href);

module.exports = declarePlugin((api) => {
  api.assertVersion(7);
  const { types: t } = api;

  const getFilename = (state) => state.file.opts.filename ?? "<unnamed>";

  const addedImports = new Map();
  const addRSDWImport = (path, state) => {
    const filename = getFilename(state);
    if (addedImports.has(filename)) {
      return addedImports.get(filename);
    }
    const id = addNamed(
      path,
      "registerServerReference",
      "react-server-dom-webpack/server"
    );
    addedImports.set(filename, id);
    return id;
  };

  const hasUseServerDirective = (path) => {
    const { body } = path.node;
    if (!t.isBlockStatement(body)) {
      return false;
    }
    if (
      !(
        body.directives.length >= 1 &&
        body.directives.some((d) => d.value.value === "use server")
      )
    ) {
      return false;
    }
    return true;
  };

  const getFreeVariables = (path) => {
    /** @type {Set<string>} */
    const freeVariablesSet = new Set();
    // Find free variables by walking through the function body.
    path.traverse({
      Identifier(innerPath) {
        const { name } = innerPath.node;
        if (freeVariablesSet.has(name)) {
          return;
        }
        if (
          !path.scope.hasOwnBinding(name) &&
          path.parentPath.scope.hasOwnBinding(name)
        ) {
          freeVariablesSet.add(name);
        }
      },
    });
    const freeVariables = [...freeVariablesSet];
    return freeVariables;
  };

  function extractInlineActionToTopLevel(path, state, { body, freeVariables }) {
    const freeVarsParam = t.objectPattern(
      freeVariables.map((variable) => {
        return t.objectProperty(t.identifier(variable), t.identifier(variable));
      })
    );

    const extractedFunctionExpr = t.arrowFunctionExpression(
      [freeVarsParam, ...path.node.params],
      t.blockStatement(body.body)
    );

    const moduleScope = path.scope.getProgramParent();
    const extractedIdentifier =
      moduleScope.generateUidIdentifier("$$INLINE_ACTION");

    const filePath = getFilename(state);
    const actionModuleId = getServerActionModuleId(filePath);

    // Create a top-level declaration for the extracted function.
    const functionDeclaration = t.exportNamedDeclaration(
      t.variableDeclaration("const", [
        t.variableDeclarator(extractedIdentifier, extractedFunctionExpr),
      ])
    );

    const registerServerReferenceId = addRSDWImport(path, state);

    const registerStmt = t.expressionStatement(
      t.callExpression(registerServerReferenceId, [
        extractedIdentifier,
        t.stringLiteral(actionModuleId),
        t.stringLiteral(extractedIdentifier.name),
      ])
    );

    // TODO: is this the best way to insert a top-level declaration...?
    moduleScope.block.body.push(functionDeclaration, registerStmt);
    return {
      extractedIdentifier,
      getReplacement: () =>
        getInlineActionReplacement({ id: extractedIdentifier, freeVariables }),
    };
  }

  const getInlineActionReplacement = ({ id, freeVariables }) => {
    return t.callExpression(t.memberExpression(id, t.identifier("bind")), [
      t.nullLiteral(),
      t.objectExpression(
        freeVariables.map((variable) => {
          return t.objectProperty(
            t.identifier(variable),
            t.identifier(variable)
          );
        })
      ),
    ]);
  };

  return {
    visitor: {
      // Find all arrow functions with the "use server" pragma in the body.
      ArrowFunctionExpression(path, state) {
        const { body } = path.node;
        if (!t.isBlockStatement(body)) {
          return;
        }
        if (!hasUseServerDirective(path)) {
          return;
        }

        const freeVariables = getFreeVariables(path);

        const { getReplacement } = extractInlineActionToTopLevel(path, state, {
          freeVariables,
          body,
        });

        path.replaceWith(getReplacement());
      },
      FunctionDeclaration(path, state) {
        if (!hasUseServerDirective(path)) {
          return;
        }

        const fnId = path.node.id;
        if (!fnId) {
          throw new Error(
            "Internal error: expected FunctionDeclaration to have a name"
          );
        }

        const freeVariables = getFreeVariables(path).filter(
          // TODO: why is `getFreeVariables` returning the function's name too?
          // TODO: if we're referencing other (named) inline actions, they'll end up in here, and do something stupid
          (name) => name !== fnId.name
        );
        const { getReplacement } = extractInlineActionToTopLevel(path, state, {
          freeVariables,
          body: path.node.body,
        });

        path.replaceWith(
          t.variableDeclaration("var", [
            t.variableDeclarator(fnId, getReplacement()),
          ])
        );
      },
      FunctionExpression(path, state) {
        if (hasUseServerDirective(path)) {
          throw new Error("TODO - inline `function () {}`");
        }
      },
    },
  };
});