Add more printing facilities and fix minor bugs

6 files changed, 120 insertions, 23 deletions

diff --git a/src/PrintPure.ml b/src/PrintPure.ml
index 54c2bedf..5e9e110a 100644
--- a/src/PrintPure.ml
+++ b/src/PrintPure.ml
@@ -68,16 +68,29 @@ let integer_type_to_string = Print.Types.integer_type_to_string
 
 let scalar_value_to_string = Print.Values.scalar_value_to_string
 
+let mk_type_formatter (type_defs : T.type_def TypeDefId.Map.t)
+    (type_params : type_var list) : type_formatter =
+  let type_var_id_to_string vid =
+    let var = T.TypeVarId.nth type_params vid in
+    type_var_to_string var
+  in
+  let type_def_id_to_string def_id =
+    let def = T.TypeDefId.Map.find def_id type_defs in
+    name_to_string def.name
+  in
+  { type_var_id_to_string; type_def_id_to_string }
+
 (* TODO: there is a bit of duplication with Print.fun_def_to_ast_formatter.
 
    TODO: use the pure defs as inputs? Note that it is a bit annoying for the
    functions (there is a difference between the forward/backward functions...)
    while we only need those definitions to lookup proper names for the def ids.
 *)
-let fun_def_to_ast_formatter (type_defs : T.type_def TypeDefId.Map.t)
-    (fun_defs : A.fun_def FunDefId.Map.t) (fdef : A.fun_def) : ast_formatter =
+let mk_ast_formatter (type_defs : T.type_def TypeDefId.Map.t)
+    (fun_defs : A.fun_def FunDefId.Map.t) (type_params : type_var list) :
+    ast_formatter =
   let type_var_id_to_string vid =
-    let var = T.TypeVarId.nth fdef.signature.type_params vid in
+    let var = T.TypeVarId.nth type_params vid in
     type_var_to_string var
   in
   let type_def_id_to_string def_id =
@@ -117,7 +130,9 @@ let type_id_to_string (fmt : type_formatter) (id : T.type_id) : string =
   | T.Tuple -> ""
   | T.Assumed aty -> (
       match aty with
-      | Box -> (* Boxes should have been eliminated *) failwith "Unreachable")
+      | Box ->
+          (* Boxes should have been eliminated *)
+          failwith "Unreachable: boxes should have been eliminated")
 
 let rec ty_to_string (fmt : type_formatter) (ty : ty) : string =
   match ty with
@@ -144,15 +159,8 @@ let variant_to_string fmt (v : variant) : string =
   ^ String.concat ", " (List.map (field_to_string fmt) v.fields)
   ^ ")"
 
-let type_def_to_string (type_def_id_to_string : TypeDefId.id -> string)
-    (def : type_def) : string =
+let type_def_to_string (fmt : type_formatter) (def : type_def) : string =
   let types = def.type_params in
-  let type_var_id_to_string id =
-    match List.find_opt (fun tv -> tv.T.index = id) types with
-    | Some tv -> type_var_to_string tv
-    | None -> failwith "Unreachable"
-  in
-  let fmt = { type_var_id_to_string; type_def_id_to_string } in
   let name = name_to_string def.name in
   let params =
     if types = [] then ""
@@ -297,6 +305,11 @@ let regular_fun_id_to_string (fmt : ast_formatter) (fun_id : A.fun_id) : string
       | A.BoxDerefMut -> "core::ops::deref::DerefMut::deref_mut"
       | A.BoxFree -> "alloc::alloc::box_free")
 
+let fun_suffix (rg_id : T.RegionGroupId.id option) : string =
+  match rg_id with
+  | None -> ""
+  | Some rg_id -> "@" ^ T.RegionGroupId.to_string rg_id
+
 let unop_to_string (unop : unop) : string =
   match unop with Not -> "¬" | Neg _ -> "-"
 
@@ -304,11 +317,9 @@ let binop_to_string = Print.CfimAst.binop_to_string
 
 let fun_id_to_string (fmt : ast_formatter) (fun_id : fun_id) : string =
   match fun_id with
-  | Regular (fun_id, rg_id) -> (
+  | Regular (fun_id, rg_id) ->
       let f = regular_fun_id_to_string fmt fun_id in
-      match rg_id with
-      | None -> f
-      | Some rg_id -> f ^ "@" ^ T.RegionGroupId.to_string rg_id)
+      f ^ fun_suffix rg_id
   | Unop unop -> unop_to_string unop
   | Binop (binop, int_ty) ->
       binop_to_string binop ^ "<" ^ integer_type_to_string int_ty ^ ">"
