From 11a2a54fdb06d56a63117e8be7cd8f687957cc47 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ulysse=20G=C3=A9rard?= <thevoodoos@gmail.com>
Date: Wed, 22 Nov 2023 16:58:35 +0100
Subject: [PATCH] shapes: dependent keep_alias

---
 src/analysis/locate.ml     |  7 ++++++-
 src/ocaml/typing/shape.ml  | 12 ++++++------
 src/ocaml/typing/shape.mli |  4 ++--
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/src/analysis/locate.ml b/src/analysis/locate.ml
index b8883d0f94..95b1ce1b99 100644
--- a/src/analysis/locate.ml
+++ b/src/analysis/locate.ml
@@ -559,8 +559,13 @@ let find_definition_uid ~config ~env ~(decl : Env_lookup.item) path =
   let shape = Env.shape_of_path ~namespace env path in
   log ~title:"shape_of_path" "initial: %a"
     Logger.fmt (Fun.flip Shape.print shape);
+  let keep_aliases =
+    if config.traverse_aliases
+    then (fun _ -> false)
+    else (fun _ -> true)
+  in
   let reduced = Shape_reduce.reduce_for_uid
-    ~keep_aliases:(not config.traverse_aliases) env shape
+    ~keep_aliases env shape
   in
   log ~title:"shape_of_path" "reduced: %a"
     Logger.fmt (fun fmt -> Shape.print_reduction_result fmt reduced);
diff --git a/src/ocaml/typing/shape.ml b/src/ocaml/typing/shape.ml
index 46b820b467..0307eb2b1b 100644
--- a/src/ocaml/typing/shape.ml
+++ b/src/ocaml/typing/shape.ml
@@ -327,7 +327,7 @@ end) = struct
 
   type env = {
     fuel: int ref;
-    keep_aliases: bool;
+    keep_aliases: t -> bool;
     global_env: Params.env;
     local_env: local_env;
     reduce_memo_table: (thunk, nf) Hashtbl.t;
@@ -455,9 +455,9 @@ end) = struct
       | Struct m ->
           let mnf = Item.Map.map (delay_reduce env) m in
           return (NStruct mnf)
-      | Alias t ->
-        let nf = reduce env t in
-        if env.keep_aliases then
+      | Alias aliased_t ->
+        let nf = reduce env aliased_t in
+        if env.keep_aliases t then
           return (NAlias nf)
         else nf
       | Error s -> return ~approximated:true (NError s)
@@ -499,7 +499,7 @@ end) = struct
   let reduce_memo_table = Hashtbl.create 42
   let read_back_memo_table = Hashtbl.create 42
 
-  let reduce ?(keep_aliases = true) global_env t =
+  let reduce ?(keep_aliases = fun _ -> true) global_env t =
     let fuel = ref Params.fuel in
     let local_env = Ident.Map.empty in
     let env = {
@@ -523,7 +523,7 @@ end) = struct
     | NError _ -> false
     | NLeaf -> false
 
-  let reduce_for_uid ?(keep_aliases = true) global_env t =
+  let reduce_for_uid ?(keep_aliases = fun _ -> true) global_env t =
     let fuel = ref Params.fuel in
     let local_env = Ident.Map.empty in
     let env = {
diff --git a/src/ocaml/typing/shape.mli b/src/ocaml/typing/shape.mli
index f6bafecd6a..18cd20331a 100644
--- a/src/ocaml/typing/shape.mli
+++ b/src/ocaml/typing/shape.mli
@@ -218,12 +218,12 @@ module Make_reduce(Context : sig
     val find_shape : env -> Ident.t -> t
   end) : sig
   val reduce :
-    ?keep_aliases:bool -> Context.env -> t -> t
+    ?keep_aliases:(t -> bool) -> Context.env -> t -> t
 
   (** Perform weak reduction and return the head's uid if any. If reduction was
     incomplete the partially reduced shape is returned. *)
   val reduce_for_uid :
-    ?keep_aliases:bool -> Context.env -> t -> reduction_result
+    ?keep_aliases:(t -> bool) -> Context.env -> t -> reduction_result
 end
 
 (** [toplevel_local_reduce] is only suitable to reduce toplevel shapes (shapes