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                 let def_map = self.def_map.borrow();
  55                 match def_map.get().find(&expr.id) {
  56                     None => fail!("path not found"),
  57                     Some(&df) => {
  58                         let mut def = df;
  59                         while i < depth {
  60                             match def {
  61                                 ast::DefUpvar(_, inner, _, _) => { def = *inner; }
  62                                 _ => break
  63                             }
  64                             i += 1;
  65                         }
  66                         if i == depth { // Made it to end of loop
  67                             let dnum = ast_util::def_id_of_def(def).node;
  68                             if !self.seen.contains(&dnum) {
  69                                 self.refs.push(@freevar_entry {
  70                                     def: def,
  71                                     span: expr.span,
  72                                 });
  73                                 self.seen.insert(dnum);
  74                             }
  75                         }
  76                     }
  77                 }
  78             }
  79             _ => visit::walk_expr(self, expr, depth)
  80         }
  81     }
  82
  83
  84 }
  85
  86 // Searches through part of the AST for all references to locals or
  87 // upvars in this frame and returns the list of definition IDs thus found.
  88 // Since we want to be able to collect upvars in some arbitrary piece
  89 // of the AST, we take a walker function that we invoke with a visitor
  90 // in order to start the search.
  91 fn collect_freevars(def_map: resolve::DefMap, blk: &ast::Block) -> freevar_info {
  92     let seen = NodeSet::new();
  93     let refs = Vec::new();
  94
  95     let mut v = CollectFreevarsVisitor {
  96         seen: seen,
  97         refs: refs,
  98         def_map: def_map,
  99     };
 100
 101     v.visit_block(blk, 1);
 102     let CollectFreevarsVisitor {
 103         refs,
 104         ..
 105     } = v;
 106     return @refs;
 107 }
 108
 109 struct AnnotateFreevarsVisitor {
 110     def_map: resolve::DefMap,
 111     freevars: freevar_map,
 112 }
 113
 114 impl Visitor<()> for AnnotateFreevarsVisitor {
 115     fn visit_fn(&mut self, fk: &visit::FnKind, fd: &ast::FnDecl,
 116                 blk: &ast::Block, s: Span, nid: ast::NodeId, _: ()) {
 117         let vars = collect_freevars(self.def_map, blk);
 118         self.freevars.insert(nid, vars);
 119         visit::walk_fn(self, fk, fd, blk, s, nid, ());
 120     }
 121 }
 122
 123 // Build a map from every function and for-each body to a set of the
 124 // freevars contained in it. The implementation is not particularly
 125 // efficient as it fully recomputes the free variables at every
 126 // node of interest rather than building up the free variables in
 127 // one pass. This could be improved upon if it turns out to matter.
 128 pub fn annotate_freevars(def_map: resolve::DefMap, krate: &ast::Crate) ->
 129    freevar_map {
 130     let mut visitor = AnnotateFreevarsVisitor {
 131         def_map: def_map,
 132         freevars: NodeMap::new(),
 133     };
 134     visit::walk_crate(&mut visitor, krate, ());
 135
 136     let AnnotateFreevarsVisitor {
 137         freevars,
 138         ..
 139     } = visitor;
 140     freevars
 141 }
 142
 143 pub fn get_freevars(tcx: &ty::ctxt, fid: ast::NodeId) -> freevar_info {
 144     let freevars = tcx.freevars.borrow();
 145     match freevars.get().find(&fid) {
 146         None => fail!("get_freevars: {} has no freevars", fid),
 147         Some(&d) => return d
 148     }
 149 }
 150
 151 pub fn has_freevars(tcx: &ty::ctxt, fid: ast::NodeId) -> bool {
 152     !get_freevars(tcx, fid).is_empty()
 153 }