normalizers.cpp

Go to the documentation of this file.

#include <type_traits>
 
#include <mim/plug/math/math.h>
 
#include "mim/plug/core/core.h"
 
namespace mim::plug::core {
 
namespace {
 
constexpr nat_t idx_shift_width(u64 size) { return std::max(1_n, Idx::size2bitwidth(size)); }
 
constexpr std::optional<unsigned> idx_shift_amount(u64 size, u64 b) {
    auto width = idx_shift_width(size);
    if (b >= width) return {};
    return static_cast<unsigned>(b);
}
 
constexpr u64 idx_unsigned_max(u64 size) { return size == 0 ? std::numeric_limits<u64>::max() : size - 1; }
 
constexpr u64 idx_signed_max(u64 size) {
    return size == 0 ? static_cast<u64>(std::numeric_limits<s64>::max()) : (size - 1) / 2;
}
 
constexpr u64 idx_signed_min_abs(u64 size) {
    return size == 0 ? static_cast<u64>(std::numeric_limits<s64>::max()) + 1_u64 : size / 2;
}
 
constexpr u64 idx_signed_abs(s64 x) { return x >= 0 ? static_cast<u64>(x) : static_cast<u64>(-(x + 1)) + 1_u64; }
 
constexpr s64 idx_neg(u64 abs) {
    if (abs == static_cast<u64>(std::numeric_limits<s64>::max()) + 1_u64) return std::numeric_limits<s64>::min();
    return -static_cast<s64>(abs);
}
 
constexpr u64 idx_pow2(unsigned k) { return 1_u64 << k; }
 
constexpr bool idx_sign(u64 size, u64 x) {
    // Pre: x is already in range.
    if (size == 0) return x > static_cast<u64>(std::numeric_limits<s64>::max()); // Idx 0 encodes 2^64.
 
    // signed representatives in [-floor(size/2), ceil(size/2)-1]
    return x > (size - 1) / 2;
}
 
constexpr s64 idx_sext(u64 size, u64 x) {
    // Pre: x is already in range.
    if (size == 0) return static_cast<s64>(x);
 
    const u64 max_pos = (size - 1) / 2;
    if (x <= max_pos) return static_cast<s64>(x);
 
    // Negative representative is -(size - x).
    return -static_cast<s64>(size - x);
}
 
constexpr u64 idx_from_signed(u64 size, s64 x) {
    if (size == 0) return static_cast<u64>(x);
    return x >= 0 ? static_cast<u64>(x) : size - static_cast<u64>(-x);
}
 
constexpr u64 idx_from_signed_mod(u64 size, s64 x) {
    if (size == 0) return static_cast<u64>(x);
    if (x >= 0) return static_cast<u64>(x) % size;
 
    auto rem = idx_signed_abs(x) % size;
    return rem == 0 ? 0 : size - rem;
}
 
constexpr bool idx_add_nuw(u64 size, u64 a, u64 b) {
    if (size == 0) return a + b < a;
    return a > size - 1 - b;
}
 
constexpr bool idx_sub_nuw(u64, u64 a, u64 b) { return a < b; }
 
constexpr bool idx_mul_nuw(u64 size, u64 a, u64 b) {
    if (a == 0 || b == 0) return false;
 
    if (size == 0) return b > std::numeric_limits<u64>::max() / a;
    return b > (size - 1) / a;
}
 
constexpr u64 idx_add(u64 size, u64 a, u64 b) {
    if (size == 0) return a + b;
    return a >= size - b ? a - (size - b) : a + b;
}
 
constexpr u64 idx_sub(u64 size, u64 a, u64 b) {
    if (size == 0) return a - b;
    return (a >= b) ? (a - b) : (size - (b - a));
}
 
constexpr u64 idx_mul(u64 size, u64 a, u64 b) {
    if (size == 0) return a * b;
 
    // Safe double-and-add modulo size, avoids overflow.
    u64 r = 0;
    while (b) {
        if (b % 2_u64 != 0) r = idx_add(size, r, a);
        b /= 2_u64;
        if (b) a = idx_add(size, a, a);
    }
    return r;
}
 
constexpr bool idx_add_nsw(u64 size, u64 a, u64 b) {
    const bool sa = idx_sign(size, a);
    const bool sb = idx_sign(size, b);
    const u64 r   = idx_add(size, a, b);
    const bool sr = idx_sign(size, r);
    return (sa == sb) && (sr != sa);
}
 
constexpr bool idx_sub_nsw(u64 size, u64 a, u64 b) {
    const bool sa = idx_sign(size, a);
    const bool sb = idx_sign(size, b);
    const u64 r   = idx_sub(size, a, b);
    const bool sr = idx_sign(size, r);
    return (sa != sb) && (sr != sa);
}
 
constexpr bool idx_mul_nsw(u64 size, u64 a, u64 b) {
    const s64 x = idx_sext(size, a);
    const s64 y = idx_sext(size, b);
 
    if (x == 0 || y == 0) return false;
 
    const s64 min_val = size == 0 ? std::numeric_limits<s64>::min() : -static_cast<s64>(size / 2);
    const s64 max_val = size == 0 ? std::numeric_limits<s64>::max() : static_cast<s64>((size - 1) / 2);
 
    if (x == -1) return size % 2 == 0 && y == min_val;
    if (y == -1) return size % 2 == 0 && x == min_val;
 
    if (x > 0)
        if (y > 0)
            return x > max_val / y;
        else
            return y < min_val / x;
    else if (y > 0)
        return x < min_val / y;
    else
        return x < max_val / y;
}
 
constexpr std::optional<u64> idx_udiv([[maybe_unused]] u64 size, u64 a, u64 b) {
    if (b == 0) return {};
    return a / b;
}
 
constexpr std::optional<u64> idx_urem([[maybe_unused]] u64 size, u64 a, u64 b) {
    if (b == 0) return {};
    return a % b;
}
 
constexpr bool idx_slt(u64 size, u64 a, u64 b) {
    const bool sa = idx_sign(size, a);
    const bool sb = idx_sign(size, b);
 
    if (a == b) return false;
    if (!sa && sb) return false;
    if (sa && !sb) return true;
    return a < b;
}
 
constexpr bool idx_sgt(u64 size, u64 a, u64 b) { return idx_slt(size, b, a); }
 
constexpr bool idx_sdivrem_ub(u64 size, u64 a, u64 b) {
    const s64 x = idx_sext(size, a);
    const s64 y = idx_sext(size, b);
 
    if (y == 0) return true;
 
    const s64 min_val = [&] {
        if (size == 0) return std::numeric_limits<s64>::min();
        return -static_cast<s64>(size / 2);
    }();
 
    return size % 2 == 0 && x == min_val && y == -1;
}
 
constexpr u64 idx_sdiv(u64 size, u64 a, u64 b) {
    const s64 x = idx_sext(size, a);
    const s64 y = idx_sext(size, b);
    return idx_from_signed(size, x / y);
}
 
constexpr u64 idx_srem(u64 size, u64 a, u64 b) {
    const s64 x = idx_sext(size, a);
    const s64 y = idx_sext(size, b);
    return idx_from_signed(size, x % y);
}
 
constexpr bool idx_shl_nuw(u64 size, u64 a, unsigned k) {
    u64 x   = a;
    u64 max = idx_unsigned_max(size);
 
    while (k--) {
        if (x > max / 2_u64) return true;
        x *= 2_u64;
    }
 
    return false;
}
 
constexpr bool idx_shl_nsw(u64 size, u64 a, unsigned k) {
    const s64 x = idx_sext(size, a);
    if (x >= 0) {
        u64 y   = static_cast<u64>(x);
        u64 max = idx_signed_max(size);
        while (k--) {
            if (y > max / 2_u64) return true;
            y *= 2_u64;
        }
    } else {
        u64 y   = idx_signed_abs(x);
        u64 min = idx_signed_min_abs(size);
        while (k--) {
            if (y > min / 2_u64) return true;
            y *= 2_u64;
        }
    }
 
    return false;
}
 
constexpr std::optional<u64> idx_shl(u64 size, u64 a, u64 b, bool nsw, bool nuw) {
    auto k = idx_shift_amount(size, b);
    if (!k) return {};
 
    if (nuw && idx_shl_nuw(size, a, *k)) return {};
    if (nsw && idx_shl_nsw(size, a, *k)) return {};
 
    return idx_mul(size, a, idx_pow2(*k));
}
 
constexpr std::optional<u64> idx_lshr(u64 size, u64 a, u64 b) {
    auto k = idx_shift_amount(size, b);
    if (!k) return {};
    return a / idx_pow2(*k);
}
 
constexpr std::optional<u64> idx_ashr(u64 size, u64 a, u64 b) {
    auto k = idx_shift_amount(size, b);
    if (!k) return {};
 
    auto divisor = idx_pow2(*k);
    auto x       = idx_sext(size, a);
    if (x >= 0) return idx_from_signed(size, static_cast<s64>(static_cast<u64>(x) / divisor));
 
    auto q = (idx_signed_abs(x) + divisor - 1_u64) / divisor;
    return idx_from_signed(size, idx_neg(q));
}
 
template<icmp id>
constexpr bool fold_icmp_idx(u64 size, u64 a, u64 b) {
    const bool su = idx_sign(size, a);
    const bool sv = idx_sign(size, b);
 
    flags_t rel = 0;
    // clang-format off
    if      (a == b)     rel = icmp_mask & flags_t(icmp::xyglE); // equal
    else if (!su &&  sv) rel = icmp_mask & flags_t(icmp::Xygle); // plus, minus
    else if ( su && !sv) rel = icmp_mask & flags_t(icmp::xYgle); // minus, plus
    else if (a > b)      rel = icmp_mask & flags_t(icmp::xyGle); // greater (same sign)
    else rel = icmp_mask & flags_t(icmp::xygLe);                 // less (same sign)
    // clang-format on
 
    return (flags_t(id) & rel) != 0;
}
 
template<class Id, Id id>
std::optional<u64> fold_idx(u64 size, u64 a, u64 b, [[maybe_unused]] bool nsw, [[maybe_unused]] bool nuw) {
    // Pre: a, b already in range for Idx size.
 
    if constexpr (std::is_same_v<Id, wrap>) {
        if constexpr (id == wrap::add) {
            if (nuw && idx_add_nuw(size, a, b)) return {};
            if (nsw && idx_add_nsw(size, a, b)) return {};
            return idx_add(size, a, b);
 
        } else if constexpr (id == wrap::sub) {
            if (nuw && idx_sub_nuw(size, a, b)) return {};
            if (nsw && idx_sub_nsw(size, a, b)) return {};
            return idx_sub(size, a, b);
 
        } else if constexpr (id == wrap::mul) {
            if (nuw && idx_mul_nuw(size, a, b)) return {};
            if (nsw && idx_mul_nsw(size, a, b)) return {};
            return idx_mul(size, a, b);
 
        } else if constexpr (id == wrap::shl) {
            return idx_shl(size, a, b, nsw, nuw);
 
        } else {
            static_assert(false, "missing wrap subtag");
        }
 
    } else if constexpr (std::is_same_v<Id, shr>) {
        if constexpr (id == shr::a)
            return idx_ashr(size, a, b);
        else if constexpr (id == shr::l)
            return idx_lshr(size, a, b);
        else
            static_assert(false, "missing shr subtag");
 
    } else if constexpr (std::is_same_v<Id, div>) {
        if constexpr (id == div::udiv) {
            return idx_udiv(size, a, b);
 
        } else if constexpr (id == div::urem) {
            return idx_urem(size, a, b);
 
        } else if constexpr (id == div::sdiv) {
            if (idx_sdivrem_ub(size, a, b)) return {};
            return idx_sdiv(size, a, b);
 
        } else if constexpr (id == div::srem) {
            if (idx_sdivrem_ub(size, a, b)) return {};
            return idx_srem(size, a, b);
 
        } else {
            static_assert(false, "missing div subtag");
        }
 
    } else if constexpr (std::is_same_v<Id, icmp>) {
        return u64(fold_icmp_idx<id>(size, a, b));
 
    } else if constexpr (std::is_same_v<Id, extrema>) {
        if constexpr (id == extrema::sm)
            return std::min(a, b);
 
        else if constexpr (id == extrema::sM)
            return std::max(a, b);
 
        else if constexpr (id == extrema::Sm)
            return idx_slt(size, a, b) ? a : b;
 
        else if constexpr (id == extrema::SM)
            return idx_sgt(size, a, b) ? a : b;
 
        else
            static_assert(false, "missing extrema subtag");
 
    } else {
        static_assert(false, "missing tag");
    }
}
 
template<class Id, Id id>
const Def* fold(World& world, const Def* type, const Def*& a, const Def*& b, const Def* mode = {}) {
    if (a->isa<Bot>() || b->isa<Bot>()) return world.bot(type);
 
    if (auto la = Lit::isa(a)) {
        if (auto lb = Lit::isa(b)) {
            assert(a->type() == b->type());
 
            auto size = Lit::as(Idx::isa(a->type()));
 
            bool nsw = false, nuw = false;
            if constexpr (std::is_same_v<Id, wrap>) {
                auto lm = Lit::isa(mode);
                if (!lm) return nullptr;
                nsw = fe::has_flag(static_cast<Mode>(*lm), Mode::nsw);
                nuw = fe::has_flag(static_cast<Mode>(*lm), Mode::nuw);
            }
 
            if (size == 1) {
                if constexpr (std::is_same_v<Id, div>) {
                    if (*lb == 0) return world.bot(type);
                }
                if constexpr (std::is_same_v<Id, icmp>)
                    return world.lit(type, u64(fold_icmp_idx<id>(1, 0, 0)));
                else
                    return world.lit(type, 0);
            }
 
            auto res = fold_idx<Id, id>(size, *la, *lb, nsw, nuw);
            return res ? world.lit(type, *res) : world.bot(type);
        }
    }
 
    if (::mim::is_commutative(id) && Def::greater(a, b)) std::swap(a, b);
    return nullptr;
}
 
template<class Id>
const Def* fold(World& world, const Def* type, const Def*& a) {
    if (a->isa<Bot>()) return world.bot(type);
 
    if (auto la = Lit::isa(a)) {
        auto size = Lit::as(Idx::isa(a->type()));
 
        if constexpr (std::is_same_v<Id, abs>) {
            auto x = idx_sext(size, *la);
            if (x >= 0) return world.lit(type, static_cast<u64>(x));
 
            auto y = idx_signed_abs(x);
            if (size % 2_u64 == 0 && y == idx_signed_min_abs(size)) return world.lit(type, *la);
 
            return world.lit(type, y);
        } else {
            static_assert(false, "missing tag");
        }
    }
 
    return nullptr;
}
 
/// Reassociates @p a and @p b according to following rules.
/// We use the following naming convention while literals are prefixed with an `l`:
/// ```
///     a    op     b
/// (x op y) op (z op w)
///
/// (1)     la    op (lz op w) -> (la op lz) op w
/// (2) (lx op y) op (lz op w) -> (lx op lz) op (y op w)
/// (3)      a    op (lz op w) ->  lz op (a op w)
/// (4) (lx op y) op      b    ->  lx op (y op b)
/// ```
template<class Id>
const Def* reassociate(Id id, World& world, const Def* a, const Def* b) {
    if (!is_associative(id)) return nullptr;
 
    auto xy     = Axm::isa<Id>(id, a);
    auto zw     = Axm::isa<Id>(id, b);
    auto la     = a->isa<Lit>();
    auto [x, y] = xy ? xy->template args<2>() : std::array<const Def*, 2>{nullptr, nullptr};
    auto [z, w] = zw ? zw->template args<2>() : std::array<const Def*, 2>{nullptr, nullptr};
    auto lx     = Lit::isa(x);
    auto lz     = Lit::isa(z);
 
    // if we reassociate, we have to forget about nsw/nuw
    auto make_op = [&world, id](const Def* a, const Def* b) {
        if constexpr (std::is_same_v<Id, nat>)
            return world.call(id, Defs{a, b});
        else
            return world.call(id, Mode::none, Defs{a, b});
    };
 
    if (la && lz) return make_op(make_op(a, z), w);             // (1)
    if (lx && lz) return make_op(make_op(x, z), make_op(y, w)); // (2)
    if (lz) return make_op(z, make_op(a, w));                   // (3)
    if (lx) return make_op(x, make_op(y, b));                   // (4)
    return nullptr;
}
 
template<class Id>
const Def* merge_cmps(std::array<std::array<u64, 2>, 2> tab, const Def* a, const Def* b) {
    static_assert(sizeof(sub_t) == 1, "if this ever changes, please adjust the logic below");
    static constexpr size_t num_bits = std::bit_width(Annex::num<Id>() - 1_u64);
 
    auto& world = a->world();
    auto a_cmp  = Axm::isa<Id>(a);
    auto b_cmp  = Axm::isa<Id>(b);
 
    if (a_cmp && b_cmp && a_cmp->arg() == b_cmp->arg()) {
        // push sub bits of a_cmp and b_cmp through truth table
        sub_t res   = 0;
        sub_t a_sub = a_cmp.sub();
        sub_t b_sub = b_cmp.sub();
        for (size_t i = 0; i != num_bits; ++i, res >>= 1, a_sub >>= 1, b_sub >>= 1)
            res |= tab[a_sub & 1][b_sub & 1] << 7_u8;
        res >>= (7_u8 - u8(num_bits));
 
        if constexpr (std::is_same_v<Id, math::cmp>)
            return world.call(math::cmp(res), /*mode*/ a_cmp->decurry()->decurry()->arg(), a_cmp->arg());
        else
            return world.call(icmp(Annex::base<icmp>() | res), a_cmp->arg());
    }
 
    return nullptr;
}
 
} // namespace
 
template<nat id>
const Def* normalize_nat(const Def* type, const Def* callee, const Def* arg) {
    auto& world = type->world();
    auto [a, b] = arg->projs<2>();
    if (is_commutative(id) && Def::greater(a, b)) std::swap(a, b);
    auto la = Lit::isa(a);
    auto lb = Lit::isa(b);
 
    if (la) {
        if (lb) {
            switch (id) {
                case nat::add: return world.lit_nat(*la + *lb);
                case nat::sub: return *la < *lb ? world.lit_nat_0() : world.lit_nat(*la - *lb);
                case nat::mul: return world.lit_nat(*la * *lb);
                case nat::div: return *lb == 0 ? world.lit_nat_0() : world.lit_nat(*la / *lb);
                case nat::rem: return *lb == 0 ? a : world.lit_nat(*la % *lb);
            }
        }
 
        if (*la == 0) {
            switch (id) {
                case nat::add: return b;
                case nat::sub: return a;                 // 0 - b = 0
                case nat::mul: return a;                 // 0 * b = 0
                case nat::div: return world.lit_nat_0(); // 0 / b = 0
                case nat::rem: return world.lit_nat_0(); // 0 % b = 0
            }
        }
 
        if (*la == 1 && id == nat::mul) return b; // 1 * b = b
    }
 
    if (lb) {
        if (*lb == 0) {
            switch (id) {
                case nat::sub: return a;                 // a - 0 = a
                case nat::div: return world.lit_nat_0(); // a / 0 = 0
                case nat::rem: return a;                 // a % 0 = a
                default: break;
            }
        }
        if (*lb == 1) {
            switch (id) {
                case nat::div: return a;                 // a / 1 = a
                case nat::rem: return world.lit_nat_0(); // a % 1 = 0
                default: break;
            }
        }
    }
 
    // (c * x) / b = (c / b) * x and (c * x) % b = 0 if c % b == 0
    if (lb && *lb != 0 && (id == nat::div || id == nat::rem)) {
        if (auto m = Axm::isa(nat::mul, a)) {
            const Def* marg = m->arg();
            auto [c, x]     = marg->projs<2>();
            if (auto lc = Lit::isa(c); lc && *lc != 0 && *lc % *lb == 0) {
                if (id == nat::rem) return world.lit_nat_0();
                return world.call(nat::mul, Defs{world.lit_nat(*lc / *lb), x});
            }
        }
    }
 
    if (a == b) {
        switch (id) {
            case nat::add: return world.call(nat::mul, Defs{world.lit_nat(2), a}); // a + a = 2 * a
            case nat::sub: return world.lit_nat(0);                                // a - a = 0
            case nat::mul: break;
            case nat::div: break;                    // 0 / 0 = 0, so we cannot fold a / a = 1 symbolically
            case nat::rem: return world.lit_nat_0(); // a % a = 0 (even for a = 0, since 0 % 0 = 0)
        }
    }
 
    if (auto res = reassociate<nat>(id, world, a, b)) return res;
 
    return world.raw_app(type, callee, {a, b});
}
 
template<ncmp id>
const Def* normalize_ncmp(const Def* type, const Def* callee, const Def* arg) {
    auto& world = type->world();
 
    if (id == ncmp::t) return world.lit_tt();
    if (id == ncmp::f) return world.lit_ff();
 
    auto [a, b] = arg->projs<2>();
    if (is_commutative(id) && Def::greater(a, b)) std::swap(a, b);
 
    if (a == b) {
        constexpr auto eq_mask = fe::to_underlying(ncmp::e) & 0xff;
        if ((fe::to_underlying(id) & eq_mask) != 0) return world.lit_tt();
        if (id == ncmp::ne) return world.lit_ff();
    }
 
    if (auto la = Lit::isa(a)) {
        if (auto lb = Lit::isa(b)) {
            // clang-format off
            switch (id) {
                case ncmp:: e: return world.lit_bool(*la == *lb);
                case ncmp::ne: return world.lit_bool(*la != *lb);
                case ncmp::l : return world.lit_bool(*la <  *lb);
                case ncmp::le: return world.lit_bool(*la <= *lb);
                case ncmp::g : return world.lit_bool(*la >  *lb);
                case ncmp::ge: return world.lit_bool(*la >= *lb);
                default: std::unreachable();
            }
            // clang-format on
        }
    }
 
    return world.raw_app(type, callee, {a, b});
}
 
template<icmp id>
const Def* normalize_icmp(const Def* type, const Def* c, const Def* arg) {
    auto& world = type->world();
    auto callee = c->as<App>();
    auto [a, b] = arg->projs<2>();
 
    if (auto result = fold<icmp, id>(world, type, a, b)) return result;
    if (id == icmp::f) return world.lit_ff();
    if (id == icmp::t) return world.lit_tt();
    if (a == b) {
        constexpr auto eq_mask = fe::to_underlying(icmp::e) & 0xff;
        if ((fe::to_underlying(id) & eq_mask) != 0) return world.lit_tt();
        if (id == icmp::ne) return world.lit_ff();
    }
 
    return world.raw_app(type, callee, {a, b});
}
 
template<extrema id>
const Def* normalize_extrema(const Def* type, const Def* c, const Def* arg) {
    auto& world = type->world();
    auto callee = c->as<App>();
    auto [a, b] = arg->projs<2>();
    if (auto result = fold<extrema, id>(world, type, a, b)) return result;
    return world.raw_app(type, callee, {a, b});
}
 
const Def* normalize_abs(const Def* type, const Def*, const Def* arg) {
    auto& world           = type->world();
    auto [mem, a]         = arg->projs<2>();
    auto [_, actual_type] = type->projs<2>();
    auto make_res         = [&, mem = mem](const Def* res) { return world.tuple({mem, res}); };
 
    if (auto result = fold<abs>(world, actual_type, a)) return make_res(result);
    return {};
}
 
template<bit1 id>
const Def* normalize_bit1(const Def* type, const Def* c, const Def* a) {
    auto& world = type->world();
    auto callee = c->as<App>();
    auto s      = callee->arg();
    // TODO cope with wrap around
 
    if constexpr (id == bit1::id) return a;
 
    if (auto ls = Lit::isa(s)) {
        switch (id) {
            case bit1::f: return world.lit_idx(*ls, 0);
            case bit1::t: return world.lit_idx(*ls, *ls - 1_u64);
            case bit1::id: std::unreachable();
            default: break;
        }
 
        assert(id == bit1::neg);
        if (auto la = Lit::isa(a)) return world.lit_idx_mod(*ls, ~*la);
    }
 
    return {};
}
 
template<bit2 id>
const Def* normalize_bit2(const Def* type, const Def* c, const Def* arg) {
    auto& world = type->world();
    auto callee = c->as<App>();
    auto [a, b] = arg->projs<2>();
    auto mode   = callee->decurry()->arg();
    auto s      = callee->arg();
    auto ls     = Lit::isa(s);
    // TODO cope with wrap around
 
    if (is_commutative(id) && Def::greater(a, b)) std::swap(a, b);
 
    auto tab = make_truth_table(id);
    if (auto res = merge_cmps<icmp>(tab, a, b)) return res;
    if (auto res = merge_cmps<math::cmp>(tab, a, b)) return res;
 
    auto la = Lit::isa(a);
    auto lb = Lit::isa(b);
 
    // clang-format off
    switch (id) {
        case bit2::    f: return world.lit(type, 0);
        case bit2::    t: if (ls) return world.lit(type, *ls-1_u64); break;
        case bit2::  fst: return a;
        case bit2::  snd: return b;
        case bit2:: nfst: return world.call(bit1::neg,  mode, a);
        case bit2:: nsnd: return world.call(bit1::neg,  mode, b);
        case bit2:: ciff: return world.call(bit2:: iff, mode, Defs{b, a});
        case bit2::nciff: return world.call(bit2::niff, mode, Defs{b, a});
        default:         break;
    }
 
    if (la && lb && ls) {
        switch (id) {
            case bit2::and_: return world.lit_idx    (*ls,   *la &  *lb);
            case bit2:: or_: return world.lit_idx    (*ls,   *la |  *lb);
            case bit2::xor_: return world.lit_idx    (*ls,   *la ^  *lb);
            case bit2::nand: return world.lit_idx_mod(*ls, ~(*la &  *lb));
            case bit2:: nor: return world.lit_idx_mod(*ls, ~(*la |  *lb));
            case bit2::nxor: return world.lit_idx_mod(*ls, ~(*la ^  *lb));
            case bit2:: iff: return world.lit_idx_mod(*ls, ~ *la |  *lb);
            case bit2::niff: return world.lit_idx    (*ls,   *la & ~*lb);
            default: std::unreachable();
        }
    }
 
    // TODO rewrite using bit2
    auto unary = [&](bool x, bool y, const Def* a) -> const Def* {
        if (!x && !y) return world.lit(type, 0);
        if ( x &&  y) return ls ? world.lit(type, *ls-1_u64) : nullptr;
        if (!x &&  y) return a;
        if ( x && !y && id != bit2::xor_) return world.call(bit1::neg, mode, a);
        return nullptr;
    };
    // clang-format on
 
    if (is_commutative(id) && a == b) {
        if (auto res = unary(tab[0][0], tab[1][1], a)) return res;
    }
 
    if (la) {
        if (*la == 0) {
            if (auto res = unary(tab[0][0], tab[0][1], b)) return res;
        } else if (ls && *la == *ls - 1_u64) {
            if (auto res = unary(tab[1][0], tab[1][1], b)) return res;
        }
    }
 
    if (lb) {
        if (*lb == 0) {
            if (auto res = unary(tab[0][0], tab[1][0], a)) return res;
        } else if (ls && *lb == *ls - 1_u64) {
            if (auto res = unary(tab[0][1], tab[1][1], a)) return res;
        }
    }
 
    if (auto res = reassociate<bit2>(id, world, a, b)) return res;
 
    return world.raw_app(type, callee, {a, b});
}
 
const Def* normalize_idx(const Def* type, const Def* c, const Def* arg) {
    auto& world = type->world();
    auto callee = c->as<App>();
    if (auto i = Lit::isa(arg)) {
        if (*i == 0) return world.lit(type, 0); // inhabits every `Idx s`, even a symbolic one
        if (auto s = Lit::isa(Idx::isa(type))) {
            if (*s == 0 || *i < *s) return world.lit_idx(*s, *i);
            if (auto m = Lit::isa(callee->decurry()->arg())) return *m ? world.bot(type) : world.lit_idx_mod(*s, *i);
        }
    }
 
    return {};
}
 
template<shr id>
const Def* normalize_shr(const Def* type, const Def* c, const Def* arg) {
    auto& world = type->world();
    auto callee = c->as<App>();
    auto [a, b] = arg->projs<2>();
    auto s      = Idx::isa(a->type());
    auto ls     = Lit::isa(s);
    auto width  = ls.transform(idx_shift_width);
 
    if (auto result = fold<shr, id>(world, type, a, b)) return result;
 
    if (auto la = Lit::isa(a); la && *la == 0) return a;
 
    if (auto lb = Lit::isa(b)) {
        if (width && *lb >= *width) return world.bot(type);
 
        if (*lb == 0) return a;
    }
 
    return world.raw_app(type, callee, {a, b});
}
 
template<wrap id>
const Def* normalize_wrap(const Def* type, const Def* c, const Def* arg) {
    auto& world = type->world();
    auto callee = c->as<App>();
    auto [a, b] = arg->projs<2>();
    auto mode   = callee->decurry()->arg();
    auto s      = Idx::isa(a->type());
    auto ls     = Lit::isa(s);
    auto width  = ls.transform(idx_shift_width);
 
    if (auto result = fold<wrap, id>(world, type, a, b, mode)) return result;
 
    // clang-format off
    if (auto la = Lit::isa(a)) {
        if (*la == 0) {
            switch (id) {
                case wrap::add: return b;    // 0  + b -> b
                case wrap::sub: break;
                case wrap::mul: return a;    // 0  * b -> 0
                case wrap::shl: return a;    // 0 << b -> 0
            }
        } else if (*la == 1) {
            switch (id) {
                case wrap::add: break;
                case wrap::sub: break;
                case wrap::mul: return b;    // 1  * b -> b
                case wrap::shl: break;
            }
        }
    }
 
    if (auto lb = Lit::isa(b)) {
        if (*lb == 0) {
            switch (id) {
                case wrap::add: return a;    // a  + 0 -> a
                case wrap::sub: return a;    // a  - 0 -> a
                case wrap::mul: return b;    // a  * 0 -> 0
                case wrap::shl: return a;    // a << 0 -> a
            }
        }
 
        if (auto lm = Lit::isa(mode); lm && ls && *lm == 0 && id == wrap::sub)
            return world.call(wrap::add, mode, Defs{a, world.lit_idx(*ls, idx_sub(*ls, 0, *lb))}); // a - lb -> a + (0 - lb)
        else if (id == wrap::shl && width && *lb >= *width)
            return world.bot(type);
    }
 
    if (a == b) {
        switch (id) {
            case wrap::add:
                if (ls && (*ls == 0 || *ls > 2)) return world.call(wrap::mul, mode, Defs{world.lit(type, 2), a}); // a + a -> 2 * a
                break;
            case wrap::sub: return world.lit(type, 0);                                       // a - a -> 0
            case wrap::mul: break;
            case wrap::shl: break;
        }
    }
    // clang-format on
 
    if (auto res = reassociate<wrap>(id, world, a, b)) return res;
 
    return world.raw_app(type, callee, {a, b});
}
 
template<div id>
const Def* normalize_div(const Def* full_type, const Def*, const Def* arg) {
    auto& world    = full_type->world();
    auto [mem, ab] = arg->projs<2>();
    auto [a, b]    = ab->projs<2>();
    auto [_, type] = full_type->projs<2>(); // peel off actual type
    auto make_res  = [&, mem = mem](const Def* res) { return world.tuple({mem, res}); };
 
    if (auto result = fold<div, id>(world, type, a, b)) return make_res(result);
 
    if (auto la = Lit::isa(a)) {
        if (*la == 0) return make_res(a); // 0 / b -> 0 and 0 % b -> 0
    }
 
    if (auto lb = Lit::isa(b)) {
        if (*lb == 0) return make_res(world.bot(type)); // a / 0 -> ⊥ and a % 0 -> ⊥
 
        if (*lb == 1) {
            switch (id) {
                case div::sdiv: return make_res(a);                  // a / 1 -> a
                case div::udiv: return make_res(a);                  // a / 1 -> a
                case div::srem: return make_res(world.lit(type, 0)); // a % 1 -> 0
                case div::urem: return make_res(world.lit(type, 0)); // a % 1 -> 0
            }
        }
    }
 
    if (a == b) {
        switch (id) {
            case div::sdiv:
            case div::udiv:
                if (auto ls = Lit::isa(Idx::isa(type)); ls && *ls != 1)
                    return make_res(world.lit(type, 1)); // a / a -> 1
                break;
            case div::srem: return make_res(world.lit(type, 0)); // a % a -> 0
            case div::urem: return make_res(world.lit(type, 0)); // a % a -> 0
        }
    }
 
    return {};
}
 
template<conv id>
const Def* normalize_conv(const Def* dst_t, const Def*, const Def* x) {
    auto& world = dst_t->world();
    auto s_t    = x->type()->as<App>();
    auto d_t    = dst_t->as<App>();
    auto s      = s_t->arg();
    auto d      = d_t->arg();
    auto ls     = Lit::isa(s);
    auto ld     = Lit::isa(d);
 
    if (s_t == d_t) return x;
    if (x->isa<Bot>()) return world.bot(d_t);
 
    // `Idx 1` has exactly one inhabitant: any conversion into it is the literal 0₁. This also keeps
    // its zero-width arithmetic out of the backend (a 1-extent loop dim would otherwise emit i0).
    if (ld && *ld == 1) return world.lit(d_t, 0);
    if (ls && *ls == 1) return world.lit(d_t, 0);
 
    if (auto l = Lit::isa(x); l && ls && ld) {
        if constexpr (id == conv::u) {
            if (*ld == 0) return world.lit(d_t, *l); // I64
            return world.lit(d_t, *l % *ld);
        }
 
        return world.lit(d_t, idx_from_signed_mod(*ld, idx_sext(*ls, *l)));
    }
 
    if (ls && ld)
        if (auto c1 = Axm::isa(id, x)) {
            auto x1 = c1->arg();
            if (auto ls1 = Lit::isa(x1->type()->as<App>()->arg()))
                if (*ls == 0 || (*ls1 != 0 && *ls > *ls1)) // the intermediate conv is widening
                    if (*ld == *ls1) return x1;            // conv(conv(x)) -> x
        }
 
    return {};
}
 
const Def* normalize_idx2nat(const Def* type, const Def*, const Def* x) {
    auto& world = type->world();
 
    if (x->isa<Bot>()) return world.bot(type);
    if (auto l = Lit::isa(x)) return world.lit_nat(*l);
    if (Idx::isa_lit(x->type()) == 1) return world.lit_nat_0();
 
    return {};
}
 
const Def* normalize_bitcast(const Def* dst_t, const Def*, const Def* src) {
    auto& world = dst_t->world();
    auto src_t  = src->type();
 
    if (src->isa<Bot>()) return world.bot(dst_t);
    if (src_t == dst_t) return src;
 
    if (auto other = Axm::isa<bitcast>(src))
        return other->arg()->type() == dst_t ? other->arg() : world.call<bitcast>(dst_t, other->arg());
 
    if (auto i2n = Axm::isa<idx2nat>(src); i2n && Idx::isa(dst_t))
        return i2n->arg()->type() == dst_t ? i2n->arg() : world.call<bitcast>(dst_t, i2n->arg());
 
    if (auto l = Lit::isa(src)) {
        if (dst_t->isa<Nat>()) return world.lit(dst_t, *l);
        if (auto size = Idx::isa(dst_t)) {
            if (auto s = Lit::isa(size)) return world.lit_idx_mod(*s, *l);
            return world.lit(dst_t, *l);
        }
    }
 
    return {};
}
 
// TODO this currently hard-codes x86_64 ABI
// TODO in contrast to C, we might want to give singleton types like 'Idx 1' or '[]' a size of 0 and simply nuke each
// and every occurance of these types in a later phase
// TODO Pi and others
template<trait id>
const Def* normalize_trait(const Def*, const Def*, const Def* type) {
    auto& world = type->world();
    if (Axm::isa<Ptr>(type)) {
        return world.lit_nat(8);
    } else if (type->isa<Pi>()) {
        return world.lit_nat(8); // Gets lowered to function ptr
    } else if (auto size = Idx::isa(type)) {
        if (auto w = Idx::size2bitwidth(size)) return world.lit_nat(std::max(1_n, std::bit_ceil(*w) / 8_n));
    } else if (auto w = math::isa_f(type)) {
        switch (*w) {
            case 16: return world.lit_nat(2);
            case 32: return world.lit_nat(4);
            case 64: return world.lit_nat(8);
            default: std::unreachable();
        }
    } else if (type->isa<Sigma>()) {
        u64 offset = 0;
        u64 align  = 1;
        for (auto t : type->ops()) {
            auto a = Lit::isa(core::op(trait::align, t));
            auto s = Lit::isa(core::op(trait::size, t));
            if (!a || !s) return {};
 
            align  = std::max(align, *a);
            offset = fe::pad(offset, *a) + *s;
        }
 
        offset   = fe::pad(offset, align);
        u64 size = std::max(1_u64, offset);
 
        switch (id) {
            case trait::align: return world.lit_nat(align);
            case trait::size: return world.lit_nat(size);
        }
    } else if (auto arr = type->isa_imm<Arr>()) {
        // One axis at a time, so a fused Arr multiplies in every extent: `«a, b; T»` is `a * size «b; T»`.
        auto elem = arr->elem();
        if (!elem) return {};
        auto align = op(trait::align, elem);
        if constexpr (id == trait::align) return align;
        auto b = op(trait::size, elem);
        if (!Axm::isa(trait::size, b)) return world.call(nat::mul, Defs{arr->arity(), b});
    }
 
    return {};
}
 
template<pe id>
const Def* normalize_pe(const Def* type, const Def*, const Def* arg) {
    auto& world = type->world();
 
    if constexpr (id == pe::is_closed) {
        if (Axm::isa(pe::hlt, arg)) return world.lit_ff();
        if (arg->is_closed()) return world.lit_tt();
    }
 
    return {};
}
 
MIM_core_NORMALIZER_IMPL
 
} // namespace mim::plug::core