WIP: Fix type-annotate action for functions

ocaml-lsp-server/src/code_actions/action_remove_type_annotation.ml

@@ -3,16 +3,17 @@ open Import
 let action_kind = "remove type annotation"
 
 let check_typeable_context pipeline pos_start =
+  let open Typedtree in
   let pos_start = Mpipeline.get_lexing_pos pipeline pos_start in
   let typer = Mpipeline.typer_result pipeline in
   let browse = Mbrowse.of_typedtree (Mtyper.get_typedtree typer) in
   let is_exp_constrained = function
-    | Typedtree.Texp_constraint _, loc, _ -> Some loc
-    | Typedtree.Texp_coerce (Some { ctyp_loc; _ }, _), _, _ -> Some ctyp_loc
+    | Texp_constraint { ctyp_loc; _ }, _, _ -> Some ctyp_loc
+    | Texp_coerce (Some { ctyp_loc; _ }, _), _, _ -> Some ctyp_loc
     | _ -> None
   in
   let is_pat_constrained = function
-    | Typedtree.Tpat_constraint _, loc, _ -> Some loc
+    | Tpat_constraint _, loc, _ -> Some loc
     | _ -> None
   in
   let is_valid loc p extras =
@@ -22,7 +23,34 @@ let check_typeable_context pipeline pos_start =
     | Some x -> `Valid (loc, x)
     | None -> `Invalid
   in
+  let rec trav_cases = function
+    | { c_lhs = { pat_desc = Tpat_var _; _ }
+      ; c_rhs = { exp_desc = Texp_function { cases; _ }; _ }
+      ; _
+      }
+      :: _ -> trav_cases cases
+    | { c_lhs = { pat_desc = Tpat_var _; pat_loc; _ }
+      ; c_rhs = { exp_extra; _ }
+      ; _
+      }
+      :: _ -> is_valid pat_loc is_exp_constrained exp_extra
+    | { c_lhs = { pat_desc = Tpat_alias _; pat_loc; pat_extra; _ }
+      ; c_rhs = { exp_extra; _ }
+      ; _
+      }
+      :: _ -> (
+      match is_valid pat_loc is_pat_constrained pat_extra with
+      | `Valid (_, loc) ->
+        is_valid (Loc.union pat_loc loc) is_exp_constrained exp_extra
+      | `Invalid -> is_valid pat_loc is_exp_constrained exp_extra)
+    | _ -> `Invalid
+  in
   match Mbrowse.enclosing pos_start [ browse ] with
+  | (_, Pattern { pat_desc = Tpat_var _; _ })
+    :: ( _
+       , Value_binding
+           { vb_expr = { exp_desc = Texp_function { cases; _ }; _ }; _ } )
+    :: _ -> trav_cases cases
   | (_, Expression e) :: _ -> is_valid e.exp_loc is_exp_constrained e.exp_extra
   | (_, Pattern { pat_desc = Typedtree.Tpat_any; pat_loc; _ })
     :: (_, Pattern { pat_desc = Typedtree.Tpat_alias _; pat_extra; _ })

ocaml-lsp-server/src/code_actions/action_type_annotate.ml

@@ -1,30 +1,66 @@
 open Import
-open Fiber.O
 
 let action_kind = "type-annotate"
 
 let check_typeable_context pipeline pos_start =
+  let open Typedtree in
   let pos_start = Mpipeline.get_lexing_pos pipeline pos_start in
   let typer = Mpipeline.typer_result pipeline in
   let browse = Mbrowse.of_typedtree (Mtyper.get_typedtree typer) in
   let is_exp_constrained = function
-    | Typedtree.Texp_constraint _, _, _ -> true
-    | Typedtree.Texp_coerce (Some _, _), _, _ -> true
+    | Texp_constraint _, _, _ -> true
+    | Texp_coerce (Some _, _), _, _ -> true
     | _ -> false
   in
   let is_pat_constrained = function
-    | Typedtree.Tpat_constraint _, _, _ -> true
+    | Tpat_constraint _, _, _ -> true
     | _ -> false
   in
-  let is_valid p extras =
-    if List.exists ~f:p extras then `Invalid else `Valid
+  let pat_constraint_loc = function
+    | Typedtree.Tpat_constraint _, loc, _ -> Some loc
+    | _ -> None
+  in
+  let is_valid env typ loc p extras =
+    if List.exists ~f:p extras then `Invalid else `Valid (env, typ, loc)
+  in
+  let rec trav_cases = function
+    | { c_lhs = { pat_desc = Tpat_var _; _ }
+      ; c_rhs = { exp_desc = Texp_function { cases; _ }; _ }
+      ; _
+      }
+      :: _ -> trav_cases cases
+    | { c_lhs = { pat_desc = Tpat_var _; pat_loc; _ }
+      ; c_rhs = { exp_extra; exp_type; exp_env; _ }
+      ; _
+      }
+      :: _ ->
+      if List.exists ~f:is_exp_constrained exp_extra then `Invalid
+      else `Valid_fun (exp_env, exp_type, pat_loc)
+    | { c_lhs = { pat_desc = Tpat_alias _; pat_loc; pat_extra; _ }
+      ; c_rhs = { exp_extra; exp_type; exp_env; _ }
+      ; _
+      }
+      :: _ -> (
+      if List.exists ~f:is_exp_constrained exp_extra then `Invalid
+      else
+        match pat_extra |> List.rev |> List.find_map ~f:pat_constraint_loc with
+        | Some loc -> `Valid_fun (exp_env, exp_type, Loc.union pat_loc loc)
+        | None -> `Valid_fun (exp_env, exp_type, pat_loc))
+    | _ -> `Invalid
   in
   match Mbrowse.enclosing pos_start [ browse ] with
-  | (_, Expression e) :: _ -> is_valid is_exp_constrained e.exp_extra
-  | (_, Pattern { pat_desc = Typedtree.Tpat_any; _ })
-    :: (_, Pattern { pat_desc = Typedtree.Tpat_alias _; pat_extra; _ })
-    :: _ -> is_valid is_pat_constrained pat_extra
-  | (_, Pattern p) :: _ -> is_valid is_pat_constrained p.pat_extra
+  | (_, Pattern { pat_desc = Tpat_var _; _ })
+    :: ( _
+       , Value_binding
+           { vb_expr = { exp_desc = Texp_function { cases; _ }; _ }; _ } )
+    :: _ -> trav_cases cases
+  | (_, Expression e) :: _ ->
+    is_valid e.exp_env e.exp_type e.exp_loc is_exp_constrained e.exp_extra
+  | (_, Pattern { pat_desc = Tpat_any; pat_loc; pat_env; pat_type; _ })
+    :: (_, Pattern { pat_desc = Tpat_alias _; pat_extra; _ })
+    :: _ -> is_valid pat_env pat_type pat_loc is_pat_constrained pat_extra
+  | (_, Pattern p) :: _ ->
+    is_valid p.pat_env p.pat_type p.pat_loc is_pat_constrained p.pat_extra
   | _ :: _ | [] -> `Invalid
 
 let get_source_text doc (loc : Loc.t) =
