trace-async-actions

trace-async-actions.ts

import { parse } from "@babel/core";
import traverse from "@babel/traverse";
import * as t from "@babel/types";
import { exec, execSync } from "child_process";
import fs from "fs";
import path from "path";

type SelectorUsage = {
  selector: string;
  consumers: string[];
};

type StateSelectors = {
  stateSlice: string;
  selectorUsages: SelectorUsage[];
};

const MERMAID_CONFIG = {
  securityLevel: "loose",
  maxTextSize: 200000,
} as const;

function findConsumersForSelector(selector: string) {
  const consumers = collectPaths(`rg "${selector}" src --json`)
    .filter((path) => !path.includes("selectors"))
    .map((path) => path.replace("src/modules/", ""));

  return dedupe(consumers);
}

function findSelectors(stateSlice: string) {
  const rootSelectors = dedupe(findSelectorsReferencing(stateSlice));
  let previousSelectors = dedupe(
    rootSelectors.flatMap((selector) => findSelectorsReferencing(selector)),
  );
  let nextSelectors = previousSelectors;

  do {
    let lastIter = nextSelectors;
    nextSelectors = dedupe(
      previousSelectors.flatMap((selector) =>
        findSelectorsReferencing(selector),
      ),
    );
    previousSelectors = lastIter;
  } while (nextSelectors.length > previousSelectors.length);

  return nextSelectors;
}

function findSelectorsReferencing(stateOrSelector: string) {
  const filesWithRootSelectors = [
    ...collectPaths(`rg "${stateOrSelector}" src -g **/selectors* --json`),
    ...collectPaths(`rg "${stateOrSelector}" src -g **/selectors/* --json`),
  ];

  return filesWithRootSelectors.flatMap((file) => {
    const selectorNames: string[] = [];
    const tree = parseFile(file);

    traverse(tree, {
      enter(path) {
        if (
          !path.isExportNamedDeclaration() ||
          !t.isVariableDeclaration(path.node.declaration)
        ) {
          return;
        }

        const serialized = JSON.stringify(path.node.declaration);
        const searchTerms = stateOrSelector.includes(".")
          ? stateOrSelector.split(".").map((term) => `"${term}"`)
          : [`"${stateOrSelector}"`];

        if (!searchTerms.every((term) => serialized.includes(term))) {
          return;
        }

        const [declaration] = path.node.declaration.declarations;
        if (
          !t.isVariableDeclarator(declaration) ||
          !t.isIdentifier(declaration.id)
        ) {
          return;
        }

        selectorNames.push(declaration.id.name);
      },
    });

    return selectorNames;
  });
}

function findStateSlices(
  actionName: string,
  stateSliceHints: string[],
  hintBehaviour: "include" | "exclude",
) {
  if (hintBehaviour === "exclude") {
    return dedupe(stateSliceHints);
  }

  const reducers = [
    ...collectPaths(`rg "${actionName}\.success" src -g **/reducer* --json`),
    ...collectPaths(`rg "${actionName}\.success" src -g **/reducer/* --json`),
  ];

  const stateSlices = reducers.flatMap((reducer) => {
    const slices: string[] = [];
    const tree = parseFile(reducer);

    traverse(tree, {
      enter(path) {
        if (!path.isSwitchCase() || !path.node.test) {
          return;
        }

        const serialized = JSON.stringify(path.node.test);
        if (!serialized.includes(`"${actionName}"`)) {
          return;
        }

        if (path.node.consequent.length <= 0) {
          path = path.getNextSibling();
          while (path.isSwitchCase() && path.node.consequent.length <= 0) {
            path = path.getNextSibling();
          }
        }

        if (!path.isSwitchCase()) {
          return;
        }

        let [consequent] = path.node.consequent;
        if (t.isBlockStatement(consequent)) {
          const returnStatement = consequent.body.find((statement) =>
            t.isReturnStatement(statement),
          );

          if (returnStatement) {
            consequent = returnStatement;
          }
        }

        if (
          !t.isReturnStatement(consequent) ||
          !t.isObjectExpression(consequent.argument)
        ) {
          return;
        }

        return consequent.argument.properties.forEach((property) => {
          if (!t.isObjectProperty(property) || !t.isIdentifier(property.key)) {
            return;
          }

          slices.push(property.key.name);
        });
      },
    });

    return slices;
  });

  return dedupe([
    ...stateSliceHints,
    ...stateSlices.filter(
      (slice) =>
        ![
          "request",
          "requests",
          "loading",
          "isLoading",
          "errors",
          "error",
        ].includes(slice),
    ),
  ]);
}

