Monomorphization properly done

compiler/driver.ml

@@ -272,8 +272,13 @@ module Passes = struct
     in
     Message.emit_debug "Retyping lambda calculus...";
     let prg = Typing.program ~leave_unresolved:true prg in
-    let prg =
-      if monomorphize_types then Lcalc.Monomorphize.program prg else prg
+    let prg, type_ordering =
+      if monomorphize_types then (
+        Message.emit_debug "Monomorphizing types...";
+        let prg, type_ordering = Lcalc.Monomorphize.program prg in
+        let prg = Typing.program ~leave_unresolved:false prg in
+        prg, type_ordering)
+      else prg, type_ordering
     in
     Program.untype prg, type_ordering
 
@@ -1090,6 +1095,11 @@ let main () =
     Message.Content.emit content Error;
     if Cli.globals.debug then Printexc.print_raw_backtrace stderr bt;
     exit Cmd.Exit.some_error
+  | exception Failure msg ->
+    let bt = Printexc.get_raw_backtrace () in
+    Message.Content.emit (Message.Content.of_string msg) Error;
+    if Printexc.backtrace_status () then Printexc.print_raw_backtrace stderr bt;
+    exit Cmd.Exit.some_error
   | exception Sys_error msg ->
     let bt = Printexc.get_raw_backtrace () in
     Message.Content.emit

compiler/lcalc/monomorphize.ml

@@ -37,117 +37,214 @@ type tuple_instance = {
   fields : (StructField.t * naked_typ) list;
 }
 
