more improvements to completions involving modules

Author: zth

analysis/src/CompletionFrontEnd.ml

@@ -1699,6 +1699,7 @@ let completionWithParser1 ~currentFile ~debug ~offset ~path ~posCursor
   in
   let module_expr (iterator : Ast_iterator.iterator)
       (me : Parsetree.module_expr) =
+    let processed = ref false in
     (match me.pmod_desc with
     | Pmod_ident lid when lid.loc |> Loc.hasPos ~pos:posBeforeCursor ->
       let lidPath = flattenLidCheckDot lid in
@@ -1710,8 +1711,17 @@ let completionWithParser1 ~currentFile ~debug ~offset ~path ~posCursor
       setResult
         (Cpath
            (CPId {loc = lid.loc; path = lidPath; completionContext = Module}))
+    | Pmod_functor (name, maybeType, body) ->
+      let oldScope = !scope in
+      scope := !scope |> Scope.addModule ~name:name.txt ~loc:name.loc;
+      (match maybeType with
+      | None -> ()
+      | Some mt -> iterator.module_type iterator mt);
+      iterator.module_expr iterator body;
+      scope := oldScope;
+      processed := true
     | _ -> ());
-    Ast_iterator.default_iterator.module_expr iterator me
+    if not !processed then Ast_iterator.default_iterator.module_expr iterator me
   in
   let module_type (iterator : Ast_iterator.iterator)
       (mt : Parsetree.module_type) =

analysis/src/ProcessCmt.ml

@@ -521,6 +521,8 @@ let rec forStructureItem ~(env : SharedTypes.Env.t) ~(exported : Exported.t)
       (fun {Typedtree.vb_pat; vb_attributes} ->
         handlePattern vb_attributes vb_pat)
       bindings;
+    bindings
+    |> List.iter (fun {Typedtree.vb_expr} -> scanLetModules ~env vb_expr);
     !items
   | Tstr_module
       {mb_id; mb_attributes; mb_loc; mb_name = name; mb_expr = {mod_desc}}
@@ -630,16 +632,14 @@ and forModule ~env mod_desc moduleName =
   | Tmod_functor (ident, argName, maybeType, resultExpr) ->
     (match maybeType with
     | None -> ()
-    | Some t -> (
-      match forTreeModuleType ~name:argName.txt ~env t with
-      | None -> ()
-      | Some kind ->
-        let stamp = Ident.binding_time ident in
-        let declared =
-          ProcessAttributes.newDeclared ~item:kind ~name:argName
-            ~extent:t.Typedtree.mty_loc ~stamp ~modulePath:NotVisible false []
-        in
-        Stamps.addModule env.stamps stamp declared));
+    | Some t ->
+      let kind = forTypeModule ~name:argName.txt ~env t.mty_type in
+      let stamp = Ident.binding_time ident in
+      let declared =
+        ProcessAttributes.newDeclared ~item:kind ~name:argName
+          ~extent:argName.loc ~stamp ~modulePath:NotVisible false []
+      in
+      Stamps.addModule env.stamps stamp declared);
     forModule ~env resultExpr.mod_desc moduleName
   | Tmod_apply (functor_, _arg, _coercion) ->
     forModule ~env functor_.mod_desc moduleName
@@ -653,6 +653,69 @@ and forModule ~env mod_desc moduleName =
     let modTypeKind = forTypeModule ~name:moduleName ~env typ in
     Constraint (modKind, modTypeKind)
 
