#include "SlidingWindow.h"

#include "Bounds.h"
#include "CompilerLogger.h"
#include "Debug.h"
#include "ExprUsesVar.h"
#include "IREquality.h"
#include "IRMatch.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "IRPrinter.h"
#include "Monotonic.h"
#include "Scope.h"
#include "Simplify.h"
#include "Solve.h"
#include "Substitute.h"
#include <list>
#include <set>
#include <utility>

namespace Halide {
namespace Internal {

using std::list;
using std::map;
using std::pair;
using std::set;
using std::string;
using std::vector;

namespace {

// Does an expression depend on a particular variable?
class ExprDependsOnVar : public IRVisitor {
    using IRVisitor::visit;

    void visit(const Variable *op) override {
        if (op->name == var) {
            result = true;
        }
    }

    void visit(const Let *op) override {
        op->value.accept(this);
        // The name might be hidden within the body of the let, in
        // which case there's no point descending.
        if (op->name != var) {
            op->body.accept(this);
        }
    }

public:
    bool result;
    string var;

    ExprDependsOnVar(string v)
        : result(false), var(std::move(v)) {
    }
};

bool expr_depends_on_var(const Expr &e, string v) {
    ExprDependsOnVar depends(std::move(v));
    e.accept(&depends);
    return depends.result;
}

class ExpandExpr : public IRMutator {
    using IRMutator::visit;
    const Scope<Expr> &scope;

    Expr visit(const Variable *var) override {
        if (scope.contains(var->name)) {
            Expr expr = scope.get(var->name);
            debug(4) << "Fully expanded " << var->name << " -> " << expr << "\n";
            return expr;
        } else {
            return var;
        }
    }

public:
    ExpandExpr(const Scope<Expr> &s)
        : scope(s) {
    }
};

// Perform all the substitutions in a scope
Expr expand_expr(const Expr &e, const Scope<Expr> &scope) {
    ExpandExpr ee(scope);
    Expr result = ee.mutate(e);
    debug(4) << "Expanded " << e << " into " << result << "\n";
    return result;
}

class FindProduce : public IRVisitor {
    const string &func;

    using IRVisitor::visit;

    void visit(const ProducerConsumer *op) override {
        if (op->is_producer && op->name == func) {
            found = true;
        } else {
            IRVisitor::visit(op);
        }
    }

public:
    bool found = false;

    FindProduce(const string &func)
        : func(func) {
    }
};

bool find_produce(const Stmt &s, const string &func) {
    FindProduce finder(func);
    s.accept(&finder);
    return finder.found;
}

// This mutator rewrites calls and provides to a particular
// func:
// - Calls and Provides are shifted to be relative to the min.
// - Provides additionally are rewritten to load values from the
//   previous iteration of the loop if they were computed in the
//   last iteration.
class RollFunc : public IRMutator {
    const Function &func;
    int dim;
    const string &loop_var;
    const Interval &old_bounds;
    const Interval &new_bounds;

    Scope<Expr> scope;

    // It helps simplify the shifted calls/provides to rebase the
    // loops that are subtracted from to have a min of 0.
    set<string> loops_to_rebase;
    bool in_produce = false;

    using IRMutator::visit;

    Stmt visit(const ProducerConsumer *op) override {
        bool produce_func = op->name == func.name() && op->is_producer;
        ScopedValue<bool> old_in_produce(in_produce, in_produce || produce_func);
        return IRMutator::visit(op);
    }

