From 6ed4f9e98e0be8651684377023270ff7b9bfc833 Mon Sep 17 00:00:00 2001
From: Maxime Buyse <maxime@cryspen.com>
Date: Wed, 4 Sep 2024 12:33:23 +0200
Subject: [PATCH] Avoid recursing twice on the arguments.

---
 engine/lib/phases/phase_specialize.ml | 36 +++++++++++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/engine/lib/phases/phase_specialize.ml b/engine/lib/phases/phase_specialize.ml
index a3bde5b6f..17ed60370 100644
--- a/engine/lib/phases/phase_specialize.ml
+++ b/engine/lib/phases/phase_specialize.ml
@@ -113,7 +113,14 @@ module Make (F : Features.T) =
 
           method! visit_expr () e =
             match e.e with
-            | App { f = { e = GlobalVar f; _ } as f'; args = l; _ } -> (
+            | App
+                {
+                  f = { e = GlobalVar f; _ } as f';
+                  args = l;
+                  trait;
+                  generic_args;
+                  bounds_impls;
+                } -> (
                 let l = List.map ~f:(self#visit_expr ()) l in
                 let matching =
                   List.filter patterns ~f:(fun { fn; args; _ } ->
@@ -139,7 +146,32 @@ module Make (F : Features.T) =
                             bounds_impls = [];
                           };
                     }
-                | [] -> super#visit_expr () e
+                | [] -> (
+                    (* In this case we need to avoid recursing again through the arguments *)
+                    let visited =
+                      super#visit_expr ()
+                        {
+                          e with
+                          e =
+                            App
+                              {
+                                f = f';
+                                args = [];
+                                trait;
+                                generic_args;
+                                bounds_impls;
+                              };
+                        }
+                    in
+                    match visited.e with
+                    | App { f; trait; generic_args; bounds_impls; _ } ->
+                        {
+                          visited with
+                          e =
+                            App
+                              { f; args = l; trait; generic_args; bounds_impls };
+                        }
+                    | _ -> super#visit_expr () e)
                 | _ ->
                     Error.assertion_failure e.span
                       "Found multiple matching patterns")