def.cpp

Go to the documentation of this file.

#include "mim/def.h"
 
#include <algorithm>
 
#include <fe/assert.h>
#include <fe/hash.h>
#include <fe/term.h>
#include <fe/worklist.h>
 
#include "mim/driver.h"
#include "mim/rule.h"
#include "mim/world.h"
 
using namespace std::literals;
 
#ifndef DOXYGEN // fe::Patricia is not part of the documented input
template void fe::Patricia<const mim::Var, mim::DefKey>::Set::dump() const;
template void fe::Patricia<mim::Def, mim::DefKey>::Set::dump() const;
#endif
 
namespace mim {
 
std::ostream& DefKey::stream(std::ostream& os, const Def* d) { return os << d->sym() << ": " << d->gid(); }
 
/*
 * constructors
 */
 
Def::Def(World* world, Node node, const Def* type, Defs ops, flags_t flags)
    : world_(world)
    , flags_(flags)
    , node_(node)
    , mut_(false)
    , external_(false)
    , annex_(false)
    , has_hole_(false)
    , num_ops_(ops.size())
    , type_(type) {
    if (node == Node::Level) {
        hash_ = fe::hash_combine(0, node_t(Node::Level));
        return;
    }
 
    // A gid is a u32, so two of them - or a Node and one - share a single hash step.
    auto h = fe::hash_combine(0, u64(node) << 32 | (type ? type->gid() : 0));
    h      = fe::hash_combine(h, flags_);
 
    auto has_hole = type && type->has_hole_;
 
    auto ptr  = ops_ptr();
    auto gids = u64(0);
    for (size_t i = 0, e = ops.size(); i != e; ++i) {
        auto op = ops[i];
        ptr[i]  = op;
        gids    = gids << 32 | op->gid();
        if (i & 1) h = fe::hash_combine(h, gids), gids = 0;
        if (node != Node::Proxy || flags & (Proxy::Vars_Bit | Proxy::Muts_Bit)) has_hole |= op->has_hole_;
    }
    if (gids) h = fe::hash_combine(h, gids); // gid 0 is never handed out, so this odd tail cannot alias a pair
 
    has_hole_ = has_hole;
    hash_     = h;
}
 
void Def::collect_locals() {
    auto vars    = Vars();
    auto muts    = Muts();
    auto collect = [&](const Def* d, bool with_vars = true, bool with_muts = true) {
        auto dvars = with_vars ? d->local_vars() : Vars();
        auto dmuts = with_muts ? d->local_muts() : Muts();
        if ((vars && dvars) || (muts && dmuts)) {
            vars = world().vars().merge(vars, dvars);
            muts = world().muts().merge(muts, dmuts);
        } else {
            if (dvars) vars = dvars;
            if (dmuts) muts = dmuts;
        }
    };
 
    if (type()) collect(type());
    for (auto op : ops())
        if (node() == Node::Proxy)
            collect(op, flags_ & Proxy::Vars_Bit, flags_ & Proxy::Muts_Bit);
        else
            collect(op);
 
    vars_ = vars;
    muts_ = muts;
}
 
Def::Def(Node n, const Def* type, Defs ops, flags_t flags)
    : Def(nullptr, n, type, ops, flags) {}
 
Def::Def(Node node, const Def* type, size_t num_ops, flags_t flags)
    : flags_(flags)
    , node_(node)
    , mut_(true)
    , external_(false)
    , annex_(false)
    , has_hole_(node == Node::Hole)
    , num_ops_(num_ops)
    , type_(type) {
    var_ = nullptr;
    std::fill_n(ops_ptr(), num_ops, nullptr);
}
 
Def::Def(Node node, Def* binder)
    : binder_(binder) // the binder is stored here instead of as an op, so a Var stays out of the operand graph
    , flags_(0)
    , node_(node)
    , mut_(false)
    , external_(false)
    , annex_(false)
    , has_hole_(false)
    , num_ops_(0)
    , type_(nullptr) {
    vars_ = Vars(as<Var>());
    hash_ = fe::hash_combine(0, u64(Node::Var) << 32 | binder->gid());
}
 
Nat::Nat(World& world)
    : Def(Node, world.type(), Defs{}, 0) {}
 
LMax::LMax(World& world, Defs ops)
    : Def(Node, world.type_level(), ops, 0) {}
 
/*
 * immutabilize
 */
 
bool Def::is_immutabilizable() {
    if (!is_set()) return false;
 
    auto v = has_var();
    for (auto op : deps()) {
        if (v && op->has_free_var(v)) return false;
        if (op->local_muts().contains(this)) return false; // recursion
    }
    return true;
}
 
/*
 * reduce
 */
 
Defs Def::reduce_(const Def* arg) const {
    if (!is_set()) fe::throwf("cannot reduce `{}`: it is not set", this);
    if (auto var = has_var()) return world().reduce(var, arg);
    auto off = reduction_offset();
    assert(off != size_t(-1) && "not a binder");
    return {ops().begin() + off, num_ops() - off};
}
 
/*
 * Def - set
 */
 
void Def::watch() const {
#ifdef MIM_ENABLE_CHECKS
    if (world().watchpoints().contains(gid())) fe::breakpoint();
#endif
}
 
Def* Def::finalize() {
    if (auto t = check()->zonk(); t != type()) type_ = t;
    return this;
}
 
Def* Def::set(Defs ops) {
    watch();
    invalidate();
 
    size_t n = ops.size();
    assert(n == num_ops() && "num ops don't match");
 
    for (size_t i = 0; i != n; ++i) {
        auto def = check(i, ops[i]);
        assert(def);
        ops_ptr()[i] = def;
    }
#ifndef NDEBUG
    curr_op_ = n;
#endif
 
    return finalize();
}
 
Def* Def::set(size_t i, const Def* def) {
    watch();
    invalidate();
    def = check(i, def);
    assert(def && !op(i) && curr_op_++ == i);
    ops_ptr()[i] = def;
 
    if (i + 1 == num_ops()) return finalize(); // set last op, so check kind
    return this;
}
 
Def* Def::set_type(const Def* type) {
    invalidate();
    type_ = type;
    return this;
}
 
Def* Def::unset() {
    invalidate();
#ifndef NDEBUG
    curr_op_ = 0;
#endif
    std::fill_n(ops_ptr(), num_ops(), nullptr);
    return this;
}
 
/*
 * free_vars
 */
 
const Def* Def::var() {
    if (var_) return var_;
    return world().var(this);
}
 
std::pair<const Var*, std::optional<nat_t>> Var::isa_tproj(const Def* def) {
    if (auto var = def->isa<Var>()) {
        if (var->num_tprojs() == 1 && var->proj(1, 0) == var) return {var, 0};
        return {var, std::nullopt};
    }
    if (auto ex = def->isa<Extract>()) {
        if (auto var = ex->tuple()->isa<Var>()) {
            // A thresholded Var's only tproj is the Var itself.
            if (auto i = Lit::isa(ex->index()); i && var->num_tprojs() == var->num_projs()) return {var, i};
            return {var, std::nullopt};
        }
    }
    return {};
}
 
// clang-format off
const Def* Def::var_type() {
    switch (mut_node()) {
        case MutNode::Lam:     return as<Lam >()->dom();
        case MutNode::Pi:      return as<Pi  >()->dom();
        case MutNode::Rule:    return as<Rule>()->dom();
        case MutNode::Arr:
        case MutNode::Pack:    return world().type_indices(as<Seq>()->shape());
        case MutNode::Sigma:   return this;
        case MutNode::Hole:
        case MutNode::Variant: return nullptr;
    }
    std::unreachable();
}
// clang-format on
 
Vars Def::free_vars() const {
    if (auto mut = isa_mut()) return mut->free_vars();
 
    auto& vars = world().vars();
    auto fvs   = local_vars();
    for (auto mut : local_muts())
        fvs = vars.merge(fvs, mut->free_vars());
 
    return fvs;
}
 
// free_vars() is a union, so any predicate over it distributes over that union: ask f for the local Vars and
// then for each local mutable's free_vars() - instead of merging them all into one throw-away Set.
static bool any_free_vars(const Def* def, std::predicate<Vars> auto f) {
    if (auto mut = def->isa_mut()) return f(mut->free_vars());
 
    if (f(def->local_vars())) return true;
    for (auto mut : def->local_muts())
        if (f(mut->free_vars())) return true;
 
    return false;
}
 
bool Def::has_free_var(const Var* var) const {
    return any_free_vars(this, [var](Vars fvs) { return fvs.contains(var); });
}
 
bool Def::has_free_vars() const {
    return any_free_vars(this, [](Vars fvs) { return !fvs.empty(); });
}
 
bool Def::has_free_vars_in(Vars vars) const {
    return vars && any_free_vars(this, [vars](Vars fvs) { return vars.has_intersection(fvs); });
}
 
Vars Def::free_vars() {
    if (mark_ == 0) {
        // fixed-point iteration to recompute free vars:
        // (run - 1) identifies the previous iteration; so make sure to offset run by 2 for the first iteration
        auto& w     = world();
        bool cyclic = false;
        w.next_run();
        free_vars<true>(w, cyclic, w.next_run());
 
        for (bool todo = cyclic; todo;) {
            todo = false;
            free_vars<false>(w, todo, w.next_run());
        }
    }
 
    return vars_;
}
 
// w is threaded through instead of re-deriving it via world() - which chases the type chain - at every node.
template<bool init>
Vars Def::free_vars(World& w, bool& todo, u32 run) {
    // If init == true : todo flag detects cycle.
    // If init == false: todo flag keeps track whether sth changed.
    //
    // Recursively recompute free vars. If
    // * mark_ == 0:        Invalid - need to recompute.
    // * mark_ == run - 1:  Previous iteration - need to recompute.
    // * mark_ == run:      We are running in cycles within the *current* iteration of our fixed-point loop.
    // * otherwise:         Valid!
    if (mark_ != 0 && mark_ != run - 1) {
        if constexpr (init) todo |= mark_ == run;
        return vars_;
    }
 
    mark_ = run;
 
    auto fvs0  = vars_;
    auto fvs   = fvs0;
    auto& muts = w.muts();
    auto& vars = w.vars();
 
    for (auto op : deps()) {
        if constexpr (init) fvs = vars.merge(fvs, op->local_vars());
 
        for (auto mut : op->local_muts()) {
            if constexpr (init) mut->muts_ = muts.insert(mut->muts_, this); // register "this" as user of local_mut
            fvs = vars.merge(fvs, mut->free_vars<init>(w, todo, run));
        }
    }
 
    if (auto var = has_var()) fvs = vars.erase(fvs, var); // FV(λx.e) = FV(e) \ {x}
 
    if constexpr (!init) todo |= fvs0 != fvs;
 
    return vars_ = fvs;
}
 
void Def::invalidate() {
    if (mark_ != 0) {
        mark_ = 0;
        // TODO optimize if vars empty?
        for (auto mut : users())
            mut->invalidate();
        vars_ = Vars();
        muts_ = Muts();
    }
}
 
bool Def::is_closed() const {
    bool closed = !has_free_vars();
    assert((!is_external() || closed) && "an external must not have free Vars");
    return closed;
}
 
Def* Def::outermost_binder() const {
    auto fvs = free_vars();
    if (fvs.empty()) return isa_mut();
    // Terminates: the binder of a free Var of `this` sits strictly further out than `this`.
    return fvs.min()->binder()->outermost_binder();
}
 
bool Def::nests(Def* mut, MutSet& checked) {
    auto var = has_var(); // `this` is fixed across the recursion, and so is its Var
    auto fvs = mut->free_vars();
    if (fvs.contains(var)) return true;
    if (auto [_, ins] = checked.emplace(mut); !ins) return false;
 
    for (auto fv : fvs)
        if (this->nests(fv->binder(), checked)) return true;
 
    return false;
}
 
bool Def::nests(Def* mut) {
    if (!this->has_var()) return false;
    auto checked = MutSet();
    return this->nests(mut, checked);
}
 
bool Def::nests(const Def* def) {
    if (auto mut = def->isa_mut()) return this->nests(mut);
    if (!this->has_var()) return false;
 
    auto checked = MutSet();
    for (auto fv : def->free_vars())
        if (this->nests(fv->binder(), checked)) return true;
 
    return false;
}
 
bool Def::in_scope(Def* mut) const {
    return std::ranges::all_of(free_vars(), [mut](const Var* fv) {
        auto binder = fv->binder();
        return binder == mut || binder->nests(mut);
    });
}
 
/*
 * Def - misc
 */
 
Driver& Def::driver() const noexcept { return world().driver(); }
fe::Error& Def::error() const noexcept { return driver().error(); }
 
void Def::bail_expected(fe::Cite what) const {
    auto loc = err_loc();
    auto& e  = error().e(loc, "expected {}", what);
    if (!loc) e.n("got `{}`", this); // no Loc means no snippet showing the offender
    e.bail();
}
 
Loc Def::err_loc() const {
    auto& w = world();
    if (auto loc = w.get_loc()) return loc;
    if (auto loc = this->loc()) return loc;
 
    // Nothing was pushed and def itself is anonymous: settle for the closest Loc among its deps.
    auto queue = fe::BFSWorklist<DefSet>{this};
 
    // Charged per inspected dep, not per dequeue: a single high-arity Def would otherwise scan - and
    // enqueue - its whole operand list before the budget got a say.
    constexpr size_t Budget = 512; // a diagnostic is not worth traversing the whole World for
    auto budget             = Budget;
 
    while (!queue.empty()) {
        for (auto dep : queue.pop()->deps()) {
            if (budget-- == 0) return {};
            if (!queue.push(dep)) continue;
            if (auto loc = dep->loc()) return loc;
        }
    }
 
    return {};
}
 
Sym Def::sym(const char* s) const { return world().sym(s); }
Sym Def::sym(std::string_view s) const { return world().sym(s); }
Sym Def::sym(std::string s) const { return world().sym(std::move(s)); }
 
Dbg Def::dbg() const { return world().driver().dbg(dbg_); }
void Def::set_dbg(Dbg d) const { dbg_ = world().driver().dbg(d); }
 
// Fills in only what is missing (unless @p ow) and interns *once* - the old form went through set(Loc) and
// set(Sym) separately, each re-reading Def::dbg_ and re-interning through the Driver's hash table.
void Def::set_dbg_(Dbg d, bool ow) const {
    if (ow) return set_dbg(d);
 
    auto mine = dbg(), merged = mine;
    if (!merged.loc()) merged.set(d.loc());
    if (!merged.sym()) merged.set(d.sym());
    if (!(merged == mine)) set_dbg(merged);
}
 
void Def::set_dbg_key_(DbgKey key, bool ow) const {
    // Nothing of ours to preserve, so adopting the index is exactly what set<Ow>(that Dbg) would compute -
    // but without materialising a Dbg or touching Driver::dbg2idx_ at all.
    if (ow || !dbg_ || dbg_ == key) return void(dbg_ = key);
    set_dbg_(world().driver().dbg(key), false); // rare: we carry a partial Dbg, so merge field-wise
}
 
const Def* Def::unfold_type() const {
    if (auto t = type()) return t;
    auto& w = world();
    if (auto t = isa<Type>()) return w.type(w.lsucc(t->level()));
    assert(isa<Level>());
    return nullptr;
}
 
std::string_view Def::node_name() const {
    static constexpr std::string_view Names[Num_Nodes] = {
#define CODE(node, _) #node,
        MIM_NODE(CODE)
#undef CODE
    };
    return Names[node_t(node())];
}
 
Defs Def::deps() const noexcept {
    // Level, Type, and Var are the only nodes built without a type - and none of them has deps.
    if (!type_) {
        assert(isa<Level>() || isa<Type>() || isa<Var>());
        return Defs();
    }
 
    // Not is_set(): its assertion scans *all* ops and deps() is an inner loop.
    bool set = num_ops_ == 0 || ops_ptr()[num_ops_ - 1];
    return Defs(ops_ptr() - 1, (set ? num_ops_ : 0) + 1);
}
 
bool Def::is_term() const {
    if (auto t = type())
        if (auto u = t->type())
            if (auto type = u->isa<Type>()) return Lit::isa(type->level()) == nat_t(0);
    return false;
}
 
#ifndef NDEBUG
const Def* Def::debug_prefix(std::string prefix) const { return set_dbg(dbg().set(sym(prefix + sym().str()))), this; }
const Def* Def::debug_suffix(std::string suffix) const { return set_dbg(dbg().set(sym(sym().str() + suffix))), this; }
#endif
 
/*
 * cmp
 */
 
Def::Cmp Def::cmp(const Def* a, const Def* b) {
    if (a == b) return Cmp::E;
 
    if (a->is_mutable() != b->is_mutable()) return a->is_mutable() ? Cmp::G : Cmp::L;
 
    // clang-format off
    if (a->node()    != b->node()   ) return a->node()    < b->node()    ? Cmp::L : Cmp::G;
    if (a->num_ops() != b->num_ops()) return a->num_ops() < b->num_ops() ? Cmp::L : Cmp::G;
    if (a->flags()   != b->flags()  ) return a->flags()   < b->flags()   ? Cmp::L : Cmp::G;
    // clang-format on
 
    if (a->is_mutable()) return Cmp::U; // two distinct mutables of the same shape are incomparable
 
    if (auto va = a->isa<Var>()) {
        auto vb = b->as<Var>();
        auto ma = va->binder();
        auto mb = vb->binder();
        if (ma->is_set() && ma->free_vars().contains(vb)) return Cmp::L;
        if (mb->is_set() && mb->free_vars().contains(va)) return Cmp::G;
        return Cmp::U;
    }
 
    // heuristic: iterate backwards as index (often a Lit) comes last and will faster find a solution
    for (size_t i = a->num_ops(); i-- != 0;)
        if (auto res = cmp(a->op(i), b->op(i)); res == Cmp::L || res == Cmp::G) return res;
 
    return cmp(a->type(), b->type());
}
 
template<Def::Cmp c>
bool Def::cmp_(const Def* a, const Def* b) {
    auto res = cmp(a, b);
    if (res == Cmp::U) {
        a->world().log().w("commute check resorts to an unstable gid-based compare");
        return c == Cmp::L ? a->gid() < b->gid() : a->gid() > b->gid();
    }
    return res == c;
}
 
// clang-format off
bool Def::less   (const Def* a, const Def* b) { return cmp_<Cmp::L>(a, b); }
bool Def::greater(const Def* a, const Def* b) { return cmp_<Cmp::G>(a, b); }
// clang-format on
 
// clang-format off
 
/*
 * dispatch
 *
 * All of these used to be `virtual`. Def has no vtable - a vptr would cost 8 bytes on *every* node in the
 * World - so each is one function switching on Def::node() with the former override inlined right here.
 */
 
const Def* Def::immutabilize() {
    auto& w = world();
    switch (mut_node()) {
        case MutNode::Pi: {
            auto pi = as<Pi>();
            return is_immutabilizable() ? w.pi(pi->dom(), pi->codom(), pi->is_implicit()) : nullptr;
        }
        case MutNode::Sigma:   return is_immutabilizable() ? w.sigma(ops()) : nullptr;
        case MutNode::Variant: return is_immutabilizable() ? w.variant(ops()) : nullptr;
        case MutNode::Rule:    return nullptr; // TODO should we ever immutabilize Rules?
        case MutNode::Lam:
        case MutNode::Hole:    return nullptr;
        case MutNode::Arr:
        case MutNode::Pack: {
            auto seq = as<Seq>();
            auto arr = node() == Node::Arr;
            if (is_immutabilizable())
                return arr ? w.arr(seq->shape(), seq->body()) : w.pack(seq->shape(), seq->body());
            // below the threshold an unrollable Seq becomes an explicit Sigma/Tuple
            if (auto n = seq->shape().is_fused() ? std::nullopt : Lit::isa(seq->arity());
                n && *n < w.flags().scalarize_threshold) {
                auto elems = DefVec(*n, [&](size_t i) { return seq->reduce(w.lit_idx(*n, i)); });
                return arr ? w.sigma(elems) : w.tuple(elems);
            }
            return nullptr;
        }
    }
    std::unreachable();
}
 
size_t Def::reduction_offset() const noexcept {
    if (!is_mut_node(node())) return size_t(-1); // only a mutable can bind a Var
    switch (mut_node()) {
        case MutNode::Lam:
        case MutNode::Sigma:   return 0;
        case MutNode::Arr:
        case MutNode::Pack:
        case MutNode::Pi:
        case MutNode::Rule:    return 1;
        case MutNode::Hole:
        case MutNode::Variant: return size_t(-1);
    }
    std::unreachable();
}
 
const Def* Def::arity() const {
    switch (node()) {
        case Node::Arr:
        case Node::Pack:  return Shape(op(0)).front();
        case Node::Sigma: return num_ops() != 1 || isa_mut() ? world().lit_nat(num_ops()) : op(0)->arity();
        default:
            if (auto t = type(); t && !t->isa<Type>()) return t->arity();
            return world().lit_nat_1();
    }
}
 
// clang-format on
 
void Def::externalize() { return world().externals().externalize(this); }
void Def::internalize() { return world().externals().internalize(this); }
 
void Def::transfer_external(Def* to) {
    assert(this->sym() == to->sym());
    internalize();
    to->externalize();
}
 
std::string Def::unique_name() const { return sym().str() + "_"s + std::to_string(gid()); }
 
static bool is_seq_arity(const Def* def) {
    if (def->isa<Seq>()) return true;
    if (auto t = def->type(); t && !t->isa<Type>()) return is_seq_arity(t);
    return false;
}
 
nat_t Def::num_tprojs() const {
    if (auto a = Lit::isa(arity()); a && (*a < world().flags().scalarize_threshold || !is_seq_arity(this))) return *a;
    return 1;
}
 
const Def* Def::proj(nat_t a, nat_t i) const {
    if (a == 1) {
        assert(i == 0 && "only inhabitant of Idx 2 is 0_1");
        auto t = type();
        if (!t) return this;
        if (!isa_mut<Sigma>() && !t->isa_mut<Sigma>()) return this;
    }
 
    if (auto seq = isa<Seq>()) {
        auto& w = world();
        if (seq->shape().is_fused()) return w.peel(seq, w.lit_idx(a, i)); // one axis off a fused shape yields a sub-Seq
        return seq->reduce(w.lit_idx(a, i));
    }
 
    if (isa<Prod>()) return op(i);
 
    return world().extract(this, a, i); // only compute world() on the path that needs it
}
 
/*
 * Idx
 */
 
const Def* Idx::isa(const Def* def) {
    if (auto app = def->isa<App>())
        if (app->callee()->isa<Idx>()) return app->arg();
 
    return nullptr;
}
 
std::optional<nat_t> Idx::isa_lit(const Def* def) {
    if (auto size = Idx::isa(def)) return Lit::isa(size);
    return {};
}
 
std::optional<nat_t> Idx::size2bitwidth(const Def* size) {
    if (auto s = Lit::isa(size)) return size2bitwidth(*s);
    return {};
}
 
} // namespace mim