Overhaul how we handle freevars.

2 files changed, 39 insertions, 26 deletions

diff --git a/src/comp/middle/freevars.rs b/src/comp/middle/freevars.rs
index 00ae8f2bf26..2b733e7ba50 100644
--- a/src/comp/middle/freevars.rs
+++ b/src/comp/middle/freevars.rs
@@ -16,13 +16,16 @@ import syntax::codemap::span;
 export annotate_freevars;
 export freevar_set;
 export freevar_map;
+export get_freevar_info;
 export get_freevars;
+export get_freevar_uses;
 export has_freevars;
 export is_freevar_of;
 export def_lookup;
 
-type freevar_set = @ast::node_id[];
-type freevar_map = hashmap[ast::node_id, freevar_set];
+type freevar_set = hashset[ast::node_id];
+type freevar_info = rec(freevar_set defs, @ast::node_id[] uses);
+type freevar_map = hashmap[ast::node_id, freevar_info];
 
 // 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.
@@ -31,7 +34,7 @@ type freevar_map = hashmap[ast::node_id, freevar_set];
 // 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) -> freevar_set {
+                    ast::node_id[] initial_decls) -> freevar_info {
     type env =
         @rec(mutable ast::node_id[] refs,
              hashset[ast::node_id] decls,
@@ -50,9 +53,9 @@ fn collect_freevars(&resolve::def_map def_map, &session::session sess,
                        "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 (ast::def_arg(?did)) { e.refs += ~[expr.id]; }
+                    case (ast::def_local(?did)) { e.refs += ~[expr.id]; }
+                    case (ast::def_binding(?did)) { e.refs += ~[expr.id]; }
                     case (_) { /* no-op */ }
                 }
             }
@@ -87,12 +90,19 @@ fn collect_freevars(&resolve::def_map def_map, &session::session sess,
     walker(*visitor);
 
     // Calculate (refs - decls). This is the set of captured upvars.
-    auto result = ~[];
+    // We build a vec of the node ids of the uses and a set of the
+    // node ids of the definitions.
+    auto uses = ~[];
+    auto defs = new_int_hash();
     for (ast::node_id ref_id_ in e.refs) {
         auto ref_id = ref_id_;
-        if (!decls.contains_key(ref_id)) { result += ~[ref_id]; }
+        auto def_id = ast::def_id_of_def(def_map.get(ref_id))._1;
+        if !decls.contains_key(def_id) {
+            uses += ~[ref_id];
+            set_add(defs, def_id);
+        }
     }
-    ret @result;
+    ret rec(defs=defs, uses=@uses);
 }
 
 // Build a map from every function and for-each body to a set of the
@@ -136,7 +146,7 @@ fn annotate_freevars(&session::session sess, &resolve::def_map def_map,
     ret e.freevars;
 }
 
-fn get_freevars(&ty::ctxt tcx, ast::node_id fid) -> freevar_set {
+fn get_freevar_info(&ty::ctxt tcx, ast::node_id fid) -> freevar_info {
     alt (tcx.freevars.find(fid)) {
         none {
             fail "get_freevars: " + int::str(fid) + " has no freevars";
@@ -144,11 +154,17 @@ fn get_freevars(&ty::ctxt tcx, ast::node_id fid) -> freevar_set {
         some(?d) { ret d; }
     }
 }
+fn get_freevars(&ty::ctxt tcx, ast::node_id fid) -> freevar_set {
+    ret get_freevar_info(tcx, fid).defs;
+}
+fn get_freevar_uses(&ty::ctxt tcx, ast::node_id fid) -> @ast::node_id[] {
+    ret get_freevar_info(tcx, fid).uses;
+}
 fn has_freevars(&ty::ctxt tcx, ast::node_id fid) -> bool {
-    ret ivec::len(*get_freevars(tcx, fid)) != 0u;
+    ret get_freevars(tcx, fid).size() != 0u;
 }
 fn is_freevar_of(&ty::ctxt tcx, ast::node_id var, ast::node_id f) -> bool {
-    ret ivec::member(var, *get_freevars(tcx, f));
+    ret get_freevars(tcx, f).contains_key(var);
 }
 fn def_lookup(&ty::ctxt tcx, ast::node_id f, ast::node_id id) ->
     option::t[ast::def] {
diff --git a/src/comp/middle/trans.rs b/src/comp/middle/trans.rs
index 52e5ac2bb75..ded2a0108f6 100644
--- a/src/comp/middle/trans.rs
+++ b/src/comp/middle/trans.rs
@@ -3792,14 +3792,12 @@ fn find_variable(&@fn_ctxt fcx, ast::node_id nid) -> ValueRef {
 // Given a block context and a list of upvars, construct a closure that
 // contains pointers to all of the upvars and all of the tydescs in
 // scope. Return the ValueRef and TypeRef corresponding to the closure.
-fn build_environment(&@block_ctxt cx, &ast::node_id[] upvars) ->
+fn build_environment(&@block_ctxt cx, &freevar_set upvars) ->
     rec(ValueRef ptr, TypeRef ptrty) {
-    auto upvar_count = std::ivec::len(upvars);
     auto has_iterbody = !option::is_none(cx.fcx.lliterbody);
-    if (has_iterbody) { upvar_count += 1u; }
     auto llbindingsptr;
 
-    if (upvar_count > 0u) {
+    if (upvars.size() > 0u || has_iterbody) {
         // Gather up the upvars.
         let ValueRef[] llbindings = ~[];
         let TypeRef[] llbindingtys = ~[];
@@ -3807,7 +3805,7 @@ fn build_environment(&@block_ctxt cx, &ast::node_id[] upvars) ->
             llbindings += ~[option::get(cx.fcx.lliterbody)];
             llbindingtys += ~[val_ty(llbindings.(0))];
         }
-        for (ast::node_id nid in upvars) {
+        for each (ast::node_id nid in upvars.keys()) {
             auto llbinding = find_variable(cx.fcx, nid);
             llbindings += ~[llbinding];
             llbindingtys += ~[val_ty(llbinding)];
@@ -3815,6 +3813,7 @@ fn build_environment(&@block_ctxt cx, &ast::node_id[] upvars) ->
 
         // Create an array of bindings and copy in aliases to the upvars.
         llbindingsptr = alloca(cx, T_struct(llbindingtys));
+        auto upvar_count = std::ivec::len(llbindings);
         auto i = 0u;
         while (i < upvar_count) {
             auto llbindingptr =
@@ -3859,7 +3858,7 @@ fn build_environment(&@block_ctxt cx, &ast::node_id[] upvars) ->
 // and a list of upvars, generate code to load and populate the environment
 // with the upvars and type descriptors.
 fn load_environment(&@block_ctxt cx, &@fn_ctxt fcx,
-                    TypeRef llenvptrty, &ast::node_id[] upvars) {
+                    TypeRef llenvptrty, &freevar_set upvars) {
     auto copy_args_bcx = new_raw_block_ctxt(fcx, fcx.llcopyargs);
 
     // Populate the upvars from the environment.
@@ -3871,22 +3870,20 @@ fn load_environment(&@block_ctxt cx, &@fn_ctxt fcx,
                                   C_int(abi::closure_elt_bindings)]);
     auto llremotebindingsptr =
         copy_args_bcx.build.Load(llremotebindingsptrptr);
-    auto base = 0u;
+
     auto i = 0u;
-    auto end = std::ivec::len(upvars);
     if (!option::is_none(cx.fcx.lliterbody)) {
-        base += 1u;
+        i += 1u;
         auto lliterbodyptr =
             copy_args_bcx.build.GEP(llremotebindingsptr,
                                     ~[C_int(0), C_int(0)]);
         auto lliterbody = copy_args_bcx.build.Load(lliterbodyptr);
         fcx.lliterbody = some(lliterbody);
     }
-    while (i < end) {
-        auto upvar_id = upvars.(i);
+    for each (ast::node_id upvar_id in upvars.keys()) {
         auto llupvarptrptr =
             copy_args_bcx.build.GEP(llremotebindingsptr,
-                                    ~[C_int(0), C_int(base+i as int)]);
+                                    ~[C_int(0), C_int(i as int)]);
         auto llupvarptr = copy_args_bcx.build.Load(llupvarptrptr);
         fcx.llupvars.insert(upvar_id, llupvarptr);
         i += 1u;
@@ -3943,7 +3940,7 @@ fn trans_for_each(&@block_ctxt cx, &@ast::local local, &@ast::expr seq,
     auto decl_id = local.node.id;
     auto upvars = get_freevars(lcx.ccx.tcx, body.node.id);
 
-    auto llenv = build_environment(cx, *upvars);
+    auto llenv = build_environment(cx, upvars);
 
     // Step 2: Declare foreach body function.
     let str s =
@@ -3964,7 +3961,7 @@ fn trans_for_each(&@block_ctxt cx, &@ast::local local, &@ast::expr seq,
 
     // Generate code to load the environment out of the
     // environment pointer.
-    load_environment(cx, fcx, llenv.ptrty, *upvars);
+    load_environment(cx, fcx, llenv.ptrty, upvars);
 
     // Add an upvar for the loop variable alias.
     fcx.llupvars.insert(decl_id, llvm::LLVMGetParam(fcx.llfn, 3u));