    Stmt visit(const Provide *op) override {
        if (!(in_produce && op->name == func.name())) {
            return IRMutator::visit(op);
        }
        vector<Expr> values = op->values;
        for (Expr &i : values) {
            i = mutate(i);
        }
        vector<Expr> args = op->args;
        for (Expr &i : args) {
            i = mutate(i);
        }
        bool sliding_up = old_bounds.max.same_as(new_bounds.max);
        Expr is_new = sliding_up ? new_bounds.min <= args[dim] : args[dim] <= new_bounds.max;
        args[dim] -= old_bounds.min;
        vector<Expr> old_args = args;
        Expr old_arg_dim = expand_expr(old_args[dim], scope);
        old_args[dim] = substitute(loop_var, Variable::make(Int(32), loop_var) - 1, old_arg_dim);
        for (int i = 0; i < (int)values.size(); i++) {
            Type t = values[i].type();
            Expr old_value =
                Call::make(t, op->name, old_args, Call::Halide, func.get_contents(), i);
            values[i] = Call::make(values[i].type(), Call::if_then_else, {is_new, values[i], likely(old_value)}, Call::PureIntrinsic);
        }
        if (const Variable *v = op->args[dim].as<Variable>()) {
            // The subtractions above simplify more easily if the loop is rebased to 0.
            loops_to_rebase.insert(v->name);
        }
        return Provide::make(func.name(), values, args, op->predicate);
    }

    Expr visit(const Call *op) override {
        if (!(op->call_type == Call::Halide && op->name == func.name())) {
            return IRMutator::visit(op);
        }
        vector<Expr> args = op->args;
        for (Expr &i : args) {
            i = mutate(i);
        }
        args[dim] -= old_bounds.min;
        return Call::make(op->type, op->name, args, Call::Halide, op->func, op->value_index, op->image, op->param);
    }

    Stmt visit(const For *op) override {
        Stmt result = IRMutator::visit(op);
        op = result.as<For>();
        internal_assert(op);
        if (loops_to_rebase.count(op->name)) {
            string new_name = op->name + ".rebased";
            Stmt body = substitute(op->name, Variable::make(Int(32), new_name) + op->min, op->body);
            // use op->name *before* the re-assignment of result, which will clobber it
            loops_to_rebase.erase(op->name);
            result = For::make(new_name, 0, op->extent, op->for_type, op->device_api, body);
        }
        return result;
    }

    Stmt visit(const LetStmt *op) override {
        ScopedBinding<Expr> bind(scope, op->name, simplify(expand_expr(op->value, scope)));
        return IRMutator::visit(op);
    }

public:
    RollFunc(const Function &func, int dim, const string &loop_var,
             const Interval &old_bounds, const Interval &new_bounds)
        : func(func), dim(dim), loop_var(loop_var), old_bounds(old_bounds), new_bounds(new_bounds) {
    }
};

// Perform sliding window optimization for a function over a
// particular serial for loop
class SlidingWindowOnFunctionAndLoop : public IRMutator {
    Function func;
    string loop_var;
    Expr loop_min;
    set<int> &slid_dimensions;
    Scope<Expr> scope;

    map<string, Expr> replacements;

    using IRMutator::visit;

    // Check if the dimension at index 'dim_idx' is always pure (i.e. equal to 'dim')
    // in the definition (including in its specializations)
    bool is_dim_always_pure(const Definition &def, const string &dim, int dim_idx) {
        const Variable *var = def.args()[dim_idx].as<Variable>();
        if ((!var) || (var->name != dim)) {
            return false;
        }

        for (const auto &s : def.specializations()) {
            bool pure = is_dim_always_pure(s.definition, dim, dim_idx);
            if (!pure) {
                return false;
            }
        }
        return true;
    }

