revamp insert missing cases by leveraging actual AST transforms to produce the missing cases and insert them into the existing switch

Author: zth

analysis/src/Cli.ml

@@ -124,6 +124,19 @@ let main () =
     Commands.codeAction ~path
       ~pos:(int_of_string line, int_of_string col)
       ~currentFile ~debug:false
+  | [_; "codemod"; path; line; col; typ; hint] ->
+    let typ =
+      match typ with
+      | "add-missing-cases" -> Codemod.AddMissingCases
+      | _ -> raise (Failure "unsupported type")
+    in
+    let res =
+      Codemod.transform ~path
+        ~pos:(int_of_string line, int_of_string col)
+        ~debug:false ~typ ~hint
+      |> Json.escape
+    in
+    Printf.printf "\"%s\"" res
   | [_; "diagnosticSyntax"; path] -> Commands.diagnosticSyntax ~path
   | _ :: "reanalyze" :: _ ->
     let len = Array.length Sys.argv in

analysis/src/Codemod.ml

@@ -0,0 +1,52 @@
+type transformType = AddMissingCases
+
+let rec collectPatterns p =
+  match p.Parsetree.ppat_desc with
+  | Ppat_or (p1, p2) -> collectPatterns p1 @ [p2]
+  | _ -> [p]
+
+let mkFailWithExp () =
+  Ast_helper.Exp.apply
+    (Ast_helper.Exp.ident {txt = Lident "failwith"; loc = Location.none})
+    [(Nolabel, Ast_helper.Exp.constant (Pconst_string ("TODO", None)))]
+
+let transform ~path ~pos ~debug ~typ ~hint =
+  let structure, printExpr, _ = Xform.parseImplementation ~filename:path in
+  match typ with
+  | AddMissingCases -> (
+    let source = "let " ^ hint ^ " = ()" in
+    let {Res_driver.parsetree = hintStructure} =
+      Res_driver.parseImplementationFromSource ~forPrinter:false
+        ~displayFilename:"<none>" ~source
+    in
+    match hintStructure with
+    | [{pstr_desc = Pstr_value (_, [{pvb_pat = pattern}])}] -> (
+      let cases =
+        collectPatterns pattern
+        |> List.map (fun (p : Parsetree.pattern) ->
+               Ast_helper.Exp.case p (mkFailWithExp ()))
+      in
+      let result = ref None in
+      let mkIterator ~pos ~result =
+        let expr (iterator : Ast_iterator.iterator) (exp : Parsetree.expression)
+            =
+          match exp.pexp_desc with
+          | Pexp_match (e, existingCases)
+            when Pos.ofLexing exp.pexp_loc.loc_start = pos ->
+            result :=
+              Some {exp with pexp_desc = Pexp_match (e, existingCases @ cases)}
+          | _ -> Ast_iterator.default_iterator.expr iterator exp
+        in
+        {Ast_iterator.default_iterator with expr}
+      in
+      let iterator = mkIterator ~pos ~result in
+      iterator.structure iterator structure;
+      match !result with
+      | None ->
+        if debug then print_endline "Found no result";
+        exit 1
+      | Some switchExpr ->
+        printExpr ~range:(Xform.rangeOfLoc switchExpr.pexp_loc) switchExpr)
+    | _ ->
+      if debug then print_endline "Mismatch in expected structure";
+      exit 1)

analysis/src/Commands.ml

@@ -392,6 +392,14 @@ let test ~path =
                                  Printf.printf "%s\nnewText:\n%s<--here\n%s%s\n"
                                    (Protocol.stringifyRange range)
                                    indent indent newText)))
+          | "c-a" ->
+            let hint = String.sub rest 3 (String.length rest - 3) in
+            print_endline
+              ("Codemod AddMissingCases" ^ path ^ " " ^ string_of_int line ^ ":"
+             ^ string_of_int col);
+            Codemod.transform ~path ~pos:(line, col) ~debug:true
+              ~typ:AddMissingCases ~hint
+            |> print_endline
           | "dia" -> diagnosticSyntax ~path
           | "hin" ->
             (* Get all inlay Hint between line 1 and n.

analysis/tests/bsconfig.json

@@ -9,7 +9,7 @@
       "subdirs": true
     }
   ],
