codeActions.ts

// This file holds code actions derived from diagnostics. There are more code
// actions available in the extension, but they are derived via the analysis
// OCaml binary.
import * as p from "vscode-languageserver-protocol";
import * as utils from "./utils";
import { fileURLToPath } from "url";

export type filesCodeActions = {
  [key: string]: { range: p.Range; codeAction: p.CodeAction }[];
};

interface findCodeActionsConfig {
  diagnostic: p.Diagnostic;
  diagnosticMessage: string[];
  file: string;
  range: p.Range;
  addFoundActionsHere: filesCodeActions;
}

let wrapRangeInText = (
  range: p.Range,
  wrapStart: string,
  wrapEnd: string
): p.TextEdit[] => {
  // We need to adjust the start of where we replace if this is a single
  // character on a single line.
  let offset =
    range.start.line === range.end.line &&
    range.start.character === range.end.character
      ? 1
      : 0;

  let startRange = {
    start: {
      line: range.start.line,
      character: range.start.character - offset,
    },
    end: {
      line: range.start.line,
      character: range.start.character - offset,
    },
  };

  let endRange = {
    start: {
      line: range.end.line,
      character: range.end.character,
    },
    end: {
      line: range.end.line,
      character: range.end.character,
    },
  };

  return [
    {
      range: startRange,
      newText: wrapStart,
    },
    {
      range: endRange,
      newText: wrapEnd,
    },
  ];
};

let insertBeforeEndingChar = (
  range: p.Range,
  newText: string
): p.TextEdit[] => {
  let beforeEndingChar = {
    line: range.end.line,
    character: range.end.character - 1,
  };

  return [
    {
      range: {
        start: beforeEndingChar,
        end: beforeEndingChar,
      },
      newText,
    },
  ];
};

let replaceText = (range: p.Range, newText: string): p.TextEdit[] => {
  return [
    {
      range,
      newText,
    },
  ];
};

let removeTrailingComma = (text: string): string => {
  let str = text.trim();
  if (str.endsWith(",")) {
    return str.slice(0, str.length - 1);
  }

  return str;
};

let extractTypename = (lines: string[]): string => {
  let arrFiltered: string[] = [];

  for (let i = 0; i <= lines.length - 1; i += 1) {
    let line = lines[i];
    if (line.includes("(defined as")) {
      let [typeStr, _] = line.split("(defined as");
      arrFiltered.push(removeTrailingComma(typeStr));
      break;
    } else {
      arrFiltered.push(removeTrailingComma(line));
    }
  }

  return arrFiltered.join("").trim();
};

let takeUntil = (array: string[], startsWith: string): string[] => {
  let res: string[] = [];
  let arr = array.slice();

  let matched = false;
  arr.forEach((line) => {
    if (matched) {
      return;
    }

    if (line.startsWith(startsWith)) {
      matched = true;
    } else {
      res.push(line);
    }
  });

  return res;
};

export let findCodeActionsInDiagnosticsMessage = ({
  diagnostic,
  diagnosticMessage,
  file,
  range,
  addFoundActionsHere: codeActions,
}: findCodeActionsConfig) => {
  diagnosticMessage.forEach((line, index, array) => {
    // Because of how actions work, there can only be one per diagnostic. So,
    // halt whenever a code action has been found.
    let codeActionEtractors = [
      simpleTypeMismatches,
      didYouMeanAction,
      addUndefinedRecordFields,
      simpleConversion,
      applyUncurried,
      simpleAddMissingCases,
    ];

    for (let extractCodeAction of codeActionEtractors) {
      let didFindAction = false;

      try {
        didFindAction = extractCodeAction({
          array,
          codeActions,
          diagnostic,
          file,
          index,
          line,
          range,
        });
      } catch (e) {
        console.error(e);
      }

      if (didFindAction) {
        break;
      }
    }
  });
};

interface codeActionExtractorConfig {
  line: string;
  index: number;
  array: string[];
  file: string;
  range: p.Range;
  diagnostic: p.Diagnostic;
  codeActions: filesCodeActions;
}

type codeActionExtractor = (config: codeActionExtractorConfig) => boolean;

// This action extracts hints the compiler emits for misspelled identifiers, and
// offers to replace the misspelled name with the correct name suggested by the
// compiler.
let didYouMeanAction: codeActionExtractor = ({
  codeActions,
  diagnostic,
  file,
  line,
  range,
}) => {
  if (line.startsWith("Hint: Did you mean")) {
    let regex = /Did you mean ([A-Za-z0-9_]*)?/;
    let match = line.match(regex);

    if (match === null) {
      return false;
    }

    let [_, suggestion] = match;

    if (suggestion != null) {
      codeActions[file] = codeActions[file] || [];
      let codeAction: p.CodeAction = {
        title: `Replace with '${suggestion}'`,
        edit: {
          changes: {
            [file]: [{ range, newText: suggestion }],
          },
        },
        diagnostics: [diagnostic],
        kind: p.CodeActionKind.QuickFix,
        isPreferred: true,
      };

      codeActions[file].push({
        range,
        codeAction,
      });

      return true;
    }
  }

  return false;
};

