src/librustc/middle/freevars.rs

   1 // Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 // A pass that annotates for each loops and functions with the free
  12 // variables that they contain.
  13
  14 #![allow(non_camel_case_types)]
  15
  16 use middle::resolve;
  17 use middle::ty;
  18 use util::nodemap::{NodeMap, NodeSet};
  19
  20 use syntax::codemap::Span;
  21 use syntax::{ast, ast_util};
  22 use syntax::visit;
  23 use syntax::visit::Visitor;
  24
  25 // A vector of defs representing the free variables referred to in a function.
  26 // (The def_upvar will already have been stripped).
  27 #[deriving(Encodable, Decodable)]
  28 pub struct freevar_entry {
  29     def: ast::Def, //< The variable being accessed free.
  30     span: Span     //< First span where it is accessed (there can be multiple)
  31 }
  32 pub type freevar_info = @Vec<@freevar_entry> ;
  33 pub type freevar_map = NodeMap<freevar_info>;
  34
  35 struct CollectFreevarsVisitor {
  36     seen: NodeSet,
  37     refs: Vec<@freevar_entry> ,
  38     def_map: resolve::DefMap,
  39 }
  40
  41 impl Visitor<int> for CollectFreevarsVisitor {
  42
  43     fn visit_item(&mut self, _: &ast::Item, _: int) {
  44         // ignore_item
  45     }
  46
  47     fn visit_expr(&mut self, expr: &ast::Expr, depth: int) {
  48         match expr.node {
  49             ast::ExprFnBlock(..) | ast::ExprProc(..) => {
  50                 visit::walk_expr(self, expr, depth + 1)
  51             }
  52             ast::ExprPath(..) => {
  53                 let mut i = 0;
  54                 match self.def_map.borrow().find(&expr.id) {
  55                     None => fail!("path not found"),
  56                     Some(&df) => {
  57                         let mut def = df;
  58                         while i < depth {
  59                             match def {
  60                                 ast::DefUpvar(_, inner, _, _) => { def = *inner; }
  61                                 _ => break
  62                             }
  63                             i += 1;
  64                         }
  65                         if i == depth { // Made it to end of loop
  66                             let dnum = ast_util::def_id_of_def(def).node;
  67                             if !self.seen.contains(&dnum) {
  68                                 self.refs.push(@freevar_entry {
  69                                     def: def,
  70                                     span: expr.span,
  71                                 });
  72                                 self.seen.insert(dnum);
  73                             }
  74                         }
  75                     }
  76                 }
  77             }
  78             _ => visit::walk_expr(self, expr, depth)
  79         }
  80     }
  81
  82
  83 }
  84
  85 // Searches through part of the AST for all references to locals or
  86 // upvars in this frame and returns the list of definition IDs thus found.
  87 // Since we want to be able to collect upvars in some arbitrary piece
  88 // of the AST, we take a walker function that we invoke with a visitor
  89 // in order to start the search.
  90 fn collect_freevars(def_map: resolve::DefMap, blk: &ast::Block) -> freevar_info {
  91     let seen = NodeSet::new();
  92     let refs = Vec::new();
  93
  94     let mut v = CollectFreevarsVisitor {
  95         seen: seen,
  96         refs: refs,
  97         def_map: def_map,
  98     };
  99
 100     v.visit_block(blk, 1);
 101     let CollectFreevarsVisitor {
 102         refs,
 103         ..
 104     } = v;
 105     return @refs;
 106 }
 107
 108 struct AnnotateFreevarsVisitor {
 109     def_map: resolve::DefMap,
 110     freevars: freevar_map,
 111 }
 112
 113 impl Visitor<()> for AnnotateFreevarsVisitor {
 114     fn visit_fn(&mut self, fk: &visit::FnKind, fd: &ast::FnDecl,
 115                 blk: &ast::Block, s: Span, nid: ast::NodeId, _: ()) {
 116         let vars = collect_freevars(self.def_map, blk);
 117         self.freevars.insert(nid, vars);
 118         visit::walk_fn(self, fk, fd, blk, s, nid, ());
 119     }
 120 }
 121
 122 // Build a map from every function and for-each body to a set of the
 123 // freevars contained in it. The implementation is not particularly
 124 // efficient as it fully recomputes the free variables at every
 125 // node of interest rather than building up the free variables in
 126 // one pass. This could be improved upon if it turns out to matter.
 127 pub fn annotate_freevars(def_map: resolve::DefMap, krate: &ast::Crate) ->
 128    freevar_map {
 129     let mut visitor = AnnotateFreevarsVisitor {
 130         def_map: def_map,
 131         freevars: NodeMap::new(),
 132     };
 133     visit::walk_crate(&mut visitor, krate, ());
 134
 135     let AnnotateFreevarsVisitor {
 136         freevars,
 137         ..
 138     } = visitor;
 139     freevars
 140 }
 141
 142 pub fn get_freevars(tcx: &ty::ctxt, fid: ast::NodeId) -> freevar_info {
 143     match tcx.freevars.borrow().find(&fid) {
 144         None => fail!("get_freevars: {} has no freevars", fid),
 145         Some(&d) => return d
 146     }
 147 }
 148
 149 pub fn has_freevars(tcx: &ty::ctxt, fid: ast::NodeId) -> bool {
 150     !get_freevars(tcx, fid).is_empty()
 151 }