Deduplicate arguments for different constructors in undefined fns

Makes the generated undefined functions smaller, easier to read, and
avoids excessive memory usage in Coq (e.g., for large AST types).

3 files changed, 143 insertions, 6 deletions

diff --git a/src/ast_util.ml b/src/ast_util.ml
index 28b27e93..78a34a33 100644
--- a/src/ast_util.ml
+++ b/src/ast_util.ml
@@ -853,6 +853,95 @@ end
 
 module BESet = Set.Make(BE)
 
+let effect_compare (Effect_aux (Effect_set l1,_)) (Effect_aux (Effect_set l2,_)) =
+  match compare (List.length l1) (List.length l2) with
+  | 0 -> Util.compare_list BE.compare l1 l2
+  | n -> n
+let order_compare (Ord_aux (o1,_)) (Ord_aux (o2,_)) =
+  match o1, o2 with
+  | Ord_var k1, Ord_var k2 -> Kid.compare k1 k2
+  | Ord_inc, Ord_inc -> 0
+  | Ord_dec, Ord_dec -> 0
+  | Ord_var _, _ -> -1 | _, Ord_var _ -> 1
+  | Ord_inc, _ -> -1   | _, Ord_inc -> 1
+let lex_ord f g x1 x2 y1 y2 =
+  match f x1 x2 with
+  | 0 -> g y1 y2
+  | n -> n
+
+module NC = struct
+  type t = n_constraint
+  let rec compare (NC_aux (nc1,_)) (NC_aux (nc2,_)) =
+    match nc1, nc2 with
+    | NC_equal (n1,n2), NC_equal (n3,n4)
+    | NC_bounded_ge (n1,n2), NC_bounded_ge (n3,n4)
+    | NC_bounded_le (n1,n2), NC_bounded_le (n3,n4)
+    | NC_not_equal (n1,n2), NC_not_equal (n3,n4)
+      -> lex_ord Nexp.compare Nexp.compare n1 n3 n2 n4
+    | NC_set (k1,s1), NC_set (k2,s2) ->
+       lex_ord Kid.compare (Util.compare_list Nat_big_num.compare) k1 k2 s1 s2
+    | NC_or (nc1,nc2), NC_or (nc3,nc4)
+    | NC_and (nc1,nc2), NC_and (nc3,nc4)
+      -> lex_ord compare compare nc1 nc3 nc2 nc4
+    | NC_true, NC_true
+    | NC_false, NC_false
+      -> 0
+    | NC_equal _, _ -> -1      | _, NC_equal _ -> 1
+    | NC_bounded_ge _, _ -> -1 | _, NC_bounded_ge _ -> 1
+    | NC_bounded_le _, _ -> -1 | _, NC_bounded_le _ -> 1
+    | NC_not_equal _, _ -> -1  | _, NC_not_equal _ -> 1
+    | NC_set _, _ -> -1        | _, NC_set _ -> 1
+    | NC_or _, _ -> -1         | _, NC_or _ -> 1
+    | NC_and _, _ -> -1        | _, NC_and _ -> 1
+    | NC_true, _ -> -1         | _, NC_true -> 1
+end
+
+module Typ = struct
+  type t = typ
+  let rec compare (Typ_aux (t1,_)) (Typ_aux (t2,_)) =
+    match t1,t2 with
+    | Typ_internal_unknown, Typ_internal_unknown -> 0
+    | Typ_id id1, Typ_id id2 -> Id.compare id1 id2
+    | Typ_var kid1, Typ_var kid2 -> Kid.compare kid1 kid2
+    | Typ_fn (t1,t2,e1), Typ_fn (t3,t4,e2) ->
+       (match compare t1 t3 with
+       | 0 -> (match compare t2 t4 with
+         | 0 -> effect_compare e1 e2
+         | n -> n)
+       | n -> n)
+    | Typ_bidir (t1,t2), Typ_bidir (t3,t4) ->
+       (match compare t1 t3 with
+       | 0 -> compare t2 t3
+       | n -> n)
+    | Typ_tup ts1, Typ_tup ts2 -> Util.compare_list compare ts1 ts2
+    | Typ_exist (ks1,nc1,t1), Typ_exist (ks2,nc2,t2) ->
+       (match Util.compare_list Kid.compare ks1 ks2 with
+       | 0 -> (match NC.compare nc1 nc2 with
+         | 0 -> compare t1 t2
+         | n -> n)
+       | n -> n)
+    | Typ_app (id1,ts1), Typ_app (id2,ts2) ->
+       (match Id.compare id1 id2 with
+       | 0 -> Util.compare_list arg_compare ts1 ts2
+       | n -> n)
+    | Typ_internal_unknown, _ -> -1 | _, Typ_internal_unknown -> 1
+    | Typ_id _, _ -> -1    | _, Typ_id _ -> 1
+    | Typ_var _, _ -> -1   | _, Typ_var _ -> 1
+    | Typ_fn _, _ -> -1    | _, Typ_fn _ -> 1
+    | Typ_bidir _, _ -> -1 | _, Typ_bidir _ -> 1
+    | Typ_tup _, _ -> -1   | _, Typ_tup _ -> 1
+    | Typ_exist _, _ -> -1 | _, Typ_exist _ -> 1
+  and arg_compare (Typ_arg_aux (ta1,_)) (Typ_arg_aux (ta2,_)) =
+    match ta1, ta2 with
+    | Typ_arg_nexp n1,  Typ_arg_nexp n2  -> Nexp.compare n1 n2
+    | Typ_arg_typ t1,   Typ_arg_typ t2   -> compare t1 t2
+    | Typ_arg_order o1, Typ_arg_order o2 -> order_compare o1 o2
+    | Typ_arg_nexp _, _ -> -1  | _, Typ_arg_nexp _ -> 1
+    | Typ_arg_typ _, _  -> -1  | _, Typ_arg_typ _ -> 1
+end
+
+module TypMap = Map.Make(Typ)
+
 let rec nexp_frees (Nexp_aux (nexp, l)) =
   match nexp with
   | Nexp_id _ -> raise (Reporting_basic.err_typ l "Unimplemented Nexp_id in nexp_frees")
