rewrite.cpp

Go to the documentation of this file.

#include "mim/rewrite.h"
 
#include <fe/assert.h>
 
#include "mim/driver.h"
#include "mim/world.h"
 
// Don't use fancy C++-lambdas; it's way too annoying stepping through them in a debugger.
 
namespace mim {
 
/*
 * Rewriter
 */
 
Rewriter::Rewriter(World& world)
    : world_(&world) {
    push(); // create root map
}
 
Rewriter::~Rewriter() = default;
 
void Rewriter::reset() {
    pop();
    assert(scopes_.empty());
    push();
}
 
const Def* Rewriter::map(const Def* old_def, const Def* new_def) {
    for (size_t i = scopes_.size(); i-- != 1;)
        if (auto& scope = scopes_[i];
            old_def->has_free_vars_in(scope.old_vars) || new_def->has_free_vars_in(scope.new_vars)) {
            // A binder rebuilt in here takes its Var - and what depends on it - along.
            if (auto var = old_def->has_var()) scope.old_vars = old_def->world().vars().insert(scope.old_vars, var);
            return scope.old2new[old_def] = new_def;
        }
    return scopes_.front().old2new[old_def] = new_def;
}
 
void Rewriter::bind_old(const Var* var) {
    if (var) scopes_.back().old_vars = var->world().vars().insert(scopes_.back().old_vars, var);
}
 
void Rewriter::bind_new(const Var* var) {
    if (var) scopes_.back().new_vars = var->world().vars().insert(scopes_.back().new_vars, var);
}
 
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_key());
}
 
// 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: std::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: std::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_Level  (const Level*    ) { return world().type_level();                                 }
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_LSucc  (const LSucc*   d) { return world().lsucc  (rewrite(d->op()),     d->offset());   }
const Def* Rewriter::rewrite_imm_LMax   (const LMax*    d) { return world().lmax   (rewrite(d->ops()));                   }
const Def* Rewriter::rewrite_imm_Single (const Single*  d) { return world().single (rewrite(d->op()));                    }
const Def* Rewriter::rewrite_imm_Wrap   (const Wrap*    d) { return world().wrap   (rewrite(d->op()));                    }
const Def* Rewriter::rewrite_imm_Nom    (const Nom*     d) { return world().nom    (d->key(),             rewrite(d->op())); }
const Def* Rewriter::rewrite_imm_Name   (const Name*    d) { return world().name   (rewrite(d->nom()), rewrite(d->op())); }
const Def* Rewriter::rewrite_imm_Struc  (const Struc*   d) { return world().struc  (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_Join   (const Join*    d) { return world().join   (rewrite(d->ops()));                   }
const Def* Rewriter::rewrite_imm_Variant(const Variant* d) { return world().variant(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());
    if (d->type()->isa<Variant>()) return world().inj(new_type, d->index(), new_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_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(), d->tracks_vars(), d->tracks_muts());
}
 
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_Tuple(const Tuple* d) {
    auto new_type = rewrite(d->type());
    auto new_ops  = rewrite(d->ops());
    return world().tuple(new_type, new_ops);
}
 
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_Pi(Pi* d) {
    return rewrite_stub(d, world().mut_pi(rewrite(d->type()), d->is_implicit()));
}
 
const Def* Rewriter::rewrite_mut_Sigma(Sigma* d) {
    return rewrite_stub(d, world().mut_sigma(rewrite(d->type()), d->num_ops()));
}
 
const Def* Rewriter::rewrite_mut_Variant(Variant* d) {
    return rewrite_stub(d, world().mut_variant(rewrite(d->type()), d->num_ops()));
}
 
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 rewrite(tuple->op(*index));
        if (auto pack = ex->tuple()->isa_imm<Pack>(); pack && pack->arity()->is_closed()) return 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_shape = rewrite(*seq->shape());
    if (auto l = Lit::isa(new_shape); l && *l == 0) return world().prod(seq->is_intro());
    return world().seq(seq->is_intro(), new_shape, rewrite(seq->body()));
}
 
const Def* Rewriter::rewrite_mut_Seq(Seq* seq) {
    if (!seq->is_set()) return rewrite_stub(seq, world().mut_seq(seq->is_intro(), rewrite(seq->type())));
 
    auto new_shape = rewrite(*seq->shape())->zonk();
    auto l         = seq->shape().is_fused() ? std::nullopt : Lit::isa(new_shape);
    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 = fe::Vector<const Def*, 32>(*l);
        for (size_t i = 0, e = *l; i != e; ++i) {
            push(Vars(var));
            map(var, world().lit_idx(e, i));
            new_ops[i] = rewrite(seq->body());
            pop();
        }
        return map(seq, world().prod(seq->is_intro(), new_ops));
    }
 
    return rewrite_stub(seq, 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()) return new_mut;
 
    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)));
 
    return seal(old_mut, new_mut);
}
 
const Def* Rewriter::rewrite_stub(Def* old_mut, Def* new_mut, fe::View<size_t> new2old) {
    assert(new_mut->num_ops() == new2old.size());
    map(old_mut, new_mut);
    if (!old_mut->is_set()) return new_mut;
 
    auto _ = enter(old_mut);
    for (size_t i = 0, e = new2old.size(); i != e; ++i)
        new_mut->set(i, rewrite(old_mut->op(new2old[i])));
 
    return seal(old_mut, new_mut);
}
 
const Def* Rewriter::seal(Def* old_mut, Def* new_mut) {
    // A stub that something already captured must stay, or the capture and its users disagree on which Def it is.
    if (new_mut->users().empty() && 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) {
    // Only Vars, mutables and non-ground immutables are ever mapped, so a ground immutable needs no lookup.
    if (!old_def->isa_mut() && old_def->local_vars().empty() && old_def->local_muts().empty()) return 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_key()) : old_mut;
 
    return has_intersection(old_def) ? rewrite_imm(old_def)->set(old_def->dbg_key()) : old_def;
}
 
const Def* VarRewriter::rewrite_mut(Def* mut) {
    bind_old(mut->has_var());
    return Rewriter::rewrite_mut(mut);
}
 
/*
 * Zonker
 */
 
const Def* Zonker::lookup(const Def* old_def) {
    // Holes only ever get set, so a cached result goes stale once a Hole within it got set since.
    // A self-map is a fixpoint, though: e.g. `binder#i` keeps a set Hole in its type until the binder is rewired.
    if (auto new_def = fe::lookup(old2new(), old_def);
        new_def && (new_def == old_def || new_def->isa_mut() || !new_def->needs_zonk()))
        return new_def;
    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  = DefVec(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;
}
 
/*
 * Importer
 */
 
const Def* Importer::rewrite(const Def* def) {
    if (auto new_def = lookup(def)) return new_def;
    if (auto flags = world().driver().library()->annexes().flags(def)) return map(def, world().annex(*flags));
    return Rewriter::rewrite(def);
}
 
} // namespace mim