Update TranslateCore and factor out some definitions in PrintPure

1 files changed, 30 insertions, 16 deletions

diff --git a/compiler/TranslateCore.ml b/compiler/TranslateCore.ml
index ba5e237b..1b1572d6 100644
--- a/compiler/TranslateCore.ml
+++ b/compiler/TranslateCore.ml
@@ -18,25 +18,37 @@ type fun_context = {
 }
 [@@deriving show]
 
+type trait_decls_context = C.trait_decls_context [@@deriving show]
+type trait_impls_context = C.trait_impls_context [@@deriving show]
 type global_context = C.global_context [@@deriving show]
 
 type trans_ctx = {
   type_context : type_context;
   fun_context : fun_context;
   global_context : global_context;
+  trait_decls_context : trait_decls_context;
+  trait_impls_context : trait_impls_context;
 }
 
 type fun_and_loops = Pure.fun_decl * Pure.fun_decl list
 type pure_fun_translation_no_loops = Pure.fun_decl * Pure.fun_decl list
 type pure_fun_translation = fun_and_loops * fun_and_loops list
 
-let type_decl_to_string (ctx : trans_ctx) (def : Pure.type_decl) : string =
-  let type_params = def.type_params in
-  let cg_params = def.const_generic_params in
+let trans_ctx_to_type_formatter (ctx : trans_ctx)
+    (type_params : Pure.type_var list)
+    (const_generic_params : Pure.const_generic_var list) :
+    PrintPure.type_formatter =
   let type_decls = ctx.type_context.type_decls in
   let global_decls = ctx.global_context.global_decls in
+  let trait_decls = ctx.trait_decls_context.trait_decls in
+  let trait_impls = ctx.trait_impls_context.trait_impls in
+  PrintPure.mk_type_formatter type_decls global_decls trait_decls trait_impls
+    type_params const_generic_params
+
+let type_decl_to_string (ctx : trans_ctx) (def : Pure.type_decl) : string =
+  let generics = def.generics in
   let fmt =
-    PrintPure.mk_type_formatter type_decls global_decls type_params cg_params
+    trans_ctx_to_type_formatter ctx generics.types generics.const_generics
   in
   PrintPure.type_decl_to_string fmt def
 
@@ -44,27 +56,29 @@ let type_id_to_string (ctx : trans_ctx) (id : Pure.TypeDeclId.id) : string =
   Print.fun_name_to_string
     (Pure.TypeDeclId.Map.find id ctx.type_context.type_decls).name
 
-let fun_sig_to_string (ctx : trans_ctx) (sg : Pure.fun_sig) : string =
-  let type_params = sg.type_params in
-  let cg_params = sg.const_generic_params in
+let trans_ctx_to_ast_formatter (ctx : trans_ctx)
+    (type_params : Pure.type_var list)
+    (const_generic_params : Pure.const_generic_var list) :
+    PrintPure.ast_formatter =
   let type_decls = ctx.type_context.type_decls in
   let fun_decls = ctx.fun_context.fun_decls in
   let global_decls = ctx.global_context.global_decls in
+  let trait_decls = ctx.trait_decls_context.trait_decls in
+  let trait_impls = ctx.trait_impls_context.trait_impls in
+  PrintPure.mk_ast_formatter type_decls fun_decls global_decls trait_decls
+    trait_impls type_params const_generic_params
+
+let fun_sig_to_string (ctx : trans_ctx) (sg : Pure.fun_sig) : string =
+  let generics = sg.generics in
   let fmt =
-    PrintPure.mk_ast_formatter type_decls fun_decls global_decls type_params
-      cg_params
+    trans_ctx_to_ast_formatter ctx generics.types generics.const_generics
   in
   PrintPure.fun_sig_to_string fmt sg
 
 let fun_decl_to_string (ctx : trans_ctx) (def : Pure.fun_decl) : string =
-  let type_params = def.signature.type_params in
-  let cg_params = def.signature.const_generic_params in
-  let type_decls = ctx.type_context.type_decls in
-  let fun_decls = ctx.fun_context.fun_decls in
-  let global_decls = ctx.global_context.global_decls in
+  let generics = def.signature.generics in
   let fmt =
-    PrintPure.mk_ast_formatter type_decls fun_decls global_decls type_params
-      cg_params
+    trans_ctx_to_ast_formatter ctx generics.types generics.const_generics
   in
   PrintPure.fun_decl_to_string fmt def