    Stmt visit(const ProducerConsumer *op) override {
        if (op->is_producer) {
            if (op->name != func.name()) {
                return IRMutator::visit(op);
            }

            // We're interested in the case where exactly one of the
            // dimensions of the buffer has a min/extent that depends
            // on the loop_var.
            string dim = "";
            int dim_idx = 0;
            Expr min_required, max_required;

            debug(3) << "Considering sliding " << func.name()
                     << " along loop variable " << loop_var << "\n"
                     << "Region provided:\n";

            string prefix = func.name() + ".s" + std::to_string(func.updates().size()) + ".";
            const std::vector<string> func_args = func.args();
            for (int i = 0; i < func.dimensions(); i++) {
                if (slid_dimensions.count(i)) {
                    debug(3) << "Already slid over dimension " << i << ", so skipping it.\n";
                    continue;
                }
                // Look up the region required of this function's last stage
                string var = prefix + func_args[i];
                internal_assert(scope.contains(var + ".min") && scope.contains(var + ".max"));
                Expr min_req = scope.get(var + ".min");
                Expr max_req = scope.get(var + ".max");
                min_req = expand_expr(min_req, scope);
                max_req = expand_expr(max_req, scope);

                debug(3) << func_args[i] << ":" << min_req << ", " << max_req << "\n";
                if (expr_depends_on_var(min_req, loop_var) ||
                    expr_depends_on_var(max_req, loop_var)) {
                    if (!dim.empty()) {
                        dim = "";
                        min_required = Expr();
                        max_required = Expr();
                        break;
                    } else {
                        dim = func_args[i];
                        dim_idx = i;
                        min_required = min_req;
                        max_required = max_req;
                    }
                } else if (!min_required.defined() &&
                           i == func.dimensions() - 1 &&
                           is_pure(min_req) &&
                           is_pure(max_req)) {
                    // The footprint doesn't depend on the loop var. Just compute everything on the first loop iteration.
                    dim = func_args[i];
                    dim_idx = i;
                    min_required = min_req;
                    max_required = max_req;
                }
            }

            if (!min_required.defined()) {
                debug(3) << "Could not perform sliding window optimization of "
                         << func.name() << " over " << loop_var << " because multiple "
                         << "dimensions of the function dependended on the loop var\n";
                return op;
            }

            // If the function is not pure in the given dimension, give up. We also
            // need to make sure that it is pure in all the specializations
            bool pure = true;
            for (const Definition &def : func.updates()) {
                pure = is_dim_always_pure(def, dim, dim_idx);
                if (!pure) {
                    break;
                }
            }
            if (!pure) {
                debug(3) << "Could not performance sliding window optimization of "
                         << func.name() << " over " << loop_var << " because the function "
                         << "scatters along the related axis.\n";
                return op;
            }

            bool can_slide_up = false;
            bool can_slide_down = false;

            Monotonic monotonic_min = is_monotonic(min_required, loop_var);
            Monotonic monotonic_max = is_monotonic(max_required, loop_var);

            if (monotonic_min == Monotonic::Increasing ||
                monotonic_min == Monotonic::Constant) {
                can_slide_up = true;
            } else if (monotonic_min == Monotonic::Unknown) {
                if (get_compiler_logger()) {
                    get_compiler_logger()->record_non_monotonic_loop_var(loop_var, min_required);
                }
            }

            if (monotonic_max == Monotonic::Decreasing ||
                monotonic_max == Monotonic::Constant) {
                can_slide_down = true;
            } else if (monotonic_max == Monotonic::Unknown) {
                if (get_compiler_logger()) {
                    get_compiler_logger()->record_non_monotonic_loop_var(loop_var, max_required);
                }
            }

            if (!can_slide_up && !can_slide_down) {
                debug(3) << "Not sliding " << func.name()
                         << " over dimension " << dim
                         << " along loop variable " << loop_var
                         << " because I couldn't prove it moved monotonically along that dimension\n"
                         << "Min is " << min_required << "\n"
                         << "Max is " << max_required << "\n";
                return op;
            }

            // Ok, we've isolated a function, a dimension to slide
            // along, and loop variable to slide over.
            debug(3) << "Sliding " << func.name()
                     << " over dimension " << dim
                     << " along loop variable " << loop_var << "\n";

            Expr loop_var_expr = Variable::make(Int(32), loop_var);

            Expr prev_max_plus_one = substitute(loop_var, loop_var_expr - 1, max_required) + 1;
            Expr prev_min_minus_one = substitute(loop_var, loop_var_expr - 1, min_required) - 1;

            // If there's no overlap between adjacent iterations, we shouldn't slide.
            if (can_prove(min_required >= prev_max_plus_one) ||
                can_prove(max_required <= prev_min_minus_one)) {
                debug(3) << "Not sliding " << func.name()
                         << " over dimension " << dim
                         << " along loop variable " << loop_var
                         << " there's no overlap in the region computed across iterations\n"
                         << "Min is " << min_required << "\n"
                         << "Max is " << max_required << "\n";
                return op;
            }

            // Update the bounds of this producer assuming the previous iteration
            // has run already.
            Expr new_min, new_max;
            if (can_slide_up) {
                new_min = prev_max_plus_one;
                new_max = max_required;
            } else {
                new_min = min_required;
                new_max = prev_min_minus_one;
            }

            // See if we can find a new min for the loop that can warm up the
            // sliding window. We're going to do this by building an equation
            // that describes the constraints we have on our new loop min. The
            // first constraint is that the new loop min is not after the
            // loop min.
            string new_loop_min_name = unique_name('x');
            Expr new_loop_min_var = Variable::make(Int(32), new_loop_min_name);
            Expr new_loop_min_eq = new_loop_min_var <= loop_min;
            Expr new_min_at_new_loop_min = substitute(loop_var, new_loop_min_var, new_min);
            Expr new_max_at_new_loop_min = substitute(loop_var, new_loop_min_var, new_max);
            if (can_slide_up) {
                // We need to find a new loop min that satisfies these constraints:
                // - The new min at the new loop min needs to be before the min
                //   required at the original min.
                // - The new max needs to be greater than the new min, both at the
                //   new loop min. This guarantees that the sliding window.
                // Together, these conditions guarantee the sliding window is warmed
                // up. The first condition checks that we reached the original loop
                // min, and the second condition checks that the iterations before
                // the original min weren't empty.
                Expr min_required_at_loop_min = substitute(loop_var, loop_min, min_required);
                new_loop_min_eq = new_loop_min_eq &&
                                  new_min_at_new_loop_min <= min_required_at_loop_min &&
                                  new_max_at_new_loop_min >= new_min_at_new_loop_min;
            } else {
                // When sliding down, the constraints are similar, just swapping
                // the roles of the min and max.
                Expr max_required_at_loop_min = substitute(loop_var, loop_min, max_required);
                new_loop_min_eq = new_loop_min_eq &&
                                  new_max_at_new_loop_min >= max_required_at_loop_min &&
                                  new_min_at_new_loop_min <= new_max_at_new_loop_min;
            }
            // Try to solve the equation.
            new_loop_min_eq = simplify(new_loop_min_eq);
            Interval solve_result = solve_for_inner_interval(new_loop_min_eq, new_loop_min_name);
            internal_assert(!new_loop_min.defined());
            if (solve_result.has_upper_bound() && !equal(solve_result.max, loop_min)) {
                new_loop_min = simplify(solve_result.max);

                // We have a new loop min, so we an assume every iteration has
                // a previous iteration. In order for this to be safe, we need
                // the new min/max at the new loop min to be less than or equal to
                // the min/max required at the original loop min.
                Expr loop_var_expr = Variable::make(Int(32), loop_var);
                Expr orig_loop_min_expr = Variable::make(Int(32), loop_var + ".loop_min.orig");
                if (can_slide_up) {
                    Expr min_required_at_loop_min = substitute(loop_var, orig_loop_min_expr, min_required);
                    new_min = max(new_min, min_required_at_loop_min);
                } else {
                    Expr max_required_at_loop_min = substitute(loop_var, orig_loop_min_expr, max_required);
                    new_max = min(new_max, max_required_at_loop_min);
                }
            } else {
                // We couldn't find a suitable new loop min, we can't assume
                // every iteration has a previous iteration. The first iteration
                // will warm up the loop.
                Expr need_explicit_warmup = loop_var_expr <= loop_min;
                if (can_slide_up) {
                    new_min = select(need_explicit_warmup, min_required, likely_if_innermost(new_min));
                } else {
                    new_max = select(need_explicit_warmup, max_required, likely_if_innermost(new_max));
                }
            }
            new_min = simplify(new_min);
            new_max = simplify(new_max);

            debug(3) << "Sliding " << func.name() << ", " << dim << "\n"
                     << "Pushing min up from " << min_required << " to " << new_min << "\n"
                     << "Shrinking max from " << max_required << " to " << new_max << "\n"
                     << "Adjusting loop_min from " << loop_min << " to " << new_loop_min << "\n"
                     << "Equation is " << new_loop_min_eq << "\n";

            slid_dimensions.insert(dim_idx);

            // If we want to slide in registers, we're done here, we just need to
            // save the updated bounds for later.
            if (func.schedule().memory_type() == MemoryType::Register) {
                this->dim_idx = dim_idx;
                old_bounds = {min_required, max_required};
                new_bounds = {new_min, new_max};
                return op;
            }

            // If we aren't sliding in registers, we need to update the bounds of
            // the producer to be only the bounds of the region newly computed.
            internal_assert(replacements.empty());
            if (can_slide_up) {
                replacements[prefix + dim + ".min"] = new_min;
            } else {
                replacements[prefix + dim + ".max"] = new_max;
            }

            for (size_t i = 0; i < func.updates().size(); i++) {
                string n = func.name() + ".s" + std::to_string(i) + "." + dim;
                replacements[n + ".min"] = Variable::make(Int(32), prefix + dim + ".min");
                replacements[n + ".max"] = Variable::make(Int(32), prefix + dim + ".max");
            }

            // Ok, we have a new min/max required and we're going to
            // rewrite all the lets that define bounds required. Now
            // we need to additionally expand the bounds required of
            // the last stage to cover values produced by stages
            // before the last one. Because, e.g., an intermediate
            // stage may be unrolled, expanding its bounds provided.
            Stmt result = op;
            if (!func.updates().empty()) {
                Box b = box_provided(op->body, func.name());
                if (can_slide_up) {
                    string n = prefix + dim + ".min";
                    Expr var = Variable::make(Int(32), n);
                    result = LetStmt::make(n, min(var, b[dim_idx].min), result);
                } else {
                    string n = prefix + dim + ".max";
                    Expr var = Variable::make(Int(32), n);
                    result = LetStmt::make(n, max(var, b[dim_idx].max), result);
                }
            }

            return result;
        } else if (!find_produce(op, func.name()) && new_loop_min.defined()) {
            // The producer might have expanded the loop before the min to warm
            // up the window. This consumer doesn't contain a producer that might
            // be part of the warmup, so guard it with an if to only run it on
            // the original loop bounds.
            Expr loop_var_expr = Variable::make(Int(32), loop_var);
            Expr orig_loop_min_expr = Variable::make(Int(32), loop_var + ".loop_min.orig");
            Expr guard = likely_if_innermost(orig_loop_min_expr <= loop_var_expr);

            // Put the if inside the consumer node, so semaphores end up outside the if.
            // TODO: This is correct, but it produces slightly suboptimal code: if we
            // didn't do this, the loop could likely be trimmed and the if simplified away.
            Stmt body = mutate(op->body);
            if (const IfThenElse *old_guard = body.as<IfThenElse>()) {
                Expr x = Variable::make(Int(32), "*");
                vector<Expr> matches;
                if (expr_match(likely_if_innermost(x <= loop_var_expr), old_guard->condition, matches)) {
                    // There's already a condition on loop_var_expr here. Since we're
                    // adding a condition at the old loop min, this if must already be
                    // guarding more than we will.
                    guard = Expr();
                }
            }
            if (guard.defined()) {
                debug(3) << "Guarding body " << guard << "\n";
                body = IfThenElse::make(guard, body);
            }
            if (body.same_as(op->body)) {
                return op;
            } else {
                return ProducerConsumer::make_consume(op->name, body);
            }
        } else {
            return IRMutator::visit(op);
        }
    }

