index.ts

import {
  tsx,
  FrontEndLanguage,
  type NapiConfig,
  type SgNode,
} from '@ast-grep/napi';

export const foo = 'bar';

export function normalizedFeatureFlagName(flagName: string): string {
  return flagName.startsWith('FEATURE_') || flagName.startsWith('FLAGS[')
    ? flagName
    : `FLAGS['${flagName}']`;
}

export function findFeatureFlagsInsideBinaryExpressionAnd(
  flagName: string,
  source: string,
): SgNode[] {
  const ast = tsx.parse(source);

  flagName = normalizedFeatureFlagName(flagName);

  const program = ast.root();
  const matcher: NapiConfig = {
    language: FrontEndLanguage.Tsx,
    rule: {
      kind: 'binary_expression',
      any: [{ pattern: `${flagName} && $A` }, { pattern: `$A && ${flagName}` }],
      // fix: '$A',
    },
  };

  return program.findAll(matcher);
}

export function findFeatureFlagsInsideTernaryExpression(
  flagName: string,
  source: string,
): SgNode[] {
  const ast = tsx.parse(source);

  flagName = normalizedFeatureFlagName(flagName);

  const program = ast.root();
  const matcher: NapiConfig = {
    language: FrontEndLanguage.Tsx,
    rule: {
      kind: 'ternary_expression',
      any: [{ pattern: `${flagName} ? $A : $B` }],
      // fix: '$A',
    },
  };

  return program.findAll(matcher);
}

export function findFeatureFlagsInsidePair(
  flagName: string,
  source: string,
): SgNode[] {
  // 只考虑 FEATURE_XXX 的情况，因为这些才会从 features.ts 中引入
  if (!flagName.startsWith('FEATURE_')) {
    return [];
  }

  const ast = tsx.parse(source);

  const program = ast.root();
  const matcher: NapiConfig = {
    language: FrontEndLanguage.Tsx,
    rule: {
      kind: 'pair',
      has: {
        field: 'key',
        regex: flagName,
      },
    },
  };

  const nodes = program.findAll(matcher);
  for (const node of nodes) {
    const nextNode = node.next();
    if (nextNode.kind() === ',') {
      const { start } = node.range();
      const { end } = nextNode.range();
      node.range = () => ({ start, end });
    }
  }
  return nodes;
}

export function findFeatureFlagsInsideIfStatement(
  flagName: string,
  source: string,
): SgNode[] {
  const ast = tsx.parse(source);

  // XXX: 这里不需要调用 normalizedFeatureFlagName
  // flagName = normalizedFeatureFlagName(flagName);

  const program = ast.root();
  const isIdentifier = flagName.startsWith('FEATURE_');
  // rule:
  //   kind: subscript_expression
  //   regex: "FLAGS\\['a.b.c'\\]"
  //   inside:
  //     kind: if_statement
  //     stopBy: end
  // ---
  // rule:
  //   kind: identifier
  //   regex: FEATURE_ENABLE_TABLE_MEMORY
  //   inside:
  //     kind: if_statement
  //     stopBy: end
  const matcher: NapiConfig = {
    language: FrontEndLanguage.Tsx,
    rule: {
      kind: isIdentifier ? 'identifier' : 'subscript_expression',
      regex: isIdentifier ? flagName : `FLAGS\\['${flagName}'\\]`,
      inside: {
        kind: 'if_statement',
        stopBy: 'end',
      },
    },
  };

  return program.findAll(matcher).map(node =>
    // if_statement
    //   'if'
    //   parenthesized_expression
    //     identifier

    // if_statement
    //   'if'
    //   subscript_expression
    //     identifier
    //     '['
    //     string
    //     ']'
    node.parent().parent(),
  );
}

export function findFeatureFlagsInsideImportStatement(
  flagName: string,
  source: string,
): SgNode[] {
  // 只考虑 FEATURE_XXX 的情况，因为这些才会从 features.ts 中引入
  if (!flagName.startsWith('FEATURE_')) {
    return [];
  }

  const ast = tsx.parse(source);

  const program = ast.root();
  const matcher: NapiConfig = {
    language: FrontEndLanguage.Tsx,
    rule: {
      kind: 'import_specifier',
      regex: flagName,
      // fix: {
      //   template: '',
      //   expandEnd: {
      //     regex: ','
      //   },
      //   expandStart: {
      //     regex: ','
      //   }
      // }
    },
  };

  // 这里需要扩展一下 import_specifier 的匹配范围，把前后的逗号包含进来
  // 1. 如果后面有逗号，那么把后面逗号包含进来
  // 2. 如果后面是'}'，那么判断前面
  // 2.1 如果前面是逗号，那么把前面逗号包含进来
  // 2.2 如果前面是'{', 那么忽略之
  const nodes = program.findAll(matcher);
  for (const node of nodes) {
    const nextNode = node.next();
    if (nextNode.kind() === ',') {
      const { start } = node.range();
      const { end } = nextNode.range();
      node.range = () => ({ start, end });
    } else if (nextNode.kind() === '}') {
      const prevNode = node.prev();
      if (prevNode.kind() === ',') {
        const { start } = prevNode.range();
        const { end } = node.range();
        node.range = () => ({ start, end });
      } else if (prevNode.kind() === '{') {
        // import { FLAGS_xxx } from 'xx';
        // import_statement
        //   'import'
        //   import_clause
        //     named_imports
        //       '{'
        //       import_specifier
        //       '}'
        const importStmtNode = node
          .parent() // named_imports
          .parent() // import_clause
          .parent(); // import_statement
        const { start, end } = importStmtNode.range();
        node.range = () => ({ start, end });
      }
    }
  }
  return nodes;
}

