functions: add uid to [Texp_newtype'] nodes

src/analysis/ast_iterators.ml

@@ -19,14 +19,14 @@ let iter_on_defs ~uid_to_locs_tbl =
     Types.Uid.Tbl.add uid_to_locs_tbl uid loc
   in
   { iter_decl with
-    expr = (fun sub ({ exp_extra; exp_env; _ } as expr) ->
+    expr = (fun sub ({ exp_extra; _ } as expr) ->
       List.iter exp_extra ~f:(fun (exp_extra, _loc, _attr) ->
         match exp_extra with
-        | Texp_newtype' (typ_id, typ_name) ->
-          log "Found definition %s (%a)\n%!" typ_name.txt
+        | Texp_newtype' (typ_id, typ_name, uid) ->
+          log "Found newtype %s wit id %a (%a)\n%!" typ_name.txt
+            Logger.fmt (Fun.flip Ident.print_with_scope typ_id)
             Logger.fmt (fun fmt -> Location.print_loc fmt typ_name.loc);
-          let decl = Env.find_type (Path.Pident typ_id) exp_env in
-          register_uid decl.type_uid typ_name;
+          register_uid uid typ_name;
           ()
         | _ -> ());
       iter_decl.expr sub expr);

src/ocaml/merlin_specific/browse_raw.ml

@@ -848,7 +848,7 @@ let expression_paths { Typedtree. exp_desc; exp_extra; _ } =
   List.fold_left ~init exp_extra
     ~f:(fun acc (extra, _, _) ->
       match extra with
-      | Texp_newtype' (id, label_loc) ->
+      | Texp_newtype' (id, label_loc, _) ->
         let path = Path.Pident id in
         let lid = Longident.Lident (label_loc.txt) in
         (mkloc path label_loc.loc, Some lid) :: acc

src/ocaml/typing/printtyped.ml

@@ -328,7 +328,7 @@ and expression_extra i ppf x attrs =
       option i core_type ppf cto;
   | Texp_newtype s ->
       line i ppf "Texp_newtype \"%s\"\n" s;
-  | Texp_newtype' (id, _) ->
+  | Texp_newtype' (id, _, _) ->
       line i ppf "Texp_newtype' \"%a\"\n" fmt_ident id;
       attributes i ppf attrs;
 

src/ocaml/typing/typecore.ml

@@ -3558,8 +3558,8 @@ and type_expect_
       re
         { exp_desc = Texp_function (params, body);
           exp_loc = loc;
-          exp_extra =
-            List.map (fun { txt; loc } -> Texp_newtype txt, loc, []) newtypes;
+          exp_extra = List.map (fun (id, txt_loc, uid) ->
+            Texp_newtype' (id, txt_loc, uid), txt_loc.loc, []) newtypes;
           exp_type;
           exp_attributes = sexp.pexp_attributes;
           exp_env = env;
@@ -4292,15 +4292,15 @@ and type_expect_
       re { exp with exp_extra =
              (Texp_poly cty, loc, sexp.pexp_attributes) :: exp.exp_extra }
   | Pexp_newtype({txt=name} as label_loc, sbody) ->
-      let body, ety, id = type_newtype loc env name (fun env ->
+      let body, ety, id, uid = type_newtype loc env name (fun env ->
         let expr = type_exp env sbody in
         expr, expr.exp_type)
       in
       (* non-expansive if the body is non-expansive, so we don't introduce
          any new extra node in the typed AST. *)
       rue { body with exp_loc = loc; exp_type = ety;
             exp_extra =
-            (Texp_newtype' (id, label_loc), loc, sexp.pexp_attributes) :: body.exp_extra }
+            (Texp_newtype' (id, label_loc, uid), loc, sexp.pexp_attributes) :: body.exp_extra }
   | Pexp_pack m ->
       let (p, fl) =
         match get_desc (Ctype.expand_head env (instance ty_expected)) with
@@ -4599,7 +4599,7 @@ and type_constraint_expect
     nodes for the newtype properly linked.
 *)
 and type_newtype
-  : type a. _ -> _ -> _ -> (Env.t -> a * type_expr) -> a * type_expr * Ident.t =
+  : type a. _ -> _ -> _ -> (Env.t -> a * type_expr) -> a * type_expr * Ident.t * Uid.t =
   fun loc env name type_body ->
   let ty =
     if Typetexp.valid_tyvar_name name then
@@ -4629,7 +4629,8 @@ and type_newtype
     in
     let ety = Subst.type_expr Subst.identity exp_type in
     replace ety;
-    (result, ety, id)
+    let uid = decl.type_uid in
+    (result, ety, id, uid)
   end
 
 and type_ident env ?(recarg=Rejected) lid =
@@ -4764,6 +4765,7 @@ and type_function
       in
       with_explanation ty_fun.explanation (fun () ->
         unify_exp_types loc env exp_type (instance ty_expected));
+      let newtype = nt_id, newtype, nt_uid in
       exp_type, params, body, newtype :: newtypes, contains_gadt
   | { pparam_desc = Pparam_val (arg_label, default_arg, pat); pparam_loc }
       :: rest
@@ -4862,7 +4864,7 @@ and type_function
           fp_arg_label = arg_label;
           fp_param;
           fp_partial = partial;
-          fp_newtypes = newtypes;
+          fp_newtypes = List.map (fun (_,v,_) -> v) newtypes;
           fp_loc = pparam_loc;
         }
       in

src/ocaml/typing/typedtree.ml

@@ -97,7 +97,7 @@ and exp_extra =
   | Texp_coerce of core_type option * core_type
   | Texp_poly of core_type option
   | Texp_newtype of string
-  | Texp_newtype' of Ident.t * label loc
+  | Texp_newtype' of Ident.t * label loc * Uid.t
 
 and expression_desc =
     Texp_ident of Path.t * Longident.t loc * Types.value_description

src/ocaml/typing/typedtree.mli

@@ -172,7 +172,7 @@ and exp_extra =
         (** Used for method bodies. *)
   | Texp_newtype of string
         (** fun (type t) ->  *)
-  | Texp_newtype' of Ident.t * label loc
+  | Texp_newtype' of Ident.t * label loc * Uid.t
   (** merlin-specific: keep enough information to correctly implement
       occurrences for local-types.
       Merlin typechecker uses [Texp_newtype'] constructor, while upstream

src/ocaml/typing/untypeast.ml

@@ -370,7 +370,7 @@ let exp_extra sub (extra, loc, attrs) sexp =
         Pexp_constraint (sexp, sub.typ sub cty)
     | Texp_poly cto -> Pexp_poly (sexp, Option.map (sub.typ sub) cto)
     | Texp_newtype s -> Pexp_newtype (mkloc s loc, sexp)
-    | Texp_newtype' (_id, label_loc) -> Pexp_newtype (label_loc, sexp)
+    | Texp_newtype' (_id, label_loc, _) -> Pexp_newtype (label_loc, sexp)
   in
   Exp.mk ~loc ~attrs desc