seo.cpp

Go to the documentation of this file.

#include "mim/plug/mem/phase/seo.h"
 
#include <mim/lam.h>
 
#include "mim/plug/mem/mem.h"
 
namespace mim::plug::mem::phase {
 
/*
 * Helpers
 */
 
// All of them track their Vars: liveness and ⊥-site visibility depend on them.
enum {
    Proxy_SCCP_Top, // proxy(var)                        <- var reached ⊤ for propagation but still awaits GVN bundling
    Proxy_Bundle,   // proxy(lam, var1, var2, ..., varn) <- GVN congruence class
    Proxy_Sloxy,    // proxy(lam, ptr)                   <- slot invoked at lam where ptr is the slot cons's ptr var
    Proxy_Phi,      // proxy(lam_var, sloxy)             <- phi we need at lam for sloxy
};
 
// vars are the tvars followed by the phis.
static size_t idx_of(Defs vars, const Def* p) {
    if (auto [_, i] = Var::isa_tproj(p); i) {
        assert(vars[*i] == p);
        return *i;
    }
    auto i = std::ranges::find(vars, p);
    assert(i != vars.end());
    return i - vars.begin();
}
 
static const Proxy* isa_bundle(const Def* def, Lam* lam);
 
void SEO::Analysis::reset() {
    Super::reset();
    visited_.clear();
    lam2sloxy2val_.clear();
    first_.clear();
    non_mem_fvs_.clear();
}
 
/*
 * Main Analysis
 */
 
// SCCP
 
const Proxy* SEO::Analysis::mk_sccp_top(const Def* var) {
    return world().proxy(var->type(), {var}, Proxy_SCCP_Top, true);
}
 
const Def* SEO::Analysis::sccp_join(Lam* lam, const Def* var, const Def* def) {
    log().d("sccp join: {} ⊔ {}", var, def);
    auto cur = lattice(var);
    if (cur == var) return var; // ⊤ is final
 
    // cur ⊔ ⊥ = cur, but a ⊥-site must see what var propagates.
    // As ⊥ never restarts the join, cur is at worst last round's value; a quiet round thus checks the final one.
    if (def->isa<Bot>()) {
        if (!cur) return lattice(var, def), def;
        if (auto caller = curr_mut(); caller && !Proxy::isa<Proxy_Sloxy>(cur) && !cur->in_scope(caller)) {
            log().d("cannot propagate {} → {}: invisible at ⊥-site {}", var, cur, caller);
            return pin(var), var;
        }
        return cur;
    }
 
    // Pin mem.M-typed vars to top: mem must stay threaded through every lam,
    // as later stages (clos conversion, ll backend) rely on each lam having its own mem var.
    if (Axm::isa<mem::M>(var->type())) return pin(var), var;
 
    // A site passing `var` itself (a backedge threading a loop-invariant argument through) contributes
    // nothing: the other call sites alone determine the value.
    if (def == var) return cur ? cur : var;
 
    // `⊥ ⊔ x` is `x`, but unusable if lam nests it.
    // A closed def can never be nested, and Def::nests allocates a fresh MutSet per call - so skip it.
    if (!def->isa<Proxy>() && !def->is_closed() && lam->nests(def)) {
        log().d("cannot propagate {} → {}: out of scope", var, def);
        return pin(var), var;
    }
 
    // Frozen: a ⊤ / this-lam's bundle wins and is kept across rounds (needs cur to exist).
    if (cur && Proxy::isa<Proxy_SCCP_Top>(cur)) return cur;
    if (cur && isa_bundle(cur, lam)) return cur;
 
    // First touch of `var` this round: restart the join here.
    // Discarding what earlier sites contributed requires a round that revisits *every* `lam` call site;
    // apply_known() taints all of lam2callers_ whenever the abstract vars change, so the next round does.
    if (auto [_, ins] = first_.emplace(var); ins) {
        log().d("first contribution; restart: {} → {}", var, def);
        lattice_force(var, def); // may descend from an earlier round's ⊤ - hence force, not lattice()
        return def;
    }
 
    // Every first_ insertion also writes the lattice, so past the first touch `cur` exists;
    // and it cannot be ⊤, as the `cur == var` check bailed out on entry.
    assert(cur && cur != var);
    if (cur == def) return cur;                         // def ⊔ def = def
    if (cur->isa<Bot>()) return lattice(var, def), def; // ⊥ ⊔ def = def
 
    log().d("cannot propagate {} → {}; cur = {}; try GVN", var, def, cur);
    auto top = mk_sccp_top(var);
    return lattice(var, top), top;
}
 
// GVN
 
static const Proxy* isa_bundle(const Def* def, Lam* lam) {
    if (auto bundle = Proxy::isa<Proxy_Bundle>(def); bundle && bundle->op(0) == lam) return bundle;
    return nullptr;
}
 
const Proxy* SEO::Analysis::mk_bundle(Lam* lam, const Def* var, Defs bundle_vars) {
    return world().proxy(var->type(), Def::cat(lam, bundle_vars), Proxy_Bundle, true)->set(var->dbg_key());
}
 
void SEO::Analysis::gvn_bundle(Lam* lam, Defs vars, Defs abstr_args, fe::Span<const Def*> abstr_vars) {
    auto n    = vars.size();
    auto idxs = fe::Vector<size_t>();
 
    for (size_t i = 0; i != n; ++i) {
        if (!Proxy::isa<Proxy_SCCP_Top>(abstr_vars[i])) continue;
 
        idxs.clear();
        idxs.emplace_back(i);
        for (size_t j = i + 1; j != n; ++j)
            if (Proxy::isa<Proxy_SCCP_Top>(abstr_vars[j]) && abstr_args[j] == abstr_args[i]) idxs.emplace_back(j);
 
        if (idxs.size() == 1) {
            pin(vars[i]);
            abstr_vars[i] = vars[i];
        } else {
            auto bundle = mk_bundle(lam, vars[i], DefVec(idxs, [&](size_t j) { return vars[j]; }));
            for (auto j : idxs)
                lattice(vars[j], abstr_vars[j] = bundle);
            log().d("bundle: {}", bundle);
        }
    }
}
 
void SEO::Analysis::gvn_split(Lam* lam, Defs vars, fe::Span<const Def*> abstr_args, fe::Span<const Def*> abstr_vars) {
    // E.g.: Say we started with `{a, b, c, d, e}` as a single bundle for all tvars of `lam`.
    // Now, we see `lam (x, y, x, y, z)`. Then we have to build:
    // a -> {a, c}
    // b -> {b, d}
    // c -> {a, c}
    // d -> {b, d}
    // e -> e      (top)
    auto split = DefVec();
 
    for (size_t i = 0, n = vars.size(); i != n; ++i) {
        if (auto bundle = isa_bundle(abstr_vars[i], lam)) {
            auto members = bundle->ops().subspan(1);
            assert(members.size() > 1 && "a bundle is only ever built for a congruence class");
 
            split.clear();
            for (auto p : members)
                if (abstr_args[idx_of(vars, p)] == abstr_args[i]) split.emplace_back(p);
 
            if (split.size() != members.size()) {
                // vars[i] is a member of its own bundle and trivially agrees with itself, so it is in split.
                if (split.size() == 1) {
                    pin(vars[i]);
                    abstr_vars[i] = vars[i];
                    log().d("GVN single: {}", vars[i]);
                } else {
                    auto new_bundle = mk_bundle(lam, vars[i], split);
                    log().d("GVN split: {}", new_bundle);
                    for (auto p : split) {
                        auto j = idx_of(vars, p);
                        lattice(vars[j], abstr_vars[j] = new_bundle);
                    }
                }
            }
        }
    }
}
 
// SSA
 
// lam's Var makes the phi only visible where lam's future param is.
static const Def* mk_phi(World& w, Lam* lam, const Def* sloxy) {
    return w.proxy(pointee(sloxy), {lam->var(), sloxy}, Proxy_Phi, true)->set(sloxy->dbg_key());
}
 
const Def* SEO::Analysis::lam2sloxy2val(Lam* lam, const Def* sloxy) {
    if (auto sloxy2val = fe::lookup(lam2sloxy2val_, lam))
        if (auto val = fe::lookup(*sloxy2val, sloxy)) return val;
 
    auto phi = mk_phi(world(), lam, sloxy);
    log().d("no value for sloxy {}; use phi {}", sloxy, phi);
    if (auto val = lattice(phi); val && !Proxy::isa<Proxy_SCCP_Top>(val)) return val;
    return nullptr;
}
 
bool SEO::Analysis::can_supply(Lam* lam, const Def* sloxy) {
    if (auto i = lam2callers_.find(lam); i != lam2callers_.end())
        for (auto caller : i->second)
            if (auto c = caller->isa_mut<Lam>(); c && !lam2sloxy2val(c, sloxy)) return false;
    return true;
}
 
static const Var* slot_var(const Def* ptr) { return ptr->as<Extract>()->tuple()->as<Var>(); }
 
/// The sloxy @p abstr_ptr points into, and the literal mem.lea indices on the way, outermost first.
static std::pair<const Proxy*, DefVec> sloxy_path(const Def* abstr_ptr) {
    auto path = DefVec();
    for (auto ptr = abstr_ptr;;) {
        if (auto sloxy = Proxy::isa<Proxy_Sloxy>(ptr)) {
            std::ranges::reverse(path);
            return {sloxy, path};
        } else if (auto lea = Axm::isa<mem::lea>(ptr)) {
            auto [base, index] = lea->args<2>();
            auto agg           = pointee(base);
            auto n             = Lit::isa(agg->arity());
            if (!index->isa<Lit>() || !n || agg->num_tprojs() != *n) return {};
            path.emplace_back(index);
            ptr = base;
        } else
            return {};
    }
}
 
Vars SEO::Analysis::reachable_slots(VarWorklist& todo) {
    if (var2slot_.empty()) return {};
 
    auto& vars = world().vars();
    auto res   = Vars();
    while (!todo.empty()) {
        auto var = todo.pop();
        if (var2slot_.contains(var)) res = vars.insert(res, var);
        if (auto abstr = lattice(var); abstr && abstr != var) todo.push(non_mem_free_vars(abstr));
        if (auto stored = fe::lookup(slot2stored_, var)) todo.push(*stored);
    }
    return res;
}
 
void SEO::Analysis::live(Lam* lam, Defs abstr_args) {
    auto todo = VarWorklist();
    todo.push(lam->free_vars());
    if (auto var = lam->has_var()) todo.push(var);
    for (auto arg : abstr_args)
        todo.push(arg->free_vars());
 
    auto& live = lam2live_[lam];
    auto grown = world().vars().merge(live, reachable_slots(todo));
    if (grown != live) live = grown, touch();
}
 
Vars SEO::Analysis::non_mem_free_vars(const Def* def) {
    if (isa_mem(def) || def->is_closed()) return {};
    if (auto var = def->isa<Var>()) return Vars{var};
    if (def->isa_mut()) return def->free_vars();
    if (auto res = fe::lookup(non_mem_fvs_, def)) return *res;
 
    auto res = Vars();
    for (auto d : def->deps())
        res = world().vars().merge(res, non_mem_free_vars(d));
    return non_mem_fvs_[def] = res;
}
 
void SEO::Analysis::pin_reachable_slots(const Def* abstr_callee, const Def* abstr_arg) {
    if (var2slot_.empty()) return;
 
    auto fvsc = non_mem_free_vars(abstr_callee);
    auto fvsa = non_mem_free_vars(abstr_arg);
    auto todo = VarWorklist();
    todo.push(world().vars().merge(fvsc, fvsa));
    for (auto var : reachable_slots(todo))
        if (auto slot = var2slot_[var]; !is_top(slot)) {
            log().d("slot var {} reaches a direct-style call", var);
            pin_slot(slot);
        }
}
 
void SEO::Analysis::pin_slot(const App* slot) {
    auto [_, ret_lam, __, ptr] = split_slot(slot);
    pin(slot);
    pin(ptr);
    // Once slot is ⊤, apply_known no longer refreshes the whole var, whose stale sloxy would resurface via ptr.
    pin(ret_lam->var());
    if (auto i = slot2resolved_.find(ret_lam->has_var()); i != slot2resolved_.end())
        for (auto concr : i->second)
            lattice_erase(concr);
}
 
void SEO::Analysis::resolve(const Def* sloxy, const Def* concr, const Def* abstr) {
    lattice(concr, abstr);
    slot2resolved_[slot_var(sloxy->op(1))].emplace(concr);
}
 
void SEO::Analysis::propagate_phis(Lam* lam, DefVec& phis, DefVec& abstr_args) {
    const auto& live = lam2live_[lam];
    for (auto ptr : slots()) {
        if (!live.contains(slot_var(ptr))) {
            profile_count("seo.phis.dead");
        } else if (auto sloxy = Proxy::isa<Proxy_Sloxy>(rewrite(ptr))) {
            if (auto value = sloxy2val(sloxy)) {
                auto phi = mk_phi(world(), lam, sloxy);
                phis.emplace_back(phi);
                abstr_args.emplace_back(value);
                log().d("propagate phi {} for sloxy {} with value {}", phi, sloxy, value);
            } else {
                log().d("no value for sloxy {}", sloxy);
            }
        }
    }
}
 
const Def* SEO::Analysis::apply_known(Lam* known, Defs abstr_targs) {
    auto n = abstr_targs.size();
    assert(n == known->num_tvars());
    log().d("known edge: {} → {}", curr_mut(), known);
    if (auto mut = curr_mut()) lam2callers_[known].emplace(mut);
    rewrite(known); // enqueue so its body is drained this round; a no-op if already scheduled
 
    if (known->is_dependent()) {
        log().d("dependent signature; keeping all vars: {}", known);
        for (auto var : known->tvars())
            pin(var);
        return abstract_app(known, world().tuple(abstr_targs));
    }
 
    auto v = version();
 
    auto all_vars = known->tvars();
    DefVec all_abstr_args(n, [&](size_t i) {
        // Strips EtaConv's per-occurrence wrappers so that call sites passing the same continuation still agree.
        if (auto f = Lam::eta_reduce(abstr_targs[i])) return profile_count("seo.eta.canonicalized"), f;
        return abstr_targs[i];
    });
 
    live(known, all_abstr_args);
    propagate_phis(known, all_vars, all_abstr_args);
 
    auto all_abstr_vars = DefVec(all_vars.size());
    for (size_t i = 0, e = all_vars.size(); i != e; ++i)
        all_abstr_vars[i] = sccp_join(known, all_vars[i], all_abstr_args[i]);
 
    gvn_bundle(known, all_vars, all_abstr_args, all_abstr_vars);
    gvn_split(known, all_vars, all_abstr_args, all_abstr_vars);
 
    lattice(known->var(), world().tuple(all_abstr_vars.span().subspan(0, n)));
 
    for (size_t i = n, e = all_vars.size(); i != e; ++i)
        lattice(all_vars[i], all_abstr_vars[i]);
 
    // Something about known's abstract vars/phis changed: the next round must re-join them from *all* call
    // sites - sccp_join's first_-restart discards any unvisited site's contribution, so re-visiting only the
    // writer would be unsound in a sparse round (e.g. it would constant-fold a loop's backedge forever).
    if (version() != v)
        for (auto caller : lam2callers_[known])
            taint(caller);
 
    return abstract_app(known, world().tuple(all_abstr_args.span().subspan(0, n)));
}
 
const Def* SEO::Analysis::abstract_app(const Def* callee, const Def* arg) {
    // Only an Axm keeps the checked, normalizing World::app: its normalizer is what folds a propagated
    // expression. Everywhere else abstract operands (⊤/⊥/proxies) need not be assignable to the domain -
    // and World::app would either throw or partially evaluate the callee.
    if (std::get<0>(Axm::next(callee))) return world().app(callee, arg);
    auto pi = callee->type()->isa<Pi>();
    return world().raw_app(pi ? pi->reduce(arg) : arg->type(), callee, arg);
}
 
const Def* SEO::Analysis::rewrite_imm_App(const App* app) {
    if (auto slot = Axm::isa<mem::slot>(app)) {
        if (!is_top(slot)) {
            if (auto [mem, ret_lam, _, ptr] = split_slot(slot); ret_lam) {
                auto abstr_mem = rewrite(mem);
                auto sloxy     = world().proxy(ptr->type(), {curr_mut(), ptr}, Proxy_Sloxy, true)->set(slot->dbg_key());
                slots_.emplace(ptr);
                var2slot_[ret_lam->has_var()] = slot;
                log().d("slot {} → sloxy {}", ptr, sloxy);
                // The slot is ptr's *defining* site: mark first_ so the ⊥ joined below cannot restart it away.
                fe::assert_emplace(first_, ptr);
                lattice(ptr, sloxy);
                sloxy2val(sloxy, world().bot(pointee(sloxy))); // a load before any store yields undef
                // Treat the slot jump like an app of `ret_lam` so mem and existing phis flow across the edge.
                // The ptr var is defined *by* the slot, so pass ⊥ (not the sloxy) as its abstract argument:
                // this keeps the sloxy out of the abstract body, so it only survives if an unresolved
                // load/store actually references it - `lattice(ptr, sloxy)` above still drives that resolution.
                return apply_known(ret_lam, {abstr_mem, world().bot(ptr->type())});
            }
        }
    } else if (auto store = Axm::isa<mem::store>(app)) {
        auto [mem, ptr, val] = store->args<3>();
        auto abstr_mem       = rewrite(mem);
        auto abstr_ptr       = rewrite(ptr);
        auto abstr_val       = rewrite(val);
 
        if (auto [sloxy, path] = sloxy_path(abstr_ptr); sloxy) {
            if (auto agg = path.empty() ? abstr_val : sloxy2val(sloxy)) {
                auto& stored = slot2stored_[slot_var(sloxy->op(1))];
                if (auto grown = world().vars().merge(stored, non_mem_free_vars(abstr_val)); grown != stored)
                    stored = grown, touch();
                if (!path.empty()) resolve(sloxy, ptr, abstr_ptr); // lets the rewrite recognize a promoted lea
                sloxy2val(sloxy, path.empty() ? abstr_val : world().insert(agg, world().tuple(path), abstr_val));
                log().d("store: {} ← {}", abstr_ptr, abstr_val);
                return abstr_mem;
            }
        }
        log().d("unresolved store: {} ← {}", abstr_ptr, abstr_val);
    } else if (auto load = Axm::isa<mem::load>(app)) {
        auto [mem, ptr] = load->args<2>();
        auto [_, val]   = load->projs<2>();
        auto abstr_mem  = rewrite(mem);
        auto abstr_ptr  = rewrite(ptr);
 
        if (auto [sloxy, path] = sloxy_path(abstr_ptr); sloxy) {
            if (auto agg = sloxy2val(sloxy)) {
                auto abstr_val = path.empty() ? agg : world().extract(agg, world().tuple(path));
                log().d("load: {} → {}", abstr_ptr, abstr_val);
                if (!path.empty()) resolve(sloxy, ptr, abstr_ptr);
                resolve(sloxy, val, abstr_val);
                return world().tuple({abstr_mem, abstr_val});
            }
            log().d("load of unknown value: {}", sloxy);
        } else {
            log().d("load from unknown ptr: {}", abstr_ptr);
        }
    } else {
        auto abstr_callee = rewrite(app->callee());
        auto abstr_arg    = rewrite(app->arg());
        // A direct-style call's effects bypass the continuation the per-Lam store state follows.
        if (has_mem_ty(app->type())) pin_reachable_slots(abstr_callee, abstr_arg);
        if (auto known = Lam::isa_rewritable(abstr_callee)) {
            auto n = known->num_tvars();
            return apply_known(known, abstr_arg->projs(n));
        }
    }
 
    // Rewrite the arg before the callee; see Rewriter::rewrite_imm_App.
    auto new_arg    = rewrite(app->arg());
    auto new_callee = rewrite(app->callee());
    return abstract_app(new_callee, new_arg);
}
 
void SEO::Analysis::leave() {
    if (auto src = curr_mut()->isa<Lam>(); src && src->is_set()) {
        auto phi_dsts       = fe::Vector<Lam*>();
        auto phi_vars       = DefVec();
        auto phi_abstr_args = DefVec();
 
        auto propagate_unknowns = [this, &phi_dsts, &phi_vars, &phi_abstr_args, src](const Def* abstr) {
            for (auto mut : abstr->local_muts())
                if (auto dst = mut->isa<Lam>(); dst && dst->is_open()) {
                    log().d("unknown edge: {} → {}", src, dst);
                    live(dst, {});
                    propagate_phis(dst, phi_vars, phi_abstr_args);
                    phi_dsts.resize(phi_vars.size(), dst);
                }
        };
 
        auto abstr = rewrite(src->body());
        if (auto app = abstr->isa<App>()) {
            if (!app->callee()->isa<Lam>()) propagate_unknowns(app->callee());
            propagate_unknowns(app->arg());
        }
 
        // Several unknown edges may reach the same phi in one round.
        for (size_t i = 0, e = phi_vars.size(); i != e; ++i)
            sccp_join(phi_dsts[i], phi_vars[i], phi_abstr_args[i]);
    }
}
 
/*
 * Post-Analysis:
 * Finds sloxies that are still present + unknown lambdas
 */
 
static bool keep(Lam* lam, const Def* old_var, const Def* abstr) {
    if (!abstr) return true;                            // no info -> keep
    if (old_var == abstr) return true;                  // ⊤       -> keep
    if (abstr->isa<Bot>()) return false;                // ⊥       -> nuke
    if (Proxy::isa<Proxy_SCCP_Top>(abstr)) return true; // pending ⊤: nothing was propagated -> keep
    if (auto bundle = isa_bundle(abstr, lam)) return bundle->op(1) == old_var; // use first in GVN bundle
    return false;                                                              // propagated value -> nuke
}
 
void SEO::Analysis::finalize() {
    for (auto def : world().externals().muts())
        analyze(def);
}
 
void SEO::Analysis::analyze(const Def* def) {
    if (def->isa<Var>()) return; // do not run escape analysis through a Var (would remap it via lookup)
    if (auto [_, ins] = visited_.emplace(def); !ins) return;
    if (auto l = lookup(def)) def = l; // get abstracted value of def
 
    if (auto proxy = def->isa<Proxy>()) {
        if (proxy->tag() == Proxy_Sloxy) {
            auto slot = var2slot_[slot_var(proxy->op(1))];
            log().d("sloxy {} survived; pinning slot {} to ⊤", proxy, slot);
            pin_slot(slot);
        }
        return; // never walk a proxy's deps (would drag in meta info)
    }
 
    // A Lam is unknown (and hence its vars must go to top) iff it is reached as a *value*.
    if (auto app = def->isa<App>()) {
        if (auto slot = Axm::isa<mem::slot>(app)) {
            // The slot jump applies its continuation, so `ret_lam` is known - not reached as a value.
            if (auto [mem, ret_lam, _, __] = split_slot(slot); ret_lam) {
                analyze(app->type());
                analyze(mem); // the ptr var has no argument - the slot itself defines it
                for (auto d : ret_lam->deps())
                    analyze(d);
                return;
            }
        }
 
        if (auto lam = Lam::isa_rewritable(app->callee())) {
            // lam is applied here, it's known: traverse its body without pinning its vars to top
            analyze(app->type());
 
            // only analyze args that we keep
            for (size_t i = 0, e = lam->num_tvars(); i != e; ++i)
                if (auto old_var = lam->tvar(i); keep(lam, old_var, lattice(old_var))) analyze(app->targ(i));
 
            for (auto d : lam->deps())
                analyze(d);
 
            return;
        }
    } else if (auto [lam, var] = def->isa_binder<Lam>(); lam) {
        log().d("unknown lam: {}", lam);
        unknowns_.emplace(lam);
        pin(var); // apply_known's whole-var tuple would otherwise resurface an earlier round's value via repr
        for (auto v : var->tprojs())
            pin(v);
    }
 
    for (auto d : def->deps())
        analyze(d);
}
 
/*
 * Transformation:
 * Apply analysis info to code
 */
 
const Def* SEO::isa_promoted(const Def* old_ptr) const {
    if (auto abstr = lattice(old_ptr))
        if (auto [sloxy, _] = sloxy_path(abstr); sloxy && lattice(sloxy->op(1)) == sloxy) return sloxy;
    return nullptr;
}
 
bool SEO::analyze() {
    if (Super::analyze()) return true;
 
    for (auto ptr : analysis_.slots())
        if (auto sloxy = isa_promoted(ptr)) sloxies_.emplace_back(sloxy);
    return false;
}
 
const SEO::Sig& SEO::sig_of(Lam* old_lam) {
    auto [i, ins] = lam2sig_.emplace(old_lam, Sig());
    auto& sig     = i->second;
    if (!ins) return sig;
 
    auto var = old_lam->has_var();
 
    sig.sieve = Sieve(old_lam->tvars(), [this, old_lam](const Def* var) { return keep(old_lam, var, lattice(var)); });
 
    for (auto sloxy : sloxies_) {
        auto phi = mk_phi(old_world(), old_lam, sloxy);
        if (var && phi->type()->has_free_var(var)) continue; // not expressible in old_lam's signature
        auto val = lattice(phi);
        if (!val || Proxy::isa<Proxy_SCCP_Top>(val)) continue;
        // build_args() needs an argument at *every* call site; else the slot has to survive.
        if (analysis_.can_supply(old_lam, sloxy)) sig.phis.emplace_back(sloxy, phi, val);
    }
 
    sig.sieve.cat(sig.phis.size(), [old_lam, &sig](size_t i) {
        auto [_, phi, val] = sig.phis[i];
        return keep(old_lam, phi, val);
    });
 
    // A new signature is needed iff some var is dropped/propagated/merged or some phi has to be threaded in -
    // except for an unknown lam, which is used as a value somewhere and hence must keep its signature as is.
    sig.todo = !sig.sieve.all() && !analysis_.unknowns().contains(old_lam);
 
    return sig;
}
 
const Def* SEO::rewrite_imm_App(const App* old_app) {
    if (auto slot = Axm::isa<mem::slot>(old_app)) {
        auto [mem, ret_lam, _, ptr] = split_slot(slot);
        if (!ret_lam) return Super::rewrite_imm_App(old_app);
 
        if (isa_promoted(ptr)) {
            // The slot was promoted away: jump straight to the (rebuilt) continuation, dropping the ptr var.
            profile_count("seo.slots.eliminated");
            assert(!analysis_.unknowns().contains(ret_lam)); // promoted -> ret_lam was never reached as a value
            auto new_lam  = build_lam(ret_lam);
            auto new_args = build_args(ret_lam, {mem, ptr});
            return map(old_app, new_world().app(new_lam, new_args));
        }
 
        // The slot survives: keep the allocation, forwarding the continuation.
        auto [T, a]      = slot->decurry()->args<2>();
        auto new_mem     = rewrite(mem);
        auto new_ret_lam = rewrite(ret_lam)->as_mut<Lam>();
        return map(old_app, mem::op_slot(rewrite(T), rewrite(a), new_mem, new_ret_lam));
    } else if (auto store = Axm::isa<mem::store>(old_app)) {
        auto [mem, ptr, val] = store->args<3>();
        if (isa_promoted(ptr)) return rewrite(mem);
    } else if (auto load = Axm::isa<mem::load>(old_app)) {
        auto [res_mem, res_val] = load->projs<2>();
        auto [mem, ptr]         = load->args<2>();
        if (auto sloxy = isa_promoted(ptr)) {
            auto abstr_val = abstracted(res_val);
            assert(abstr_val && "a promoted slot implies every load from it resolved");
            log().d("load from {} resolved to {}", sloxy, abstr_val);
            auto new_mem = rewrite(mem);
            return new_world().tuple({new_mem, rewrite(abstr_val)});
        }
    } else {
        auto old_lam = old_app->callee()->isa_mut<Lam>();
        if (!old_lam) {
            // The callee may fold to a rebuilt lam in the new world only,
            // e.g. a branch `(f, t)#cond` whose cond becomes constant after GVN merged vars.
            if (auto new_lam = rewrite(old_app->callee())->isa_mut<Lam>())
                if (auto ol = fe::lookup(lam_new2old_, new_lam)) old_lam = ol;
        }
 
        if (old_lam) {
            log().d("in {}: app of {}", curr_mut(), old_lam);
 
            if (sig_of(old_lam).todo) {
                log().d("needs seo: {}", old_lam);
                auto new_lam  = build_lam(old_lam);
                auto new_args = build_args(old_lam, old_app->targs());
                return map(old_app, new_world().app(new_lam, new_args));
            }
        }
    }
 
    return Super::rewrite_imm_App(old_app);
}
 
const Def* SEO::rewrite_mut_Lam(Lam* old_lam) {
    // A lam that gets a new signature must never be rebuilt generically:
    // otherwise two new versions of old_lam exist and their (hash-consed, cached) body defs
    // reference whichever version was rewritten first - leaving free vars in the other one.
    if (sig_of(old_lam).todo) return build_lam(old_lam);
    return Super::rewrite_mut_Lam(old_lam);
}
 
const Def* SEO::rewrite_imm_Var(const Var* var) {
    // The generic Var rewrite fabricates var(new_binder) - the wrong arity for a signature build_lam()
    // narrowed. Reachable before build_lam() maps the whole var, e.g. from rewriting a propagated value.
    if (auto old_lam = var->binder()->isa_mut<Lam>())
        if (sig_of(old_lam).todo) return var_of(old_lam);
    return Super::rewrite_imm_Var(var);
}
 
const Def* SEO::var_of(Lam* old_lam) {
    auto new_lam    = build_lam(old_lam);
    const auto& sig = sig_of(old_lam);
    size_t n        = old_lam->num_tvars();
    auto elems      = DefVec(n);
 
    for (size_t i = 0; i != n; ++i)
        if (auto j = sig.sieve[i]; j != Sieve::Gone)
            elems[i] = new_lam->var(sig.sieve.num_new(), j);
        else if (auto abstr = lattice(old_lam->tvar(i)); sloxy_path(abstr).first)
            elems[i] = new_world().bot(rewrite(old_lam->tdom(i))); // a promoted slot carries no value
        else
            elems[i] = rewrite(abstr); // SCCP propagate
 
    return new_world().tuple(elems);
}
 
Lam* SEO::build_lam(Lam* old_lam) {
    assert(!old_lam->is_dependent() && "apply_known pins a dependent signature to \u22a4");
    if (auto memo = fe::lookup(lam_old2new_, old_lam)) return memo;
 
    log().d("build new lam for {}", old_lam);
    invalidate();
    const auto& sig = sig_of(old_lam);
    size_t n        = old_lam->num_tvars();
    auto new_doms   = DefVec();
 
    for (size_t i = 0; i != n; ++i)
        if (sig.sieve[i] != Sieve::Gone)
            new_doms.emplace_back(rewrite(old_lam->tdom(i)));
        else if (auto abstr = lattice(old_lam->tvar(i)); isa_bundle(abstr, old_lam))
            profile_count("seo.gvn.vars_merged");
        else if (!sloxy_path(abstr).first)
            profile_count("seo.sccp.vars_eliminated");
    for (size_t i = 0, e = sig.phis.size(); i != e; ++i)
        if (sig.sieve.cat(i) != Sieve::Gone) new_doms.emplace_back(rewrite(sig.phis[i].phi->type()));
 
    auto new_lam          = new_world().mut_lam(new_doms, rewrite(old_lam->codom()))->set(old_lam->dbg_key());
    lam_old2new_[old_lam] = new_lam;
    lam_new2old_[new_lam] = old_lam;
 
    // Map all *kept* vars/phis before rewriting any propagated value below: that rewrite may recursively
    // re-enter old_lam and must be able to resolve them - see var_of().
    for (size_t i = 0; i != n; ++i)
        if (auto j = sig.sieve[i]; j != Sieve::Gone) {
            auto old_var = old_lam->tvar(i);
            auto v       = new_lam->var(sig.sieve.num_new(), j)->set(old_var->dbg_key());
            map(old_var, v);
            if (auto bundle = isa_bundle(lattice(old_var), old_lam)) map(bundle, v); // GVN bundle
        }
 
    for (size_t i = 0, e = sig.phis.size(); i != e; ++i)
        if (auto j = sig.sieve.cat(i); j != Sieve::Gone) {
            const auto& p = sig.phis[i];
            auto v        = new_lam->var(sig.sieve.num_new(), j);
            profile_count("phis.materialized");
            log().d("map phi {} → {}", p.phi, v);
            map(p.phi, v);
            if (p.val != p.phi) map(p.val, v); // phi is part of a GVN bundle
        }
 
    // Map the whole var *before* rewriting a dropped phi's value below: such a value may project a *dropped*
    // var of old_lam (e.g. its now-removed empty closure env), which only the whole var resolves.
    map(old_lam->var(), var_of(old_lam));
 
    for (size_t i = 0, e = sig.phis.size(); i != e; ++i)
        if (sig.sieve.cat(i) == Sieve::Gone) {
            const auto& p = sig.phis[i];
            log().d("map phi {} → its propagated value {}", p.phi, p.val);
            map(p.phi, rewrite(p.val));
        }
 
    auto _ = enter(old_lam);
    new_lam->set(rewrite(old_lam->filter()), rewrite(old_lam->body()));
    return new_lam;
}
 
DefVec SEO::build_args(Lam* old_lam, Defs old_targs) {
    const auto& sig = sig_of(old_lam);
    size_t n        = old_lam->num_tvars();
    assert(old_targs.size() == n);
    auto new_args = DefVec();
 
    for (size_t i = 0; i != n; ++i)
        if (sig.sieve[i] != Sieve::Gone) new_args.emplace_back(rewrite(old_targs[i]));
 
    for (size_t i = 0, e = sig.phis.size(); i != e; ++i)
        if (sig.sieve.cat(i) != Sieve::Gone) {
            auto sloxy = sig.phis[i].sloxy;
            auto arg   = analysis_.lam2sloxy2val(curr_mut<Lam>(), sloxy);
            assert(arg);
            log().d("wire up phi argument {} for sloxy {}", arg, sloxy);
            new_args.emplace_back(rewrite(arg));
        }
 
    return new_args;
}
 
} // namespace mim::plug::mem::phase