    Stmt visit(const For *op) override {
        // It's not safe to enter an inner loop whose bounds depend on
        // the var we're sliding over.
        Expr min = expand_expr(op->min, scope);
        Expr extent = expand_expr(op->extent, scope);
        if (is_const_one(extent)) {
            // Just treat it like a let
            Stmt s = LetStmt::make(op->name, min, op->body);
            s = mutate(s);
            // Unpack it back into the for
            const LetStmt *l = s.as<LetStmt>();
            internal_assert(l);
            return For::make(op->name, op->min, op->extent, op->for_type, op->device_api, l->body);
        } else if (is_monotonic(min, loop_var) != Monotonic::Constant ||
                   is_monotonic(extent, loop_var) != Monotonic::Constant) {
            debug(3) << "Not entering loop over " << op->name
                     << " because the bounds depend on the var we're sliding over: "
                     << min << ", " << extent << "\n";
            return op;
        } else {
            return IRMutator::visit(op);
        }
    }

    Stmt visit(const LetStmt *op) override {
        ScopedBinding<Expr> bind(scope, op->name, simplify(expand_expr(op->value, scope)));
        Stmt new_body = mutate(op->body);

        Expr value = op->value;

        map<string, Expr>::iterator iter = replacements.find(op->name);
        if (iter != replacements.end()) {
            value = iter->second;
            replacements.erase(iter);
        }

        if (new_body.same_as(op->body) && value.same_as(op->value)) {
            return op;
        } else {
            return LetStmt::make(op->name, value, new_body);
        }
    }

public:
    SlidingWindowOnFunctionAndLoop(Function f, string v, Expr v_min, set<int> &slid_dimensions)
        : func(std::move(f)), loop_var(std::move(v)), loop_min(std::move(v_min)), slid_dimensions(slid_dimensions) {
    }

