Simplify monomorphisation, and preserve type annotations

compiler/lcalc/monomorphize.ml

@@ -212,112 +212,96 @@ let is_some c =
   (assert (EnumConstructor.equal Expr.none_constr c);
    false)
 
-(* We output a typed expr but the types in the output are wrong, it should be
-   untyped and re-typed later. *)
 let rec monomorphize_expr
     (monomorphized_instances : monomorphized_instances)
-    (e : typed expr) : typed expr boxed =
+    (e0 : typed expr) : typed expr boxed =
-  let typ = Expr.ty e in
+  let ty0 = Expr.ty e0 in
-  match Mark.remove e with
+  (* Keys in [monomorphized_instances] are before monomorphization, so collect
+     this top-down *)
+  let f_expr = monomorphize_expr monomorphized_instances in
+  let f_ty = monomorphize_typ monomorphized_instances in
+  (* Proceed bottom-up: apply first to the sub-terms *)
+  let e = Expr.map ~f:f_expr ~typ:f_ty ~op:Fun.id e0 in
+  let m = Mark.get e in
+  let map_box f = Expr.Box.app1 e (fun e -> f (Mark.remove e)) m in
+  map_box
+  @@ function
   | ETuple args ->
-    let new_args = List.map (monomorphize_expr monomorphized_instances) args in
+    let tuple_instance = Type.Map.find ty0 monomorphized_instances.tuples in
-    let tuple_instance = Type.Map.find typ monomorphized_instances.tuples in
+    EStruct
-    let fields =
+      {
-      StructField.Map.of_list
+        name = tuple_instance.name;
-      @@ List.map2
+        fields =
-           (fun new_arg (tuple_field, _) -> tuple_field, new_arg)
+          StructField.Map.of_list
-           new_args tuple_instance.fields
+          @@ List.map2
+               (fun (tuple_field, _) arg -> tuple_field, arg)
+               tuple_instance.fields args;
+      }
+  | ETupleAccess { e; index; _ } ->
+    (* The type of the tuple needs to be recovered from the untransformed
+       expr *)
+    let tup_ty =
+      match e0 with ETupleAccess { e; _ }, _ -> Expr.ty e | _ -> assert false
     in
-    Expr.estruct ~name:tuple_instance.name ~fields (Mark.get e)
-  | ETupleAccess { e = e1; index; _ } ->
     let tuple_instance =
-      Type.Map.find (Expr.ty e1) monomorphized_instances.tuples
+      Type.Map.find tup_ty monomorphized_instances.tuples
-    in
-    let new_e1 = monomorphize_expr monomorphized_instances e1 in
-    Expr.estructaccess ~name:tuple_instance.name
-      ~field:(fst (List.nth tuple_instance.fields index))
-      ~e:new_e1 (Mark.get e)
-  | EMatch { name; e = e1; cases } when EnumName.equal name Expr.option_enum ->
-    let new_e1 = monomorphize_expr monomorphized_instances e1 in
-    let new_cases =
-      EnumConstructor.Map.bindings
-        (EnumConstructor.Map.map
-           (monomorphize_expr monomorphized_instances)
-           cases)
     in
+    EStructAccess
+      {
+        name = tuple_instance.name;
+        e;
+        field = fst (List.nth tuple_instance.fields index);
+      }
+  | EMatch { name; e; cases } when EnumName.equal name Expr.option_enum ->
+    let option_instance = Type.Map.find ty0 monomorphized_instances.options in
+    EMatch
+      {
+        name = option_instance.name;
+        e;
+        cases =
+          EnumConstructor.Map.fold
+            (fun c ->
+              EnumConstructor.Map.add
+                (if is_some c then option_instance.some_cons
+                 else option_instance.none_cons))
+            cases EnumConstructor.Map.empty;
+      }
+  | EInj { name; e; cons } when EnumName.equal name Expr.option_enum ->
     let option_instance =
       Type.Map.find
-        (match Mark.remove (Expr.ty e1) with
+        (match Mark.remove ty0 with TOption t -> t | _ -> assert false)
-        | TOption t -> t
-        | _ -> failwith "should not happen")
         monomorphized_instances.options
     in
-    let new_cases =
+    EInj
-      match new_cases with
+      {
-      | [(n1, e1); (n2, e2)] -> (
+        name = option_instance.name;
-        match is_some n1, is_some n2 with
+        e;
-        | true, false ->
+        cons =
-          [option_instance.some_cons, e1; option_instance.none_cons, e2]
+          (if is_some cons then option_instance.some_cons
-        | false, true ->
+           else option_instance.none_cons);
-          [option_instance.some_cons, e2; option_instance.none_cons, e1]
+      }
-        | _ -> failwith "should not happen")
+  | EArray elts as e ->
-      | _ -> failwith "should not happen"
+    let elt_ty =
+      match Mark.remove ty0 with TArray t -> t | _ -> assert false
     in