// This action handles when the compiler errors on certain fields of a record
// being undefined. We then offers an action that inserts all of the record
// fields, with an `assert false` dummy value. `assert false` is so applying the
// code action actually compiles.
let addUndefinedRecordFields: codeActionExtractor = ({
  array,
  codeActions,
  diagnostic,
  file,
  index,
  line,
  range,
}) => {
  if (line.startsWith("Some record fields are undefined:")) {
    let recordFieldNames = line
      .trim()
      .split("Some record fields are undefined: ")[1]
      ?.split(" ");

    // This collects the rest of the fields if fields are printed on
    // multiple lines.
    array.slice(index + 1).forEach((line) => {
      recordFieldNames.push(...line.trim().split(" "));
    });

    if (recordFieldNames != null) {
      codeActions[file] = codeActions[file] || [];

      // The formatter outputs trailing commas automatically if the record
      // definition is on multiple lines, and no trailing comma if it's on a
      // single line. We need to adapt to this so we don't accidentally
      // insert an invalid comma.
      let multilineRecordDefinitionBody = range.start.line !== range.end.line;

      // Let's build up the text we're going to insert.
      let newText = "";

      if (multilineRecordDefinitionBody) {
        // If it's a multiline body, we know it looks like this:
        // ```
        // let someRecord = {
        //   atLeastOneExistingField: string,
        // }
        // ```
        // We can figure out the formatting from the range the code action
        // gives us. We'll insert to the direct left of the ending brace.

        // The end char is the closing brace, and it's always going to be 2
        // characters back from the record fields.
        let paddingCharacters = multilineRecordDefinitionBody
          ? range.end.character + 2
          : 0;
        let paddingContentRecordField = Array.from({
          length: paddingCharacters,
        }).join(" ");
        let paddingContentEndBrace = Array.from({
          length: range.end.character,
        }).join(" ");

        recordFieldNames.forEach((fieldName, index) => {
          if (index === 0) {
            // This adds spacing from the ending brace up to the equivalent
            // of the last record field name, needed for the first inserted
            // record field name.
            newText += "  ";
          } else {
            // The rest of the new record field names will start from a new
            // line, so they need left padding all the way to the same level
            // as the rest of the record fields.
            newText += paddingContentRecordField;
          }

          newText += `${fieldName}: assert false,\n`;
        });

        // Let's put the end brace back where it was (we still have it to the direct right of us).
        newText += `${paddingContentEndBrace}`;
      } else {
        // A single line record definition body is a bit easier - we'll just add the new fields on the same line.
        newText += ", ";
        newText += recordFieldNames
          .map((fieldName) => `${fieldName}: assert false`)
          .join(", ");
      }

      let codeAction: p.CodeAction = {
        title: `Add missing record fields`,
        edit: {
          changes: {
            [file]: insertBeforeEndingChar(range, newText),
          },
        },
        diagnostics: [diagnostic],
        kind: p.CodeActionKind.QuickFix,
        isPreferred: true,
      };

      codeActions[file].push({
        range,
        codeAction,
      });

      return true;
    }
  }

  return false;
};

// This action detects suggestions of converting between mismatches in types
// that the compiler tells us about.
let simpleConversion: codeActionExtractor = ({
  line,
  codeActions,
  file,
  range,
  diagnostic,
}) => {
  if (line.startsWith("You can convert ")) {
    let regex = /You can convert (\w*) to (\w*) with ([\w.]*).$/;
    let match = line.match(regex);

    if (match === null) {
      return false;
    }

    let [_, from, to, fn] = match;

    if (from != null && to != null && fn != null) {
      codeActions[file] = codeActions[file] || [];

      let codeAction: p.CodeAction = {
        title: `Convert ${from} to ${to} with ${fn}`,
        edit: {
          changes: {
            [file]: wrapRangeInText(range, `${fn}(`, `)`),
          },
        },
        diagnostics: [diagnostic],
        kind: p.CodeActionKind.QuickFix,
        isPreferred: true,
      };

      codeActions[file].push({
        range,
        codeAction,
      });

      return true;
    }
  }

  return false;
};

