From 2aa759d18bc4e08e687e8c9564e27aa50140dae7 Mon Sep 17 00:00:00 2001
From: Andrew Adams <andrew.b.adams@gmail.com>
Date: Thu, 13 Jul 2017 10:52:40 -0700
Subject: [PATCH] Fix print_loop_nest

---
 src/PrintLoopNest.cpp | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/PrintLoopNest.cpp b/src/PrintLoopNest.cpp
index bb4fa211efc5..0299a3cfe4fb 100644
--- a/src/PrintLoopNest.cpp
+++ b/src/PrintLoopNest.cpp
@@ -157,8 +157,7 @@ string print_loop_nest(const vector<Function> &output_funcs) {
     // Compute an environment
     map<string, Function> env;
     for (Function f : output_funcs) {
-        map<string, Function> more_funcs = find_transitive_calls(f);
-        env.insert(more_funcs.begin(), more_funcs.end());
+        populate_environment(f, env);
     }
 
     // Create a deep-copy of the entire graph of Funcs.
@@ -170,6 +169,11 @@ string print_loop_nest(const vector<Function> &output_funcs) {
         Func(f).compute_root().store_root();
     }
 
+    // Ensure that all ScheduleParams become well-defined constant Exprs.
+    for (auto &f : env) {
+        f.second.substitute_schedule_param_exprs();
+    }
+
     // Substitute in wrapper Funcs
     env = wrap_func_calls(env);