rustc: add a closure depth to DefUpvar.

Author: eddyb

src/librustc/middle/astencode.rs

@@ -472,9 +472,10 @@ impl tr for def::Def {
           def::DefPrimTy(p) => def::DefPrimTy(p),
           def::DefTyParam(s, did, v) => def::DefTyParam(s, did.tr(dcx), v),
           def::DefUse(did) => def::DefUse(did.tr(dcx)),
-          def::DefUpvar(nid1, def, nid2, nid3) => {
+          def::DefUpvar(nid1, def, depth, nid2, nid3) => {
             def::DefUpvar(dcx.tr_id(nid1),
                            box(GC) (*def).tr(dcx),
+                           depth,
                            dcx.tr_id(nid2),
                            dcx.tr_id(nid3))
           }

src/librustc/middle/def.rs

@@ -30,10 +30,12 @@ pub enum Def {
     DefPrimTy(ast::PrimTy),
     DefTyParam(ParamSpace, ast::DefId, uint),
     DefUse(ast::DefId),
-    DefUpvar(ast::NodeId,  // id of closed over var
-             Gc<Def>,     // closed over def
+    DefUpvar(ast::NodeId,  // id of closed over local
+             Gc<Def>,      // closed over def
+             u32,          // number of closures implicitely capturing this local
              ast::NodeId,  // expr node that creates the closure
-             ast::NodeId), // id for the block/body of the closure expr
+             ast::NodeId), // block node for the closest enclosing proc
+                           // or unboxed closure, DUMMY_NODE_ID otherwise
 
     /// Note that if it's a tuple struct's definition, the node id of the ast::DefId
     /// may either refer to the item definition's id or the StructDef.ctor_id.
@@ -68,7 +70,7 @@ impl Def {
             }
             DefLocal(id) |
             DefSelfTy(id) |
-            DefUpvar(id, _, _, _) |
+            DefUpvar(id, _, _, _,  _) |
             DefRegion(id) |
             DefTyParamBinder(id) |
             DefLabel(id) => {

src/librustc/middle/freevars.rs

@@ -26,7 +26,7 @@ use syntax::visit;
 
 #[deriving(Clone, Decodable, Encodable, Show)]
 pub enum CaptureMode {
-    /// Copy/move the value from this llvm ValueRef into the environment.
+    /// Copy/move the value into the environment.
     CaptureByValue,
 
     /// Access by reference (used for stack closures).
@@ -45,23 +45,22 @@ pub type freevar_map = NodeMap<Vec<freevar_entry>>;
 
 pub type CaptureModeMap = NodeMap<CaptureMode>;
 
-struct CollectFreevarsVisitor<'a> {
+struct CollectFreevarsVisitor<'a, 'b:'a> {
+    node_id: ast::NodeId,
     seen: NodeSet,
-    refs: Vec<freevar_entry>,
-    def_map: &'a resolve::DefMap,
-    capture_mode_map: &'a mut CaptureModeMap,
-    depth: uint
+    cx: &'a mut AnnotateFreevarsVisitor<'b>,
+    depth: u32
 }
 
-impl<'a, 'v> Visitor<'v> for CollectFreevarsVisitor<'a> {
+impl<'a, 'b, 'v> Visitor<'v> for CollectFreevarsVisitor<'a, 'b> {
     fn visit_item(&mut self, _: &ast::Item) {
         // ignore_item
     }
 
     fn visit_expr(&mut self, expr: &ast::Expr) {
         match expr.node {
             ast::ExprProc(..) => {
-                self.capture_mode_map.insert(expr.id, CaptureByValue);
+                self.cx.capture_mode_map.insert(expr.id, CaptureByValue);
                 self.depth += 1;
                 visit::walk_expr(self, expr);
                 self.depth -= 1;
@@ -74,7 +73,7 @@ impl<'a, 'v> Visitor<'v> for CollectFreevarsVisitor<'a> {
                 //    ast::CaptureByRef => CaptureByRef,
                 //};
                 let capture_mode = CaptureByRef;
-                self.capture_mode_map.insert(expr.id, capture_mode);
+                self.cx.capture_mode_map.insert(expr.id, capture_mode);
                 self.depth += 1;
                 visit::walk_expr(self, expr);
                 self.depth -= 1;
@@ -84,60 +83,45 @@ impl<'a, 'v> Visitor<'v> for CollectFreevarsVisitor<'a> {
                     ast::CaptureByValue => CaptureByValue,
                     ast::CaptureByRef => CaptureByRef,
                 };
-                self.capture_mode_map.insert(expr.id, capture_mode);
+                self.cx.capture_mode_map.insert(expr.id, capture_mode);
                 self.depth += 1;
                 visit::walk_expr(self, expr);
                 self.depth -= 1;
             }
             ast::ExprPath(..) => {
-                let mut def = *self.def_map.borrow().find(&expr.id)
-                                                    .expect("path not found");
-                let mut i = 0;
-                while i < self.depth {
-                    match def {
-                        def::DefUpvar(_, inner, _, _) => { def = *inner; }
-                        _ => break
-                    }
-                    i += 1;
-                }
-                if i == self.depth { // Made it to end of loop
-                    let dnum = def.def_id().node;
-                    if !self.seen.contains(&dnum) {
-                        self.refs.push(freevar_entry {
-                            def: def,
-                            span: expr.span,
-                        });
-                        self.seen.insert(dnum);
-                    }
+                let def = *self.cx.def_map.borrow().find(&expr.id)
+                                                   .expect("path not found");
+                let dnum = def.def_id().node;
+                if self.seen.contains(&dnum) {
+                    return;
                 }
+                let def = match def {
+                    def::DefUpvar(_, _, depth, _, _) => {
+                        if depth < self.depth {
+                            return;
+                        }
+                        let mut def = def;
+                        for _ in range(0, depth - self.depth) {
+                            match def {
+                                def::DefUpvar(_, inner, _, _, _) => { def = *inner; }
+                                _ => unreachable!()
+                            }
+                        }
+                        def
+                    },
+                    _ => return
+                };
+                self.cx.freevars.find_or_insert(self.node_id, vec![]).push(freevar_entry {
+                    def: def,
+                    span: expr.span,
+                });
+                self.seen.insert(dnum);
             }
             _ => visit::walk_expr(self, expr)
         }
     }
 }
 
-// Searches through part of the AST for all references to locals or
-// upvars in this frame and returns the list of definition IDs thus found.
-// Since we want to be able to collect upvars in some arbitrary piece
-// of the AST, we take a walker function that we invoke with a visitor
-// in order to start the search.
-fn collect_freevars(def_map: &resolve::DefMap,
-                    blk: &ast::Block,
-                    capture_mode_map: &mut CaptureModeMap)
-                    -> Vec<freevar_entry> {
-    let mut v = CollectFreevarsVisitor {
-        seen: NodeSet::new(),
-        refs: Vec::new(),
-        def_map: def_map,
-        capture_mode_map: &mut *capture_mode_map,
-        depth: 1
-    };
-
-    v.visit_block(blk);
-
-    v.refs
-}
-
 struct AnnotateFreevarsVisitor<'a> {
     def_map: &'a resolve::DefMap,
     freevars: freevar_map,
@@ -147,10 +131,12 @@ struct AnnotateFreevarsVisitor<'a> {
 impl<'a, 'v> Visitor<'v> for AnnotateFreevarsVisitor<'a> {
     fn visit_fn(&mut self, fk: visit::FnKind<'v>, fd: &'v ast::FnDecl,
                 blk: &'v ast::Block, s: Span, nid: ast::NodeId) {
-        let vars = collect_freevars(self.def_map,
-                                    blk,
-                                    &mut self.capture_mode_map);
-        self.freevars.insert(nid, vars);
+        CollectFreevarsVisitor {
+            node_id: nid,
+            seen: NodeSet::new(),
+            cx: self,
+            depth: 0
+        }.visit_block(blk);
         visit::walk_fn(self, fk, fd, blk, s);
     }
 }
@@ -168,13 +154,7 @@ pub fn annotate_freevars(def_map: &resolve::DefMap, krate: &ast::Crate)
         capture_mode_map: NodeMap::new(),
     };
     visit::walk_crate(&mut visitor, krate);
-
-    let AnnotateFreevarsVisitor {
-        freevars,
-        capture_mode_map,
-        ..
-    } = visitor;
-    (freevars, capture_mode_map)
+    (visitor.freevars, visitor.capture_mode_map)
 }
 
 pub fn with_freevars<T>(tcx: &ty::ctxt, fid: ast::NodeId, f: |&[freevar_entry]| -> T) -> T {

src/librustc/middle/mem_categorization.rs

@@ -547,7 +547,7 @@ impl<'t,'tcx,TYPER:Typer<'tcx>> MemCategorizationContext<'t,TYPER> {
               }))
           }
 
-          def::DefUpvar(var_id, _, fn_node_id, _) => {
+          def::DefUpvar(var_id, _, _, fn_node_id, _) => {
               let ty = if_ok!(self.node_ty(fn_node_id));
               match ty::get(ty).sty {
                   ty::ty_closure(ref closure_ty) => {

src/librustc/middle/resolve.rs

@@ -277,9 +277,13 @@ enum RibKind {
     // No translation needs to be applied.
     NormalRibKind,
 
-    // We passed through a function scope at the given node ID. Translate
-    // upvars as appropriate.
-    FunctionRibKind(NodeId /* func id */, NodeId /* body id */),
+    // We passed through a closure scope at the given node ID.
+    // Translate upvars as appropriate.
+    ClosureRibKind(NodeId /* func id */),
+
+    // We passed through a proc or unboxed closure scope at the given node ID.
+    // Translate upvars as appropriate.
+    ProcRibKind(NodeId /* func id */, NodeId /* body id */),
 
     // We passed through an impl or trait and are now in one of its
     // methods. Allow references to ty params that impl or trait
@@ -3859,12 +3863,22 @@ impl<'a> Resolver<'a> {
                 NormalRibKind => {
                     // Nothing to do. Continue.
                 }
-                FunctionRibKind(function_id, body_id) => {
+                ClosureRibKind(function_id) => {
                     if !is_ty_param {
-                        def = DefUpvar(def.def_id().node,
-                                       box(GC) def,
-                                       function_id,
-                                       body_id);
+                        let (depth, block_id) = match def {
+                            DefUpvar(_, _, depth, _, block_id) => (depth + 1, block_id),
+                            _ => (0, ast::DUMMY_NODE_ID)
+                        };
+                        def = DefUpvar(def.def_id().node, box(GC) def, depth, function_id, block_id);
+                    }
+                }
+                ProcRibKind(function_id, block_id) => {
+                    if !is_ty_param {
+                        let depth = match def {
+                            DefUpvar(_, _, depth, _, _) => depth + 1,
+                            _ => 0
+                        };
+                        def = DefUpvar(def.def_id().node, box(GC) def, depth, function_id, block_id);
                     }
                 }
                 MethodRibKind(item_id, _) => {
@@ -5758,10 +5772,14 @@ impl<'a> Resolver<'a> {
                 visit::walk_expr(self, expr);
             }
 
-            ExprFnBlock(_, ref fn_decl, ref block) |
+            ExprFnBlock(_, ref fn_decl, ref block) => {
+                self.resolve_function(ClosureRibKind(expr.id),
+                                      Some(&**fn_decl), NoTypeParameters,
+                                      &**block);
+            }
             ExprProc(ref fn_decl, ref block) |
             ExprUnboxedFn(_, _, ref fn_decl, ref block) => {
-                self.resolve_function(FunctionRibKind(expr.id, block.id),
+                self.resolve_function(ProcRibKind(expr.id, block.id),
                                       Some(&**fn_decl), NoTypeParameters,
                                       &**block);
             }

src/librustc/middle/save/mod.rs

@@ -232,7 +232,7 @@ impl <'l, 'tcx> DxrVisitor<'l, 'tcx> {
             def::DefStatic(_, _) |
             def::DefLocal(_) |
             def::DefVariant(_, _, _) |
-            def::DefUpvar(_, _, _, _) => Some(recorder::VarRef),
+            def::DefUpvar(..) => Some(recorder::VarRef),
 
             def::DefFn(_, _) => Some(recorder::FnRef),
 

src/librustc/middle/trans/_match.rs

@@ -1226,7 +1226,7 @@ pub fn trans_match<'blk, 'tcx>(bcx: Block<'blk, 'tcx>,
 fn is_discr_reassigned(bcx: Block, discr: &ast::Expr, body: &ast::Expr) -> bool {
     match discr.node {
         ast::ExprPath(..) => match bcx.def(discr.id) {
-            def::DefLocal(vid) | def::DefUpvar(vid, _, _, _) => {
+            def::DefLocal(vid) | def::DefUpvar(vid, _, _, _, _) => {
                 let mut rc = ReassignmentChecker {
                     node: vid,
                     reassigned: false

src/librustc/middle/trans/expr.rs

@@ -1176,7 +1176,7 @@ pub fn trans_local_var<'blk, 'tcx>(bcx: Block<'blk, 'tcx>,
     let _icx = push_ctxt("trans_local_var");
 
     match def {
-        def::DefUpvar(nid, _, _, _) => {
+        def::DefUpvar(nid, _, _, _, _) => {
             // Can't move upvars, so this is never a ZeroMemLastUse.
             let local_ty = node_id_type(bcx, nid);
             match bcx.fcx.llupvars.borrow().find(&nid) {

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

@@ -5027,7 +5027,7 @@ pub fn polytype_for_def(fcx: &FnCtxt,
                         defn: def::Def)
                         -> Polytype {
     match defn {
-      def::DefLocal(nid) => {
+      def::DefLocal(nid) | def::DefUpvar(nid, _, _, _, _) => {
           let typ = fcx.local_ty(sp, nid);
           return no_params(typ);
       }
@@ -5036,9 +5036,6 @@ pub fn polytype_for_def(fcx: &FnCtxt,
       def::DefStruct(id) => {
         return ty::lookup_item_type(fcx.ccx.tcx, id);
       }
-      def::DefUpvar(_, inner, _, _) => {
-        return polytype_for_def(fcx, sp, *inner);
-      }
       def::DefTrait(_) |
       def::DefTy(..) |
       def::DefAssociatedTy(..) |

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

@@ -245,10 +245,11 @@ fn region_of_def(fcx: &FnCtxt, def: def::Def) -> ty::Region {
         def::DefLocal(node_id) => {
             tcx.region_maps.var_region(node_id)
         }
-        def::DefUpvar(_, subdef, closure_id, body_id) => {
-            match ty::ty_closure_store(fcx.node_ty(closure_id)) {
-                ty::RegionTraitStore(..) => region_of_def(fcx, *subdef),
-                ty::UniqTraitStore => ReScope(body_id)
+        def::DefUpvar(node_id, _, _, _, body_id) => {
+            if body_id == ast::DUMMY_NODE_ID {
+                tcx.region_maps.var_region(node_id)
+            } else {
+                ReScope(body_id)
             }
         }
         _ => {
@@ -1029,7 +1030,7 @@ fn check_expr_fn_block(rcx: &mut Rcx,
             // determining the final borrow_kind) and propagate that as
             // a constraint on the outer closure.
             match freevar.def {
-                def::DefUpvar(var_id, _, outer_closure_id, _) => {
+                def::DefUpvar(var_id, _, _, outer_closure_id, _) => {
                     // thing being captured is itself an upvar:
                     let outer_upvar_id = ty::UpvarId {
                         var_id: var_id,