-    let new_cases = EnumConstructor.Map.of_list new_cases in
+    let array_instance = Type.Map.find elt_ty monomorphized_instances.arrays in
-    Expr.ematch ~name:option_instance.name ~e:new_e1 ~cases:new_cases
+    EStruct
-      (Mark.get e)
+      {
-  | EInj { name; e = e1; cons } when EnumName.equal name Expr.option_enum ->
+        name = array_instance.name;
-    let option_instance =
+        fields =
-      Type.Map.find
+          StructField.Map.of_list
-        (match Mark.remove (Expr.ty e) with
+            [
-        | TOption t -> t
+              ( array_instance.len_field,
-        | _ -> failwith "should not happen")
+                ( ELit (LInt (Runtime.integer_of_int (List.length elts))),
-        monomorphized_instances.options
+                  Expr.with_ty m (TLit TInt, Expr.mark_pos m) ) );
-    in
+              ( array_instance.content_field,
-    let new_e1 = monomorphize_expr monomorphized_instances e1 in
+                (e, Expr.with_ty m (TArray (f_ty elt_ty), Expr.mark_pos m)) );
-    let new_cons =
+            ];
-      if is_some cons then option_instance.some_cons
+      }
-      else option_instance.none_cons
+  | e -> e
-    in
-    Expr.einj ~name:option_instance.name ~e:new_e1 ~cons:new_cons (Mark.get e)
-    (* We do not forget to tweak types stored directly in the AST in the nodes
-       of kind [EAbs], [EApp] and [EAppOp]. *)
-  | EAbs { binder; tys } ->
-    let new_tys = List.map (monomorphize_typ monomorphized_instances) tys in
-    let vars, body = Bindlib.unmbind binder in
-    let new_body = monomorphize_expr monomorphized_instances body in
-    Expr.make_abs vars new_body new_tys (Expr.pos e)
-  | EApp { f; args; tys } ->
-    let new_f = monomorphize_expr monomorphized_instances f in
-    let new_args = List.map (monomorphize_expr monomorphized_instances) args in
-    let new_tys = List.map (monomorphize_typ monomorphized_instances) tys in
-    Expr.eapp ~f:new_f ~args:new_args ~tys:new_tys (Mark.get e)
-  | EAppOp { op; args; tys } ->
-    let new_args = List.map (monomorphize_expr monomorphized_instances) args in
-    let new_tys = List.map (monomorphize_typ monomorphized_instances) tys in
-    Expr.eappop ~op ~args:new_args ~tys:new_tys (Mark.get e)
-  | EArray args ->
-    let new_args = List.map (monomorphize_expr monomorphized_instances) args in
-    let array_instance =
-      Type.Map.find
-        (match Mark.remove (Expr.ty e) with
-        | TArray t -> t
-        | _ -> failwith "should not happen")
-        monomorphized_instances.arrays
-    in
-    Expr.estruct ~name:array_instance.name
-      ~fields:
-        (StructField.Map.add array_instance.content_field
-           (Expr.earray new_args (Mark.get e))
-           (StructField.Map.singleton array_instance.len_field
-              (Expr.elit
-                 (LInt (Runtime.integer_of_int (List.length args)))
-                 (Mark.get e))))
-      (Mark.get e)
-  | _ -> Expr.map ~f:(monomorphize_expr monomorphized_instances) e
 
 let program (prg : typed program) :
-    untyped program * Scopelang.Dependency.TVertex.t list =
+    typed program * Scopelang.Dependency.TVertex.t list =
   let monomorphized_instances = collect_monomorphized_instances prg in
   (* First we remove the polymorphic option type *)
   let prg =
@@ -417,7 +401,6 @@ let program (prg : typed program) :
                      scope_body)))
          ~varf:Fun.id prg.code_items
   in
-  let prg = Program.untype { prg with code_items } in
+  ( { prg with code_items },
-  ( prg,
     Scopelang.Dependency.check_type_cycles prg.decl_ctx.ctx_structs
       prg.decl_ctx.ctx_enums )

compiler/lcalc/monomorphize.mli

@@ -18,7 +18,7 @@ open Shared_ast
 open Ast
 
 val program :
-  typed program -> untyped program * Scopelang.Dependency.TVertex.t list
+  typed program -> typed program * Scopelang.Dependency.TVertex.t list
 (** This function performs type monomorphization in a Catala program with two
     main actions: {ul
     {- transforms tuples into named structs.}

compiler/shared_ast/print.ml

@@ -474,13 +474,20 @@ end
 
 module ExprGen (C : EXPR_PARAM) = struct
   let rec expr_aux :
-      type a.
+      type a t.
       Bindlib.ctxt ->
       Ocolor_types.color4 list ->
       Format.formatter ->
-      (a, 't) gexpr ->
+      (a, t) gexpr ->
       unit =
    fun bnd_ctx colors fmt e ->
+    (* (* Uncomment for type annotations everywhere *)
+     * (fun f ->
+     *    Format.fprintf fmt "@[<hv 1>(%a:@ %a)@]"
+     *      f e
+     *      typ_debug
+     *      (match Mark.get e with Typed {ty; _} -> ty | _ -> TAny,Pos.no_pos))
+     * @@ fun fmt e -> *)
     let exprb bnd_ctx colors e = expr_aux bnd_ctx colors e in
     let exprc colors e = exprb bnd_ctx colors e in
     let expr e = exprc colors e in