Connect renamed members of an applied functor to their declarations

Codex review on #9084: the members of an applied functor are registered
under fresh stamps, so references to `IntSet.add` were no longer
references to `MakeSet.add`; Find References missed them and renaming the
declaration left them unchanged. Each renamed stamp now maps to the stamp
it was renamed from: a use of the member is also recorded as a use of
the original, and references starting from the member start from the
original.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Kbf3gLMawVgtmiKy3TMBc3
Signed-off-by: Christoph Knittel <ck@cca.io>

Author: cknitt

analysis/src/process_cmt.ml

@@ -189,6 +189,39 @@ and for_type_module ~name ~env module_type =
    in a typed tree, so that they never meet one of them. *)
 let renamed_stamp_floor = 1 lsl 40
 
+(* [Subst] renames the items of a signature in place, so the original and
+   the renamed signature list the same items in the same order. Records
+   each renamed stamp's original, nested signatures included, for
+   references to find the declarations in the functor body. *)
+let rec record_renamed_stamps table (original : Types.module_type)
+    (renamed : Types.module_type) =
+  match (original, renamed) with
+  | Mty_signature items, Mty_signature renamed_items
+    when List.length items = List.length renamed_items ->
+    List.iter2
+      (fun (item : Types.signature_item) (renamed_item : Types.signature_item)
+         ->
+        let ids =
+          match (item, renamed_item) with
+          | Sig_value (id, _), Sig_value (id', _)
+          | Sig_type (id, _, _), Sig_type (id', _, _)
+          | Sig_typext (id, _, _), Sig_typext (id', _, _)
+          | Sig_modtype (id, _), Sig_modtype (id', _) ->
+            Some (id, id')
+          | Sig_module (id, md, _), Sig_module (id', md', _) ->
+            record_renamed_stamps table md.md_type md'.md_type;
+            Some (id, id')
+          | _ -> None
+        in
+        match ids with
+        | Some (id, id') ->
+          Hashtbl.replace table (Ident.binding_time id') (Ident.binding_time id)
+        | None -> ())
+      items renamed_items
+  | Mty_functor (_, _, result), Mty_functor (_, _, renamed_result) ->
+    record_renamed_stamps table result renamed_result
+  | _ -> ()
+
 let get_module_type_path mod_desc =
   match mod_desc with
   | Typedtree.Tmty_ident (path, _) | Tmty_alias (path, _) -> Some path
@@ -685,8 +718,9 @@ and for_module ~env (mod_expr : Typedtree.module_expr) module_name =
     in
     if Ident.current_time () < renamed_stamp_floor then
       Ident.set_current_time renamed_stamp_floor;
-    for_type_module ~name:module_name ~env
-      (Subst.modtype Subst.identity mod_expr.mod_type)
+    let renamed = Subst.modtype Subst.identity mod_expr.mod_type in
+    record_renamed_stamps env.renamed_stamps mod_expr.mod_type renamed;
+    for_type_module ~name:module_name ~env renamed
   | Tmod_unpack (_expr, module_type) ->
     let env = env |> Env.add_module ~name:module_name in
     for_type_module ~name:module_name ~env module_type
@@ -742,7 +776,11 @@ and for_structure ~name ~env str_items =
 let file_for_cmt_infos ~module_name ~uri
     ({cmt_modname; cmt_annots} : Cmt_format.cmt_infos) =
   let env =
-    {Env.stamps = Stamps.init (); module_path = File (uri, module_name)}
+    {
+      Env.stamps = Stamps.init ();
+      module_path = File (uri, module_name);
+      renamed_stamps = Hashtbl.create 1;
+    }
   in
   match cmt_annots with
   | Partial_implementation parts ->
@@ -756,7 +794,13 @@ let file_for_cmt_infos ~module_name ~uri
       |> List.concat
     in
     let structure = for_structure ~name:module_name ~env items in
-    {File.uri; module_name = cmt_modname; stamps = env.stamps; structure}
+    {
+      File.uri;
+      module_name = cmt_modname;
+      stamps = env.stamps;
+      structure;
+      renamed_stamps = env.renamed_stamps;
+    }
   | Partial_interface parts ->
     let items =
       parts |> Array.to_list
@@ -768,13 +812,31 @@ let file_for_cmt_infos ~module_name ~uri
       |> List.concat
     in
     let structure = for_signature ~name:module_name ~env items in
-    {uri; module_name = cmt_modname; stamps = env.stamps; structure}
+    {
+      uri;
+      module_name = cmt_modname;
+      stamps = env.stamps;
+      structure;
+      renamed_stamps = env.renamed_stamps;
+    }
   | Implementation structure ->
     let structure = for_structure ~name:module_name ~env structure.str_items in
-    {uri; module_name = cmt_modname; stamps = env.stamps; structure}
+    {
+      uri;
+      module_name = cmt_modname;
+      stamps = env.stamps;
+      structure;
+      renamed_stamps = env.renamed_stamps;
+    }
   | Interface signature ->
     let structure = for_signature ~name:module_name ~env signature.sig_items in
-    {uri; module_name = cmt_modname; stamps = env.stamps; structure}
+    {
+      uri;
+      module_name = cmt_modname;
+      stamps = env.stamps;
+      structure;
+      renamed_stamps = env.renamed_stamps;
+    }
 
 let file_for_cmt ~state ~module_name ~cmt ~uri =
   match Hashtbl.find_opt state.cmt_cache cmt with

analysis/src/process_extra.ml

@@ -135,13 +135,21 @@ let extra_for_cmt ~(iterator : Tast_iterator.iterator)
     extra_for_signature_items ~iterator items;
     extra_for_parts parts
 
+(* A use of an applied functor's member is also a use of the declaration
+   in the functor's body that it was renamed from. *)
+let add_reference_with_original ~(env : Query_env.t) ~extra stamp loc =
+  add_reference ~extra stamp loc;
+  match Hashtbl.find_opt env.file.renamed_stamps stamp with
+  | Some original -> add_reference ~extra original loc
+  | None -> ()
+
 let add_for_path ~env ~extra path lident loc typ tip =
   let ident_name = Longident.last lident in
   let ident_loc = Utils.end_of_location loc (String.length ident_name) in
   let loc_type =
     match Resolve_path.from_compiler_path ~env path with
     | Stamp stamp ->
-      add_reference ~extra stamp ident_loc;
+      add_reference_with_original ~env ~extra stamp ident_loc;
       LocalReference (stamp, tip)
     | NotFound -> NotFound
     | Global (module_name, path) ->
@@ -154,7 +162,7 @@ let add_for_path ~env ~extra path lident loc typ tip =
         | _ -> Exported.find env.exported Exported.Value name
       with
       | Some stamp ->
-        add_reference ~extra stamp ident_loc;
+        add_reference_with_original ~env ~extra stamp ident_loc;
         LocalReference (stamp, tip)
       | None -> NotFound)
     | GlobalMod _ -> NotFound

analysis/src/references.ml

@@ -440,6 +440,11 @@ type references = {
 
 let for_local_stamp ~state ~full:{file; extra; package} stamp (tip : Tip.t) =
   let env = Query_env.from_file file in
+  let stamp =
+    match Hashtbl.find_opt file.renamed_stamps stamp with
+    | Some original -> original
+    | None -> stamp
+  in
   match
     match tip with
     | Constructor name ->

analysis/src/shared_types.ml

@@ -274,13 +274,15 @@ module File = struct
     stamps: Stamps.t;
     module_name: string;
     structure: Module.structure;
+    renamed_stamps: (int, int) Hashtbl.t;
   }
 
   let create module_name uri =
     {
       uri;
       stamps = Stamps.init ();
       module_name;
+      renamed_stamps = Hashtbl.create 1;
       structure =
         {
           name = module_name;
@@ -404,7 +406,13 @@ and completion_type =
     }
 
 module Env = struct
-  type t = {stamps: Stamps.t; module_path: Module_path.t}
+  type t = {
+    stamps: Stamps.t;
+    module_path: Module_path.t;
+    renamed_stamps: (int, int) Hashtbl.t;
+        (** the stamps of an applied functor's renamed members, mapped to
+            those of the functor body's declarations they come from *)
+  }
   let add_exported_module ~name ~is_type env =
     {
       env with

tests/analysis_tests/tests/src/AppliedFunctor.res

@@ -6,6 +6,7 @@ module type Ord = {
 module MakeSet = (O: Ord) => {
   type t = array<O.t>
   let add = (s: t, x: O.t) => s->Array.concat([x])
+  //   ^ref
 }
 
 module IntSet = MakeSet({
@@ -16,6 +17,12 @@ module IntSet = MakeSet({
 let s = IntSet.add([], 1)
 //             ^hov
 
+let s3 = IntSet.add([], 3)
+//              ^ref
+
+let s4 = IntSet.add([], 4)
+//              ^ren insert
+
 let t: IntSet.t = []
 //            ^hov
 

tests/analysis_tests/tests/src/expected/AppliedFunctor.res.txt

@@ -1,22 +1,123 @@
-Hover src/AppliedFunctor.res 15:15
+References src/AppliedFunctor.res 7:7
+[
+  {
+    "range": {
+      "end": { "character": 9, "line": 7 },
+      "start": { "character": 6, "line": 7 }
+    },
+    "uri": "file:///AppliedFunctor.res"
+  },
+  {
+    "range": {
+      "end": { "character": 18, "line": 16 },
+      "start": { "character": 15, "line": 16 }
+    },
+    "uri": "file:///AppliedFunctor.res"
+  },
+  {
+    "range": {
+      "end": { "character": 19, "line": 19 },
+      "start": { "character": 16, "line": 19 }
+    },
+    "uri": "file:///AppliedFunctor.res"
+  },
+  {
+    "range": {
+      "end": { "character": 19, "line": 22 },
+      "start": { "character": 16, "line": 22 }
+    },
+    "uri": "file:///AppliedFunctor.res"
+  }
+]
+
+Hover src/AppliedFunctor.res 16:15
 {
   "contents": {
     "kind": "markdown",
     "value": "```rescript\n(IntSet.t, int) => array<int>\n```\n\n---\n\n```\n \n```\n```rescript\ntype IntSet.t = array<int>\n```\nGo to: [Type definition](AppliedFunctor.res#L7,3)\n"
   }
 }
 
-Hover src/AppliedFunctor.res 18:14
+References src/AppliedFunctor.res 19:16
+[
+  {
+    "range": {
+      "end": { "character": 9, "line": 7 },
+      "start": { "character": 6, "line": 7 }
+    },
+    "uri": "file:///AppliedFunctor.res"
+  },
+  {
+    "range": {
+      "end": { "character": 18, "line": 16 },
+      "start": { "character": 15, "line": 16 }
+    },
+    "uri": "file:///AppliedFunctor.res"
+  },
+  {
+    "range": {
+      "end": { "character": 19, "line": 19 },
+      "start": { "character": 16, "line": 19 }
+    },
+    "uri": "file:///AppliedFunctor.res"
+  },
+  {
+    "range": {
+      "end": { "character": 19, "line": 22 },
+      "start": { "character": 16, "line": 22 }
+    },
+    "uri": "file:///AppliedFunctor.res"
+  }
+]
+
+Rename src/AppliedFunctor.res 22:16 insert
+[
+  {
+    "edits": [
+      {
+        "newText": "insert",
+        "range": {
+          "end": { "character": 9, "line": 7 },
+          "start": { "character": 6, "line": 7 }
+        }
+      },
+      {
+        "newText": "insert",
+        "range": {
+          "end": { "character": 18, "line": 16 },
+          "start": { "character": 15, "line": 16 }
+        }
+      },
+      {
+        "newText": "insert",
+        "range": {
+          "end": { "character": 19, "line": 19 },
+          "start": { "character": 16, "line": 19 }
+        }
+      },
+      {
+        "newText": "insert",
+        "range": {
+          "end": { "character": 19, "line": 22 },
+          "start": { "character": 16, "line": 22 }
+        }
+      }
+    ],
+    "textDocument": { "uri": "file:///AppliedFunctor.res", "version": 0 }
+  }
+]
+
+Hover src/AppliedFunctor.res 25:14
 {
   "contents": {
     "kind": "markdown",
     "value": "```rescript\nIntSet.t\n```\n\n---\n\n```\n \n```\n```rescript\ntype IntSet.t = array<int>\n```\nGo to: [Type definition](AppliedFunctor.res#L7,3)\n"
   }
 }
 
-Complete src/AppliedFunctor.res 21:19
-pos_cursor:[21:19] pos_no_white:[21:18] Found expr:[21:12->21:19]
-Pexp_ident IntSet.:[21:12->21:19]
+Complete src/AppliedFunctor.res 28:19
+pos_cursor:[28:19] pos_no_white:[28:18] Found expr:[28:12->28:19]
+Pexp_ident IntSet.:[28:12->28:19]
 Completable: Cpath Value[IntSet, ""]
 Package opens Stdlib.place holder Pervasives.JsxModules.place holder
 Resolved opens 1 Stdlib