function findInitiators(actionName: string) {
  const initiators = collectPaths(
    `rg "${actionName}\.request" src --json`,
  ).filter((path) => !path.includes("reducer"));
  return initiators;
}

function buildGraph(
  actionName: string,
  initiators: string[],
  stateSelectors: StateSelectors[],
) {
  let graph = `
flowchart LR
  `;

  // Add initiators to graph
  graph = `${graph}\nsubgraph initiators`;
  initiators.forEach((path, i) => {
    graph = `${graph}\nI${i}[${path}]`;
  });
  graph = `${graph}\nend`;

  // Add action to graph
  graph = `
${graph}
subgraph action
A[${actionName}]
end
`;

  // Link initiators to action
  initiators.forEach((_, i) => {
    graph = `${graph}\nI${i}-->A`;
  });

  // Link action to state slices
  Array.from({ length: stateSelectors.length }).forEach((_, i) => {
    graph = `${graph}\nA-->S${i}`;
  });

  // Add state slices to graph
  graph = `${graph}\nsubgraph state`;
  stateSelectors.forEach(({ stateSlice }, i) => {
    graph = `${graph}\nS${i}(${stateSlice})`;
  });
  graph = `${graph}\nend`;

  // Add + link selectors to state slices
  graph = `${graph}\nsubgraph selectors`;
  stateSelectors.forEach(({ selectorUsages }, i) => {
    selectorUsages.forEach(({ selector }) => {
      // Add the selector
      graph = `${graph}\nSE_${selector}(${selector})`;

      // Link the selector to state
      graph = `${graph}\nS${i}-->SE_${selector}`;
    });
  });
  graph = `${graph}\nend`;

  // Add + link consumers to selectors
  graph = `${graph}\nsubgraph consumers`;
  stateSelectors.forEach(({ selectorUsages }) => {
    selectorUsages.forEach(({ consumers, selector }) => {
      consumers.forEach((consumer) => {
        // Add consumer
        graph = `${graph}\nC_${consumer}[${consumer}]`;

        // Link selector to consumer
        graph = `${graph}\nSE_${selector}-->C_${consumer}`;
      });
    });
  });
  graph = `${graph}\nend`;

  return graph;
}

function saveAndOpenGraph(actionName: string, graph: string, quiet: boolean) {
  const mermaidDir = path.join(
    process.cwd(),
    "output",
    "trace-async-action",
    "mermaid",
  );
  if (!fs.existsSync(mermaidDir)) {
    fs.mkdirSync(mermaidDir, { recursive: true });
  }

  // Write mermaid config if needed
  const mermaidConfigPath = path.join(mermaidDir, "mermaidConfig.json");
  if (!fs.existsSync(mermaidConfigPath)) {
    fs.writeFileSync(mermaidConfigPath, JSON.stringify(MERMAID_CONFIG));
  }

  // Save mermaid graph
  const fileName = `${actionName}--${new Date().toISOString()}`;
  const mermaidPath = path.join(mermaidDir, `${fileName}.mmd`);
  fs.writeFileSync(mermaidPath, graph);

  const svgDir = path.join(
    process.cwd(),
    "output",
    "trace-async-action",
    "svg",
  );
  if (!fs.existsSync(svgDir)) {
    fs.mkdirSync(svgDir, { recursive: true });
  }

  // Save mermaid svg
  const svgPath = path.join(svgDir, `${fileName}.svg`);
  execSync(`mmdc -i ${mermaidPath} -o ${svgPath} -c ${mermaidConfigPath}`);

  if (!quiet) {
    exec(`open ${svgPath}`);
  } else {
    console.log(">> Saved graph to", svgPath);
  }
}

