Add extract code actions (#870)

* work on extract action

* add changes

* add documentation and tests

* update test

* formatting

Author: jfeser

CHANGES.md

@@ -49,6 +49,7 @@
 
 - Add "Remove type annotation" code action. (#1039)
 - Support settings through `didChangeConfiguration` notification (#1103)
+- Add "Extract local" and "Extract function" code actions. (#870)
 - Depend directly on `merlin-lib` 4.9 (#1070)
 
 # 1.15.1

ocaml-lsp-server/src/code_actions.ml

@@ -92,6 +92,8 @@ let compute server (params : CodeActionParams.t) =
             ; Action_mark_remove_unused.mark
             ; Action_mark_remove_unused.remove
             ; Action_inline.t
+            ; Action_extract.local
+            ; Action_extract.function_
             ]
         in
         List.concat

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

@@ -0,0 +1,217 @@
+open Import
+module H = Ocaml_parsing.Ast_helper
+
+let range_contains_loc range loc =
+  match Range.of_loc_opt loc with
+  | Some range' -> Range.contains range range'
+  | None -> false
+
+let range_contained_by_loc range loc =
+  match Range.of_loc_opt loc with
+  | Some range' -> Range.contains range' range
+  | None -> false
+
+let largest_enclosed_expression typedtree range =
+  let exception Found of Typedtree.expression in
+  let module I = Ocaml_typing.Tast_iterator in
+  let expr_iter (iter : I.iterator) (expr : Typedtree.expression) =
+    if range_contains_loc range expr.exp_loc then raise (Found expr)
+    else I.default_iterator.expr iter expr
+  in
+  let iterator = { I.default_iterator with expr = expr_iter } in
+  try
+    iterator.structure iterator typedtree;
+    None
+  with Found e -> Some e
+
+let enclosing_structure_item typedtree range =
+  let exception Found of Typedtree.structure_item in
+  let module I = Ocaml_typing.Tast_iterator in
+  let structure_item_iter (iter : I.iterator) (item : Typedtree.structure_item)
+      =
+    if range_contained_by_loc range item.str_loc then
+      match item.str_desc with
+      | Tstr_value _ -> raise (Found item)
+      | _ -> I.default_iterator.structure_item iter item
+  in
+  let iterator =
+    { I.default_iterator with structure_item = structure_item_iter }
+  in
+  try
+    iterator.structure iterator typedtree;
+    None
+  with Found e -> Some e
+
+let tightest_enclosing_binder_position typedtree range =
+  let exception Found of Position.t in
+  let module I = Ocaml_typing.Tast_iterator in
+  let found_loc loc =
+    Position.of_lexical_position loc
+    |> Option.iter ~f:(fun p -> raise (Found p))
+  in
+  let found_if_expr_contains (expr : Typedtree.expression) =
+    let loc = expr.exp_loc in
+    if range_contained_by_loc range loc then found_loc loc.loc_start
+  in
+  let found_if_case_contains cases =
+    List.iter cases ~f:(fun (case : _ Typedtree.case) ->
+        found_if_expr_contains case.c_rhs)
+  in
+  let expr_iter (iter : I.iterator) (expr : Typedtree.expression) =
+    if range_contained_by_loc range expr.exp_loc then (
+      I.default_iterator.expr iter expr;
+      match expr.exp_desc with
+      | Texp_let (_, _, body)
+      | Texp_while (_, body)
+      | Texp_for (_, _, _, _, _, body)
+      | Texp_letmodule (_, _, _, _, body)
+      | Texp_letexception (_, body)
+      | Texp_open (_, body) -> found_if_expr_contains body
+      | Texp_letop { body; _ } -> found_if_case_contains [ body ]
+      | Texp_function { cases; _ } -> found_if_case_contains cases
+      | Texp_match (_, cases, _) -> found_if_case_contains cases
+      | Texp_try (_, cases) -> found_if_case_contains cases
+      | _ -> ())
+  in
+  let structure_item_iter (iter : I.iterator) (item : Typedtree.structure_item)
+      =
+    if range_contained_by_loc range item.str_loc then (
+      I.default_iterator.structure_item iter item;
+      match item.str_desc with
+      | Tstr_value (_, bindings) ->
+        List.iter bindings ~f:(fun (binding : Typedtree.value_binding) ->
+            found_if_expr_contains binding.vb_expr)
+      | _ -> ())
+  in
+  let iterator =
+    { I.default_iterator with
+      expr = expr_iter
+    ; structure_item = structure_item_iter
+    }
+  in
+  try
+    iterator.structure iterator typedtree;
+    None
+  with Found e -> Some e
+
+module LongidentSet = Set.Make (struct
+  type t = Longident.t
+
+  let compare = compare
+end)
+
+(** [free expr] returns the free variables in [expr]. *)
+let free (expr : Typedtree.expression) =
+  let module I = Ocaml_typing.Tast_iterator in
+  let idents = ref [] in
+  let expr_iter (iter : I.iterator) (expr : Typedtree.expression) =
+    match expr.exp_desc with
+    | Texp_ident (path, { txt = ident; _ }, _) ->
+      idents := (ident, path) :: !idents
+    | _ ->
+      I.default_iterator.expr iter expr;
+
+      (* if a variable was bound but is no longer, it must be associated with a
+         binder inside the expression *)
+      idents :=
+        List.filter !idents ~f:(fun (ident, path) ->
+            match Env.find_value_by_name ident expr.exp_env with
+            | path', _ -> Path.same path path'
+            | exception Not_found -> false)
+  in
+  let iter = { I.default_iterator with expr = expr_iter } in
+  iter.expr iter expr;
+  !idents
+
+let must_pass expr env =
+  List.filter (free expr) ~f:(fun (ident, path) ->
+      match Env.find_value_by_name ident env with
+      | path', _ ->
+        (* new environment binds ident to a different path than the old one *)
+        not (Path.same path path')
+      | exception Not_found -> true)
+  |> List.map ~f:fst
+
+let extract_local doc typedtree range =
+  let open Option.O in
+  let* to_extract = largest_enclosed_expression typedtree range in
+  let* extract_range = Range.of_loc_opt to_extract.exp_loc in
+  let* edit_pos = tightest_enclosing_binder_position typedtree range in
+  let new_name = "var_name" in
+  let* local_text = Document.substring doc extract_range in
+  let newText = sprintf "let %s = %s in\n" new_name local_text in
+  let insert_range = { Range.start = edit_pos; end_ = edit_pos } in
+  Some
+    [ TextEdit.create ~newText ~range:insert_range
+    ; TextEdit.create ~newText:new_name ~range:extract_range
+    ]
+
+let extract_function doc typedtree range =
+  let open Option.O in
+  let* to_extract = largest_enclosed_expression typedtree range in
+  let* extract_range = Range.of_loc_opt to_extract.exp_loc in
+  let* parent_item = enclosing_structure_item typedtree range in
+  let* edit_pos = Position.of_lexical_position parent_item.str_loc.loc_start in
+  let new_name = "fun_name" in
+  let* args_str =
+    let free_vars = must_pass to_extract parent_item.str_env in
+    let+ args =
+      List.map free_vars ~f:(function
+          | Longident.Lident id -> Some id
+          | _ -> None)
+      |> Option.List.all
+    in
+    let s = String.concat ~sep:" " args in
+    if String.is_empty s then "()" else s
+  in
+  let* func_text = Document.substring doc extract_range in
+  let new_function = sprintf "let %s %s = %s\n\n" new_name args_str func_text in
+  let new_call = sprintf "%s %s" new_name args_str in
+  let insert_range = { Range.start = edit_pos; end_ = edit_pos } in
+  Some
+    [ TextEdit.create ~newText:new_function ~range:insert_range
+    ; TextEdit.create ~newText:new_call ~range:extract_range
+    ]
+
+let run_extract_local doc (params : CodeActionParams.t) =
+  let open Option.O in
+  Document.Merlin.with_pipeline_exn
+    (Document.merlin_exn doc)
+    Mpipeline.typer_result
+  |> Fiber.map ~f:(fun typer ->
+         let* typedtree =
+           match Mtyper.get_typedtree typer with
+           | `Interface _ -> None
+           | `Implementation x -> Some x
+         in
+         let+ edits = extract_local doc typedtree params.range in
+         CodeAction.create
+           ~title:"Extract local"
+           ~kind:CodeActionKind.RefactorExtract
+           ~edit:(Document.edit doc edits)
+           ~isPreferred:false
+           ())
+
+let run_extract_function doc (params : CodeActionParams.t) =
+  let open Option.O in
+  Document.Merlin.with_pipeline_exn
+    (Document.merlin_exn doc)
+    Mpipeline.typer_result
+  |> Fiber.map ~f:(fun typer ->
+         let* typedtree =
+           match Mtyper.get_typedtree typer with
+           | `Interface _ -> None
+           | `Implementation x -> Some x
+         in
+         let+ edits = extract_function doc typedtree params.range in
+         CodeAction.create
+           ~title:"Extract function"
+           ~kind:CodeActionKind.RefactorExtract
+           ~edit:(Document.edit doc edits)
+           ~isPreferred:false
+           ())
+
+let local = { Code_action.kind = RefactorExtract; run = run_extract_local }
+
+let function_ =
+  { Code_action.kind = RefactorExtract; run = run_extract_function }

ocaml-lsp-server/src/code_actions/action_extract.mli

@@ -0,0 +1,3 @@
+val local : Code_action.t
+
+val function_ : Code_action.t

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

@@ -1,11 +1,5 @@
 open Import
 
-let slice doc (range : Range.t) =
-  let src = Document.source doc in
-  let (`Offset start) = Msource.get_offset src @@ Position.logical range.start
-  and (`Offset end_) = Msource.get_offset src @@ Position.logical range.end_ in
-  String.sub (Msource.text src) ~pos:start ~len:(end_ - start)
-
 (* Return contexts enclosing `pos` in order from most specific to most
    general. *)
 let enclosing_pos pipeline pos =
@@ -72,14 +66,14 @@ let rec mark_value_unused_edit name contexts =
 let code_action_mark_value_unused doc (diagnostic : Diagnostic.t) =
   let open Option.O in
   Document.Merlin.with_pipeline_exn (Document.merlin_exn doc) (fun pipeline ->
-      let var_name = slice doc diagnostic.range in
+      let* var_name = Document.substring doc diagnostic.range in
       let pos = diagnostic.range.start in
       let+ text_edit =
         enclosing_pos pipeline pos
         |> List.rev_map ~f:(fun (_, x) -> x)
         |> mark_value_unused_edit var_name
       in
-      let edit = Document.edit doc text_edit in
+      let edit = Document.edit doc [ text_edit ] in
       CodeAction.create
         ~diagnostics:[ diagnostic ]
         ~title:"Mark as unused"
@@ -114,7 +108,7 @@ let enclosing_value_binding_range name =
 
 (* Create a code action that removes [range] and refers to [diagnostic]. *)
 let code_action_remove_range doc (diagnostic : Diagnostic.t) range =
-  let edit = Document.edit doc { range; newText = "" } in
+  let edit = Document.edit doc [ { range; newText = "" } ] in
   CodeAction.create
     ~diagnostics:[ diagnostic ]
     ~title:"Remove unused"
@@ -125,8 +119,9 @@ let code_action_remove_range doc (diagnostic : Diagnostic.t) range =
 
 (* Create a code action that removes the value mentioned in [diagnostic]. *)
 let code_action_remove_value doc pos (diagnostic : Diagnostic.t) =
+  let open Option.O in
   Document.Merlin.with_pipeline_exn (Document.merlin_exn doc) (fun pipeline ->
-      let var_name = slice doc diagnostic.range in
+      let* var_name = Document.substring doc diagnostic.range in
       enclosing_pos pipeline pos |> List.map ~f:snd
       |> enclosing_value_binding_range var_name
       |> Option.map ~f:(fun range ->

ocaml-lsp-server/src/document.ml

@@ -314,13 +314,15 @@ module Merlin = struct
     with_pipeline_exn doc (fun pipeline -> doc_comment pipeline pos)
 end
 
-let edit t text_edit =
+let edit t text_edits =
   let version = version t in
   let textDocument =
     OptionalVersionedTextDocumentIdentifier.create ~uri:(uri t) ~version ()
   in
   let edit =
-    TextDocumentEdit.create ~textDocument ~edits:[ `TextEdit text_edit ]
+    TextDocumentEdit.create
+      ~textDocument
+      ~edits:(List.map text_edits ~f:(fun text_edit -> `TextEdit text_edit))
   in
   WorkspaceEdit.create ~documentChanges:[ `TextDocumentEdit edit ] ()
 
@@ -379,3 +381,9 @@ let get_impl_intf_counterparts uri =
     | to_switch_to -> to_switch_to
   in
   List.map ~f:Uri.of_path files_to_switch_to
+
+let substring doc range =
+  let start, end_ = Text_document.absolute_range (tdoc doc) range in
+  let text = text doc in
+  if start < 0 || start > end_ || end_ > String.length text then None
+  else Some (String.sub text ~pos:start ~len:(end_ - start))

ocaml-lsp-server/src/document.mli

@@ -99,4 +99,12 @@ val close : t -> unit Fiber.t
     For instance, the counterparts of the file [/file.ml] are [/file.mli]. *)
 val get_impl_intf_counterparts : Uri.t -> Uri.t list
 
-val edit : t -> TextEdit.t -> WorkspaceEdit.t
+(** [edits t edits] creates a [WorkspaceEdit.t] that applies edits [edits] to
+    the document [t]. *)
+val edit : t -> TextEdit.t list -> WorkspaceEdit.t
+
+(** [substring t range] returns the substring of the document [t] that
+    corresponds to the range [range].
+
+    Returns [None] when there is no corresponding substring. *)
+val substring : t -> Range.t -> string option

ocaml-lsp-server/src/range.ml

@@ -10,6 +10,12 @@ let to_dyn { start; end_ } =
   Dyn.record
     [ ("start", Position.to_dyn start); ("end_", Position.to_dyn end_) ]
 
+let contains (x : t) (y : t) =
+  let open Ordering in
+  match (Position.compare x.start y.start, Position.compare x.end_ y.end_) with
+  | (Lt | Eq), (Gt | Eq) -> true
+  | _ -> false
+
 (* Compares ranges by their lengths*)
 let compare_size (x : t) (y : t) =
   let dx = Position.(x.end_ - x.start) in

ocaml-lsp-server/src/range.mli

@@ -6,6 +6,9 @@ include module type of Lsp.Types.Range with type t = Lsp.Types.Range.t
     positions. *)
 val compare : t -> t -> Ordering.t
 
+(** [contains r1 r2] returns true if [r1] contains [r2]. *)
+val contains : t -> t -> bool
+
 val to_dyn : t -> Dyn.t
 
 val compare_size : t -> t -> Ordering.t