diff --git a/src/ast_util.mli b/src/ast_util.mli
index 57e1a7cc..b4ef27ee 100644
--- a/src/ast_util.mli
+++ b/src/ast_util.mli
@@ -251,6 +251,12 @@ module BE : sig
   val compare : base_effect -> base_effect -> int
 end
 
+(* NB: the comparison function does not expand synonyms *)
+module Typ : sig
+  type t = typ
+  val compare : typ -> typ -> int
+end
+
 module IdSet : sig
   include Set.S with type elt = id
 end
@@ -279,6 +285,10 @@ module Bindings : sig
   include Map.S with type key = id
 end
 
+module TypMap : sig
+  include Map.S with type key = typ
+end
+
 val nexp_frees : nexp -> KidSet.t
 val nexp_identical : nexp -> nexp -> bool
 val is_nexp_constant : nexp -> bool
diff --git a/src/initial_check.ml b/src/initial_check.ml
index af3d5fd5..60c22afc 100644
--- a/src/initial_check.ml
+++ b/src/initial_check.ml
@@ -1155,14 +1155,52 @@ let generate_undefineds vs_ids (Defs defs) =
                             (mk_exp (E_record (mk_fexps (List.map (fun (_, id) -> mk_fexp id (mk_lit_exp L_undef)) fields))))]]
     | TD_variant (id, _, typq, tus, _) when not (IdSet.mem (prepend_id "undefined_" id) vs_ids) ->
        let pat = mk_pat (P_tup (quant_items typq |> List.map quant_item_param |> List.concat |> List.map (fun id -> mk_pat (P_id id)))) in
+       let body =
+         if !opt_fast_undefined && List.length tus > 0 then
+           undefined_tu (List.hd tus)
+         else
+           (* Deduplicate arguments for each constructor to keep definitions
+              manageable. *)
+           let extract_tu = function
+             | Tu_aux (Tu_ty_id (Typ_aux (Typ_tup typs, _), id), _) -> (id, typs)
+             | Tu_aux (Tu_ty_id (typ, id), _) -> (id, [typ])
+           in
+           let record_arg_typs m (_,typs) =
+             let m' = 
+               List.fold_left (fun m typ ->
+                 TypMap.add typ (1 + try TypMap.find typ m with Not_found -> 0) m) TypMap.empty typs in
+             TypMap.merge (fun _ x y -> match x,y with Some m, Some n -> Some (max m n)
+             | None, x -> x
+             | x, None -> x) m m'
+           in
+           let make_undef_var typ n (i,lbs,m) =
+             let j = i+n in
+             let rec aux k =
+               if k = j then [] else
+                 let v = mk_id ("u_" ^ string_of_int k) in
+                 (mk_letbind (mk_pat (P_typ (typ,mk_pat (P_id v)))) (mk_lit_exp L_undef))::
+                   (aux (k+1))
+             in
+             (j, aux i @ lbs, TypMap.add typ i m)
+           in
+           let make_constr m (id,typs) =
+             let args, _ = List.fold_right (fun typ (acc,m) ->
+               let i = TypMap.find typ m in
+               (mk_exp (E_id (mk_id ("u_" ^ string_of_int i)))::acc,
+                TypMap.add typ (i+1) m)) typs ([],m) in
+             mk_exp (E_app (id, args))
+           in
+           let constr_args = List.map extract_tu tus in
+           let typs_needed = List.fold_left record_arg_typs TypMap.empty constr_args in
+           let (_,letbinds,typ_to_var) = TypMap.fold make_undef_var typs_needed (0,[],TypMap.empty) in
+           List.fold_left (fun e lb -> mk_exp (E_let (lb,e)))
+             (mk_exp (E_app (mk_id "internal_pick",
+                             [mk_exp (E_list (List.map (make_constr typ_to_var) constr_args))]))) letbinds
+       in
        [mk_val_spec (VS_val_spec (undefined_typschm id typq, prepend_id "undefined_" id, (fun _ -> None), false));
         mk_fundef [mk_funcl (prepend_id "undefined_" id)
-                            pat
-                            (if !opt_fast_undefined && List.length tus > 0 then
-                               undefined_tu (List.hd tus)
-                             else
-                               mk_exp (E_app (mk_id "internal_pick",
-                                            [mk_exp (E_list (List.map undefined_tu tus))])))]]
+                      pat
+                      body]]
     | _ -> []
   in
   let rec undefined_defs = function