sort_methods.ts

import {
  type MethodDeclaration,
  Project,
  SyntaxKind,
} from "jsr:@ts-morph/ts-morph@28.0.0";

const paths = Deno.args.length === 0
  ? [
    "client/1_client_generic.ts",
    "client/6_client_dispatcher.ts",
    "client/6_client.ts",
  ]
  : Deno.args;

const sectionPattern =
  /^(\s*\/\/\s*\n\s*\/\/ ========================= (.+) ========================= \/\/\s*\n\s*\/\/\s*\n)/gm;

const sectionMethodTags = new Map([
  ["CONNECTION", "cn"],
  ["AUTHORIZATION", "au"],
  ["LOW-LEVEL", "ll"],
  ["ACCOUNT", "ac"],
  ["MESSAGES", "ms"],
  ["POLLS", "pl"],
  ["CHECKLISTS", "cl"],
  ["FILES", "fs"],
  ["CHATS", "ch"],
  ["CALLBACK QUERIES", "cq"],
  ["INLINE QUERIES", "iq"],
  ["BOTS", "bo"],
  ["REACTIONS", "re"],
  ["STORIES", "st"],
  ["STORY ALBUMS", "sa"],
  ["NETWORK STATISTICS", "ns"],
  ["VIDEO CHATS", "vc"],
  ["PAYMENTS", "pa"],
  ["CONTACTS", "co"],
  ["COMMUNITIES", "ct"],
  ["TRANSLATIONS", "ta"],
  ["GIFTS", "gf"],
  ["GIFT COLLECTIONS", "gc"],
  ["TAKEOUTS", "to"],
  ["STICKER SETS", "ss"],
  ["MANAGED BOTS", "mb"],
  ["GUEST QUERIES", "gq"],
  ["SECRET CHATS", "sc"],
]);

const project = new Project();

function getMainClass(sourceFile: ReturnType<Project["addSourceFileAtPath"]>) {
  const classDeclaration = sourceFile.getClasses()[0];
  if (classDeclaration === undefined) {
    throw new Error(`No class found in ${sourceFile.getFilePath()}`);
  }
  return classDeclaration;
}

function getSectionBoundaries(text: string) {
  const sections = [...text.matchAll(sectionPattern)].map((match) => ({
    contentStart: match.index! + match[1].length,
    sectionStart: match.index!,
    tag: sectionMethodTags.get(match[2]),
    title: match[2],
  }));

  return sections.map((section, index) => ({
    contentEnd: sections[index + 1]?.sectionStart ?? text.length,
    contentStart: section.contentStart,
    tag: section.tag,
    title: section.title,
  }));
}

function getMethodSegmentStart(
  method: MethodDeclaration,
  sectionStart: number,
): number {
  return Math.max(method.getFullStart(), sectionStart);
}

function normalizeMethodSegment(text: string): string {
  return `  ${text.trimStart()}`;
}

function normalizeSortedRun(
  sortedText: string,
  replacementStart: number,
  sectionStart: number,
): string {
  return replacementStart === sectionStart ? sortedText : `\n\n${sortedText}`;
}

function isSortableMethod(method: MethodDeclaration) {
  return !method.getName().startsWith("#");
}

function getMethodTag(text: string, method: MethodDeclaration) {
  return text.slice(method.getFullStart(), method.getStart()).match(
    /@method\s+(\w+)/,
  )?.[1];
}

function getSectionForMember(
  sections: ReturnType<typeof getSectionBoundaries>,
  member: { getEnd(): number; getStart(): number },
) {
  return sections.find((section) =>
    member.getStart() >= section.contentStart &&
    member.getEnd() <= section.contentEnd
  );
}

for (const path of paths) {
  const sourceFile = project.addSourceFileAtPath(path);
  const classDeclaration = getMainClass(sourceFile);
  const members = classDeclaration.getMembers();
  const text = sourceFile.getFullText();
  const sections = getSectionBoundaries(text);
  const replacements = [];
  const sortableMethods = members.flatMap((member) => {
    if (
      member.getKind() !== SyntaxKind.MethodDeclaration ||
      !isSortableMethod(member as MethodDeclaration)
    ) {
      return [];
    }

    const method = member as MethodDeclaration;
    const section = getSectionForMember(sections, method);
    if (section === undefined) {
      return [];
    }

    return [{
      index: members.indexOf(method),
      method,
      section,
      tag: getMethodTag(text, method) ?? section.tag,
    }];
  });

  for (const section of sections) {
    const sectionMethods = sortableMethods
      .filter((entry) => entry.section === section)
      .map((entry) => entry.method);
    const targetMethods = sortableMethods
      .filter((entry) => entry.tag === section.tag)
      .toSorted((left, right) => {
        const byName = left.method.getName().localeCompare(
          right.method.getName(),
        );
        return byName === 0 ? left.index - right.index : byName;
      });

    if (sectionMethods.length === 0 && targetMethods.length === 0) {
      continue;
    }

    const replacementStart = sectionMethods.length === 0
      ? section.contentStart
      : getMethodSegmentStart(sectionMethods[0], section.contentStart);
    const replacementEnd = sectionMethods.at(-1)?.getEnd() ??
      section.contentStart;
    const sortedText = targetMethods
      .map((entry) =>
        normalizeMethodSegment(
          text.slice(
            getMethodSegmentStart(
              entry.method,
              entry.section.contentStart,
            ),
            entry.method.getEnd(),
          ),
        )
      )
      .join("\n\n");

    replacements.push({
      end: replacementEnd,
      start: replacementStart,
      text: normalizeSortedRun(
        sortedText,
        replacementStart,
        section.contentStart,
      ),
    });
  }

  for (const replacement of replacements.toReversed()) {
    sourceFile.replaceText(
      [replacement.start, replacement.end],
      replacement.text,
    );
  }

  sourceFile.saveSync();
}