From d14b9a452945ea7a59ead320446eb9fa0149456b Mon Sep 17 00:00:00 2001
From: Cristiano Calcagno <cristianoc@users.noreply.github.com>
Date: Wed, 18 Dec 2024 10:20:37 +0100
Subject: [PATCH] Move arity decoding to ast conversion.

---
 compiler/ml/ast_mapper_from0.ml | 12 +++++++++++-
 compiler/ml/ast_uncurried.ml    | 14 ++++----------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/compiler/ml/ast_mapper_from0.ml b/compiler/ml/ast_mapper_from0.ml
index 91cb6a809b..60f94866eb 100644
--- a/compiler/ml/ast_mapper_from0.ml
+++ b/compiler/ml/ast_mapper_from0.ml
@@ -109,7 +109,17 @@ module T = struct
       | Ptyp_constr
           (lid, [({ptyp_desc = Ptyp_arrow (lbl, t1, t2, _)} as fun_t); t_arity])
         when lid.txt = Lident "function$" ->
-        let arity = Ast_uncurried.arity_from_type t_arity in
+        let decode_arity_string arity_s =
+          int_of_string
+            ((String.sub [@doesNotRaise]) arity_s 9 (String.length arity_s - 9))
+        in
+        let arity_from_type (typ : Parsetree.core_type) =
+          match typ.ptyp_desc with
+          | Ptyp_variant ([Rtag ({txt}, _, _, _)], _, _) ->
+            decode_arity_string txt
+          | _ -> assert false
+        in
+        let arity = arity_from_type t_arity in
         let fun_t =
           {fun_t with ptyp_desc = Ptyp_arrow (lbl, t1, t2, Some arity)}
         in
diff --git a/compiler/ml/ast_uncurried.ml b/compiler/ml/ast_uncurried.ml
index f464026fbd..92900e615d 100644
--- a/compiler/ml/ast_uncurried.ml
+++ b/compiler/ml/ast_uncurried.ml
@@ -1,20 +1,12 @@
 (* Uncurried AST *)
 
 let encode_arity_string arity = "Has_arity" ^ string_of_int arity
-let decode_arity_string arity_s =
-  int_of_string
-    ((String.sub [@doesNotRaise]) arity_s 9 (String.length arity_s - 9))
 
 let arity_type ~loc arity =
   Ast_helper.Typ.variant ~loc
     [Rtag ({txt = encode_arity_string arity; loc}, [], true, [])]
     Closed None
 
-let arity_from_type (typ : Parsetree.core_type) =
-  match typ.ptyp_desc with
-  | Ptyp_variant ([Rtag ({txt}, _, _, _)], _, _) -> decode_arity_string txt
-  | _ -> assert false
-
 let uncurried_type ~loc ~arity (t_arg : Parsetree.core_type) =
   let t_arg =
     match t_arg.ptyp_desc with
@@ -52,8 +44,10 @@ let core_type_is_uncurried_fun (typ : Parsetree.core_type) =
 
 let core_type_extract_uncurried_fun (typ : Parsetree.core_type) =
   match typ.ptyp_desc with
-  | Ptyp_constr ({txt = Lident "function$"}, [t_arg; t_arity]) ->
-    (arity_from_type t_arity, t_arg)
+  | Ptyp_constr
+      ( {txt = Lident "function$"},
+        [({ptyp_desc = Ptyp_arrow (_, _, _, Some arity)} as t_arg); _] ) ->
+    (arity, t_arg)
   | _ -> assert false
 
 let type_is_uncurried_fun = Ast_uncurried_utils.type_is_uncurried_fun