#include "Closure.h"
#include "Debug.h"
#include "ExprUsesVar.h"
#include "IRMutator.h"
#include "IROperator.h"

namespace Halide {
namespace Internal {

using std::string;

namespace {
constexpr int DBG = 3;
} // namespace

void Closure::include(const Stmt &s, const string &loop_variable) {
 if (!loop_variable.empty()) {
 ignore.push(loop_variable);
 }
 s.accept(this);
 if (!loop_variable.empty()) {
 ignore.pop(loop_variable);
 }
}

void Closure::visit(const Let *op) {
 op->value.accept(this);
 ScopedBinding<> p(ignore, op->name);
 op->body.accept(this);
}

void Closure::visit(const LetStmt *op) {
 op->value.accept(this);
 ScopedBinding<> p(ignore, op->name);
 op->body.accept(this);
}

void Closure::visit(const For *op) {
 ScopedBinding<> p(ignore, op->name);
 op->min.accept(this);
 op->extent.accept(this);
 op->body.accept(this);
}

void Closure::found_buffer_ref(const string &name, Type type,
 bool read, bool written, const Halide::Buffer<> &image) {
 if (!ignore.contains(name)) {
 debug(DBG) << "Adding buffer " << name << " to closure:\n";
 Buffer &ref = buffers[name];
 ref.type = type.element_of(); // TODO: Validate type is the same as existing refs?
 ref.read = ref.read || read;
 ref.write = ref.write || written;

 // If reading an image/buffer, compute the size.
 if (image.defined()) {
 ref.size = image.size_in_bytes();
 ref.dimensions = image.dimensions();
 }
 debug(DBG) << " "
 << " t=" << ref.type
 << " d=" << (int)ref.dimensions
 << " r=" << ref.read
 << " w=" << ref.write
 << " mt=" << (int)ref.memory_type
 << " sz=" << ref.size << "\n";
 } else {
 debug(DBG) << "Not adding buffer " << name << " to closure\n";
 }
}

void Closure::visit(const Load *op) {
 op->predicate.accept(this);
 op->index.accept(this);
 found_buffer_ref(op->name, op->type, true, false, op->image);
}

void Closure::visit(const Store *op) {
 op->predicate.accept(this);
 op->index.accept(this);
 op->value.accept(this);
 found_buffer_ref(op->name, op->value.type(), false, true, Halide::Buffer<>());
}

void Closure::visit(const Allocate *op) {
 if (op->new_expr.defined()) {
 op->new_expr.accept(this);
 }
 ScopedBinding<> p(ignore, op->name);
 for (const auto &extent : op->extents) {
 extent.accept(this);
 }
 op->condition.accept(this);
 op->body.accept(this);
}

void Closure::visit(const Variable *op) {
 if (ignore.contains(op->name)) {
 debug(DBG) << "Not adding var " << op->name << " to closure\n";
 } else {
 debug(DBG) << "Adding var " << op->name << " to closure\n";
 vars[op->name] = op->type;
 }
}

void Closure::visit(const Atomic *op) {
 if (!op->mutex_name.empty()) {
 found_buffer_ref(op->mutex_name, type_of<void *>(), true, true, Halide::Buffer<>());
 }
 op->body.accept(this);
}

Expr Closure::pack_into_struct() const {
 std::vector<Expr> elements;

 for (const auto &b : buffers) {
 Expr ptr_var = Variable::make(type_of<void *>(), b.first);
 elements.emplace_back(ptr_var);
 }
 for (const auto &v : vars) {
 Expr var = Variable::make(v.second, v.first);
 elements.emplace_back(var);
 }

 // Sort by decreasing size, to guarantee the struct is densely packed in
 // memory. We don't actually rely on this, it's just nice to have.
 std::stable_sort(elements.begin(), elements.end(),
 [&](const Expr &a, const Expr &b) {
 return a.type().bytes() > b.type().bytes();
 });

 Expr result = Call::make(Handle(),
 Call::make_struct, elements, Call::Intrinsic);
 return result;
}

Stmt Closure::unpack_from_struct(const Expr &e, const Stmt &s) const {
 // Use the struct-packing code just to make sure the order of elements is
 // the same.
 Expr packed = pack_into_struct();

 // Make a prototype of the packed struct
 class ReplaceCallArgsWithZero : public IRMutator {
 public:
 using IRMutator::mutate;
 Expr mutate(const Expr &e) override {
 if (!e.as<Call>()) {
 return make_zero(e.type());
 } else {
 return IRMutator::mutate(e);
 }
 }
 } replacer;
 string prototype_name = unique_name("closure_prototype");
 Expr prototype = replacer.mutate(packed);
 Expr prototype_var = Variable::make(Handle(), prototype_name);

 const Call *c = packed.as<Call>();

 Stmt result = s;
 for (int idx = (int)c->args.size() - 1; idx >= 0; idx--) {
 Expr arg = c->args[idx];
 const Variable *var = arg.as<Variable>();
 Expr val = Call::make(var->type,
 Call::load_typed_struct_member,
 {e, prototype_var, idx},
 Call::Intrinsic);
 if (stmt_uses_var(result, var->name)) {
 // If a closure is generated for multiple consuming blocks of IR,
 // then some of those blocks might only need some of the field.
 result = LetStmt::make(var->name, val, result);
 }
 }
 result = LetStmt::make(prototype_name, prototype, result);

 return result;
}

} // namespace Internal
} // namespace Halide