Add freevars finding pass

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, _));

src/comp/middle/freevars.rs

@@ -0,0 +1,138 @@
+// 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::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_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,
+             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) {
+        for (ast::arg a in f.decl.inputs) { e.decls.insert(a.id, ()); }
+    }
+    fn walk_expr(env e, &@ast::expr expr) {
+        alt (expr.node) {
+            case (ast::expr_path(?path)) {
+                if (! e.def_map.contains_key(expr.id)) {
+                    e.sess.span_fatal(expr.span,
+                       "internal error in collect_freevars");
+                }
+                alt (e.def_map.get(expr.id)) {
+                    case (ast::def_arg(?did)) { e.refs += ~[did._1]; }
+                    case (ast::def_local(?did)) { e.refs += ~[did._1]; }
+                    case (ast::def_binding(?did)) { e.refs += ~[did._1]; }
+                    case (_) { /* no-op */ }
+                }
+            }
+            case (_) { }
+        }
+    }
+    fn walk_local(env e, &@ast::local local) {
+        e.decls.insert(local.node.id, ());
+    }
+    fn walk_pat(env e, &@ast::pat p) {
+        alt (p.node) {
+            case (ast::pat_bind(_)) {
+                e.decls.insert(p.id, ());
+            }
+            case (_) {}
+        }
+    }
+    let hashmap[ast::node_id, ()] decls = new_int_hash[()]();
+    for (ast::node_id decl in initial_decls) { decls.insert(decl, ()); }
+
+    let env e =
+        @rec(mutable refs=~[],
+             decls=decls,
+             def_map=def_map,
+             sess=sess);
+    auto visitor =
+        @rec(visit_fn_pre=bind walk_fn(e, _, _, _, _, _),
+             visit_local_pre=bind walk_local(e, _),
+             visit_expr_pre=bind walk_expr(e, _),
+             visit_pat_pre=bind walk_pat(e, _)
+             with walk::default_visitor());
+    walker(*visitor);
+
+    // 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_;
+        if (!decls.contains_key(ref_id)) { result += ~[ref_id]; }
+    }
+    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;
+// indent-tabs-mode: nil
+// c-basic-offset: 4
+// buffer-file-coding-system: utf-8-unix
+// compile-command: "make -k -C $RBUILD 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
+// End:

src/comp/middle/trans.rs

@@ -26,6 +26,7 @@ import syntax::ast;
 import syntax::walk;
 import driver::session;
 import middle::ty;
+import middle::freevars;
 import back::link;
 import back::x86;
 import back::abi;
@@ -4080,76 +4081,6 @@ fn trans_for(&@block_ctxt cx, &@ast::local local, &@ast::expr seq,
 
 // Iterator translation
 
-// 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(&@block_ctxt cx, &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,
-             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) {
-        for (ast::arg a in f.decl.inputs) { e.decls.insert(a.id, ()); }
-    }
-    fn walk_expr(env e, &@ast::expr expr) {
-        alt (expr.node) {
-            case (ast::expr_path(?path)) {
-                if (! e.def_map.contains_key(expr.id)) {
-                    e.sess.span_fatal(expr.span,
-                       "internal error in collect_upvars");
-                }
-                alt (e.def_map.get(expr.id)) {
-                    case (ast::def_arg(?did)) { e.refs += ~[did._1]; }
-                    case (ast::def_local(?did)) { e.refs += ~[did._1]; }
-                    case (ast::def_binding(?did)) { e.refs += ~[did._1]; }
-                    case (_) { /* no-op */ }
-                }
-            }
-            case (_) { }
-        }
-    }
-    fn walk_local(env e, &@ast::local local) {
-        e.decls.insert(local.node.id, ());
-    }
-    fn walk_pat(env e, &@ast::pat p) {
-        alt (p.node) {
-            case (ast::pat_bind(_)) {
-                e.decls.insert(p.id, ());
-            }
-            case (_) {}
-        }
-    }
-    let hashmap[ast::node_id, ()] decls = new_int_hash[()]();
-    for (ast::node_id decl in initial_decls) { decls.insert(decl, ()); }
-
-    let env e =
-        @rec(mutable refs=~[],
-             decls=decls,
-             def_map=cx.fcx.lcx.ccx.tcx.def_map,
-             sess=cx.fcx.lcx.ccx.tcx.sess);
-    auto visitor =
-        @rec(visit_fn_pre=bind walk_fn(e, _, _, _, _, _),
-             visit_local_pre=bind walk_local(e, _),
-             visit_expr_pre=bind walk_expr(e, _),
-             visit_pat_pre=bind walk_pat(e, _)
-             with walk::default_visitor());
-    walker(*visitor);
-    // 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_;
-        if (!decls.contains_key(ref_id)) { result += ~[ref_id]; }
-    }
-    ret result;
-}
-
 // Finds the ValueRef associated with a variable in a function
 // context. It checks locals, upvars, and args.
 fn find_variable(&@fn_ctxt fcx, ast::node_id nid) -> ValueRef {
@@ -4318,16 +4249,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 = collect_upvars(cx, 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;

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(),

src/comp/rustc.rc

@@ -26,6 +26,7 @@ mod middle {
     mod resolve;
     mod typeck;
     mod alias;
+    mod freevars;
 
     mod tstate {
         mod ck;