    Expr new_loop_min;
    int dim_idx;
    Interval old_bounds;
    Interval new_bounds;

    Stmt translate_loop(const Stmt &s) {
        return RollFunc(func, dim_idx, loop_var, old_bounds, new_bounds).mutate(s);
    }
};

// In Stmt s, does the production of b depend on a?
// We can't use produce/consume nodes to determine this, because they're "loose".
// For example, we get this:
//
//  produce a {
//   a(...) = ...
//  }
//  consume a {
//   produce b {
//    b(...) = ... // not depending on a
//   }
//   consume b {
//    c(...) = a(...) + b(...)
//   }
//  }
//
// When we'd rather see this:
//
//  produce a {
//   a(...) = ...
//  }
//  produce b {
//   b(...) = ... // not depending on a
//  }
//  consume a {
//   consume b {
//    c(...) = a(...) + b(...)
//   }
//  }
//
// TODO: We might also need to figure out transitive dependencies...? If so, it
// would be best to just fix the produce/consume relationships as above. We would
// just be able to look for produce b inside produce a.
class Dependencies : public IRVisitor {
    using IRVisitor::visit;

    const string &producer;
    bool in_producer = false;

    void visit(const ProducerConsumer *op) override {
        ScopedValue<bool> old_finding_a(in_producer, in_producer || (op->is_producer && op->name == producer));
        return IRVisitor::visit(op);
    }

