Add a pass that finds all of the free variables.

4 files changed, 69 insertions, 18 deletions

diff --git a/src/comp/driver/rustc.rs b/src/comp/driver/rustc.rs
index b810fbecc72..b24dedffdd0 100644
--- a/src/comp/driver/rustc.rs
+++ b/src/comp/driver/rustc.rs
@@ -10,6 +10,7 @@ import syntax::codemap;
 import front::attr;
 import middle::trans;
 import middle::resolve;
+import middle::freevars;
 import middle::ty;
 import middle::typeck;
 import middle::tstate::ck;
@@ -131,7 +132,10 @@ fn compile_input(session::session sess, ast::crate_cfg cfg, str input,
     auto d =
         time(time_passes, "resolution",
              bind resolve::resolve_crate(sess, ast_map, crate));
-    auto ty_cx = ty::mk_ctxt(sess, d._0, d._1, ast_map);
+    auto freevars =
+        time(time_passes, "freevar finding",
+             bind freevars::annotate_freevars(sess, d._0, crate));
+    auto ty_cx = ty::mk_ctxt(sess, d._0, d._1, ast_map, freevars);
     time[()](time_passes, "typechecking",
              bind typeck::check_crate(ty_cx, crate));
     if (sess.get_opts().run_typestate) {
@@ -196,7 +200,8 @@ fn pretty_print_input(session::session sess, ast::crate_cfg cfg,
         case (ppm_typed) {
             auto amap = middle::ast_map::map_crate(*crate);
             auto d = resolve::resolve_crate(sess, amap, crate);
-            auto ty_cx = ty::mk_ctxt(sess, d._0, d._1, amap);
+            auto freevars = freevars::annotate_freevars(sess, d._0, crate);
+            auto ty_cx = ty::mk_ctxt(sess, d._0, d._1, amap, freevars);
             typeck::check_crate(ty_cx, crate);
             ann = rec(pre=ann_paren_for_expr,
                       post=bind ann_typed_post(ty_cx, _));
diff --git a/src/comp/middle/freevars.rs b/src/comp/middle/freevars.rs
index bf2a8f7e2c9..4b88f2945c8 100644
--- a/src/comp/middle/freevars.rs
+++ b/src/comp/middle/freevars.rs
@@ -1,23 +1,29 @@
-// A pass that annotates for each loops with the free variables that
-// they contain.
+// A pass that annotates for each loops and functions with the free
+// variables that they contain.
 
 import std::map;
 import std::map::*;
 import syntax::ast;
 import syntax::walk;
 import driver::session;
-import middle::ty;
+import middle::resolve;
 import syntax::codemap::span;
 
+export annotate_freevars;
+export freevar_set;
+export freevar_map;
 
+type freevar_set = ast::node_id[];
+type freevar_map = hashmap[ast::node_id, freevar_set];
 
 // 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_upvars(&ty::ctxt tcx, &fn (&walk::ast_visitor) walker,
-                  ast::node_id[] initial_decls) -> ast::node_id[] {
+fn collect_freevars(&resolve::def_map def_map, &session::session sess,
+                    &fn (&walk::ast_visitor) walker,
+                    ast::node_id[] initial_decls) -> ast::node_id[] {
     type env =
         @rec(mutable ast::node_id[] refs,
              hashmap[ast::node_id, ()] decls,
@@ -33,7 +39,7 @@ fn collect_upvars(&ty::ctxt tcx, &fn (&walk::ast_visitor) walker,
             case (ast::expr_path(?path)) {
                 if (! e.def_map.contains_key(expr.id)) {
                     e.sess.span_fatal(expr.span,
-                       "internal error in collect_upvars");
+                       "internal error in collect_freevars");
                 }
                 alt (e.def_map.get(expr.id)) {
                     case (ast::def_arg(?did)) { e.refs += ~[did._1]; }
@@ -62,8 +68,8 @@ fn collect_upvars(&ty::ctxt tcx, &fn (&walk::ast_visitor) walker,
     let env e =
         @rec(mutable refs=~[],
              decls=decls,
-             def_map=tcx.def_map,
-             sess=tcx.sess);
+             def_map=def_map,
+             sess=sess);
     auto visitor =
         @rec(visit_fn_pre=bind walk_fn(e, _, _, _, _, _),
              visit_local_pre=bind walk_local(e, _),
@@ -71,8 +77,8 @@ fn collect_upvars(&ty::ctxt tcx, &fn (&walk::ast_visitor) walker,
              visit_pat_pre=bind walk_pat(e, _)
              with walk::default_visitor());
     walker(*visitor);
-    // Calculate (refs - decls). This is the set of captured upvars.
 
+    // Calculate (refs - decls). This is the set of captured upvars.
     let ast::node_id[] result = ~[];
     for (ast::node_id ref_id_ in e.refs) {
         auto ref_id = ref_id_;
@@ -81,6 +87,47 @@ fn collect_upvars(&ty::ctxt tcx, &fn (&walk::ast_visitor) walker,
     ret result;
 }
 
+// Build a map from every function and for-each body to a set of the
+// freevars contained in it. The implementation is not particularly
+// efficient as it fully recomputes the free variables at every
+// node of interest rather than building up the free variables in
+// one pass. This could be improved upon if it turns out to matter.
+fn annotate_freevars(&session::session sess, &resolve::def_map def_map,
+                     &@ast::crate crate) -> freevar_map {
+    type env =
+        rec(freevar_map freevars,
+            resolve::def_map def_map,
+            session::session sess);
+
+    fn walk_fn(env e, &ast::_fn f, &ast::ty_param[] tps, &span sp,
+               &ast::fn_ident i, ast::node_id nid) {
+        auto walker = bind walk::walk_fn(_, f, tps, sp, i, nid);
+        auto vars = collect_freevars(e.def_map, e.sess, walker, ~[]);
+        e.freevars.insert(nid, vars);
+    }
+    fn walk_expr(env e, &@ast::expr expr) {
+        alt (expr.node) {
+            ast::expr_for_each(?local, _, ?body) {
+                auto vars = collect_freevars(e.def_map, e.sess,
+                                             bind walk::walk_block(_, body),
+                                             ~[local.node.id]);
+                e.freevars.insert(body.node.id, vars);
+            }
+            _ {}
+        }
+    }
+
+    let env e =
+        rec(freevars = new_int_hash(), def_map=def_map, sess=sess);
+    auto visitor =
+        rec(visit_fn_pre=bind walk_fn(e, _, _, _, _, _),
+            visit_expr_pre=bind walk_expr(e, _)
+            with walk::default_visitor());
+    walk::walk_crate(visitor, *crate);
+
+    ret e.freevars;
+}
+
 // Local Variables:
 // mode: rust
 // fill-column: 78;
diff --git a/src/comp/middle/trans.rs b/src/comp/middle/trans.rs
index 710b91a18a9..ee290dc18d5 100644
--- a/src/comp/middle/trans.rs
+++ b/src/comp/middle/trans.rs
@@ -4276,17 +4276,14 @@ fn trans_for_each(&@block_ctxt cx, &@ast::local local, &@ast::expr seq,
      *
      */
 
-    // Step 1: walk body and figure out which references it makes
-    // escape. This could be determined upstream, and probably ought
-    // to be so, eventualy.
+    // Step 1: Generate code to build an environment containing pointers
+    // to all of the upvars
     auto lcx = cx.fcx.lcx;
 
     // FIXME: possibly support alias-mode here?
     auto decl_ty = node_id_type(lcx.ccx, local.node.id);
     auto decl_id = local.node.id;
-    auto upvars = freevars::collect_upvars(cx.fcx.lcx.ccx.tcx,
-                                           bind walk::walk_block(_, body),
-                                           ~[decl_id]);
+    auto upvars = cx.fcx.lcx.ccx.tcx.freevars.get(body.node.id);
 
     auto environment_data = build_environment(cx, upvars);
     auto llenvptr = environment_data._0;
diff --git a/src/comp/middle/ty.rs b/src/comp/middle/ty.rs
index b5242900faf..30c04a74922 100644
--- a/src/comp/middle/ty.rs
+++ b/src/comp/middle/ty.rs
@@ -210,6 +210,7 @@ type ctxt =
         resolve::def_map def_map,
         node_type_table node_types,
         ast_map::map items,
+        freevars::freevar_map freevars,
 
         constr_table fn_constrs,
         type_cache tcache,
@@ -392,7 +393,7 @@ fn mk_rcache() -> creader_cache {
 }
 
 fn mk_ctxt(session::session s, resolve::def_map dm, constr_table cs,
-           ast_map::map amap) -> ctxt {
+           ast_map::map amap, freevars::freevar_map freevars) -> ctxt {
     let node_type_table ntt =
         @smallintmap::mk[ty::ty_param_substs_opt_and_ty]();
     auto tcache = new_def_hash[ty::ty_param_count_and_ty]();
@@ -403,6 +404,7 @@ fn mk_ctxt(session::session s, resolve::def_map dm, constr_table cs,
             def_map=dm,
             node_types=ntt,
             items=amap,
+            freevars=freevars,
             fn_constrs=cs,
             tcache=tcache,
             rcache=mk_rcache(),