-  "bsc-flags": ["-w -33-44"],
+  "bsc-flags": ["-w -33-44-8"],
   "bs-dependencies": ["@rescript/react"],
   "jsx": { "version": 3 }
 }

analysis/tests/src/Codemod.res

@@ -0,0 +1,9 @@
+type someTyp = [#valid | #invalid]
+
+let ff = (v1: someTyp, v2: someTyp) => {
+  let x = switch (v1, v2) {
+  //      ^c-a (#valid, #valid) | (#invalid, _)
+  | (#valid, #invalid) => ()
+  }
+  x
+}

analysis/tests/src/expected/Codemod.res.txt

@@ -0,0 +1,8 @@
+Codemod AddMissingCasessrc/Codemod.res 3:10
+switch (v1, v2) {
+          //      ^c-a (#valid, #valid) | (#invalid, _)
+          | (#valid, #invalid) => ()
+          | (#valid, #valid) => failwith("TODO")
+          | (#invalid, _) => failwith("TODO")
+          }
+

server/src/codeActions.ts

@@ -2,6 +2,8 @@
 // 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 }[];
@@ -82,6 +84,15 @@ let insertBeforeEndingChar = (
   ];
 };
 
+let replaceText = (range: p.Range, newText: string): p.TextEdit[] => {
+  return [
+    {
+      range,
+      newText,
+    },
+  ];
+};
+
 let removeTrailingComma = (text: string): string => {
   let str = text.trim();
   if (str.endsWith(",")) {
@@ -514,67 +525,30 @@ let simpleAddMissingCases: codeActionExtractor = ({
   if (
     line.startsWith("You forgot to handle a possible case here, for example:")
   ) {
-    let cases: string[] = [];
-
     // This collects the rest of the fields if fields are printed on
     // multiple lines.
     let allCasesAsOneLine = array
       .slice(index + 1)
       .join("")
       .trim();
 
-    // We only handle the simplest possible cases until the compiler actually
-    // outputs ReScript. This means bailing on anything that's not a
-    // variant/polyvariant, with one payload (or no payloads at all).
-    let openParensCount = allCasesAsOneLine.split("(").length - 1;
-
-    if (openParensCount > 1 || allCasesAsOneLine.includes("{")) {
-      return false;
-    }
-
-    // Remove surrounding braces if they exist
-    if (allCasesAsOneLine[0] === "(") {
-      allCasesAsOneLine = allCasesAsOneLine.slice(
-        1,
-        allCasesAsOneLine.length - 1
-      );
-    }
-
-    cases.push(
-      ...(allCasesAsOneLine
-        .split("|")
-        .map(transformMatchPattern)
-        .filter(Boolean) as string[])
-    );
-
-    if (cases.length === 0) {
-      return false;
-    }
-
-    // The end char is the closing brace. In switches, the leading `|` always
-    // has the same left padding as the end brace.
-    let paddingContentSwitchCase = Array.from({
-      length: range.end.character,
-    }).join(" ");
-
-    let newText = cases
-      .map((variantName, index) => {
-        // The first case will automatically be padded because we're inserting
-        // it where the end brace is currently located.
-        let padding = index === 0 ? "" : paddingContentSwitchCase;
-        return `${padding}| ${variantName} => assert false`;
-      })
-      .join("\n");
+    let filePath = fileURLToPath(file);
 
-    // Let's put the end brace back where it was (we still have it to the direct right of us).
-    newText += `\n${paddingContentSwitchCase}`;
+    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]: insertBeforeEndingChar(range, newText),
+          [file]: replaceText(range, newSwitchCode),
         },
       },
       diagnostics: [diagnostic],