function parseFile(filePath: string) {
  const fileContents = fs.readFileSync(filePath);
  const tree = parse(fileContents.toString(), {
    filename: filePath,
    plugins: ["@babel/plugin-syntax-typescript"],
  });

  return tree as t.Node;
}

function collectPaths(rgCommand: string) {
  try {
    const result = execSync(rgCommand).toString();

    return dedupe(
      result
        .trim()
        .split("\n")
        .map((line) => JSON.parse(line))
        .filter(
          (res) => res.type === "match" && !res.data.path.text.includes("spec"),
        )
        .map((match) => match.data.path.text),
    );
  } catch {
    return [];
  }
}

function dedupe(list: string[]) {
  return [...new Set(list)];
}

function maybe(cb: () => any) {
  try {
    return ["success", cb()] as const;
  } catch (error) {
    return ["error", error] as const;
  }
}

function parseParams() {
  const helpText =
    "yarn trace-async-action [actionName] -s [state slice hints] -x [state slices to exclude] -q [quiet mode (don't open browser)]";
  const [_, __, param2, param3, param4] = process.argv;

  if (!param2 || param2.length <= 0) {
    throw new Error(helpText);
  }

  let actionName = param2;
  let stateSliceHints: string[] = [];
  let hintBehaviour = "include";
  let quiet = false;

  if (process.argv.includes("-h") || process.argv.includes("--help")) {
    throw new Error(helpText);
  }

  // Parse arguments
  switch (param3) {
    // Add state slice hints (ie, if the entire payload is spread into state, let the tool know what the slice(s) are called)
    case "-s":
    case "--state":
      if (!param4 || param4.length <= 0) {
        throw new Error(helpText);
      } else {
        stateSliceHints = param4
          .split(",")
          .filter((hint) => hint.trim().length > 0);
      }
      break;

    // Add state slices to exclude (ie, if the state slice is not relevant to the action, let the tool know to exclude it)
    case "-x":
    case "--only":
      if (!param4 || param4.length <= 0) {
        throw new Error(helpText);
      } else {
        stateSliceHints = param4
          .split(",")
          .filter((hint) => hint.trim().length > 0);
        hintBehaviour = "exclude";
      }
      break;
  }

  // Add quiet mode (don't open browser)
  if (process.argv.includes("-q") || process.argv.includes("--quiet")) {
    quiet = true;
  }

  return {
    actionName,
    stateSliceHints,
    hintBehaviour,
    quiet,
  };
}

(function () {
  const [hasRipgrep] = maybe(() => execSync("which rg"));
  if (hasRipgrep === "error") {
    console.error(
      "Command `rg` not found, `brew install ripgrep` and try again",
    );
    return;
  }

  const [hasMermaidCli] = maybe(() => execSync("which mmdc"));
  if (hasMermaidCli === "error") {
    console.error(
      "Command `mmdc` not found, `npm install -g @mermaid-js/mermaid-cli` and try again",
    );
    return;
  }

  const [type, value] = maybe(parseParams);
  if (type === "error") {
    console.log((value as Error).message);
    return;
  }

  const { actionName, stateSliceHints, hintBehaviour, quiet } = value;

  console.log(">> Finding initiators...");
  const initiators = findInitiators(actionName);

  console.log(">> Finding state slices...");
  const stateSlices = findStateSlices(
    actionName,
    stateSliceHints,
    hintBehaviour,
  );

  if (!stateSlices.length) {
    console.error(
      '>> No state modification found from running "',
      actionName,
      '". Exiting.',
    );
    return;
  }

  console.log(">> Finding selectors and consumers...");
  const stateSelectors = stateSlices.map((stateSlice) => ({
    stateSlice,
    selectorUsages: findSelectors(stateSlice).map((selector) => ({
      selector,
      consumers: findConsumersForSelector(selector),
    })),
  }));

  console.log(">> Building graph...");
  const graph = buildGraph(actionName, initiators, stateSelectors);
  saveAndOpenGraph(actionName, graph, quiet);
})();