    void visit(const Call *op) override {
        if (in_producer && op->call_type == Call::Halide) {
            if (op->name != producer) {
                dependencies.insert(op->name);
            }
        }
        IRVisitor::visit(op);
    }

public:
    set<string> dependencies;

    Dependencies(const string &producer)
        : producer(producer) {
    }
};

bool depends_on(const string &a, const string &b, const Stmt &s, map<string, bool> &cache) {
    if (a == b) {
        return true;
    }
    auto cached = cache.find(b);
    if (cached != cache.end()) {
        return cached->second;
    }
    Dependencies deps(b);
    s.accept(&deps);
    // Recursively search for dependencies.
    for (const string &i : deps.dependencies) {
        if (depends_on(a, i, s, cache)) {
            cache[b] = true;
            return true;
        }
    }
    cache[b] = false;
    return false;
}

bool depends_on(const string &a, const string &b, const Stmt &s) {
    map<string, bool> cache;
    return depends_on(a, b, s, cache);
}

// Update the loop variable referenced by prefetch directives.
class SubstitutePrefetchVar : public IRMutator {
    const string &old_var;
    const string &new_var;

    using IRMutator::visit;

    Stmt visit(const Prefetch *op) override {
        Stmt new_body = mutate(op->body);
        if (op->prefetch.at == old_var || op->prefetch.from == old_var) {
            PrefetchDirective p = op->prefetch;
            if (op->prefetch.at == old_var) {
                p.at = new_var;
            }
            if (op->prefetch.from == old_var) {
                p.from = new_var;
            }
            return Prefetch::make(op->name, op->types, op->bounds, p, op->condition, std::move(new_body));
        } else if (!new_body.same_as(op->body)) {
            return Prefetch::make(op->name, op->types, op->bounds, op->prefetch, op->condition, std::move(new_body));
        } else {
            return op;
        }
    }

public:
    SubstitutePrefetchVar(const string &old_var, const string &new_var)
        : old_var(old_var), new_var(new_var) {
    }
};

// Perform sliding window optimization for all functions
class SlidingWindow : public IRMutator {
    const map<string, Function> &env;