export function findFeatureFlagsInsideBinaryExpressionOr(
  flagName: string,
  source: string,
): SgNode[] {
  const ast = tsx.parse(source);

  flagName = normalizedFeatureFlagName(flagName);

  const program = ast.root();
  const matcher: NapiConfig = {
    language: FrontEndLanguage.Tsx,
    rule: {
      kind: 'binary_expression',
      any: [{ pattern: `${flagName} || $A` }, { pattern: `$A || ${flagName}` }],
      // fix: 'true',
    },
  };

  return program.findAll(matcher);
}

/**
 * 对于 FLAG_foo 替换的时候，需要考虑如下的情况
 *
 * 1. if (FLAG_foo) {} else {}
 * 2. if (x) {} else if (FLAG_foo) {} else {}
 * 3. FLAG_foo ? 'yes' : 'no'
 * 4. x && FLAG_foo ? 'yes' : 'no'
 * 5. x || FLAG_foo ? 'yes' : 'no'
 * 6. {x && FLAG_foo}
 *
 * @param flagName The feature flag name.
 * @param source The source code.
 * @return The source code after removing the feature flag.
 */
export function removeFeatureFlag(flagName: string, source: string): string {
  let nodes: SgNode[];

  // rule:
  //   kind: binary_expression
  //   any:
  //     - pattern: $A && FEATURE_ENABLE_TABLE_MEMORY
  //     - pattern: FEATURE_ENABLE_TABLE_MEMORY && $A
  // fix: "$A"
  nodes = findFeatureFlagsInsideBinaryExpressionAnd(flagName, source);
  source = copySource(source, nodes, node => node.getMatch('A').text());

  // rule:
  //   kind: binary_expression
  //   any:
  //     - pattern: $A || FEATURE_ENABLE_TABLE_MEMORY
  //     - pattern: FEATURE_ENABLE_TABLE_MEMORY || $A
  // fix: "true"
  nodes = findFeatureFlagsInsideBinaryExpressionOr(flagName, source);
  source = copySource(source, nodes, () => 'true');

  // rule:
  //   kind: ternary_expression
  //   any:
  //     - pattern: "FEATURE_ENABLE_TABLE_MEMORY ? $A : $B"
  // fix: "$A"
  nodes = findFeatureFlagsInsideTernaryExpression(flagName, source);
  source = copySource(source, nodes, node => node.getMatch('A').text());

  if (flagName.startsWith('FEATURE_')) {
    // rule:
    //   kind: import_specifier
    //   regex: FEATURE_ENABLE_TABLE_MEMORY
    // fix: ""
    nodes = findFeatureFlagsInsideImportStatement(flagName, source);
    source = copySource(source, nodes, () => '');

    // rule:
    //   kind: pair
    //   has:
    //     field: key
    //     regex: FEATURE_ENABLE_TABLE_MEMORY
    // fix: ''
    nodes = findFeatureFlagsInsidePair(flagName, source);
    source = copySource(source, nodes, () => '');
  }

  // rule:
  //   kind: if_statement
  //   has:
  //     regex: FEATURE_ENABLE_TABLE_MEMORY
  nodes = findFeatureFlagsInsideIfStatement(flagName, source);
  source = copySource(source, nodes, ifStmtNode => {
    // if_statement
    //   'if'
    //   parenthesized_expression
    //   statement_block
    //     '{'
    //     expression_statement                 <---- FIND IT (expression_statement)
    //     '}'
    //   else_clause
    //     'else'
    //     statement_block

    // if_statement
    //   'if'
    //   parenthesized_expression
    //   expression_statement                   <---- FIND IT (expression_statement)
    //   else_clause
    //     'else'
    //     statement_block

    // if_statement
    //   'if'
    //   parenthesized_expression
    //   statement_block
    //   else_clause
    //     if_statement                        <---- parent().kind() is 'else_clause'
    //       'if'
    //       parenthesized_expression
    //       statement_block                   <---- FIND IT (statement_block)

    const consequenceNode = ifStmtNode.field('consequence');
    if (!consequenceNode) {
      return '';
    }

    if (
      /**
       * if (FLAG_a)
       *  console.log(10);
       */
      consequenceNode.kind() === 'expression_statement' ||
      /**
       * if (x) {
       * } else if (FLAG_a) {
       *  console.log(10);
       * }
       */
      ifStmtNode.parent().kind() === 'else_clause'
    ) {
      return consequenceNode.text();
    } else if (consequenceNode.kind() === 'statement_block') {
      /**
       * if (FLAG_a) {
       *  console.log(10);
       * }
       */
      for (const childNode of consequenceNode.children()) {
        if (childNode.kind() === 'expression_statement') {
          return childNode.text();
        }
      }
    }

    return '';
  });

  return source;
}

// xxxxxxxxxxxx[++++xxxx++++]xxxxxxxxx[+++xxxx++++]xxxxxxx[++xx+]xxxxx
//             ^    s  e    ^         ^   s  e    ^       ^  se ^
//             S            E         S           E       S     E
export function copySource(
  source: string,
  nodes: SgNode[],
  callback: (node: SgNode) => string,
): string {
  const target = [];

  let start = 0;
  let end = 0;
  for (const node of nodes) {
    const range = node.range();
    end = range.start.index;

    target.push(source.slice(start, end));
    target.push(callback(node));

    start = range.end.index;
  }

  target.push(source.slice(start));

  return target.join('');
}