-type monorphized_instances = {
+type monomorphized_instances = {
   (* The keys are the types inside the [TOption] (before monomorphization). *)
   options : option_instance TypMap.t;
   (* The keys are the [TTuple] types themselves (before monomorphization). *)
   tuples : tuple_instance TypMap.t;
 }
 
-let program (prg : typed program) : typed program =
-  let monomorphized_instances : monorphized_instances =
-    let option_instances_counter = ref 0 in
-    let tuple_instances_counter = ref 0 in
-    let rec monomorphize_typ acc typ =
-      match Mark.remove typ with
-      | TStruct _ | TEnum _ | TAny | TClosureEnv | TLit _ -> acc
-      | TTuple args ->
-        let new_acc =
-          {
-            acc with
-            tuples =
-              TypMap.update (Mark.remove typ)
-                (fun monomorphized_name ->
-                  match monomorphized_name with
-                  | Some e -> Some e
-                  | None ->
-                    incr tuple_instances_counter;
-                    Some
-                      {
-                        fields =
-                          List.mapi
-                            (fun i arg ->
-                              ( StructField.fresh
-                                  ("elt_" ^ string_of_int i, Pos.no_pos),
-                                Mark.remove arg ))
-                            args;
-                        name =
-                          StructName.fresh []
-                            ( "tuple_" ^ string_of_int !option_instances_counter,
-                              Pos.no_pos );
-                      })
-                acc.tuples;
-          }
-        in
-        List.fold_left monomorphize_typ new_acc args
-      | TArray t | TDefault t -> monomorphize_typ acc t
-      | TArrow (args, ret) ->
-        List.fold_left monomorphize_typ (monomorphize_typ acc ret) args
-      | TOption t ->
-        let new_acc =
-          {
-            acc with
-            options =
-              TypMap.update (Mark.remove t)
-                (fun monomorphized_name ->
-                  match monomorphized_name with
-                  | Some e -> Some e
-                  | None ->
-                    incr option_instances_counter;
-                    Some
-                      {
-                        some_cons =
-                          EnumConstructor.fresh
-                            ( "Some_" ^ string_of_int !option_instances_counter,
-                              Pos.no_pos );
-                        none_cons =
-                          EnumConstructor.fresh
-                            ( "None_" ^ string_of_int !option_instances_counter,
-                              Pos.no_pos );
-                        some_typ = Mark.remove t;
-                        name =
-                          EnumName.fresh []
-                            ( "option_" ^ string_of_int !option_instances_counter,
-                              Pos.no_pos );
-                      })
-                acc.options;
-          }
-        in
-        monomorphize_typ new_acc t
-    in
-    let rec monomorphize_expr acc e =
-      let acc = monomorphize_typ acc (Expr.maybe_ty (Mark.get e)) in
-      Expr.shallow_fold (fun e acc -> monomorphize_expr acc e) e acc
-    in
-    let acc =
-      Scope.fold_left
-        ~init:{ options = TypMap.empty; tuples = TypMap.empty }
-        ~f:(fun acc item _ ->
-          match item with
-          | Topdef (_, typ, e) -> monomorphize_typ (monomorphize_expr acc e) typ
-          | ScopeDef (_, body) ->
-            let _, body = Bindlib.unbind body.scope_body_expr in
-            Scope.fold_left_lets ~init:acc
-              ~f:(fun acc { scope_let_typ; scope_let_expr; _ } _ ->
-                monomorphize_typ
-                  (monomorphize_expr acc scope_let_expr)
-                  scope_let_typ)
-              body)
-        prg.code_items
-    in
-    EnumName.Map.fold
-      (fun _ constructors acc ->
-        EnumConstructor.Map.fold
-          (fun _ t acc -> monomorphize_typ acc t)
-          constructors acc)
-      prg.decl_ctx.ctx_enums
-      (StructName.Map.fold
-         (fun _ fields acc ->
-           StructField.Map.fold
-             (fun _ t acc -> monomorphize_typ acc t)
-             fields acc)
-         prg.decl_ctx.ctx_structs acc)
+let collect_monomorphized_instances (prg : typed program) :
+    monomorphized_instances =
+  let option_instances_counter = ref 0 in
+  let tuple_instances_counter = ref 0 in
+  let rec collect_typ acc typ =
+    match Mark.remove typ with
+    | TStruct _ | TEnum _ | TAny | TClosureEnv | TLit _ -> acc
+    | TTuple args ->
+      let new_acc =
+        {
+          acc with
+          tuples =
+            TypMap.update (Mark.remove typ)
+              (fun monomorphized_name ->
+                match monomorphized_name with
+                | Some e -> Some e
+                | None ->
+                  incr tuple_instances_counter;
+                  Some
+                    {
+                      fields =
+                        List.mapi
+                          (fun i arg ->
+                            ( StructField.fresh
+                                ("elt_" ^ string_of_int i, Pos.no_pos),
+                              Mark.remove arg ))
+                          args;
+                      name =
+                        StructName.fresh []
+                          ( "tuple_" ^ string_of_int !option_instances_counter,
+                            Pos.no_pos );
+                    })
+              acc.tuples;
+        }
+      in
+      List.fold_left collect_typ new_acc args
+    | TArray t | TDefault t -> collect_typ acc t
+    | TArrow (args, ret) ->
+      List.fold_left collect_typ (collect_typ acc ret) args
+    | TOption t ->
+      let new_acc =
+        {
+          acc with
+          options =
+            TypMap.update (Mark.remove t)
+              (fun monomorphized_name ->
+                match monomorphized_name with
+                | Some e -> Some e
+                | None ->
+                  incr option_instances_counter;
+                  Some
+                    {
+                      some_cons =
+                        EnumConstructor.fresh
+                          ( "Some_" ^ string_of_int !option_instances_counter,
+                            Pos.no_pos );
+                      none_cons =
+                        EnumConstructor.fresh
+                          ( "None_" ^ string_of_int !option_instances_counter,
+                            Pos.no_pos );
+                      some_typ = Mark.remove t;
+                      name =
+                        EnumName.fresh []
+                          ( "option_" ^ string_of_int !option_instances_counter,
+                            Pos.no_pos );
+                    })
+              acc.options;
+        }
+      in
+      collect_typ new_acc t
   in
+  let rec collect_expr acc e =
+    let acc = collect_typ acc (Expr.maybe_ty (Mark.get e)) in
+    Expr.shallow_fold (fun e acc -> collect_expr acc e) e acc
+  in
+  let acc =
+    Scope.fold_left
+      ~init:{ options = TypMap.empty; tuples = TypMap.empty }
+      ~f:(fun acc item _ ->
+        match item with
+        | Topdef (_, typ, e) -> collect_typ (collect_expr acc e) typ
+        | ScopeDef (_, body) ->
+          let _, body = Bindlib.unbind body.scope_body_expr in
+          Scope.fold_left_lets ~init:acc
+            ~f:(fun acc { scope_let_typ; scope_let_expr; _ } _ ->
+              collect_typ (collect_expr acc scope_let_expr) scope_let_typ)
+            body)
+      prg.code_items
+  in
+  EnumName.Map.fold
+    (fun _ constructors acc ->
+      EnumConstructor.Map.fold
+        (fun _ t acc -> collect_typ acc t)
+        constructors acc)
+    prg.decl_ctx.ctx_enums
+    (StructName.Map.fold
+       (fun _ fields acc ->
+         StructField.Map.fold (fun _ t acc -> collect_typ acc t) fields acc)
+       prg.decl_ctx.ctx_structs acc)
+
+let rec monomorphize_typ
+    (monomorphized_instances : monomorphized_instances)
+    (typ : typ) : typ =
+  match Mark.remove typ with
+  | TStruct _ | TEnum _ | TAny | TClosureEnv | TLit _ -> typ
+  | TArray t1 ->
+    TArray (monomorphize_typ monomorphized_instances t1), Mark.get typ
+  | TDefault t1 ->
+    TDefault (monomorphize_typ monomorphized_instances t1), Mark.get typ
+  | TArrow (t1s, t2) ->
+    ( TArrow
+        ( List.map (monomorphize_typ monomorphized_instances) t1s,
+          monomorphize_typ monomorphized_instances t2 ),
+      Mark.get typ )
+  | TTuple _ ->
+    ( TStruct (TypMap.find (Mark.remove typ) monomorphized_instances.tuples).name,
+      Mark.get typ )
+  | TOption t1 ->
+    ( TEnum (TypMap.find (Mark.remove t1) monomorphized_instances.options).name,
+      Mark.get typ )
+
+(* 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 =
+  let typ = Expr.maybe_ty (Mark.get e) in
+  match Mark.remove e with
+  | ETuple args ->
+    let new_args = List.map (monomorphize_expr monomorphized_instances) args in
+    let tuple_instance =
+      TypMap.find (Mark.remove typ) monomorphized_instances.tuples
+    in
+    let fields =
+      StructField.Map.of_list
+      @@ List.map2
+           (fun new_arg (tuple_field, _) -> tuple_field, new_arg)
+           new_args tuple_instance.fields
+    in
+    Expr.estruct ~name:tuple_instance.name ~fields (Mark.get e)
+  | ETupleAccess { e = e1; index; _ } ->
+    let tuple_instance =
+      TypMap.find
+        (Mark.remove (Expr.maybe_ty (Mark.get e1)))
+        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
+    let option_instance =
+      TypMap.find
+        (match Mark.remove (Expr.maybe_ty (Mark.get e1)) with
+        | TOption t -> Mark.remove t
+        | _ -> failwith "should not happen")
+        monomorphized_instances.options
+    in
+    let new_cases =
+      match new_cases with
+      | [(n1, e1); (n2, e2)] -> (
+        match
+          ( Mark.remove (EnumConstructor.get_info n1),
+            Mark.remove (EnumConstructor.get_info n2) )
+        with
+        | "ESome", "ENone" ->
+          [option_instance.some_cons, e1; option_instance.none_cons, e2]
+        | "ENone", "ESome" ->
+          [option_instance.some_cons, e2; option_instance.none_cons, e1]
+        | _ -> failwith "should not happen")
+      | _ -> failwith "should not happen"
+    in
+    let new_cases = EnumConstructor.Map.of_list new_cases in
+    Expr.ematch ~name:option_instance.name ~e:new_e1 ~cases:new_cases
+      (Mark.get e)
+  | EInj { name; e = e1; cons } when EnumName.equal name Expr.option_enum ->
+    let option_instance =
+      TypMap.find
+        (Mark.remove (Expr.maybe_ty (Mark.get e1)))
+        monomorphized_instances.options
+    in
+    let new_e1 = monomorphize_expr monomorphized_instances e1 in
+    let new_cons =
+      match Mark.remove (EnumConstructor.get_info cons) with
+      | "ESome" -> option_instance.some_cons
+      | "ENone" -> option_instance.none_cons
+      | __ -> failwith "should not happen"
+    in
+    Expr.einj ~name:option_instance.name ~e:new_e1 ~cons:new_cons (Mark.get e)
+  | _ -> Expr.map ~f:(monomorphize_expr monomorphized_instances) e
+
+let program (prg : typed program) :
+    untyped program * Scopelang.Dependency.TVertex.t list =
+  let monomorphized_instances = collect_monomorphized_instances prg in
   (* First we augment the [decl_ctx] with the option instances *)
   let prg =
     {
@@ -179,5 +276,48 @@ let program (prg : typed program) : typed program =
         };
     }
   in
-  (* TODO replace types in exprs *)
-  prg
+  (* Then we replace all hardcoded types and expressions with the monomorphized
+     instances *)
+  let prg =
+    {
+      prg with
+      decl_ctx =
+        {
+          prg.decl_ctx with
+          ctx_enums =
+            EnumName.Map.map
+              (EnumConstructor.Map.map
+                 (monomorphize_typ monomorphized_instances))
+              prg.decl_ctx.ctx_enums;
+          ctx_structs =
+            StructName.Map.map
+              (StructField.Map.map (monomorphize_typ monomorphized_instances))
+              prg.decl_ctx.ctx_structs;
+        };
+    }
+  in
+  let prg =
+    Bindlib.unbox
+    @@ Bindlib.box_apply
+         (fun code_items -> { prg with code_items })
+         (Scope.map
+            ~f:(fun code_item ->
+              match code_item with
+              | Topdef (name, typ, e) -> Bindlib.box (Topdef (name, typ, e))
+              | ScopeDef (name, body) ->
+                let s_var, scope_body = Bindlib.unbind body.scope_body_expr in
+                Bindlib.box_apply
+                  (fun scope_body_expr ->
+                    ScopeDef (name, { body with scope_body_expr }))
+                  (Bindlib.bind_var s_var
+                     (Scope.map_exprs_in_lets ~varf:Fun.id
+                        ~transform_types:
+                          (monomorphize_typ monomorphized_instances)
+                        ~f:(monomorphize_expr monomorphized_instances)
+                        scope_body)))
+            ~varf:Fun.id prg.code_items)
+  in
+  let prg = Program.untype prg in
+  ( prg,
+    Scopelang.Dependency.check_type_cycles prg.decl_ctx.ctx_structs
+      prg.decl_ctx.ctx_enums )

compiler/lcalc/monomorphize.mli

@@ -17,9 +17,11 @@
 open Shared_ast
 open Ast
 
-val program : typed program -> typed program
+val program :
+  typed program -> untyped 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.}
     {- creates monomorphized instances of TOption for every occurence of the type.}}
