Isolate while-header bug to minimal testcase, fix in rustboot, remove workaround in rustc.

3 files changed, 70 insertions, 33 deletions

diff --git a/src/boot/me/semant.ml b/src/boot/me/semant.ml
index 4f02455a..c4930c12 100644
--- a/src/boot/me/semant.ml
+++ b/src/boot/me/semant.ml
@@ -132,6 +132,7 @@ type ctxt =
 
       (* Typestate-y stuff. *)
       ctxt_stmt_is_init: (node_id,unit) Hashtbl.t;
+      ctxt_while_header_slots: (node_id,node_id list) Hashtbl.t;
       ctxt_post_stmt_slot_drops: (node_id,node_id list) Hashtbl.t;
       ctxt_post_block_slot_drops: (node_id,node_id list) Hashtbl.t;
 
@@ -239,6 +240,7 @@ let new_ctxt sess abi crate =
     ctxt_required_syms = crate.Ast.crate_required_syms;
 
     ctxt_stmt_is_init = Hashtbl.create 0;
+    ctxt_while_header_slots = Hashtbl.create 0;
     ctxt_post_stmt_slot_drops = Hashtbl.create 0;
     ctxt_post_block_slot_drops = Hashtbl.create 0;
 
diff --git a/src/boot/me/trans.ml b/src/boot/me/trans.ml
index 1fdd40ad..c3d5240b 100644
--- a/src/boot/me/trans.ml
+++ b/src/boot/me/trans.ml
@@ -2614,27 +2614,24 @@ let trans_visitor
         | Ast.EXPR_atom a ->
             trans_atom a
 
+  and drop_slot_by_id (depth:int) (slot_id:node_id) : unit =
+    let slot = get_slot cx slot_id in
+    let k = Hashtbl.find cx.ctxt_slot_keys slot_id in
+      iflog (fun _ ->
+               annotate
+                 (Printf.sprintf
+                    "drop_slot %d = %s "
+                    (int_of_node slot_id)
+                    (Fmt.fmt_to_str Ast.fmt_slot_key k)));
+      drop_slot_in_current_frame
+        (cell_of_block_slot
+           ~access_depth:(Some depth) slot_id) slot
+
   and drop_slots_after_block bid : unit =
+    let depth = Hashtbl.find cx.ctxt_block_loop_depths bid in
     match htab_search cx.ctxt_post_block_slot_drops bid with
         None -> ()
-      | Some slots ->
-          List.iter
-            begin
-              fun slot_id ->
-                let slot = get_slot cx slot_id in
-                let k = Hashtbl.find cx.ctxt_slot_keys slot_id in
-                let depth = Hashtbl.find cx.ctxt_block_loop_depths bid in
-                  iflog (fun _ ->
-                           annotate
-                             (Printf.sprintf
-                                "post-block, drop_slot %d = %s "
-                                (int_of_node slot_id)
-                                (Fmt.fmt_to_str Ast.fmt_slot_key k)));
-                  drop_slot_in_current_frame
-                    (cell_of_block_slot
-                       ~access_depth:(Some depth) slot_id) slot
-            end
-            slots
+      | Some slots -> List.iter (drop_slot_by_id depth) slots
 
   and trans_block (block:Ast.block) : unit =
     flush_emitter_size_cache();
@@ -5260,6 +5257,36 @@ let trans_visitor
           trans_log_int a
       | _ -> unimpl (Some id) "logging type"
 
+  and trans_while (id:node_id) (sw:Ast.stmt_while) : unit =
+    let (head_stmts, head_expr) = sw.Ast.while_lval in
+    let fwd_jmp = mark () in
+      emit (Il.jmp Il.JMP Il.CodeNone);
+      let block_begin = mark () in
+        Stack.push (Stack.create()) simple_break_jumps;
+        trans_block sw.Ast.while_body;
+        patch fwd_jmp;
+        Array.iter trans_stmt head_stmts;
+        check_interrupt_flag ();
+        let flag = next_vreg_cell (Il.ValTy Il.Bits8) in
+          mov flag imm_true;
+          let true_jmps = trans_cond false head_expr in
+            mov flag imm_false;
+            List.iter patch true_jmps;
+            begin
+              begin
+                match htab_search cx.ctxt_while_header_slots id with
+                    None -> ()
+                  | Some slots ->
+                      let depth = get_stmt_depth cx id in
+                        List.iter (drop_slot_by_id depth) slots
+              end;
+              let back_jmps =
+                trans_compare_simple Il.JE (Il.Cell flag) imm_true
+              in
+                List.iter (fun j -> patch_existing j block_begin) back_jmps;
+        end;
+        Stack.iter patch (Stack.pop simple_break_jumps);
+
 
   and trans_stmt_full (stmt:Ast.stmt) : unit =
     match stmt.node with
@@ -5378,20 +5405,7 @@ let trans_visitor
           trans_block block
 
       | Ast.STMT_while sw ->
-          let (head_stmts, head_expr) = sw.Ast.while_lval in
-          let fwd_jmp = mark () in
-            emit (Il.jmp Il.JMP Il.CodeNone);
-            let block_begin = mark () in
-              Stack.push (Stack.create()) simple_break_jumps;
-              trans_block sw.Ast.while_body;
-              patch fwd_jmp;
-              Array.iter trans_stmt head_stmts;
-              check_interrupt_flag ();
-              begin
-                let back_jmps = trans_cond false head_expr in
-                  List.iter (fun j -> patch_existing j block_begin) back_jmps;
-              end;
-              Stack.iter patch (Stack.pop simple_break_jumps);
+          trans_while stmt.id sw
 
       | Ast.STMT_if si ->
           let skip_thn_jmps = trans_cond true si.Ast.if_test in
diff --git a/src/boot/me/typestate.ml b/src/boot/me/typestate.ml
index 4ca0caf2..ea0204f3 100644
--- a/src/boot/me/typestate.ml
+++ b/src/boot/me/typestate.ml
@@ -1467,7 +1467,28 @@ let lifecycle_visitor
                 f.Ast.for_each_body.id
                 [ (fst f.Ast.for_each_slot).id ]
 
-          | Ast.STMT_while _ ->
+          | Ast.STMT_while sw ->
+              (* Collect any header-locals. *)
+              Array.iter
+                begin
+                  fun stmt ->
+                    match stmt.node with
+                        Ast.STMT_decl (Ast.DECL_slot (_, slot)) ->
+                          begin
+                            match
+                              htab_search cx.ctxt_while_header_slots s.id
+                            with
+                                None ->
+                                  Hashtbl.add cx.ctxt_while_header_slots
+                                    s.id [slot.id]
+                              | Some slots ->
+                                  Hashtbl.replace cx.ctxt_while_header_slots
+                                    s.id (slot.id :: slots)
+                          end
+                      | _ -> ()
+                end
+                (fst sw.Ast.while_lval);
+
               iflog cx (fun _ -> log cx "entering a loop");
               Stack.push (Some (Stack.create ()))  loop_blocks;