Make capture-by-ref unboxed closures work

src/librustc/middle/mem_categorization.rs

@@ -604,17 +604,21 @@ impl<'t,'tcx,TYPER:Typer<'tcx>> MemCategorizationContext<'t,TYPER> {
                           ty::FnMutUnboxedClosureKind => ast::Many,
                           ty::FnOnceUnboxedClosureKind => ast::Once,
                       };
-                      Ok(Rc::new(cmt_ {
-                          id: id,
-                          span: span,
-                          cat: cat_copied_upvar(CopiedUpvar {
-                              upvar_id: var_id,
-                              onceness: onceness,
-                              capturing_proc: fn_node_id,
-                          }),
-                          mutbl: MutabilityCategory::from_local(self.tcx(), var_id),
-                          ty: expr_ty
-                      }))
+                      if self.typer.capture_mode(fn_node_id) == ast::CaptureByRef {
+                          self.cat_upvar(id, span, var_id, fn_node_id)
+                      } else {
+                          Ok(Rc::new(cmt_ {
+                              id: id,
+                              span: span,
+                              cat: cat_copied_upvar(CopiedUpvar {
+                                  upvar_id: var_id,
+                                  onceness: onceness,
+                                  capturing_proc: fn_node_id,
+                              }),
+                              mutbl: MutabilityCategory::from_local(self.tcx(), var_id),
+                              ty: expr_ty
+                          }))
+                      }
                   }
                   _ => {
                       self.tcx().sess.span_bug(

src/librustc/middle/trans/closure.rs

@@ -301,6 +301,7 @@ fn load_environment<'blk, 'tcx>(bcx: Block<'blk, 'tcx>,
 fn load_unboxed_closure_environment<'blk, 'tcx>(
                                     bcx: Block<'blk, 'tcx>,
                                     arg_scope_id: ScopeId,
+                                    freevar_mode: ast::CaptureClause,
                                     freevars: &Vec<ty::Freevar>,
                                     closure_id: ast::DefId)
                                     -> Block<'blk, 'tcx> {
@@ -326,11 +327,14 @@ fn load_unboxed_closure_environment<'blk, 'tcx>(
     };
 
     for (i, freevar) in freevars.iter().enumerate() {
-        let upvar_ptr = GEPi(bcx, llenv, [0, i]);
+        let mut upvar_ptr = GEPi(bcx, llenv, [0, i]);
+        if freevar_mode == ast::CaptureByRef {
+            upvar_ptr = Load(bcx, upvar_ptr);
+        }
         let def_id = freevar.def.def_id();
         bcx.fcx.llupvars.borrow_mut().insert(def_id.node, upvar_ptr);
 
-        if kind == ty::FnOnceUnboxedClosureKind {
+        if kind == ty::FnOnceUnboxedClosureKind && freevar_mode == ast::CaptureByValue {
             bcx.fcx.schedule_drop_mem(arg_scope_id,
                                       upvar_ptr,
                                       node_id_type(bcx, def_id.node))
@@ -477,6 +481,7 @@ pub fn trans_unboxed_closure<'blk, 'tcx>(
     let freevars: Vec<ty::Freevar> =
         ty::with_freevars(bcx.tcx(), id, |fv| fv.iter().map(|&fv| fv).collect());
     let freevars_ptr = &freevars;
+    let freevar_mode = bcx.tcx().capture_mode(id);
 
     trans_closure(bcx.ccx(),
                   decl,
@@ -493,6 +498,7 @@ pub fn trans_unboxed_closure<'blk, 'tcx>(
                   |bcx, arg_scope| {
                       load_unboxed_closure_environment(bcx,
                                                        arg_scope,
+                                                       freevar_mode,
                                                        freevars_ptr,
                                                        closure_id)
                   });
@@ -518,7 +524,14 @@ pub fn trans_unboxed_closure<'blk, 'tcx>(
                                                    dest_addr,
                                                    0,
                                                    i);
-        bcx = datum.store_to(bcx, upvar_slot_dest);
+        match freevar_mode {
+            ast::CaptureByValue => {
+                bcx = datum.store_to(bcx, upvar_slot_dest);
+            }
+            ast::CaptureByRef => {
+                Store(bcx, datum.to_llref(), upvar_slot_dest);
+            }
+        }
     }
     adt::trans_set_discr(bcx, &*repr, dest_addr, 0);
 

src/librustc/middle/ty.rs

@@ -4668,15 +4668,27 @@ pub struct UnboxedClosureUpvar {
 pub fn unboxed_closure_upvars(tcx: &ctxt, closure_id: ast::DefId)
                               -> Vec<UnboxedClosureUpvar> {
     if closure_id.krate == ast::LOCAL_CRATE {
+        let capture_mode = tcx.capture_modes.borrow().get_copy(&closure_id.node);
         match tcx.freevars.borrow().find(&closure_id.node) {
             None => vec![],
             Some(ref freevars) => {
                 freevars.iter().map(|freevar| {
                     let freevar_def_id = freevar.def.def_id();
+                    let mut freevar_ty = node_id_to_type(tcx, freevar_def_id.node);
+                    if capture_mode == ast::CaptureByRef {
+                        let borrow = tcx.upvar_borrow_map.borrow().get_copy(&ty::UpvarId {
+                            var_id: freevar_def_id.node,
+                            closure_expr_id: closure_id.node
+                        });
+                        freevar_ty = mk_rptr(tcx, borrow.region, ty::mt {
+                            ty: freevar_ty,
+                            mutbl: borrow.kind.to_mutbl_lossy()
+                        });
+                    }
                     UnboxedClosureUpvar {
                         def: freevar.def,
                         span: freevar.span,
-                        ty: node_id_to_type(tcx, freevar_def_id.node),
+                        ty: freevar_ty
                     }
                 }).collect()
             }

src/librustc/middle/typeck/check/regionck.rs

@@ -844,7 +844,7 @@ fn check_expr_fn_block(rcx: &mut Rcx,
                     // has static lifetime.
                 } else {
                     // Variables being referenced must outlive closure.
-                    constrain_free_variables_in_stack_closure(
+                    constrain_free_variables_in_by_ref_closure(
                         rcx, bounds.region_bound, expr, freevars);
 
                     // Closure is stack allocated and hence cannot
@@ -856,20 +856,17 @@ fn check_expr_fn_block(rcx: &mut Rcx,
             });
         }
         ty::ty_unboxed_closure(_, region) => {
-            ty::with_freevars(tcx, expr.id, |freevars| {
-                // No free variables means that there is no environment and
-                // hence the closure has static lifetime. Otherwise, the
-                // closure must not outlive the variables it closes over
-                // by-reference.
-                //
-                // NDM -- this seems wrong, discuss with pcwalton, should
-                // be straightforward enough.
-                if !freevars.is_empty() {
-                    let bounds = ty::region_existential_bound(region);
-                    ensure_free_variable_types_outlive_closure_bound(
-                        rcx, bounds, expr, freevars);
-                }
-            })
+            let bounds = ty::region_existential_bound(region);
+            if tcx.capture_modes.borrow().get_copy(&expr.id) == ast::CaptureByRef {
+                ty::with_freevars(tcx, expr.id, |freevars| {
+                    if !freevars.is_empty() {
+                        // Variables being referenced must be constrained and registered
+                        // in the upvar borrow map
+                        constrain_free_variables_in_by_ref_closure(
+                            rcx, bounds.region_bound, expr, freevars);
+                    }
+                })
+            }
         }
         _ => { }
     }
@@ -884,6 +881,13 @@ fn check_expr_fn_block(rcx: &mut Rcx,
                 propagate_upupvar_borrow_kind(rcx, expr, freevars);
             })
         }
+        ty::ty_unboxed_closure(..) => {
+            if tcx.capture_modes.borrow().get_copy(&expr.id) == ast::CaptureByRef {
+                ty::with_freevars(tcx, expr.id, |freevars| {
+                    propagate_upupvar_borrow_kind(rcx, expr, freevars);
+                });
+            }
+        }
         _ => {}
     }
 
@@ -893,6 +897,12 @@ fn check_expr_fn_block(rcx: &mut Rcx,
                 ensure_free_variable_types_outlive_closure_bound(rcx, bounds, expr, freevars);
             })
         }
+        ty::ty_unboxed_closure(_, region) => {
+            ty::with_freevars(tcx, expr.id, |freevars| {
+                let bounds = ty::region_existential_bound(region);
+                ensure_free_variable_types_outlive_closure_bound(rcx, bounds, expr, freevars);
+            })
+        }
         _ => {}
     }
 
@@ -959,7 +969,7 @@ fn check_expr_fn_block(rcx: &mut Rcx,
         }
     }
 
-    fn constrain_free_variables_in_stack_closure(
+    fn constrain_free_variables_in_by_ref_closure(
         rcx: &mut Rcx,
         region_bound: ty::Region,
         expr: &ast::Expr,

src/test/compile-fail/unboxed-closure-region.rs

@@ -0,0 +1,20 @@
+// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#![feature(unboxed_closures)]
+
+// Test that an unboxed closure that captures a free variable by
+// reference cannot escape the region of that variable.
+fn main() {
+    let _f = {
+        let x = 0u;
+        |:| x //~ ERROR cannot infer an appropriate lifetime due to conflicting requirements
+    };
+}

src/test/compile-fail/unboxed-closures-borrow-conflict.rs

@@ -0,0 +1,20 @@
+// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#![feature(unboxed_closures)]
+
+// Test that an unboxed closure that mutates a free variable will
+// cause borrow conflicts.
+
+fn main() {
+    let mut x = 0u;
+    let f = |:| x += 1;
+    let _y = x; //~ ERROR cannot use `x` because it was mutably borrowed
+}

src/test/run-pass/capture-clauses-unboxed-closures.rs

@@ -8,13 +8,9 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
-// ignore-test
-//
-// This is ignored because it depends on #16122.
-
 #![feature(overloaded_calls, unboxed_closures)]
 
-fn each<'a,T,F:|&mut: &'a T|>(x: &'a [T], mut f: F) {
+fn each<'a,T,F:FnMut(&'a T)>(x: &'a [T], mut f: F) {
     for val in x.iter() {
         f(val)
     }
@@ -23,7 +19,6 @@ fn each<'a,T,F:|&mut: &'a T|>(x: &'a [T], mut f: F) {
 fn main() {
     let mut sum = 0u;
     let elems = [ 1u, 2, 3, 4, 5 ];
-    each(elems, ref |&mut: val: &uint| sum += *val);
+    each(elems, |&mut: val: &uint| sum += *val);
     assert_eq!(sum, 15);
 }
-

src/test/run-pass/unboxed-closures-by-ref.rs

@@ -0,0 +1,35 @@
+// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#![feature(overloaded_calls, unboxed_closures)]
+
+// Test by-ref capture of environment in unboxed closure types
+
+fn call_fn<F: Fn()>(f: F) {
+    f()
+}
+
+fn call_fn_mut<F: FnMut()>(mut f: F) {
+    f()
+}
+
+fn call_fn_once<F: FnOnce()>(f: F) {
+    f()
+}
+
+fn main() {
+    let mut x = 0u;
+    let y = 2u;
+
+    call_fn(|&:| x += y);
+    call_fn_mut(|&mut:| x += y);
+    call_fn_once(|:| x += y);
+    assert_eq!(x, y * 3);
+}