Make the scrutinee in Pure.Switch an expression rather than a value

4 files changed, 25 insertions, 21 deletions

diff --git a/src/PrintPure.ml b/src/PrintPure.ml
index 4f7b653a..44c0be24 100644
--- a/src/PrintPure.ml
+++ b/src/PrintPure.ml
@@ -361,7 +361,7 @@ let rec expression_to_string (fmt : ast_formatter) (indent : string)
   | Call call -> call_to_string fmt indent indent_incr call
   | Let (monadic, lv, re, e) ->
       let_to_string fmt indent indent_incr monadic lv re e
-  | Switch (scrutinee, _, body) ->
+  | Switch (scrutinee, body) ->
       switch_to_string fmt indent indent_incr scrutinee body
   | Meta (meta, e) ->
       let meta = meta_to_string fmt meta in
@@ -394,10 +394,13 @@ and let_to_string (fmt : ast_formatter) (indent : string) (indent_incr : string)
   else "let " ^ lv ^ " = " ^ re ^ " in\n" ^ indent ^ e
 
 and switch_to_string (fmt : ast_formatter) (indent : string)
-    (indent_incr : string) (scrutinee : typed_rvalue) (body : switch_body) :
+    (indent_incr : string) (scrutinee : expression) (body : switch_body) :
     string =
-  let scrut = typed_rvalue_to_string fmt scrutinee in
   let indent1 = indent ^ indent_incr in
+  (* Printing can mess up on the scrutinee, because it is an expression - but
+   * in most situations it will be a value or a function call, so it should be
+   * ok*)
+  let scrut = expression_to_string fmt indent1 indent_incr scrutinee in
   match body with
   | If (e_true, e_false) ->
       let e_true = expression_to_string fmt indent1 indent_incr e_true in