    // A map of which dimensions we've already slid over, by Func name.
    map<string, set<int>> slid_dimensions;

    // Keep track of realizations we want to slide, from innermost to
    // outermost.
    list<Function> sliding;

    using IRMutator::visit;

    Stmt visit(const Realize *op) override {
        // Find the args for this function
        map<string, Function>::const_iterator iter = env.find(op->name);

        // If it's not in the environment it's some anonymous
        // realization that we should skip (e.g. an inlined reduction)
        if (iter == env.end()) {
            return IRMutator::visit(op);
        }

        // If the Function in question has the same compute_at level
        // as its store_at level, skip it.
        const FuncSchedule &sched = iter->second.schedule();
        if (sched.compute_level() == sched.store_level()) {
            return IRMutator::visit(op);
        }

        // We want to slide innermost first, so put it on the front of
        // the list.
        sliding.push_front(iter->second);
        Stmt new_body = mutate(op->body);
        sliding.pop_front();
        // Remove tracking of slid dimensions when we're done realizing
        // it in case a realization appears elsewhere.
        auto slid_it = slid_dimensions.find(iter->second.name());
        if (slid_it != slid_dimensions.end()) {
            slid_dimensions.erase(slid_it);
        }

        if (new_body.same_as(op->body)) {
            return op;
        } else {
            return Realize::make(op->name, op->types, op->memory_type,
                                 op->bounds, op->condition, new_body);
        }
    }

