rewrite.cpp

Go to the documentation of this file.

#include "mim/rewrite.h"
 
#include <absl/container/fixed_array.h>
 
#include "mim/world.h"
 
#include "fe/assert.h"
 
// Don't use fancy C++-lambdas; it's way too annoying stepping through them in a debugger.
 
namespace mim {
 
/*
 * Rewriter
 */
 
Rewriter::Rewriter(std::unique_ptr<World>&& ptr)
    : ptr_(std::move(ptr))
    , world_(ptr_.get()) {
    push(); // create root map
}
 
Rewriter::Rewriter(World& world)
    : world_(&world) {
    push(); // create root map
}
 
Rewriter::~Rewriter() = default;
 
void Rewriter::reset(std::unique_ptr<World>&& ptr) {
    ptr_   = std::move(ptr);
    world_ = ptr_.get();
    reset();
}
 
void Rewriter::reset() {
    pop();
    assert(old2news_.empty());
    push();
}
 
const Def* Rewriter::map(const Def* old_def, Defs new_defs) {
    auto new_tuple = world().tuple(new_defs);
    return map(old_def, new_tuple);
}
const Def* Rewriter::map(Defs old_defs, const Def* new_def) {
    auto old_tuple = world().tuple(old_defs);
    return map(old_tuple, new_def);
}
const Def* Rewriter::map(Defs old_defs, Defs new_defs) {
    auto old_tuple = world().tuple(old_defs);
    auto new_tuple = world().tuple(new_defs);
    return map(old_tuple, new_tuple);
}
 
const Def* Rewriter::rewrite(const Def* old_def) {
    if (auto new_def = lookup(old_def)) return new_def;
 
    auto new_def = old_def->isa_mut() ? rewrite_mut((Def*)old_def) : rewrite_imm(old_def);
    return new_def->set(old_def->dbg());
}
 
// clang-format off
#define CODE_MUT(N) case Node::N: new_def = rewrite_mut_##N(old_mut->as<N>()); break;
#define CODE_IMM(N) case Node::N: new_def = rewrite_imm_##N(old_def->as<N>()); break;
// clang-format on
 
const Def* Rewriter::rewrite_imm(const Def* old_def) {
    const Def* new_def;
    switch (old_def->node()) {
        MIM_IMM_NODE(CODE_IMM)
        default: fe::unreachable();
    }
    return map(old_def, new_def);
}
 
const Def* Rewriter::rewrite_mut(Def* old_mut) {
    const Def* new_def;
    switch (old_mut->node()) {
        MIM_MUT_NODE(CODE_MUT)
        default: fe::unreachable();
    }
    return new_def;
}
 
#undef CODE_MUT
#undef CODE_IMM
 
DefVec Rewriter::rewrite(Defs ops) {
    auto new_ops = DefVec(ops.size());
    for (size_t i = 0, e = ops.size(); i != e; ++i)
        new_ops[i] = rewrite(ops[i]);
    return new_ops;
}
 
#ifndef DOXYGEN
// clang-format off
const Def* Rewriter::rewrite_imm_Idx   (const Idx*     ) { return world().type_idx(); }
const Def* Rewriter::rewrite_imm_Nat   (const Nat*     ) { return world().type_nat(); }
const Def* Rewriter::rewrite_imm_Univ  (const Univ*    ) { return world().univ();     }
const Def* Rewriter::rewrite_imm_Lit   (const Lit*    d) { return world().lit   (rewrite(d->type()),   d->get());    }
const Def* Rewriter::rewrite_imm_Match (const Match*  d) { return world().match (rewrite(d->ops()));                 }
const Def* Rewriter::rewrite_imm_Reform(const Reform* d) { return world().reform(rewrite(d->dom()));                 }
const Def* Rewriter::rewrite_imm_Sigma (const Sigma*  d) { return world().sigma (rewrite(d->ops()));                 }
const Def* Rewriter::rewrite_imm_Type  (const Type*   d) { return world().type  (rewrite(d->level()));               }
const Def* Rewriter::rewrite_imm_UInc  (const UInc*   d) { return world().uinc  (rewrite(d->op()),     d->offset()); }
const Def* Rewriter::rewrite_imm_UMax  (const UMax*   d) { return world().umax  (rewrite(d->ops()));                 }
const Def* Rewriter::rewrite_imm_Uniq  (const Uniq*   d) { return world().uniq  (rewrite(d->op()));                  }
const Def* Rewriter::rewrite_imm_Var   (const Var*    d) { return world().var   (rewrite(d->binder())->as_mut());       }
const Def* Rewriter::rewrite_imm_Top   (const Top*    d) { return world().top   (rewrite(d->type()));                }
const Def* Rewriter::rewrite_imm_Bot   (const Bot*    d) { return world().bot   (rewrite(d->type()));                }
const Def* Rewriter::rewrite_imm_Meet  (const Meet*   d) { return world().meet  (rewrite(d->ops()));                 }
const Def* Rewriter::rewrite_imm_Join  (const Join*   d) { return world().join  (rewrite(d->ops()));                 }
 
const Def* Rewriter::rewrite_imm_Arr (const Arr*  d) { return rewrite_imm_Seq(d); }
const Def* Rewriter::rewrite_imm_Pack(const Pack* d) { return rewrite_imm_Seq(d); }
const Def* Rewriter::rewrite_mut_Arr (      Arr*  d) { return rewrite_mut_Seq(d); }
const Def* Rewriter::rewrite_mut_Pack(      Pack* d) { return rewrite_mut_Seq(d); }
// clang-format on
 
const Def* Rewriter::rewrite_imm_App(const App* d) {
    // Rewrite the arg before the callee:
    // the callee may be a recursive mutable that, when rewritten first, eagerly expands its body before the concrete
    // argument is available to specialize it, breaking partial-evaluation termination
    auto new_arg    = rewrite(d->arg());
    auto new_callee = rewrite(d->callee());
    return world().app(new_callee, new_arg);
}
 
const Def* Rewriter::rewrite_imm_Inj(const Inj* d) {
    auto new_type  = rewrite(d->type());
    auto new_value = rewrite(d->value());
    return world().inj(new_type, new_value);
}
 
const Def* Rewriter::rewrite_imm_Insert(const Insert* d) {
    auto new_tuple = rewrite(d->tuple());
    auto new_index = rewrite(d->index());
    auto new_value = rewrite(d->value());
    return world().insert(new_tuple, new_index, new_value);
}
 
const Def* Rewriter::rewrite_imm_Lam(const Lam* d) {
    auto new_type   = rewrite(d->type())->as<Pi>();
    auto new_filter = rewrite(d->filter());
    auto new_body   = rewrite(d->body());
    return world().lam(new_type, new_filter, new_body);
}
 
const Def* Rewriter::rewrite_imm_Merge(const Merge* d) {
    auto new_type = rewrite(d->type());
    auto new_ops  = rewrite(d->ops());
    return world().merge(new_type, new_ops);
}
 
const Def* Rewriter::rewrite_imm_Pi(const Pi* d) {
    auto new_dom   = rewrite(d->dom());
    auto new_codom = rewrite(d->codom());
    return world().pi(new_dom, new_codom, d->is_implicit());
}
 
const Def* Rewriter::rewrite_imm_Proxy(const Proxy* d) {
    auto new_type = rewrite(d->type());
    auto new_ops  = rewrite(d->ops());
    return world().proxy(new_type, new_ops, d->tag());
}
 
const Def* Rewriter::rewrite_imm_Rule(const Rule* d) {
    auto new_type  = rewrite(d->type())->as<Reform>();
    auto new_lhs   = rewrite(d->lhs());
    auto new_rhs   = rewrite(d->rhs());
    auto new_guard = rewrite(d->guard());
    return world().rule(new_type, new_lhs, new_rhs, new_guard);
}
 
const Def* Rewriter::rewrite_imm_Split(const Split* d) {
    auto new_type  = rewrite(d->type());
    auto new_value = rewrite(d->value());
    return world().split(new_type, new_value);
}
 
const Def* Rewriter::rewrite_imm_Tuple(const Tuple* d) {
    auto new_type = rewrite(d->type());
    auto new_ops  = rewrite(d->ops());
    return world().tuple(new_type, new_ops);
}
 
// clang-format on
const Def* Rewriter::rewrite_mut_Pi(Pi* pi) {
    if (pi->is_immutabilizable()) return rewrite_imm_Pi(pi);
    return rewrite_stub(pi, world().mut_pi(rewrite(pi->type()), pi->is_implicit()));
}
const Def* Rewriter::rewrite_mut_Lam(Lam* d) { return rewrite_stub(d, world().mut_lam(rewrite(d->type())->as<Pi>())); }
const Def* Rewriter::rewrite_mut_Rule(Rule* d) {
    return rewrite_stub(d, world().mut_rule(rewrite(d->type())->as<Reform>()));
}
 
const Def* Rewriter::rewrite_mut_Sigma(Sigma* sigma) {
    if (sigma->is_immutabilizable()) return rewrite_imm_Sigma(sigma);
    return rewrite_stub(sigma, world().mut_sigma(rewrite(sigma->type()), sigma->num_ops()));
}
 
const Def* Rewriter::rewrite_mut_Global(Global* d) {
    return rewrite_stub(d, world().global(rewrite(d->type()), d->is_mutable()));
}
 
const Def* Rewriter::rewrite_imm_Axm(const Axm* a) {
    if (&a->world() != &world()) {
        auto type = rewrite(a->type());
        return world().axm(a->normalizer(), a->curry(), a->trip(), type, a->plugin(), a->tag(), a->sub());
    }
    return a;
}
 
const Def* Rewriter::rewrite_imm_Extract(const Extract* ex) {
    auto new_index = rewrite(ex->index());
    if (auto index = Lit::isa(new_index)) {
        if (auto tuple = ex->tuple()->isa<Tuple>()) return map(ex, rewrite(tuple->op(*index)));
        if (auto pack = ex->tuple()->isa_imm<Pack>(); pack && pack->arity()->is_closed())
            return map(ex, rewrite(pack->body()));
    }
 
    auto new_tuple = rewrite(ex->tuple());
    return world().extract(new_tuple, new_index);
}
 
const Def* Rewriter::rewrite_mut_Hole(Hole* hole) {
    auto [last, op] = hole->find();
    return op ? rewrite(op) : rewrite_stub(last, world().mut_hole(rewrite(last->type())));
}
 
#endif
 
const Def* Rewriter::rewrite_imm_Seq(const Seq* seq) {
    auto new_arity = rewrite(seq->arity());
    if (auto l = Lit::isa(new_arity); l && *l == 0) return world().prod(seq->is_intro());
    return world().seq(seq->is_intro(), new_arity, rewrite(seq->body()));
}
 
const Def* Rewriter::rewrite_mut_Seq(Seq* seq) {
    if (seq->is_immutabilizable()) return rewrite_imm_Seq(seq);
 
    if (!seq->is_set()) {
        auto new_seq = seq->as_mut<Seq>()->stub(world(), rewrite(seq->type()));
        return map(seq, new_seq);
    }
 
    auto new_arity = rewrite(seq->arity())->zonk();
    auto l         = Lit::isa(new_arity);
    if (l && *l == 0) return world().prod(seq->is_intro());
 
    if (auto var = seq->has_var(); var && l && *l <= world().flags().scalarize_threshold) {
        auto new_ops = absl::FixedArray<const Def*>(*l);
        for (size_t i = 0, e = *l; i != e; ++i) {
            push();
            map(var, world().lit_idx(e, i));
            new_ops[i] = rewrite(seq->body());
            pop();
        }
        return map(seq, world().prod(seq->is_intro(), new_ops));
    }
 
    if (!seq->has_var()) return map(seq, world().seq(seq->is_intro(), new_arity, rewrite(seq->body())));
    return rewrite_stub(seq->as_mut(), world().mut_seq(seq->is_intro(), rewrite(seq->type())));
}
 
const Def* Rewriter::rewrite_stub(Def* old_mut, Def* new_mut) {
    map(old_mut, new_mut);
 
    if (old_mut->is_set()) {
        auto _ = enter(old_mut);
        for (size_t i = 0, e = old_mut->num_ops(); i != e; ++i)
            new_mut->set(i, rewrite(old_mut->op(i)));
 
        // Immutabilize the *new* binder in hindsight:
        // even when the old binder was not immutabilizable, rewriting may have made it vacuous.
        if (new_mut->is_immutabilizable())
            if (auto new_imm = new_mut->immutabilize()) return map(old_mut, new_imm);
    }
 
    return new_mut;
}
 
/*
 * VarRewriter
 */
 
const Def* VarRewriter::rewrite(const Def* old_def) {
    if (auto new_def = lookup(old_def)) return new_def;
 
    if (auto old_mut = old_def->isa_mut())
        return has_intersection(old_mut) ? rewrite_mut(old_mut)->set(old_mut->dbg()) : old_mut;
 
    if (old_def->local_vars().empty() && old_def->local_muts().empty()) return old_def; // safe to skip
 
    return has_intersection(old_def) ? rewrite_imm(old_def)->set(old_def->dbg()) : old_def;
}
 
const Def* VarRewriter::rewrite_mut(Def* mut) {
    if (auto var = mut->has_var()) {
        auto& vars = vars_.back();
        vars       = world().vars().insert(vars, var);
    }
 
    return Rewriter::rewrite_mut(mut);
}
 
/*
 * Zonker
 */
 
const Def* Zonker::map(const Def* old_def, const Def* new_def) {
    auto repr = lookup(new_def); // always normalize new_def to its representative
    if (!repr) repr = new_def;
    return old2news_.back()[old_def] = repr;
}
 
const Def* Zonker::lookup(const Def* old_def) {
    for (auto& old2new : old2news_ | std::views::reverse) {
        const Def* repr;
        auto path = DefVec();
        while (true) {
            repr = get(old_def);
 
            if (repr == nullptr) break;
 
            path.emplace_back(repr);
            if (repr == old_def) break; // explicit self-map
 
            old_def = repr;
        }
 
        if (path.empty()) continue;
 
        // path compression: flatten all visited nodes
        for (auto def : path)
            old2new[def] = repr;
 
        return repr;
    }
 
    return nullptr;
}
 
const Def* Zonker::rewrite(const Def* def) {
    if (auto hole = def->isa_mut<Hole>()) {
        auto [last, op] = hole->find();
        def             = op ? op : last;
    }
 
    return def->needs_zonk() ? Rewriter::rewrite(def) : def;
}
 
const Def* Zonker::rewire_mut(Def* mut) {
    map(mut, mut);
 
    auto old_type = mut->type();
    auto old_ops  = absl::FixedArray<const Def*>(mut->ops().begin(), mut->ops().end());
 
    mut->unset()->set_type(rewrite(old_type));
 
    for (size_t i = 0, e = mut->num_ops(); i != e; ++i)
        mut->set(i, rewrite(old_ops[i]));
 
    if (auto new_imm = mut->immutabilize()) return map(mut, new_imm);
 
    return mut;
}
 
} // namespace mim