+    It also returns the new type ordering for the program.
 *)

compiler/shared_ast/scope.ml

@@ -34,12 +34,12 @@ let rec fold_right_lets ~f ~init scope_body_expr =
     f scope_let var next_result
 
 let map_exprs_in_lets :
-    ?reset_types:bool ->
+    ?transform_types:(typ -> typ) ->
     f:('expr1 -> 'expr2 boxed) ->
     varf:('expr1 Var.t -> 'expr2 Var.t) ->
     'expr1 scope_body_expr ->
     'expr2 scope_body_expr Bindlib.box =
- fun ?(reset_types = false) ~f ~varf scope_body_expr ->
+ fun ?(transform_types = Fun.id) ~f ~varf scope_body_expr ->
   fold_right_lets
     ~f:(fun scope_let var_next acc ->
       Bindlib.box_apply2
@@ -49,9 +49,7 @@ let map_exprs_in_lets :
               scope_let with
               scope_let_next;
               scope_let_expr;
-              scope_let_typ =
-                (if reset_types then Mark.copy scope_let.scope_let_typ TAny
-                 else scope_let.scope_let_typ);
+              scope_let_typ = transform_types scope_let.scope_let_typ;
             })
         (Bindlib.bind_var (varf var_next) acc)
         (Expr.Box.lift (f scope_let.scope_let_expr)))

