Reorganize branching symbolic expansions to prepare for the join operation

12 files changed, 184 insertions, 104 deletions

diff --git a/compiler/ExtractToFStar.ml b/compiler/ExtractToFStar.ml
index 7267101c..6d680984 100644
--- a/compiler/ExtractToFStar.ml
+++ b/compiler/ExtractToFStar.ml
@@ -105,7 +105,9 @@ let fstar_keywords =
       "with";
       "assert";
       "assert_norm";
+      "assume";
       "Type0";
+      "Type";
       "unit";
       "not";
       "scalar_cast";
@@ -144,7 +146,7 @@ let fstar_assumed_functions :
     (VecIndexMut, rg0, "vec_index_mut_back");
   ]
 
-let fstar_names_map_init =
+let fstar_names_map_init : names_map_init =
   {
     keywords = fstar_keywords;
     assumed_adts = fstar_assumed_adts;
@@ -414,6 +416,7 @@ let mk_formatter (ctx : trans_ctx) (crate_name : string)
     char_name = "char";
     int_name;
     str_name = "string";
+    assert_name = "massert";
     field_name;
     variant_name;
     struct_constructor;
@@ -968,12 +971,17 @@ and extract_function_call (ctx : extraction_ctx) (fmt : F.formatter)
       ctx.fmt.extract_binop
         (extract_texpression ctx fmt)
         fmt inside binop int_ty arg0 arg1
-  | Regular (fun_id, rg_id), _ ->
+  | Regular (_, _), _ | Assert, [ _ ] ->
       if inside then F.pp_print_string fmt "(";
       (* Open a box for the function call *)
       F.pp_open_hovbox fmt ctx.indent_incr;
       (* Print the function name *)
-      let fun_name = ctx_get_function fun_id rg_id ctx in
+      let fun_name =
+        match fid with
+        | Regular (fun_id, rg_id) -> ctx_get_function fun_id rg_id ctx
+        | Assert -> ctx.fmt.assert_name
+        | _ -> raise (Failure "Unreachable")
+      in
       F.pp_print_string fmt fun_name;
       (* Print the type parameters *)
       List.iter
@@ -991,7 +999,7 @@ and extract_function_call (ctx : extraction_ctx) (fmt : F.formatter)
       F.pp_close_box fmt ();
       (* Return *)
       if inside then F.pp_print_string fmt ")"
-  | _ ->
+  | (Unop _ | Binop _ | Assert), _ ->
       raise
         (Failure
            ("Unreachable:\n" ^ "Function: " ^ show_fun_id fid
diff --git a/compiler/InterpreterExpansion.ml b/compiler/InterpreterExpansion.ml
index a0fe5ecb..fe9ccbf4 100644
--- a/compiler/InterpreterExpansion.ml
+++ b/compiler/InterpreterExpansion.ml
@@ -192,12 +192,6 @@ let replace_symbolic_values (at_most_once : bool)
   (* Return *)
   ctx
 
-(** Apply a symbolic expansion to a context, by replacing the original
-    symbolic value with its expanded value. Is valid only if the expansion
-    is not a borrow (i.e., an adt...).
-    
-    This function does update the synthesis.
-*)
 let apply_symbolic_expansion_non_borrow (config : C.config)
     (original_sv : V.symbolic_value) (expansion : V.symbolic_expansion)
     (ctx : C.eval_ctx) : C.eval_ctx =
@@ -459,25 +453,43 @@ let expand_symbolic_value_borrow (config : C.config)
     branches. Note that the expansion is optional for every branch (this is
     used for integer expansion: see {!expand_symbolic_int}).
     
-    [see_cf_l]: list of pairs (optional symbolic expansion, continuation)
+    [see_cf_l]: list of pairs (optional symbolic expansion, continuation).
+    We use [None] for the symbolic expansion for the [_] (default) case of the
+    integer expansions.
+    The continuation are used to execute the content of the branches, but not
+    what comes after.
+
+    [cf_after_join]: this continuation is called *after* the branches have been evaluated.
+    We need this continuation separately (i.e., we can't compose it with the
+    continuations in [see_cf_l]) because we perform a join *before* calling it.
 *)
 let apply_branching_symbolic_expansions_non_borrow (config : C.config)
     (sv : V.symbolic_value) (sv_place : SA.mplace option)
-    (see_cf_l : (V.symbolic_expansion option * m_fun) list) : m_fun =
+    (see_cf_l : (V.symbolic_expansion option * st_cm_fun) list)
+    (cf_after_join : st_m_fun) : m_fun =
  fun ctx ->
   assert (see_cf_l <> []);
-  (* Apply the symbolic expansion in in the context and call the continuation *)
+  (* Apply the symbolic expansion in the context and call the continuation *)
   let resl =
     List.map
-      (fun (see_opt, cf) ->
+      (fun (see_opt, cf_br) ->
+        (* Remember the initial context for printing purposes *)
+        let ctx0 = ctx in
         (* Expansion *)
         let ctx =
           match see_opt with
           | None -> ctx
           | Some see -> apply_symbolic_expansion_non_borrow config sv see ctx
         in
+        (* Debug *)
+        log#ldebug
+          (lazy
+            ("apply_branching_symbolic_expansions_non_borrow: "
+            ^ symbolic_value_to_string ctx0 sv
+            ^ "\n\n- original context:\n" ^ eval_ctx_to_string ctx0
+            ^ "\n\n- new context:\n" ^ eval_ctx_to_string ctx ^ "\n"));
         (* Continuation *)
-        cf ctx)
+        cf_br cf_after_join ctx)
       see_cf_l
   in
   (* Collect the result: either we computed no subterm, or we computed all
@@ -495,7 +507,8 @@ let apply_branching_symbolic_expansions_non_borrow (config : C.config)
   S.synthesize_symbolic_expansion sv sv_place seel subterms
 
 let expand_symbolic_bool (config : C.config) (sv : V.symbolic_value)
-    (sv_place : SA.mplace option) (cf_true : m_fun) (cf_false : m_fun) : m_fun =
+    (sv_place : SA.mplace option) (cf_true : st_cm_fun) (cf_false : st_cm_fun)
+    (cf_after_join : st_m_fun) : m_fun =
  fun ctx ->
   (* Compute the expanded value *)
   let original_sv = sv in
@@ -508,7 +521,7 @@ let expand_symbolic_bool (config : C.config) (sv : V.symbolic_value)
   let seel = [ (Some see_true, cf_true); (Some see_false, cf_false) ] in
   (* Apply the symbolic expansion (this also outputs the updated symbolic AST) *)
   apply_branching_symbolic_expansions_non_borrow config original_sv
-    original_sv_place seel ctx
+    original_sv_place seel cf_after_join ctx
 
 let expand_symbolic_value_no_branching (config : C.config)
     (sv : V.symbolic_value) (sv_place : SA.mplace option) : cm_fun =
@@ -572,54 +585,38 @@ let expand_symbolic_value_no_branching (config : C.config)
   cc cf ctx
 
 let expand_symbolic_adt (config : C.config) (sv : V.symbolic_value)
-    (sv_place : SA.mplace option) : cm_fun =
- fun cf ctx ->
+    (sv_place : SA.mplace option) (cf_branches : st_cm_fun)
+    (cf_after_join : st_m_fun) : m_fun =
+ fun ctx ->
   (* Debug *)
   log#ldebug (lazy ("expand_symbolic_adt:" ^ symbolic_value_to_string ctx sv));
-  (* Remember the initial context for printing purposes *)
-  let ctx0 = ctx in
   (* Compute the expanded value - note that when doing so, we may introduce
    * fresh symbolic values in the context (which thus gets updated) *)
   let original_sv = sv in
   let original_sv_place = sv_place in
   let rty = original_sv.V.sv_ty in
-  let cc : cm_fun =
-   fun cf ctx ->
-    match rty with
-    (* ADTs *)
-    | T.Adt (adt_id, regions, types) ->
-        let allow_branching = true in
-        (* Compute the expanded value *)
-        let seel =
-          compute_expanded_symbolic_adt_value allow_branching sv.sv_kind adt_id
-            regions types ctx
-        in
-        (* Apply *)
-        let seel = List.map (fun see -> (Some see, cf)) seel in
-        apply_branching_symbolic_expansions_non_borrow config original_sv
-          original_sv_place seel ctx
-    | _ ->
-        raise
-          (Failure ("expand_symbolic_adt: unexpected type: " ^ T.show_rty rty))
-  in
-  (* Debug *)
-  let cc =
-    comp_unit cc (fun ctx ->
-        log#ldebug
-          (lazy
-            ("expand_symbolic_adt: "
-            ^ symbolic_value_to_string ctx0 sv
-            ^ "\n\n- original context:\n" ^ eval_ctx_to_string ctx0
-            ^ "\n\n- new context:\n" ^ eval_ctx_to_string ctx ^ "\n"));
-        (* Sanity check: the symbolic value has disappeared *)
-        assert (not (symbolic_value_id_in_ctx original_sv.V.sv_id ctx)))
-  in
-  (* Continue *)
-  cc cf ctx
+  (* Execute *)
+  match rty with
+  (* ADTs *)
+  | T.Adt (adt_id, regions, types) ->
+      let allow_branching = true in
+      (* Compute the expanded value *)
+      let seel =
+        compute_expanded_symbolic_adt_value allow_branching sv.sv_kind adt_id
+          regions types ctx
+      in
+      (* Apply *)
+      let seel = List.map (fun see -> (Some see, cf_branches)) seel in
+      apply_branching_symbolic_expansions_non_borrow config original_sv
+        original_sv_place seel cf_after_join ctx
+  | _ ->
+      raise
+        (Failure ("expand_symbolic_adt: unexpected type: " ^ T.show_rty rty))
 
 let expand_symbolic_int (config : C.config) (sv : V.symbolic_value)
     (sv_place : SA.mplace option) (int_type : T.integer_type)
-    (tgts : (V.scalar_value * m_fun) list) (otherwise : m_fun) : m_fun =
+    (tgts : (V.scalar_value * st_cm_fun) list) (otherwise : st_cm_fun)
+    (cf_after_join : st_m_fun) : m_fun =
   (* Sanity check *)
   assert (sv.V.sv_ty = T.Integer int_type);
   (* For all the branches of the switch, we expand the symbolic value
@@ -631,12 +628,13 @@ let expand_symbolic_int (config : C.config) (sv : V.symbolic_value)
    * First, generate the list of pairs:
    * (optional expansion, statement to execute)
    *)
-  let tgts =
+  let seel =
     List.map (fun (v, cf) -> (Some (V.SePrimitive (PV.Scalar v)), cf)) tgts
   in
-  let tgts = List.append tgts [ (None, otherwise) ] in
+  let seel = List.append seel [ (None, otherwise) ] in
   (* Then expand and evaluate - this generates the proper symbolic AST *)
-  apply_branching_symbolic_expansions_non_borrow config sv sv_place tgts
+  apply_branching_symbolic_expansions_non_borrow config sv sv_place seel
+    cf_after_join
 
 (** Expand all the symbolic values which contain borrows.
     Allows us to restrict ourselves to a simpler model for the projectors over
diff --git a/compiler/InterpreterExpansion.mli b/compiler/InterpreterExpansion.mli
index 54f9036f..a75f88fd 100644
--- a/compiler/InterpreterExpansion.mli
+++ b/compiler/InterpreterExpansion.mli
@@ -13,17 +13,60 @@ open InterpreterBorrows
 
 type proj_kind = LoanProj | BorrowProj
 
-(** Expand a symbolic boolean *)
+(** Apply a symbolic expansion to a context, by replacing the original
+    symbolic value with its expanded value. Is valid only if the expansion
+    is *not a borrow* (i.e., an adt...).
+
+    This function does *not* update the synthesis.
+*)
+val apply_symbolic_expansion_non_borrow :
+  C.config ->
+  V.symbolic_value ->
+  V.symbolic_expansion ->
+  C.eval_ctx ->
+  C.eval_ctx
+
+(** Expand a symhbolic value, without branching *)
+val expand_symbolic_value_no_branching :
+  C.config -> V.symbolic_value -> SA.mplace option -> cm_fun
+
+(** Expand a symbolic enumeration (leads to branching if the enumeration has
+    more than one variant).
+
+    Parameters:
+    - [config]
+    - [sv]
+    - [sv_place]
+    - [cf_branches]: the continuation to evaluate the branches. This continuation
+      typically evaluates a [match] statement *after* we have performed the symbolic
+      expansion (in particular, we can have one continuation for all the branches).
+    - [cf_after_join]: the continuation for *after* the match (we perform a join
+      then call it).
+ *)
+val expand_symbolic_adt :
+  C.config ->
+  V.symbolic_value ->
+  SA.mplace option ->
+  st_cm_fun ->
+  st_m_fun ->
+  m_fun
+
+(** Expand a symbolic boolean.
+
+    Parameters: see {!expand_symbolic_adt}.
+    The two parameters of type [st_cm_fun] correspond to the [cf_branches]
+    parameter (here, there are exactly two branches).
+ *)
 val expand_symbolic_bool :
-  C.config -> V.symbolic_value -> SA.mplace option -> m_fun -> m_fun -> m_fun
+  C.config ->
+  V.symbolic_value ->
+  SA.mplace option ->
+  st_cm_fun ->
+  st_cm_fun ->
+  st_m_fun ->
+  m_fun
 
-(** Symbolic integers are expanded upon evaluating a [switch], when the integer
-    is not an enumeration discriminant.
-    Note that a discriminant is never symbolic: we evaluate discriminant values
-    upon evaluating [eval_discriminant], which always generates a concrete value
-    (because if we call it on a symbolic enumeration, we expand the enumeration
-    *then* evaluate the discriminant). This is how we can spot "regular" switches
-    over integers.
+(** Symbolic integers are expanded upon evaluating a [SwitchInt].
 
     When expanding a boolean upon evaluating an [if ... then ... else ...],
     or an enumeration just before matching over it, we can simply expand the
@@ -47,20 +90,11 @@ val expand_symbolic_int :
   V.symbolic_value ->
   SA.mplace option ->
   T.integer_type ->
-  (V.scalar_value * m_fun) list ->
-  m_fun ->
+  (V.scalar_value * st_cm_fun) list ->
+  st_cm_fun ->
+  st_m_fun ->
   m_fun
 
-(** See {!expand_symbolic_value} *)
-val expand_symbolic_value_no_branching :
-  C.config -> V.symbolic_value -> SA.mplace option -> cm_fun
-
-(** Expand a symbolic enumeration (leads to branching if the enumeration has
-    more than one variant).
- *)
-val expand_symbolic_adt :
-  C.config -> V.symbolic_value -> SA.mplace option -> cm_fun
-
 (** If this mode is activated through the [config], greedily expand the symbolic
     values which need to be expanded. See {!C.config} for more information.
  *)
diff --git a/compiler/InterpreterStatements.ml b/compiler/InterpreterStatements.ml
index 4dd03c1a..15962ee3 100644
--- a/compiler/InterpreterStatements.ml
+++ b/compiler/InterpreterStatements.ml
@@ -185,16 +185,18 @@ let eval_assertion (config : C.config) (assertion : A.assertion) : st_cm_fun =
     | Symbolic sv ->
         assert (config.mode = C.SymbolicMode);
         assert (sv.V.sv_ty = T.Bool);
-        (* Expand the symbolic value and call the proper continuation functions
-         * for the true and false cases - TODO: call an "assert" function instead *)
-        let cf_true : m_fun = fun ctx -> cf Unit ctx in
-        let cf_false : m_fun = fun ctx -> cf Panic ctx in
-        let expand =
-          expand_symbolic_bool config sv
-            (S.mk_opt_place_from_op assertion.cond ctx)
-            cf_true cf_false
+        (* We continue the execution as if the test had succeeded, and thus
+         * perform the symbolic expansion: sv ~~> true.
+         * We will of course synthesize an assertion in the generated code
+         * (see below). *)
+        let ctx =
+          apply_symbolic_expansion_non_borrow config sv
+            (V.SePrimitive (PV.Bool true)) ctx
         in
-        expand ctx
+        (* Continue *)
+        let expr = cf Unit ctx in
+        (* Add the synthesized assertion *)
+        S.synthesize_assertion v expr
     | _ ->
         raise
           (Failure ("Expected a boolean, got: " ^ typed_value_to_string ctx v))
@@ -949,11 +951,11 @@ and eval_switch (config : C.config) (switch : A.switch) : st_cm_fun =
           | V.Symbolic sv ->
               (* Expand the symbolic boolean, and continue by evaluating
                * the branches *)
-              let cf_true : m_fun = eval_statement config st1 cf in
-              let cf_false : m_fun = eval_statement config st2 cf in
+              let cf_true : st_cm_fun = eval_statement config st1 in
+              let cf_false : st_cm_fun = eval_statement config st2 in
               expand_symbolic_bool config sv
                 (S.mk_opt_place_from_op op ctx)
-                cf_true cf_false ctx
+                cf_true cf_false cf ctx
           | _ -> raise (Failure "Inconsistent state")
         in
         (* Compose *)
@@ -982,7 +984,7 @@ and eval_switch (config : C.config) (switch : A.switch) : st_cm_fun =
                * proper branches *)
               let stgts =
                 List.map
-                  (fun (cv, tgt_st) -> (cv, eval_statement config tgt_st cf))
+                  (fun (cv, tgt_st) -> (cv, eval_statement config tgt_st))
                   stgts
               in
               (* Several branches may be grouped together: every branch is described
@@ -996,11 +998,11 @@ and eval_switch (config : C.config) (switch : A.switch) : st_cm_fun =
                      stgts)
               in
               (* Translate the otherwise branch *)
-              let otherwise = eval_statement config otherwise cf in
+              let otherwise = eval_statement config otherwise in
               (* Expand and continue *)
               expand_symbolic_int config sv
                 (S.mk_opt_place_from_op op ctx)
-                int_ty stgts otherwise ctx
+                int_ty stgts otherwise cf ctx
           | _ -> raise (Failure "Inconsistent state")
         in
         (* Compose *)
@@ -1034,7 +1036,7 @@ and eval_switch (config : C.config) (switch : A.switch) : st_cm_fun =
               in
               (* Re-evaluate the switch - the value is not symbolic anymore,
                  which means we will go to the other branch *)
-              comp cf_expand (eval_switch config switch) cf ctx
+              cf_expand (eval_switch config switch) cf ctx
           | _ -> raise (Failure "Inconsistent state")
         in
         (* Compose *)
diff --git a/compiler/PrintPure.ml b/compiler/PrintPure.ml
index 64b527a2..0d1288d7 100644
--- a/compiler/PrintPure.ml
+++ b/compiler/PrintPure.ml
@@ -428,6 +428,7 @@ let fun_id_to_string (fmt : ast_formatter) (fun_id : fun_id) : string =
   | Unop unop -> unop_to_string unop
   | Binop (binop, int_ty) ->
       binop_to_string binop ^ "<" ^ integer_type_to_string int_ty ^ ">"
+  | Assert -> "assert"
 
 (** [inside]: controls the introduction of parentheses *)
 let rec texpression_to_string (fmt : ast_formatter) (inside : bool)
diff --git a/compiler/Pure.ml b/compiler/Pure.ml
index 62657fc7..f11397e9 100644
--- a/compiler/Pure.ml
+++ b/compiler/Pure.ml
@@ -292,9 +292,11 @@ type fun_id =
           works only if we unfold all the monadic let-bindings, and we
           then replace the content of the occurrences of [Return] to also
           return the state (which is really super ugly).
+          TODO: also add Assert...
        *)
   | Unop of unop
   | Binop of E.binop * integer_type
+  | Assert
 [@@deriving show, ord]
 
 (** An identifier for an ADT constructor *)
@@ -394,8 +396,9 @@ type expression =
       
           TODO: the boolean should be replaced by an enum: sometimes we use
           the error-monad, sometimes we use the state-error monad (and we
-          do this an a per-function basis! For instance, arithmetic functions
-          are always in the error monad).
+          should do this an a per-function basis! For instance, arithmetic
+          functions are always in the error monad, they shouldn't use the
+          state-error monad).
 
           The boolean controls whether the let is monadic or not.
           For instance, in F*:
diff --git a/compiler/PureToExtract.ml b/compiler/PureToExtract.ml
index e38a92db..860949a7 100644
--- a/compiler/PureToExtract.ml
+++ b/compiler/PureToExtract.ml
@@ -51,6 +51,7 @@ type formatter = {
   char_name : string;
   int_name : integer_type -> string;
   str_name : string;
+  assert_name : string;
   field_name : name -> FieldId.id -> string option -> string;
       (** Inputs:
           - type name
@@ -632,11 +633,20 @@ type names_map_init = {
 
 (** Initialize a names map with a proper set of keywords/names coming from the
     target language/prover. *)
-let initialize_names_map (init : names_map_init) : names_map =
+let initialize_names_map (fmt : formatter) (init : names_map_init) : names_map =
+  let int_names = List.map fmt.int_name T.all_int_types in
+  let keywords =
+    List.concat
+      [
+        [ fmt.bool_name; fmt.char_name; fmt.str_name; fmt.assert_name ];
+        int_names;
+        init.keywords;
+      ]
+  in
+  let names_set = StringSet.of_list keywords in
   let name_to_id =
-    StringMap.of_list (List.map (fun x -> (x, UnknownId)) init.keywords)
+    StringMap.of_list (List.map (fun x -> (x, UnknownId)) keywords)
   in
-  let names_set = StringSet.of_list init.keywords in
   (* We fist initialize [id_to_name] as empty, because the id of a keyword is [UnknownId].
    * Also note that we don't need this mapping for keywords: we insert keywords only
    * to check collisions. *)
diff --git a/compiler/PureUtils.ml b/compiler/PureUtils.ml
index 9d364dc7..e292576c 100644
--- a/compiler/PureUtils.ml
+++ b/compiler/PureUtils.ml
@@ -350,6 +350,10 @@ let mk_typed_pattern_from_var (v : var) (mp : mplace option) : typed_pattern =
   let ty = v.ty in
   { value; ty }
 
+let mk_dummy_pattern (ty : ty) : typed_pattern =
+  let value = PatDummy in
+  { value; ty }
+
 let mk_meta (m : meta) (e : texpression) : texpression =
   let ty = e.ty in
   let e = Meta (m, e) in
diff --git a/compiler/SymbolicAst.ml b/compiler/SymbolicAst.ml
index 39709a13..528d8255 100644
--- a/compiler/SymbolicAst.ml
+++ b/compiler/SymbolicAst.ml
@@ -67,6 +67,7 @@ type expression =
   | EndAbstraction of V.abs * expression
   | EvalGlobal of A.GlobalDeclId.id * V.symbolic_value * expression
       (** Evaluate a global to a fresh symbolic value *)
+  | Assertion of V.typed_value * expression  (** An assertion *)
   | Expansion of mplace option * V.symbolic_value * expansion
       (** Expansion of a symbolic value.
     
diff --git a/compiler/SymbolicToPure.ml b/compiler/SymbolicToPure.ml
index 3636d4b8..8329d80e 100644
--- a/compiler/SymbolicToPure.ml
+++ b/compiler/SymbolicToPure.ml
@@ -1084,6 +1084,7 @@ let rec translate_expression (config : config) (e : S.expression) (ctx : bs_ctx)
   | FunCall (call, e) -> translate_function_call config call e ctx
   | EndAbstraction (abs, e) -> translate_end_abstraction config abs e ctx
   | EvalGlobal (gid, sv, e) -> translate_global_eval config gid sv e ctx
+  | Assertion (v, e) -> translate_assertion config v e ctx
   | Expansion (p, sv, exp) -> translate_expansion config p sv exp ctx
   | Meta (meta, e) -> translate_meta config meta e ctx
 
@@ -1480,6 +1481,18 @@ and translate_global_eval (config : config) (gid : A.GlobalDeclId.id)
   let e = translate_expression config e ctx in
   mk_let false (mk_typed_pattern_from_var var None) gval e
 
+and translate_assertion (config : config) (v : V.typed_value) (e : S.expression)
+    (ctx : bs_ctx) : texpression =
+  let next_e = translate_expression config e ctx in
+  let monadic = true in
+  let v = typed_value_to_texpression ctx v in
+  let args = [ v ] in
+  let func = { id = Func Assert; type_args = [] } in
+  let func_ty = mk_arrow Bool mk_unit_ty in
+  let func = { e = Qualif func; ty = func_ty } in
+  let assertion = mk_apps func args in
+  mk_let monadic (mk_dummy_pattern mk_unit_ty) assertion next_e
+
 and translate_expansion (config : config) (p : S.mplace option)
     (sv : V.symbolic_value) (exp : S.expansion) (ctx : bs_ctx) : texpression =
   (* Translate the scrutinee *)
diff --git a/compiler/SynthesizeSymbolic.ml b/compiler/SynthesizeSymbolic.ml
index 504fcdb5..c74a831e 100644
--- a/compiler/SynthesizeSymbolic.ml
+++ b/compiler/SynthesizeSymbolic.ml
@@ -155,3 +155,6 @@ let synthesize_assignment (lplace : mplace) (rvalue : V.typed_value)
   match expr with
   | None -> None
   | Some expr -> Some (Meta (Assignment (lplace, rvalue, rplace), expr))
+
+let synthesize_assertion (v : V.typed_value) (expr : expression option) =
+  match expr with None -> None | Some expr -> Some (Assertion (v, expr))
diff --git a/compiler/Translate.ml b/compiler/Translate.ml
index 46f8172a..d7cc9155 100644
--- a/compiler/Translate.ml
+++ b/compiler/Translate.ml
@@ -634,15 +634,18 @@ let translate_module (filename : string) (dest_dir : string) (config : config)
       config.use_state crate
   in
 
-  (* Initialize the extraction context - for now we extract only to F* *)
-  let names_map =
-    PureToExtract.initialize_names_map ExtractToFStar.fstar_names_map_init
-  in
+  (* Initialize the extraction context - for now we extract only to F*.
+   * We initialize the names map by registering the keywords used in the
+   * language, as well as some primitive names ("u32", etc.) *)
   let variant_concatenate_type_name = true in
   let fstar_fmt =
     ExtractToFStar.mk_formatter trans_ctx crate.name
       variant_concatenate_type_name
   in
+  let names_map =
+    PureToExtract.initialize_names_map fstar_fmt
+      ExtractToFStar.fstar_names_map_init
+  in
   let ctx =
     { PureToExtract.trans_ctx; names_map; fmt = fstar_fmt; indent_incr = 2 }
   in