// This action will apply a curried function (essentially inserting a dot in the
// correct place).
let applyUncurried: codeActionExtractor = ({
  line,
  codeActions,
  file,
  range,
  diagnostic,
}) => {
  if (
    line.startsWith(
      "This is an uncurried ReScript function. It must be applied with a dot."
    )
  ) {
    const locOfOpenFnParens = {
      line: range.end.line,
      character: range.end.character + 1,
    };

    codeActions[file] = codeActions[file] || [];
    let codeAction: p.CodeAction = {
      title: `Apply uncurried function call with dot`,
      edit: {
        changes: {
          [file]: [
            {
              range: {
                start: locOfOpenFnParens,
                end: locOfOpenFnParens,
              },
              /*
               * Turns `fn(123)` into `fn(. 123)`.
               */
              newText: `. `,
            },
          ],
        },
      },
      diagnostics: [diagnostic],
      kind: p.CodeActionKind.QuickFix,
      isPreferred: true,
    };

    codeActions[file].push({
      range,
      codeAction,
    });

    return true;
  }

  return false;
};

// This action detects missing cases for exhaustive pattern matches, and offers
// to insert dummy branches (using `failwith("TODO")`) for those branches.
let simpleAddMissingCases: codeActionExtractor = ({
  line,
  codeActions,
  file,
  range,
  diagnostic,
  array,
  index,
}) => {
  if (
    line.startsWith("You forgot to handle a possible case here, for example:")
  ) {
    // This collects the rest of the fields if fields are printed on
    // multiple lines.
    let allCasesAsOneLine = array
      .slice(index + 1)
      .join("")
      .trim();

    let filePath = fileURLToPath(file);

    let newSwitchCode = utils.runAnalysisAfterSanityCheck(filePath, [
      "codemod",
      filePath,
      range.start.line,
      range.start.character,
      "add-missing-cases",
      allCasesAsOneLine,
    ]);

    codeActions[file] = codeActions[file] || [];
    let codeAction: p.CodeAction = {
      title: `Insert missing cases`,
      edit: {
        changes: {
          [file]: replaceText(range, newSwitchCode),
        },
      },
      diagnostics: [diagnostic],
      kind: p.CodeActionKind.QuickFix,
      isPreferred: true,
    };

    codeActions[file].push({
      range,
      codeAction,
    });

    return true;
  }

  return false;
};

// This detects concrete variables or values put in a position which expects an
// optional of that same type, and offers to wrap the value/variable in
// `Some()`.
let simpleTypeMismatches: codeActionExtractor = ({
  line,
  codeActions,
  file,
  range,
  diagnostic,
  array,
  index,
}) => {
  // Examples:
  //
  // 46 │ let as_ = {
  // 47 │   someProp: "123",
  // 48 │   another: "123",
  // 49 │ }
  // 50 │
  // This has type: string
  // Somewhere wanted: option<string>
  //
  // ...but types etc can also be on multilines, so we need a good
  // amount of cleanup.

  let lookFor = "This has type:";

  if (line.startsWith(lookFor)) {
    let thisHasTypeArr = takeUntil(
      [line.slice(lookFor.length), ...array.slice(index + 1)],
      "Somewhere wanted:"
    );
    let somewhereWantedArr = array
      .slice(index + thisHasTypeArr.length)
      .map((line) => line.replace("Somewhere wanted:", ""));

    let thisHasType = extractTypename(thisHasTypeArr);
    let somewhereWanted = extractTypename(somewhereWantedArr);

    // Switching over an option
    if (thisHasType === `option<${somewhereWanted}>`) {
      codeActions[file] = codeActions[file] || [];

      // We can figure out default values for primitives etc.
      let defaultValue = "assert false";

      switch (somewhereWanted) {
        case "string": {
          defaultValue = `"-"`;
          break;
        }
        case "bool": {
          defaultValue = `false`;
          break;
        }
        case "int": {
          defaultValue = `-1`;
          break;
        }
        case "float": {
          defaultValue = `-1.`;
          break;
        }
      }

      let codeAction: p.CodeAction = {
        title: `Unwrap optional value`,
        edit: {
          changes: {
            [file]: wrapRangeInText(
              range,
              "switch ",
              ` { | None => ${defaultValue} | Some(v) => v }`
            ),
          },
        },
        diagnostics: [diagnostic],
        kind: p.CodeActionKind.QuickFix,
        isPreferred: true,
      };

      codeActions[file].push({
        range,
        codeAction,
      });

      return true;
    }

    // Wrapping a non-optional in Some
    if (`option<${thisHasType}>` === somewhereWanted) {
      codeActions[file] = codeActions[file] || [];

      let codeAction: p.CodeAction = {
        title: `Wrap value in Some`,
        edit: {
          changes: {
            [file]: wrapRangeInText(range, "Some(", ")"),
          },
        },
        diagnostics: [diagnostic],
        kind: p.CodeActionKind.QuickFix,
        isPreferred: true,
      };

      codeActions[file].push({
        range,
        codeAction,
      });

      return true;
    }
  }

  return false;
};