@@ -416,3 +427,14 @@ and switch_to_string (fmt : ast_formatter) (indent : string)
       in
       let branches = List.map branch_to_string branches in
       indent ^ "match " ^ scrut ^ " with\n" ^ String.concat "\n" branches
+
+let fun_def_to_string (fmt : ast_formatter) (def : fun_def) : string =
+  let type_fmt = ast_to_type_formatter fmt in
+  let name = name_to_string def.basename ^ fun_suffix def.back_id in
+  let signature = fun_sig_to_string fmt def.signature in
+  let inputs = List.map (var_to_string type_fmt) def.inputs in
+  let inputs =
+    if inputs = [] then "" else "  fun " ^ String.concat " " inputs ^ " ->\n"
+  in
+  let body = expression_to_string fmt "  " "  " def.body in
+  "let " ^ name ^ " : " ^ signature ^ " =\n" ^ inputs ^ body
diff --git a/src/Pure.ml b/src/Pure.ml
index b7a61ab0..0016aedb 100644
--- a/src/Pure.ml
+++ b/src/Pure.ml
@@ -227,6 +227,7 @@ type inst_fun_sig = { inputs : ty list; outputs : ty list }
 
 type fun_def = {
   def_id : FunDefId.id;
+  back_id : T.RegionGroupId.id option;
   basename : name;
       (** The "base" name of the function.
   
diff --git a/src/SymbolicToPure.ml b/src/SymbolicToPure.ml
index e246e804..63c40f20 100644
--- a/src/SymbolicToPure.ml
+++ b/src/SymbolicToPure.ml
@@ -213,6 +213,29 @@ type bs_ctx = {
 (** Body synthesis context *)
 
 (* TODO: move *)
+let type_def_to_string (ctx : bs_ctx) (def : type_def) : string =
+  let type_params = def.type_params in
+  let type_defs = ctx.type_context.cfim_type_defs in
+  let fmt = PrintPure.mk_type_formatter type_defs type_params in
+  PrintPure.type_def_to_string fmt def
+
+(* TODO: move *)
+let fun_sig_to_string (ctx : bs_ctx) (sg : fun_sig) : string =
+  let type_params = sg.type_params in
+  let type_defs = ctx.type_context.cfim_type_defs in
+  let fun_defs = ctx.fun_context.cfim_fun_defs in
+  let fmt = PrintPure.mk_ast_formatter type_defs fun_defs type_params in
+  PrintPure.fun_sig_to_string fmt sg
+
+(* TODO: move *)
+let fun_def_to_string (ctx : bs_ctx) (def : Pure.fun_def) : string =
+  let type_params = def.signature.type_params in
+  let type_defs = ctx.type_context.cfim_type_defs in
+  let fun_defs = ctx.fun_context.cfim_fun_defs in
+  let fmt = PrintPure.mk_ast_formatter type_defs fun_defs type_params in
+  PrintPure.fun_def_to_string fmt def
+
+(* TODO: move *)
 let bs_ctx_to_value_formatter (ctx : bs_ctx) : Print.CfimAst.ast_formatter =
   Print.CfimAst.fun_def_to_ast_formatter ctx.type_context.cfim_type_defs
     ctx.fun_context.cfim_fun_defs ctx.fun_def
@@ -340,14 +363,21 @@ let translate_type_def (def : T.type_def) : type_def =
 let rec translate_fwd_ty (types_infos : TA.type_infos) (ty : 'r T.ty) : ty =
   let translate = translate_fwd_ty types_infos in
   match ty with
-  | T.Adt (type_id, regions, tys) ->
+  | T.Adt (type_id, regions, tys) -> (
       (* Can't translate types with regions for now *)
       assert (regions = []);
       (* No general parametricity for now *)
       assert (not (List.exists (TypesUtils.ty_has_borrows types_infos) tys));
       (* Translate the type parameters *)
       let tys = List.map translate tys in
-      Adt (type_id, tys)
+      (* Eliminate boxes *)
+      match type_id with
+      | T.AdtId _ | Tuple -> Adt (type_id, tys)
+      | T.Assumed T.Box -> (
+          match tys with
+          | [ bty ] -> bty
+          | _ ->
+              failwith "Unreachable: boxes receive exactly one type parameter"))
   | TypeVar vid -> TypeVar vid
   | Bool -> Bool
   | Char -> Char
@@ -381,13 +411,21 @@ let rec translate_back_ty (types_infos : TA.type_infos)
   match ty with
   | T.Adt (type_id, _, tys) -> (
       match type_id with
-      | T.AdtId _ | Assumed _ ->
+      | T.AdtId _ ->
           (* Don't accept ADTs (which are not tuples) with borrows for now *)
           assert (not (TypesUtils.ty_has_borrows types_infos ty));
           if inside_mut then
             let tys_t = List.filter_map translate tys in
             Some (Adt (type_id, tys_t))
           else None
+      | Assumed T.Box -> (
+          (* Don't accept ADTs (which are not tuples) with borrows for now *)
+          assert (not (TypesUtils.ty_has_borrows types_infos ty));
+          (* Eliminate the box *)
+          match tys with
+          | [ bty ] -> translate bty
+          | _ ->
+              failwith "Unreachable: boxes receive exactly one type parameter")
       | T.Tuple -> (
           (* Tuples can contain borrows (which we eliminated) *)
           let tys_t = List.filter_map translate tys in
@@ -1214,7 +1252,14 @@ let translate_fun_def (ctx : bs_ctx) (body : S.expression) : fun_def =
     List.for_all
       (fun (var, ty) -> (var : var).ty = ty)
       (List.combine inputs signature.inputs));
-  { def_id; basename; signature; inputs; body }
+  let def = { def_id; back_id = bid; basename; signature; inputs; body } in
+  (* Debugging *)
+  log#ldebug
+    (lazy
+      ("SymbolicToPure.translate_fun_def: translated:\n"
+     ^ fun_def_to_string ctx def));
+  (* return *)
+  def
 
 let translate_type_defs (type_defs : T.type_def list) : type_def TypeDefId.Map.t
     =
diff --git a/src/Translate.ml b/src/Translate.ml
index ce9f085c..58847d42 100644
--- a/src/Translate.ml
+++ b/src/Translate.ml
@@ -9,6 +9,28 @@ module SA = SymbolicAst
 (** The local logger *)
 let log = L.translate_log
 
+type trans_ctx = { type_context : C.type_context; fun_context : C.fun_context }
+
+let type_def_to_string (ctx : trans_ctx) (def : Pure.type_def) : string =
+  let type_params = def.type_params in
+  let type_defs = ctx.type_context.type_defs in
+  let fmt = PrintPure.mk_type_formatter type_defs type_params in
+  PrintPure.type_def_to_string fmt def
+
+let fun_sig_to_string (ctx : trans_ctx) (sg : Pure.fun_sig) : string =
+  let type_params = sg.type_params in
+  let type_defs = ctx.type_context.type_defs in
+  let fun_defs = ctx.fun_context.fun_defs in
+  let fmt = PrintPure.mk_ast_formatter type_defs fun_defs type_params in
+  PrintPure.fun_sig_to_string fmt sg
+
+let fun_def_to_string (ctx : trans_ctx) (def : Pure.fun_def) : string =
+  let type_params = def.signature.type_params in
+  let type_defs = ctx.type_context.type_defs in
+  let fun_defs = ctx.fun_context.fun_defs in
+  let fmt = PrintPure.mk_ast_formatter type_defs fun_defs type_params in
+  PrintPure.fun_def_to_string fmt def
+
 type symbolic_fun_translation = V.symbolic_value list * SA.expression
 (** The result of running the symbolic interpreter on a function:
     - the list of symbolic values used for the input values
diff --git a/src/TypesUtils.ml b/src/TypesUtils.ml
index 1eac5cee..e6e59a51 100644
--- a/src/TypesUtils.ml
+++ b/src/TypesUtils.ml
@@ -12,10 +12,15 @@ let type_def_get_fields (def : type_def) (opt_variant_id : VariantId.id option)
   | Enum variants, Some variant_id -> (VariantId.nth variants variant_id).fields
   | Struct fields, None -> fields
   | _ ->
+      let opt_variant_id =
+        match opt_variant_id with None -> "None" | Some _ -> "Some"
+      in
       raise
         (Invalid_argument
-           "The variant id should be [Some] if and only if the definition is \
-            an enumeration")
+           ("The variant id should be [Some] if and only if the definition is \
+             an enumeration:\n\
+             - def: " ^ show_type_def def ^ "\n- opt_variant_id: "
+          ^ opt_variant_id))
 
 (** Return [true] if a [ty] is actually `unit` *)
 let ty_is_unit (ty : 'r ty) : bool =
diff --git a/src/main.ml b/src/main.ml
index 46b54d55..a4de3d7a 100644
--- a/src/main.ml
+++ b/src/main.ml
@@ -59,7 +59,9 @@ let () =
   (* Load the module *)
   let json = Yojson.Basic.from_file !filename in
   match cfim_module_of_json json with
-  | Error s -> main_log#error "error: %s\n" s
+  | Error s ->
+      main_log#error "error: %s\n" s;
+      exit 1
   | Ok m ->
       (* Print the module *)
       main_log#ldebug (lazy ("\n" ^ Print.Module.module_to_string m ^ "\n"));