Source file enclosing.ml

(*
  This computes the "enclosing" ranges a given location. Those ranges do not
  necessarily correspond to a node in the AST.

  The ranges are purely syntactic, however the AST used here is the
  typedtree. This creates several issues, where some information has been lost
  during a translation phase. We worked around those limitations, but for the
  record, here is a (potentially non-exhaustive) list of things that would be
  improved by switching to the (non-preprocessed) parsetree instead:

   - [let%map ... in ...] could be recogized without having to recognize the
     rewritten expression

   - Parenthesized expression could recover their inner/outer locations without
     [node_loc_stack] (which we could get rid of)

   - Let bindings turned into match node in the typedtree would not need special
     handling.

   On the other hand, a "parsetree" version of [MBrowse] does not exist yet, and
   we use it to get the node at cursor, as well as some of its helper
   functions. So, reader, whenever the parsetree version get added, rewrite this
   module!
*)

open Std

let ( << ) (x : Location.t) (y : Location.t) =
  x.loc_end.pos_cnum <= y.loc_start.pos_cnum

let contains (loc1 : Location.t) (loc2 : Location.t) =
  if Location.is_none loc1 then false
  else
    loc1.loc_start.pos_cnum <= loc2.loc_start.pos_cnum
    && loc2.loc_end.pos_cnum <= loc1.loc_end.pos_cnum

let ( <= ) x y = contains y x

let merge (loc1 : Location.t) (loc2 : Location.t) =
  let resolve_non_pos (pos1 : Lexing.position) (pos2 : Lexing.position) f =
    if pos1.pos_cnum < 0 then pos2 else if pos2.pos_cnum < 0 then pos1 else f ()
  in
  let min_pos pos1 pos2 =
    resolve_non_pos pos1 pos2 @@ fun () ->
    if pos1.pos_cnum < pos2.pos_cnum then pos1 else pos2
  in
  let max_pos pos1 pos2 =
    resolve_non_pos pos1 pos2 @@ fun () ->
    if pos1.pos_cnum < pos2.pos_cnum then pos2 else pos1
  in
  let loc_start = min_pos loc1.loc_start loc2.loc_start in
  let loc_end = max_pos loc1.loc_end loc2.loc_end in
  (* I don't think loc_ghosts are a problem, as this is generated by Mbrowse
     module it should not contain any ghost location (apart from the initial
     [Location.none]). *)
  { Location.loc_start; loc_end; loc_ghost = false }

let ( ++ ) = merge

(* In some cases, we don't want to go "up" to a full node, but include new
   location little by little. For instance:

   {[
   let a = 1 in
   let b = 2 in
   let c = 3 in ...
   ]}

   We don't want to go go up from [2] to [let b = 2 in let c 3 in ...], we
   want intermediate stages that do not correspond to actual nodes:
   [let b = 2 in], then [let b = 2 in let c = 3 in], ...

   Similar example can be done with chains of [;] sequences, and other
   cases. All those cases share a similar setting: they look like a sequence (of
   [let ... in ...], of [;], etc) that are actually represented with binary
   operators, with "right" priority for parentheses.

   In all those cases, we split the location in the left part and the right part
   (using [merlin_loc] to split at the right point), include the left part
   first, and continue "going down-right".
*)

let rec go_down_right acc ~current_loc (expr : Typedtree.expression) =
  if Mbrowse.node_loc (Expression expr) <= current_loc then (current_loc, acc)
  else
    let go right_expr =
      (* The location without [begin ... end] or parenthesis delimiters. *)
      let undelimited_loc =
        match Mbrowse.node_loc_stack (Expression expr) with
        | [] -> Some (Mbrowse.node_loc (Expression expr))
        | innermost :: _ ->
          (* If we are in a delimited environment, and we don't come from below,
             we don't want to continue. Example, where brackets are current_loc
             and curlies are innermost:

             {[
               [let x = 1 in
               {let y = 2 in]
                x + y}
             ]}*)
          if current_loc <= innermost then Some innermost else None
      in
      match undelimited_loc with
      | None -> (current_loc, acc)
      | Some full_loc ->
        let left_loc =
          let right_merlin_loc =
            Mbrowse.node_merlin_loc (Expression right_expr)
          in
          { full_loc with loc_end = right_merlin_loc.loc_start }
        in
        let current_loc = current_loc ++ left_loc in
        let acc = current_loc :: acc in
        let is_splittable =
          Mbrowse.node_loc_stack (Expression right_expr) |> List.is_empty
        in
        if not is_splittable then (current_loc, acc)
        else go_down_right acc ~current_loc right_expr
    in
    match expr with
    (* [let x = exp1 in][ exp2] *)
    | { exp_desc = Texp_let (_, _, right_expr); _ } -> go right_expr
    (* [exp1 ;][ exp2] *)
    | { exp_desc = Texp_sequence (_, right_expr); _ } -> go right_expr
    (* [let%map x = exp1 in][ exp2] *)
    | { exp_desc =
          Texp_apply
            ( _map_or_bind,
              [ (Nolabel, Arg value_binded);
                ( Labelled "f",
                  Arg
                    { exp_desc =
                        Texp_function
                          ( [ { fp_kind = Tparam_pat pat; _ } ],
                            Tfunction_body right_expr );
                      _
                    } )
              ] );
        _
      }
      when pat.pat_loc << value_binded.exp_loc
           && value_binded.exp_loc << right_expr.exp_loc -> go right_expr
    (* [let () = exp1 in][ exp2] *)
    | { exp_desc =
          Texp_match
            ( exp,
              [ { c_lhs; c_rhs = right_expr; c_guard = None; c_cont = _ } ],
              [],
              _ );
        _
      }
      when c_lhs.pat_loc << exp.exp_loc -> go right_expr
    (* [let* x = exp1 in][ exp2] *)
    | { exp_desc = Texp_letop { body = { c_rhs = right_expr; _ }; _ }; _ } ->
      go right_expr
    | _ -> (current_loc, acc)