    Stmt visit(const For *op) override {
        if (!(op->for_type == ForType::Serial || op->for_type == ForType::Unrolled)) {
            return IRMutator::visit(op);
        }
        debug(3) << "Doing sliding window analysis on loop " << op->name << "\n";

        string name = op->name;
        Stmt body = op->body;
        Expr loop_min = op->min;
        Expr loop_extent = op->extent;
        Expr loop_max = Variable::make(Int(32), op->name + ".loop_max");

        list<pair<string, Expr>> prev_loop_mins;
        list<pair<string, Expr>> new_lets;
        for (const Function &func : sliding) {
            debug(3) << "Doing sliding window analysis on function " << func.name() << "\n";

            // Figure out where we should start sliding from. If no
            // other func needs this func, we can just start at the
            // original loop min.
            Expr prev_loop_min = op->min;
            // If a previously slid func needs this func to be warmed
            // up, then we need to back up the loop to warm up this
            // func before the already slid func starts warming up.
            for (const auto &i : prev_loop_mins) {
                if (depends_on(func.name(), i.first, body)) {
                    prev_loop_min = i.second;
                    break;
                }
            }

            SlidingWindowOnFunctionAndLoop slider(func, name, prev_loop_min, slid_dimensions[func.name()]);
            body = slider.mutate(body);

            if (func.schedule().memory_type() == MemoryType::Register &&
                slider.old_bounds.has_lower_bound()) {
                body = slider.translate_loop(body);
            }

            if (slider.new_loop_min.defined()) {
                Expr new_loop_min = slider.new_loop_min;
                if (!prev_loop_min.same_as(loop_min)) {
                    // If we didn't start sliding from the previous
                    // loop min, we the old loop min might already
                    // be further back than this new one.
                    new_loop_min = min(new_loop_min, loop_min);
                }

                // Put this at the front of the list, so we find it first
                // when checking subsequent funcs.
                prev_loop_mins.emplace_front(func.name(), new_loop_min);

                // Update the loop body to use the adjusted loop min.
                string new_name = name + ".$n";
                loop_min = Variable::make(Int(32), new_name + ".loop_min");
                loop_extent = Variable::make(Int(32), new_name + ".loop_extent");
                body = substitute({
                                      {name, Variable::make(Int(32), new_name)},
                                      {name + ".loop_min", loop_min},
                                      {name + ".loop_extent", loop_extent},
                                  },
                                  body);
                body = SubstitutePrefetchVar(name, new_name).mutate(body);

                name = new_name;

                // The new loop interval is the new loop min to the loop max.
                new_lets.emplace_front(name + ".loop_min", new_loop_min);
                new_lets.emplace_front(name + ".loop_min.orig", loop_min);
                new_lets.emplace_front(name + ".loop_extent", (loop_max - loop_min) + 1);
            }
        }

        body = mutate(body);

        if (body.same_as(op->body) && loop_min.same_as(op->min) && loop_extent.same_as(op->extent) && name == op->name) {
            return op;
        } else {
            Stmt result = For::make(name, loop_min, loop_extent, op->for_type, op->device_api, body);
            if (!new_lets.empty()) {
                result = LetStmt::make(name + ".loop_max", loop_max, result);
            }
            for (const auto &i : new_lets) {
                result = LetStmt::make(i.first, i.second, result);
            }
            return result;
        }
    }

    Stmt visit(const IfThenElse *op) override {
        // Don't let specializations corrupt the tracking of which
        // dimensions have been slid.
        map<string, set<int>> old_slid_dimensions = slid_dimensions;
        Stmt then_case = mutate(op->then_case);
        slid_dimensions = old_slid_dimensions;
        Stmt else_case = mutate(op->else_case);
        slid_dimensions = old_slid_dimensions;
        if (then_case.same_as(op->then_case) && else_case.same_as(op->else_case)) {
            return op;
        } else {
            return IfThenElse::make(op->condition, then_case, else_case);
        }
    }

public:
    SlidingWindow(const map<string, Function> &e)
        : env(e) {
    }
};

// It is convenient to be able to assume that loops have a .loop_min.orig
// let in addition to .loop_min. Most of these will get simplified away.
class AddLoopMinOrig : public IRMutator {
    using IRMutator::visit;

    Stmt visit(const For *op) override {
        Stmt body = mutate(op->body);
        Expr min = mutate(op->min);
        Expr extent = mutate(op->extent);

        Stmt result;
        if (body.same_as(op->body) && min.same_as(op->min) && extent.same_as(op->extent)) {
            result = op;
        } else {
            result = For::make(op->name, min, extent, op->for_type, op->device_api, body);
        }
        return LetStmt::make(op->name + ".loop_min.orig", Variable::make(Int(32), op->name + ".loop_min"), result);
    }
};

}  // namespace

Stmt sliding_window(const Stmt &s, const map<string, Function> &env) {
    return SlidingWindow(env).mutate(AddLoopMinOrig().mutate(s));
}

}  // namespace Internal
}  // namespace Halide