diff --git a/src/Pure.ml b/src/Pure.ml
index f20cdf50..dc35a181 100644
--- a/src/Pure.ml
+++ b/src/Pure.ml
@@ -479,7 +479,7 @@ type expression =
           ...
           ```
        *)
-  | Switch of typed_rvalue * mplace option * switch_body
+  | Switch of expression * switch_body
   | Meta of meta * expression  (** Meta-information *)
 
 and call = {
diff --git a/src/PureMicroPasses.ml b/src/PureMicroPasses.ml
index 23ae22af..99c8341c 100644
--- a/src/PureMicroPasses.ml
+++ b/src/PureMicroPasses.ml
@@ -201,7 +201,7 @@ let compute_pretty_names (def : fun_def) : fun_def =
     | Value (v, mp) -> update_value v mp ctx
     | Call call -> update_call call ctx
     | Let (monadic, lb, re, e) -> update_let monadic lb re e ctx
-    | Switch (scrut, mp, body) -> update_switch_body scrut mp body ctx
+    | Switch (scrut, body) -> update_switch_body scrut body ctx
     | Meta (meta, e) -> update_meta meta e ctx
   (* *)
   and update_value (v : typed_rvalue) (mp : mplace option) (ctx : pn_ctx) :
@@ -226,9 +226,9 @@ let compute_pretty_names (def : fun_def) : fun_def =
     let lv = update_typed_lvalue ctx lv in
     (ctx, Let (monadic, lv, re, e))
   (* *)
-  and update_switch_body (scrut : typed_rvalue) (mp : mplace option)
-      (body : switch_body) (ctx : pn_ctx) : pn_ctx * expression =
-    let ctx = add_opt_right_constraint mp scrut ctx in
+  and update_switch_body (scrut : expression) (body : switch_body)
+      (ctx : pn_ctx) : pn_ctx * expression =
+    let ctx, scrut = update_expression scrut ctx in
 
     let ctx, body =
       match body with
@@ -264,7 +264,7 @@ let compute_pretty_names (def : fun_def) : fun_def =
           let ctx = merge_ctxs_ls ctxs in
           (ctx, Match branches)
     in
-    (ctx, Switch (scrut, mp, body))
+    (ctx, Switch (scrut, body))
   (* *)
   and update_meta (meta : meta) (e : expression) (ctx : pn_ctx) :
       pn_ctx * expression =
@@ -460,7 +460,7 @@ let expression_contains_child_call_in_all_paths (ctx : trans_ctx) (call0 : call)
               self#visit_expression env re () && self#visit_expression env e ()
         | Call call1 -> fun () -> check_call call1
         | Meta (_, e) -> self#visit_expression env e
-        | Switch (_, _, body) -> self#visit_switch_body env body
+        | Switch (_, body) -> self#visit_switch_body env body
       (** We need to reimplement the way we compose the booleans *)
 
       method! visit_switch_body env body =
@@ -542,7 +542,7 @@ let filter_unused (filter_monadic_calls : bool) (ctx : trans_ctx)
 
       method! visit_expression env e =
         match e with
-        | Value (_, _) | Call _ | Switch (_, _, _) | Meta (_, _) ->
+        | Value (_, _) | Call _ | Switch (_, _) | Meta (_, _) ->
             super#visit_expression env e
         | Let (monadic, lv, re, e) ->
             (* Compute the set of values used in the next expression *)
diff --git a/src/SymbolicToPure.ml b/src/SymbolicToPure.ml
index ada64720..f6f610dd 100644
--- a/src/SymbolicToPure.ml
+++ b/src/SymbolicToPure.ml
@@ -39,13 +39,13 @@ type fun_sig_named_outputs = {
   sg : fun_sig;  (** A function signature *)
   output_names : string option list;
       (** In case the signature is for a backward function, we may provides names
-    for the outputs. The reason is that the outputs of backward functions
-    come from (in case there are no nested borrows) borrows present in the
-    inputs of the original rust function. In this situation, we can use the
-    names of those inputs to name the outputs. Those names are very useful
-    to generate beautiful codes (we may need to introduce temporary variables
-    in the bodies of the backward functions to store the returned values, in
-    which case we use those names).
+          for the outputs. The reason is that the outputs of backward functions
+          come from (in case there are no nested borrows) borrows present in the
+          inputs of the original rust function. In this situation, we can use the
+          names of those inputs to name the outputs. Those names are very useful
+          to generate beautiful codes (we may need to introduce temporary variables
+          in the bodies of the backward functions to store the returned values, in
+          which case we use those names).
  *)
 }
 
@@ -1229,13 +1229,13 @@ and translate_expansion (p : S.mplace option) (sv : V.symbolic_value)
           let branches =
             List.map (fun (vid, svl, e) -> translate_branch vid svl e) branches
           in
-          Switch (scrutinee, scrutinee_mplace, Match branches))
+          Switch (Value (scrutinee, scrutinee_mplace), Match branches))
   | ExpandBool (true_e, false_e) ->
       (* We don't need to update the context: we don't introduce any
        * new values/variables *)
       let true_e = translate_expression true_e ctx in
       let false_e = translate_expression false_e ctx in
-      Switch (scrutinee, scrutinee_mplace, If (true_e, false_e))
+      Switch (Value (scrutinee, scrutinee_mplace), If (true_e, false_e))
   | ExpandInt (int_ty, branches, otherwise) ->
       let translate_branch ((v, branch_e) : V.scalar_value * S.expression) :
           scalar_value * expression =
@@ -1247,7 +1247,8 @@ and translate_expansion (p : S.mplace option) (sv : V.symbolic_value)
       let branches = List.map translate_branch branches in
       let otherwise = translate_expression otherwise ctx in
       Switch
-        (scrutinee, scrutinee_mplace, SwitchInt (int_ty, branches, otherwise))
+        ( Value (scrutinee, scrutinee_mplace),
+          SwitchInt (int_ty, branches, otherwise) )
 
 and translate_meta (meta : S.meta) (e : S.expression) (ctx : bs_ctx) :
     expression =