Support returning from loop blocks

The code is somewhat invasive, but it seems hard to do this in a
clean way, since the design itself involves a bunch of 'action
at a distance'.

Issue #1819

Author: marijnh

src/rustc/middle/check_loop.rs

@@ -25,7 +25,8 @@ fn check_crate(tcx: ty::ctxt, crate: @crate) {
                 v.visit_block(b, {in_loop: false, can_ret: false}, v);
               }
               expr_loop_body(@{node: expr_fn_block(_, b), _}) {
-                v.visit_block(b, {in_loop: true, can_ret: false}, v);
+                let blk = is_blockish(ty::ty_fn_proto(ty::expr_ty(tcx, e)));
+                v.visit_block(b, {in_loop: true, can_ret: blk}, v);
               }
               expr_break {
                 if !cx.in_loop {

src/rustc/middle/trans/base.rs

@@ -2516,15 +2516,43 @@ fn trans_cast(cx: block, e: @ast::expr, id: ast::node_id,
     ret store_in_dest(e_res.bcx, newval, dest);
 }
 
+fn trans_loop_body(bcx: block, e: @ast::expr, ret_flag: option<ValueRef>,
+                   dest: dest) -> block {
+    alt check e.node {
+      ast::expr_loop_body(b@@{node: ast::expr_fn_block(decl, body), _}) {
+        alt check ty::get(expr_ty(bcx, e)).struct {
+          ty::ty_fn({proto, _}) {
+            closure::trans_expr_fn(bcx, proto, decl, body, e.span, b.id,
+                                   {copies: [], moves: []}, some(ret_flag),
+                                   dest)
+          }
+        }
+      }
+    }
+}
+
 // temp_cleanups: cleanups that should run only if failure occurs before the
 // call takes place:
 fn trans_arg_expr(cx: block, arg: ty::arg, lldestty: TypeRef, e: @ast::expr,
-                  &temp_cleanups: [ValueRef]) -> result {
+                  &temp_cleanups: [ValueRef], ret_flag: option<ValueRef>)
+    -> result {
     let _icx = cx.insn_ctxt("trans_arg_expr");
     let ccx = cx.ccx();
     let e_ty = expr_ty(cx, e);
     let is_bot = ty::type_is_bot(e_ty);
-    let lv = trans_temp_lval(cx, e);
+    let lv = alt ret_flag {
+      // If there is a ret_flag, this *must* be a loop body
+      some(ptr) {
+        alt check e.node {
+          ast::expr_loop_body(blk) {
+            let scratch = alloc_ty(cx, expr_ty(cx, blk));
+            let bcx = trans_loop_body(cx, e, ret_flag, save_in(scratch));
+            {bcx: bcx, val: scratch, kind: temporary}
+          }
+        }
+      }
+      none { trans_temp_lval(cx, e) }
+    };
     let mut bcx = lv.bcx;
     let mut val = lv.val;
     let arg_mode = ty::resolved_mode(ccx.tcx, arg.mode);
@@ -2595,7 +2623,7 @@ enum call_args {
 //  - new_fn_ctxt
 //  - trans_args
 fn trans_args(cx: block, llenv: ValueRef, args: call_args, fn_ty: ty::t,
-              dest: dest)
+              dest: dest, ret_flag: option<ValueRef>)
     -> {bcx: block, args: [ValueRef], retslot: ValueRef} {
     let _icx = cx.insn_ctxt("trans_args");
     let mut temp_cleanups = [];
@@ -2630,13 +2658,13 @@ fn trans_args(cx: block, llenv: ValueRef, args: call_args, fn_ty: ty::t,
     alt args {
       arg_exprs(es) {
         let llarg_tys = type_of_explicit_args(ccx, arg_tys);
-        let mut i = 0u;
-        for e: @ast::expr in es {
+        let last = es.len() - 1u;
+        vec::iteri(es) {|i, e|
             let r = trans_arg_expr(bcx, arg_tys[i], llarg_tys[i],
-                                   e, temp_cleanups);
+                                   e, temp_cleanups, if i == last { ret_flag }
+                                                     else { none });
             bcx = r.bcx;
             llargs += [r.val];
-            i += 1u;
         }
       }
       arg_vals(vs) {
@@ -2664,14 +2692,44 @@ fn trans_call(in_cx: block, f: @ast::expr,
                      {|cx| trans_callee(cx, f)}, args, dest)
 }
 
+fn body_contains_ret(body: ast::blk) -> bool {
+    let cx = {mut found: false};
+    visit::visit_block(body, cx, visit::mk_vt(@{
+        visit_item: {|_i, _cx, _v|},
+        visit_expr: {|e: @ast::expr, cx: {mut found: bool}, v|
+            if !cx.found {
+                alt e.node {
+                  ast::expr_ret(_) { cx.found = true; }
+                  _ { visit::visit_expr(e, cx, v); }
+                }
+            }
+        } with *visit::default_visitor()
+    }));
+    cx.found
+}
+
 fn trans_call_inner(in_cx: block, fn_expr_ty: ty::t, ret_ty: ty::t,
                     get_callee: fn(block) -> lval_maybe_callee,
                     args: call_args, dest: dest)
     -> block {
+    let ret_in_loop = alt args {
+      arg_exprs(args) { args.len() > 0u && alt vec::last(args).node {
+        ast::expr_loop_body(@{node: ast::expr_fn_block(_, body), _}) {
+          body_contains_ret(body)
+        }
+        _ { false }
+      } }
+      _ { false }
+    };
     with_scope(in_cx, "call") {|cx|
         let f_res = get_callee(cx);
         let mut bcx = f_res.bcx;
         let ccx = cx.ccx();
+        let ret_flag = if ret_in_loop {
+            let flag = alloca(in_cx, T_bool());
+            Store(cx, C_bool(false), flag);
+            some(flag)
+        } else { none };
 
         let mut faddr = f_res.val;
         let llenv = alt f_res.env {
@@ -2695,7 +2753,7 @@ fn trans_call_inner(in_cx: block, fn_expr_ty: ty::t, ret_ty: ty::t,
         };
 
         let args_res = {
-            trans_args(bcx, llenv, args, fn_expr_ty, dest)
+            trans_args(bcx, llenv, args, fn_expr_ty, dest, ret_flag)
         };
         bcx = args_res.bcx;
         let mut llargs = args_res.args;
@@ -2718,7 +2776,19 @@ fn trans_call_inner(in_cx: block, fn_expr_ty: ty::t, ret_ty: ty::t,
             *cell = Load(bcx, llretslot);
           }
         }
-        if ty::type_is_bot(ret_ty) { Unreachable(bcx); }
+        if ty::type_is_bot(ret_ty) {
+            Unreachable(bcx);
+        } else if ret_in_loop {
+            bcx = with_cond(bcx, Load(bcx, option::get(ret_flag))) {|bcx|
+                option::may(bcx.fcx.loop_ret) {|lret|
+                    Store(bcx, C_bool(true), lret.flagptr);
+                    Store(bcx, C_bool(false), bcx.fcx.llretptr);
+                }
+                cleanup_and_leave(bcx, none, some(bcx.fcx.llreturn));
+                Unreachable(bcx);
+                bcx
+            }
+        }
         bcx
     }
 }
@@ -2991,25 +3061,20 @@ fn trans_expr(bcx: block, e: @ast::expr, dest: dest) -> block {
       ast::expr_addr_of(_, x) { ret trans_addr_of(bcx, x, dest); }
       ast::expr_fn(proto, decl, body, cap_clause) {
         ret closure::trans_expr_fn(bcx, proto, decl, body, e.span, e.id,
-                                   *cap_clause, false, dest);
+                                   *cap_clause, none, dest);
       }
       ast::expr_fn_block(decl, body) {
         alt check ty::get(expr_ty(bcx, e)).struct {
           ty::ty_fn({proto, _}) {
             #debug("translating fn_block %s with type %s",
                    expr_to_str(e), ty_to_str(tcx, expr_ty(bcx, e)));
             ret closure::trans_expr_fn(bcx, proto, decl, body, e.span, e.id,
-                                       {copies: [], moves: []}, false, dest);
+                                       {copies: [], moves: []}, none, dest);
           }
         }
       }
-      ast::expr_loop_body(b@@{node: ast::expr_fn_block(decl, body), _}) {
-        alt check ty::get(expr_ty(bcx, e)).struct {
-          ty::ty_fn({proto, _}) {
-            ret closure::trans_expr_fn(bcx, proto, decl, body, e.span, b.id,
-                                       {copies: [], moves: []}, true, dest);
-          }
-        }
+      ast::expr_loop_body(blk) {
+        ret trans_loop_body(bcx, e, none, dest);
       }
       ast::expr_bind(f, args) {
         ret closure::trans_bind(
@@ -3406,8 +3471,25 @@ fn trans_cont(cx: block) -> block {
 fn trans_ret(bcx: block, e: option<@ast::expr>) -> block {
     let _icx = bcx.insn_ctxt("trans_ret");
     let mut bcx = bcx;
+    let retptr = alt bcx.fcx.loop_ret {
+      some({flagptr, retptr}) {
+        // This is a loop body return. Must set continue flag (our retptr)
+        // to false, return flag to true, and then store the value in the
+        // parent's retptr.
+        Store(bcx, C_bool(true), flagptr);
+        Store(bcx, C_bool(false), bcx.fcx.llretptr);
+        alt e {
+          some(x) { PointerCast(bcx, retptr,
+                                T_ptr(type_of(bcx.ccx(), expr_ty(bcx, x)))) }
+          none { retptr }
+        }
+      }
+      none { bcx.fcx.llretptr }
+    };
     alt e {
-      some(x) { bcx = trans_expr_save_in(bcx, x, bcx.fcx.llretptr); }
+      some(x) {
+        bcx = trans_expr_save_in(bcx, x, retptr);
+      }
       _ {}
     }
     cleanup_and_leave(bcx, none, some(bcx.fcx.llreturn));
@@ -3793,6 +3875,7 @@ fn new_fn_ctxt_w_id(ccx: @crate_ctxt, path: path,
           mut llreturn: llbbs.rt,
           mut llself: none,
           mut personality: none,
+          mut loop_ret: none,
           llargs: int_hash::<local_val>(),
           lllocals: int_hash::<local_val>(),
           llupvars: int_hash::<ValueRef>(),

src/rustc/middle/trans/closure.rs

@@ -287,7 +287,8 @@ fn store_environment(bcx: block,
 fn build_closure(bcx0: block,
                  cap_vars: [capture::capture_var],
                  ck: ty::closure_kind,
-                 id: ast::node_id) -> closure_result {
+                 id: ast::node_id,
+                 include_ret_handle: option<ValueRef>) -> closure_result {
     let _icx = bcx0.insn_ctxt("closure::build_closure");
     // If we need to, package up the iterator body to call
     let mut env_vals = [];
@@ -324,6 +325,16 @@ fn build_closure(bcx0: block,
           }
         }
     }
+    option::may(include_ret_handle) {|flagptr|
+        let our_ret = alt bcx.fcx.loop_ret {
+          some({retptr, _}) { retptr }
+          none { bcx.fcx.llretptr }
+        };
+        let nil_ret = PointerCast(bcx, our_ret, T_ptr(T_nil()));
+        env_vals +=
+            [env_ref(flagptr, ty::mk_mut_ptr(tcx, ty::mk_bool(tcx)), owned),
+             env_ref(nil_ret, ty::mk_nil_ptr(tcx), owned)];
+    }
     ret store_environment(bcx, env_vals, ck);
 }
 
@@ -333,6 +344,7 @@ fn build_closure(bcx0: block,
 fn load_environment(fcx: fn_ctxt,
                     cdata_ty: ty::t,
                     cap_vars: [capture::capture_var],
+                    load_ret_handle: bool,
                     ck: ty::closure_kind) {
     let _icx = fcx.insn_ctxt("closure::load_environment");
     let bcx = raw_block(fcx, fcx.llloadenv);
@@ -341,23 +353,30 @@ fn load_environment(fcx: fn_ctxt,
     let llcdata = base::opaque_box_body(bcx, cdata_ty, fcx.llenv);
 
     // Populate the upvars from the environment.
-    let mut i = 0u;
+    let mut i = 0;
     vec::iter(cap_vars) { |cap_var|
         alt cap_var.mode {
           capture::cap_drop { /* ignore */ }
           _ {
             let mut upvarptr =
-                GEPi(bcx, llcdata, [0, abi::closure_body_bindings, i as int]);
+                GEPi(bcx, llcdata, [0, abi::closure_body_bindings, i]);
             alt ck {
               ty::ck_block { upvarptr = Load(bcx, upvarptr); }
               ty::ck_uniq | ty::ck_box { }
             }
             let def_id = ast_util::def_id_of_def(cap_var.def);
             fcx.llupvars.insert(def_id.node, upvarptr);
-            i += 1u;
+            i += 1;
           }
         }
     }
+    if load_ret_handle {
+        let flagptr = Load(bcx, GEPi(bcx, llcdata,
+                                     [0, abi::closure_body_bindings, i]));
+        let retptr = Load(bcx, GEPi(bcx, llcdata,
+                                    [0, abi::closure_body_bindings, i+1]));
+        fcx.loop_ret = some({flagptr: flagptr, retptr: retptr});
+    }
 }
 
 fn trans_expr_fn(bcx: block,
@@ -367,7 +386,7 @@ fn trans_expr_fn(bcx: block,
                  sp: span,
                  id: ast::node_id,
                  cap_clause: ast::capture_clause,
-                 is_loop_body: bool,
+                 is_loop_body: option<option<ValueRef>>,
                  dest: dest) -> block {
     let _icx = bcx.insn_ctxt("closure::trans_expr_fn");
     if dest == ignore { ret bcx; }
@@ -382,12 +401,17 @@ fn trans_expr_fn(bcx: block,
     let trans_closure_env = fn@(ck: ty::closure_kind) -> ValueRef {
         let cap_vars = capture::compute_capture_vars(
             ccx.tcx, id, proto, cap_clause);
-        let {llbox, cdata_ty, bcx} = build_closure(bcx, cap_vars, ck, id);
+        let ret_handle = alt is_loop_body { some(x) { x } none { none } };
+        let {llbox, cdata_ty, bcx} = build_closure(bcx, cap_vars, ck, id,
+                                                   ret_handle);
         trans_closure(ccx, sub_path, decl, body, llfn, no_self,
                       bcx.fcx.param_substs, id, {|fcx|
-            load_environment(fcx, cdata_ty, cap_vars, ck);
+            load_environment(fcx, cdata_ty, cap_vars,
+                             option::is_some(ret_handle), ck);
         }, {|bcx|
-            if is_loop_body { Store(bcx, C_bool(true), bcx.fcx.llretptr); }
+            if option::is_some(is_loop_body) {
+                Store(bcx, C_bool(true), bcx.fcx.llretptr);
+            }
         });
         llbox
     };

src/rustc/middle/trans/common.rs

@@ -165,6 +165,9 @@ type fn_ctxt = @{
     // The a value alloca'd for calls to upcalls.rust_personality. Used when
     // outputting the resume instruction.
     mut personality: option<ValueRef>,
+    // If this is a for-loop body that returns, this holds the pointers needed
+    // for that
+    mut loop_ret: option<{flagptr: ValueRef, retptr: ValueRef}>,
 
     // Maps arguments to allocas created for them in llallocas.
     llargs: hashmap<ast::node_id, local_val>,

src/rustc/middle/trans/impl.rs

@@ -36,7 +36,7 @@ fn trans_self_arg(bcx: block, base: @ast::expr) -> result {
     let mut temp_cleanups = [];
     let result = trans_arg_expr(bcx, {mode: m_by_ref, ty: basety},
                                 T_ptr(type_of::type_of(bcx.ccx(), basety)),
-                                base, temp_cleanups);
+                                base, temp_cleanups, none);
 
     // by-ref self argument should not require cleanup in the case of
     // other arguments failing: