#include "RemoveTrivialForLoops.h"
#include "CodeGen_GPU_Dev.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "Simplify.h"
#include "Substitute.h"

namespace Halide {
namespace Internal {

class RemoveTrivialForLoops : public IRMutator2 {
    using IRMutator2::visit;

    Stmt visit(const For *for_loop) override {

        Stmt body = mutate(for_loop->body);

        if (is_one(for_loop->extent) &&
            for_loop->device_api != DeviceAPI::None) {
            // Don't assume any device API loops are trivial, but we
            // can substitute in the min value if the loop is size 1.
            body = substitute(for_loop->name, for_loop->min, body);
            return For::make(for_loop->name, for_loop->min, for_loop->extent,
                             for_loop->for_type, for_loop->device_api, body);
        } else if (is_one(for_loop->extent)) {
            if (for_loop->for_type == ForType::Parallel) {
                std::cerr << "Warning: Parallel for loop over "
                          << for_loop->name << " has extent one. "
                          << "Can't do one piece of work in parallel.\n";
            } else if (for_loop->for_type == ForType::Vectorized) {
                std::cerr << "Warning: Vectorized for loop over "
                          << for_loop->name << " has extent one. "
                          << "Not vectorizing.\n";
            }
            return LetStmt::make(for_loop->name, for_loop->min, body);
        } else if (is_zero(for_loop->extent)) {
            return Evaluate::make(0);
        } else if (can_prove(for_loop->extent <= 1)) {
            // Loop has at most one iteration
            Stmt stmt = LetStmt::make(for_loop->name, for_loop->min, body);
            return IfThenElse::make(for_loop->extent > 0, stmt, Stmt());
        } else if (body.same_as(for_loop->body)) {
            return for_loop;
        } else {
            return For::make(for_loop->name, for_loop->min, for_loop->extent,
                             for_loop->for_type, for_loop->device_api, body);
        }
    }
};

// Turn for loops of size one into let statements
Stmt remove_trivial_for_loops(Stmt s) {
    return RemoveTrivialForLoops().mutate(s);
}

}  // namespace Internal
}  // namespace Halide