+(*
+  Walk a typed expression and register any `let module M = ...` bindings as local
+  modules in stamps. This makes trailing-dot completion work for aliases like `M.`
+  that are introduced inside expression scopes. The declared module is marked as
+  NotVisible (non-exported) and the extent is the alias identifier location so
+  scope lookups match precisely.
+*)
+and scanLetModules ~env (e : Typedtree.expression) =
+  match e.exp_desc with
+  | Texp_letmodule (id, name, mexpr, body) ->
+    let stamp = Ident.binding_time id in
+    let item = forModule ~env mexpr.mod_desc name.txt in
+    let declared =
+      ProcessAttributes.newDeclared ~item ~extent:name.loc ~name ~stamp
+        ~modulePath:NotVisible false []
+    in
+    Stamps.addModule env.stamps stamp declared;
+    scanLetModules ~env body
+  | Texp_let (_rf, bindings, body) ->
+    List.iter (fun {Typedtree.vb_expr} -> scanLetModules ~env vb_expr) bindings;
+    scanLetModules ~env body
+  | Texp_apply {funct; args; _} ->
+    scanLetModules ~env funct;
+    args
+    |> List.iter (function
+         | _, Some e -> scanLetModules ~env e
+         | _, None -> ())
+  | Texp_tuple exprs -> List.iter (scanLetModules ~env) exprs
+  | Texp_sequence (e1, e2) ->
+    scanLetModules ~env e1;
+    scanLetModules ~env e2
+  | Texp_match (e, cases, exn_cases, _) ->
+    scanLetModules ~env e;
+    let scan_case {Typedtree.c_lhs = _; c_guard; c_rhs} =
+      (match c_guard with
+      | Some g -> scanLetModules ~env g
+      | None -> ());
+      scanLetModules ~env c_rhs
+    in
+    List.iter scan_case cases;
+    List.iter scan_case exn_cases
+  | Texp_function {case; _} ->
+    let {Typedtree.c_lhs = _; c_guard; c_rhs} = case in
+    (match c_guard with
+    | Some g -> scanLetModules ~env g
+    | None -> ());
+    scanLetModules ~env c_rhs
+  | Texp_try (e, cases) ->
+    scanLetModules ~env e;
+    cases
+    |> List.iter (fun {Typedtree.c_lhs = _; c_guard; c_rhs} ->
+           (match c_guard with
+           | Some g -> scanLetModules ~env g
+           | None -> ());
+           scanLetModules ~env c_rhs)
+  | Texp_ifthenelse (e1, e2, e3Opt) -> (
+    scanLetModules ~env e1;
+    scanLetModules ~env e2;
+    match e3Opt with
+    | Some e3 -> scanLetModules ~env e3
+    | None -> ())
+  | _ -> ()
+
 and forStructure ~name ~env strItems =
   let exported = Exported.init () in
   let items =

tests/analysis_tests/tests/src/FirstClassModules.res

@@ -16,4 +16,55 @@ let someFn = (~ctx: {"someModule": module(SomeModule)}) => {
 
   let _ff = SomeModule.doStuff
   //    ^hov
+
+  module M = CompletionFromModule.SomeModule
+  // M.
+  //   ^com
+
+  // M.g
+  //    ^com
+  ()
+}
+
+// Module type alias + unpack
+module type S2 = SomeModule
+let testAliasUnpack = (~ctx: {"someModule": module(SomeModule)}) => {
+  let module(S2) = ctx["someModule"]
+  // S2.
+  //    ^com
+  ()
+}
+// Functor param completion
+module Functor = (X: SomeModule) => {
+  // X.
+  //   ^com
+  let _u = X.doStuff
 }
+// First-class type hover without binding via module pattern
+let typeHover = (~ctx: {"someModule": module(SomeModule)}) => {
+  let v: module(SomeModule) = ctx["someModule"]
+  //   ^hov
+  ()
+}
+// Nested unpack inside nested module
+module Outer = {
+  let nested = (~ctx: {"someModule": module(SomeModule)}) => {
+    let module(SomeModule) = ctx["someModule"]
+    //SomeModule.
+    //           ^com
+    ()
+  }
+}
+// Shadowing: inner binding should be used for completion
+let shadowing = (
+  ~ctx1: {"someModule": module(SomeModule)},
+  ~ctx2: {"someModule": module(SomeModule)},
+) => {
+  let module(SomeModule) = ctx1["someModule"]
+  {
+    let module(SomeModule) = ctx2["someModule"]
+    //SomeModule.
+    //           ^com
+    ()
+  }
+}