@@ -36,10 +72,22 @@ let get_source_text doc (loc : Loc.t) =
   let (`Offset end_) = Msource.get_offset source (Position.logical end_) in
   String.sub (Msource.text source) ~pos:start ~len:(end_ - start)
 
-let code_action_of_type_enclosing uri doc (loc, typ) =
+let code_action uri doc str_fmt (env, typ, loc) =
   let open Option.O in
   let+ original_text = get_source_text doc loc in
-  let newText = Printf.sprintf "(%s : %s)" original_text typ in
+  let typ_str =
+    let buffer = Buffer.create 16 in
+    let ppf = Format.formatter_of_buffer buffer in
+    let pp_type env ppf ty =
+      let open Merlin_analysis in
+      let module Printtyp = Type_utils.Printtyp in
+      Printtyp.wrap_printing_env env ~verbosity:(Lvl 0) (fun () ->
+          Printtyp.type_expr ppf ty)
+    in
+    Format.fprintf ppf "%a%!" (pp_type env) typ;
+    Buffer.contents buffer
+  in
+  let newText = Printf.sprintf str_fmt original_text typ_str in
   let edit : WorkspaceEdit.t =
     let textedit : TextEdit.t = { range = Range.of_loc loc; newText } in
     let version = Document.version doc in
@@ -62,29 +110,16 @@ let code_action_of_type_enclosing uri doc (loc, typ) =
 let code_action doc (params : CodeActionParams.t) =
   match Document.kind doc with
   | `Other -> Fiber.return None
-  | `Merlin merlin -> (
+  | `Merlin merlin ->
     let pos_start = Position.logical params.range.start in
-    let+ res =
-      Document.Merlin.with_pipeline_exn merlin (fun pipeline ->
-          let context = check_typeable_context pipeline pos_start in
-          match context with
-          | `Invalid -> None
-          | `Valid ->
-            let command =
-              Query_protocol.Type_enclosing (None, pos_start, None)
-            in
-            let config = Mpipeline.final_config pipeline in
-            let config =
-              { config with query = { config.query with verbosity = Lvl 0 } }
-            in
-            let pipeline = Mpipeline.make config (Document.source doc) in
-            Some (Query_commands.dispatch pipeline command))
+    let action fmt_str data =
+      code_action params.textDocument.uri doc fmt_str data
     in
-    match res with
-    | None | Some [] | Some ((_, `Index _, _) :: _) -> None
-    | Some ((location, `String value, _) :: _) ->
-      code_action_of_type_enclosing params.textDocument.uri doc (location, value)
-    )
+    Document.Merlin.with_pipeline_exn merlin (fun pipeline ->
+        match check_typeable_context pipeline pos_start with
+        | `Invalid -> None
+        | `Valid_fun x -> action "%s : %s" x
+        | `Valid x -> action "(%s : %s)" x)
 
 let t =
   { Code_action.kind = CodeActionKind.Other action_kind; run = code_action }