rustc: Don't rebuild the AST when writing back locals

src/comp/middle/ty.rs

@@ -1671,17 +1671,7 @@ fn block_ty(&ctxt cx, &node_type_table ntt, &ast::block b) -> t {
 // Returns the type of a pattern as a monotype. Like @expr_ty, this function
 // doesn't provide type parameter substitutions.
 fn pat_ty(&ctxt cx, &node_type_table ntt, &@ast::pat pat) -> t {
-    alt (pat.node) {
-        case (ast::pat_wild(?ann))      { ret ann_to_monotype(cx, ntt, ann); }
-        case (ast::pat_lit(_, ?ann))    { ret ann_to_monotype(cx, ntt, ann); }
-        case (ast::pat_bind(_, _, ?ann)) {
-            ret ann_to_monotype(cx, ntt, ann);
-        }
-        case (ast::pat_tag(_, _, ?ann)) {
-            ret ann_to_monotype(cx, ntt, ann);
-        }
-    }
-    fail;   // not reached
+    ret ann_to_monotype(cx, ntt, pat_ann(pat));
 }
 
 fn item_ann(&@ast::item it) -> ast::ann {
@@ -1768,6 +1758,26 @@ fn expr_has_ty_params(&node_type_table ntt, &@ast::expr expr) -> bool {
     ret ann_has_type_params(ntt, expr_ann(expr));
 }
 
+fn stmt_ann(&@ast::stmt s) -> ast::ann {
+    alt (s.node) {
+        case (ast::stmt_decl(_, ?a)) { ret a; }
+        case (ast::stmt_expr(_, ?a)) { ret a; }
+        case (ast::stmt_crate_directive(_)) {
+            log_err "ty::stmt_ann(): crate directive found";
+            fail;
+        }
+    }
+}
+
+fn pat_ann(&@ast::pat p) -> ast::ann {
+    alt (p.node) {
+        case (ast::pat_wild(?a))        { ret a; }
+        case (ast::pat_bind(_, _, ?a))  { ret a; }
+        case (ast::pat_lit(_, ?a))      { ret a; }
+        case (ast::pat_tag(_, _, ?a))   { ret a; }
+    }
+}
+
 
 // Expression utilities
 

src/comp/middle/typeck.rs

@@ -1495,65 +1495,88 @@ fn pushdown_block(&@fn_ctxt fcx, &ty::t expected, &ast::block bloc) {
 // Local variable resolution: the phase that finds all the types in the AST
 // and replaces opaque "ty_local" types with the resolved local types.
 
-fn writeback_local(&option::t[@fn_ctxt] env, &span sp, &@ast::local local)
-        -> @ast::decl {
-    auto fcx = option::get[@fn_ctxt](env);
-
-    auto local_ty;
-    alt (fcx.locals.find(local.id)) {
-        case (none[ty::t]) {
-            fcx.ccx.sess.span_err(sp, "unable to determine type of local: "
-                                  + local.ident);
-            fail;
-        }
-        case (some[ty::t](?lt)) {
-            local_ty = lt;
+mod writeback {
+    fn wb_local(&@fn_ctxt fcx, &span sp, &@ast::local local) {
+        auto local_ty;
+        alt (fcx.locals.find(local.id)) {
+            case (none[ty::t]) {
+                fcx.ccx.sess.span_err(sp,
+                    "unable to determine type of local: " + local.ident);
+                fail;
+            }
+            case (some[ty::t](?lt)) {
+                local_ty = lt;
+            }
         }
-    }
 
-    auto local_wb = @rec(ann=triv_ann(local.ann.id, local_ty) with *local);
-    write_type_only(fcx.ccx.node_types, local.ann.id, local_ty);
-    ret @fold::respan[ast::decl_](sp, ast::decl_local(local_wb));
-}
+        write_type_only(fcx.ccx.node_types, local.ann.id, local_ty);
+    }
 
-fn resolve_local_types_in_annotation(&option::t[@fn_ctxt] env, &ast::ann ann)
-        -> ast::ann {
-    fn resolver(@fn_ctxt fcx, ty::t typ) -> ty::t {
-        alt (struct(fcx.ccx.tcx, typ)) {
-            case (ty::ty_local(?lid)) { ret fcx.locals.get(lid); }
-            case (_)                 { ret typ; }
+    fn resolve_local_types(&@fn_ctxt fcx, &ast::ann ann) {
+        fn resolver(@fn_ctxt fcx, ty::t typ) -> ty::t {
+            alt (struct(fcx.ccx.tcx, typ)) {
+                case (ty::ty_local(?lid))   { ret fcx.locals.get(lid); }
+                case (_)                    { ret typ; }
+            }
         }
+
+        auto tpot = ty::ann_to_ty_param_substs_opt_and_ty(fcx.ccx.node_types,
+                                                          ann);
+        auto tt = tpot._1;
+        if (!ty::type_contains_locals(fcx.ccx.tcx, tt)) { ret; }
+
+        auto f = bind resolver(fcx, _);
+        auto new_type = ty::fold_ty(fcx.ccx.tcx, f, tt);
+        write_type(fcx.ccx.node_types, ann.id, tup(tpot._0, new_type));
     }
 
-    auto fcx = option::get[@fn_ctxt](env);
+    fn visit_stmt_pre(@fn_ctxt fcx, &@ast::stmt s) {
+        resolve_local_types(fcx, ty::stmt_ann(s));
+    }
 
-    auto tt = ann_to_type(fcx.ccx.node_types, ann);
-    if (!ty::type_contains_locals(fcx.ccx.tcx, tt)) { ret ann; }
+    fn visit_expr_pre(@fn_ctxt fcx, &@ast::expr e) {
+        resolve_local_types(fcx, ty::expr_ann(e));
+    }
 
-    auto f = bind resolver(fcx, _);
-    auto new_type = ty::fold_ty(fcx.ccx.tcx, f,
-                                ann_to_type(fcx.ccx.node_types, ann));
-    write_type(fcx.ccx.node_types, ann.id, tup(ann.tps, new_type));
-    ret mk_ann_type(ann.id, new_type, ann.tps);
-}
+    fn visit_block_pre(@fn_ctxt fcx, &ast::block b) {
+        resolve_local_types(fcx, b.node.a);
+    }
 
-fn resolve_local_types_in_block(&@fn_ctxt fcx, &ast::block block)
-        -> ast::block {
-    fn update_env_for_item(&option::t[@fn_ctxt] env, &@ast::item i)
-            -> option::t[@fn_ctxt] {
-        ret none[@fn_ctxt];
+    fn visit_arm_pre(@fn_ctxt fcx, &ast::arm a) {
+        // FIXME: Need a visit_pat_pre
+        resolve_local_types(fcx, ty::pat_ann(a.pat));
     }
-    fn keep_going(&option::t[@fn_ctxt] env) -> bool {
-        ret !option::is_none[@fn_ctxt](env);
+
+    fn visit_decl_pre(@fn_ctxt fcx, &@ast::decl d) {
+        alt (d.node) {
+            case (ast::decl_local(?l)) { wb_local(fcx, d.span, l); }
+            case (ast::decl_item(_)) { /* no annotation */ }
+        }
     }
 
-    auto fld = fold::new_identity_fold[option::t[@fn_ctxt]]();
-    fld = @rec(fold_decl_local = writeback_local,
-               fold_ann = resolve_local_types_in_annotation,
-               update_env_for_item = update_env_for_item,
-               keep_going = keep_going
-               with *fld);
-    ret fold::fold_block[option::t[@fn_ctxt]](some(fcx), fld, block);
+    fn resolve_local_types_in_block(&@fn_ctxt fcx, &ast::block block) {
+        // A trick to ignore any contained items.
+        auto ignore = @mutable false;
+        fn visit_item_pre(@mutable bool ignore, &@ast::item item) {
+            *ignore = true;
+        }
+        fn visit_item_post(@mutable bool ignore, &@ast::item item) {
+            *ignore = false;
+        }
+        fn keep_going(@mutable bool ignore) -> bool { ret !*ignore; }
+
+        auto fld = fold::new_identity_fold[option::t[@fn_ctxt]]();
+        auto visit = rec(keep_going=bind keep_going(ignore),
+                         visit_item_pre=bind visit_item_pre(ignore, _),
+                         visit_item_post=bind visit_item_post(ignore, _),
+                         visit_stmt_pre=bind visit_stmt_pre(fcx, _),
+                         visit_expr_pre=bind visit_expr_pre(fcx, _),
+                         visit_block_pre=bind visit_block_pre(fcx, _),
+                         visit_arm_pre=bind visit_arm_pre(fcx, _),
+                         visit_decl_pre=bind visit_decl_pre(fcx, _)
+                         with walk::default_visitor());
+        walk::walk_block(visit, block);
+    }
 }
 
 
@@ -2699,11 +2722,9 @@ fn check_fn(&@crate_ctxt ccx, &ast::fn_decl decl, ast::proto proto,
         case (_) {}
     }
 
-    auto block_wb = resolve_local_types_in_block(fcx, block_t);
+    writeback::resolve_local_types_in_block(fcx, block_t);
 
-    auto fn_t = rec(decl=decl,
-                    proto=proto,
-                    body=block_wb);
+    auto fn_t = rec(decl=decl, proto=proto, body=block_t);
     ret fn_t;
 }
 

src/comp/middle/walk.rs

@@ -4,6 +4,7 @@
 import std::option::some;
 import std::option::none;
 
+// FIXME: Should visit patterns as well.
 type ast_visitor =
     rec(fn () -> bool                  keep_going,
         fn () -> bool                  want_crate_directives,