let add_intermediate_locs (current_loc, acc) locs =
  List.fold_left
    ~f:(fun (current_loc, acc) loc ->
      let current_loc = current_loc ++ loc in
      (current_loc, current_loc :: acc))
    ~init:(current_loc, acc) locs

let with_extra ~current_loc acc extra node =
  let loc_stack = Mbrowse.node_loc_stack node in
  let node_loc = Mbrowse.node_loc node in
  let node_locs = loc_stack @ [ node_loc ] in
  (* Exclude locations not fully containing current_loc. This is to avoid when
       the cursor is eg on a closing parenthesis, including the innermost
       location would not take the closing parenthesis. *)
  let node_locs = List.filter ~f:(fun loc -> current_loc <= loc) node_locs in
  let acc = add_intermediate_locs (current_loc, acc) node_locs in
  add_intermediate_locs acc (List.rev_map ~f:(fun (_, loc, _) -> loc) extra)

let expr_locs acc ~current_loc (exp : Typedtree.expression) =
  match exp with
  | { exp_extra; _ } -> with_extra ~current_loc acc exp_extra (Expression exp)

let pat_locs acc ~current_loc pat =
  match pat with
  | { Typedtree.pat_extra; _ } ->
    with_extra ~current_loc acc pat_extra (Pattern pat)

let typ_locs acc ~current_loc typ =
  let loc_stack =
    Mbrowse.node_loc_stack (Core_type typ)
    |> List.filter ~f:(fun loc -> current_loc <= loc)
  in
  let typ_loc = Mbrowse.node_loc (Core_type typ) in
  add_intermediate_locs (current_loc, acc) (loc_stack @ [ typ_loc ])

let rec expand_node acc ~current_loc (nodes : Browse_raw.node list) =
  match nodes with
  | [] -> List.rev acc
  (* Skipping node already included *)
  | node :: nodes when Mbrowse.node_loc node <= current_loc ->
    expand_node acc ~current_loc nodes
  (* skipping "[let x = 5] in exp2" location.
     TODO: would we want to add "let [x = 5] in exp2"? *)
  | Value_binding _ :: (Expression _ :: _ as nodes) ->
    expand_node acc ~current_loc nodes
  | Expression exp :: nodes ->
    let current_loc, acc = go_down_right acc ~current_loc exp in
    let current_loc, acc = expr_locs acc ~current_loc exp in
    expand_node acc ~current_loc nodes
  | Pattern pat :: nodes ->
    let current_loc, acc = pat_locs acc ~current_loc pat in
    expand_node acc ~current_loc nodes
  | Core_type typ :: nodes ->
    let current_loc, acc = typ_locs acc ~current_loc typ in
    expand_node acc ~current_loc nodes
  | node :: q ->
    let loc = Mbrowse.node_loc node in
    let current_loc = current_loc ++ loc in
    let acc = current_loc :: acc in
    expand_node acc ~current_loc q

let locs current_loc (mbrowse : Mbrowse.t) =
  mbrowse |> List.map ~f:snd
  |> expand_node [] ~current_loc
  |> List.fold_left ~init:[] ~f:(fun acc loc ->
      match acc with
      | hd :: _ as acc when Location_aux.compare hd loc = 0 -> acc
      | _ -> loc :: acc)
  |> List.rev