compiler/shared_ast/scope.mli

@@ -44,7 +44,7 @@ val fold_right_lets :
     scope lets to be examined (which are before in the program order). *)
 
 val map_exprs_in_lets :
-  ?reset_types:bool ->
+  ?transform_types:(typ -> typ) ->
   f:('expr1 -> 'expr2 boxed) ->
   varf:('expr1 Var.t -> 'expr2 Var.t) ->
   'expr1 scope_body_expr ->

compiler/shared_ast/typing.ml

@@ -292,8 +292,6 @@ let polymorphic_op_type (op : Operator.polymorphic A.operator Mark.pos) :
   let it = lazy (UnionFind.make (TLit TInt, pos)) in
   let cet = lazy (UnionFind.make (TClosureEnv, pos)) in
   let array a = lazy (UnionFind.make (TArray (Lazy.force a), pos)) in
-  let option a = lazy (UnionFind.make (TOption (Lazy.force a), pos)) in
-  let _default a = lazy (UnionFind.make (TDefault (Lazy.force a), pos)) in
   let ( @-> ) x y =
     lazy (UnionFind.make (TArrow (List.map Lazy.force x, Lazy.force y), pos))
   in
@@ -309,8 +307,11 @@ let polymorphic_op_type (op : Operator.polymorphic A.operator Mark.pos) :
     | Log _ -> [any] @-> any
     | Length -> [array any] @-> it
     | HandleDefault -> [array ([ut] @-> any); [ut] @-> bt; [ut] @-> any] @-> any
-    | HandleDefaultOpt ->
-      [array (option any); [ut] @-> bt; [ut] @-> option any] @-> option any
+    (* The [any] in this type definition should be [option any] but when
+       retyping after option monomorphization it would not work, so we let
+       unification instantiate a monomorphized option type at each new operator
+       instead. *)
+    | HandleDefaultOpt -> [array any; [ut] @-> bt; [ut] @-> any] @-> any
     | ToClosureEnv -> [any] @-> cet
     | FromClosureEnv -> [cet] @-> any
   in