# %% [markdown]
# # 0307 - Aeneas semantics: Rust (safe subset) -> LLBC -> pure Lean, the Result monad, and what a proof of `not P` certifies
# Companion to `notes/0307 - Aeneas Semantics.pdf`.  Pure Python, exhaustive enumeration over `u8`; no Lean needed.
# The "Rust abstract machine" below is a small stateful interpreter (store, borrows, panics, fuel).  The "translation" is a pure function
# into `Result` with explicit threading of borrowed values (backward functions).  Every check compares them on the **whole** finite domain.
# This is a finite test of the refinement statement, not a proof of Aeneas' soundness theorem (that is a meta-theorem about the tool, see the note).
# %%
import os
os.environ["OMP_NUM_THREADS"] = "1"
import itertools, math
from dataclasses import dataclass
CHECKS = []
def check(name, ok, detail=""):
    CHECKS.append((name, bool(ok), detail))
    assert ok, (name, detail)
    print(f"[ok] {name}  {detail}")
# %% [markdown]
# ## 1. The Result monad: ok / fail(e) / div
# `Result a = ok a | fail Error | div`.  `bind` propagates `fail` and `div`.  Errors we use: `integerOverflow`, `arrayOutOfBounds`, `panic`.
# Monad laws are checked on a finite carrier (all values of `u8`, three error tags, `div`).
# %%
@dataclass(frozen=True)
class R:
    tag: str            # 'ok' | 'fail' | 'div'
    v: object = None    # value for ok, error name for fail
def ok(v): return R('ok', v)
def fail(e): return R('fail', e)
DIV = R('div')
def bind(m, f):
    return f(m.v) if m.tag == 'ok' else m
U8 = range(256)
carrier = [ok(v) for v in U8] + [fail('integerOverflow'), fail('arrayOutOfBounds'), fail('panic'), DIV]
f = lambda x: ok(x + 1) if x < 255 else fail('integerOverflow')   # a fallible step
g = lambda x: ok(x // 2) if x % 3 else fail('panic')
left  = all(bind(ok(v), f) == f(v) for v in U8)
right = all(bind(m, ok) == m for m in carrier)
assoc = all(bind(bind(m, f), g) == bind(m, lambda x: bind(f(x), g)) for m in carrier)
check("1. monad laws on 260 values: left identity, right identity, associativity", left and right and assoc)
# %% [markdown]
# ## 2. Naive translation fails: integers without overflow
# Rust (debug / `overflow-checks`): `fn avg(a: u8, b: u8) -> u8 { (a + b) / 2 }`.  The abstract machine panics iff `a + b > 255`.
# NAIVE translation: `a, b` as mathematical integers, no failure case.  CORRECT translation: `let s <- a + b; s / 2` with checked `+` in `Result`.
# %%
def rust_avg(a, b):                       # abstract machine, checked arithmetic
    s = a + b
    if s > 255: return ('panic', 'integerOverflow')
    return ('ret', s // 2)
def tr_naive_avg(a, b):  return ok((a + b) // 2)            # no Result
def u8_add(a, b):  return ok(a + b) if a + b <= 255 else fail('integerOverflow')
def tr_avg(a, b):  return bind(u8_add(a, b), lambda s: ok(s // 2))
def obs(r):        return ('ret', r.v) if r.tag == 'ok' else (('panic', r.v) if r.tag == 'fail' else ('div', None))
DOM = [(a, b) for a in U8 for b in U8]
bad_naive = sum(1 for a, b in DOM if obs(tr_naive_avg(a, b)) != rust_avg(a, b))
bad_good  = sum(1 for a, b in DOM if obs(tr_avg(a, b)) != rust_avg(a, b))
panics = sum(1 for a, b in DOM if rust_avg(a, b)[0] == 'panic')
check("2. naive translation disagrees with the machine on exactly the panicking inputs", bad_naive == panics == 255 * 256 // 2, f"disagree={bad_naive}, panics={panics}, 255*256/2={255*256//2}")
check("2. Result translation agrees with the machine on all 65536 inputs", bad_good == 0, f"disagree={bad_good}")
# %% [markdown]
# ## 3. Borrow elimination: `&mut` becomes "return the new value" (forward) and, when a borrow is returned, a backward function
# Machine: a store maps addresses to values; `&mut x` is an address.  `fn bump(x: &mut u8) { *x += 1 }`.
# Translation: `bump (x : U8) : Result U8` returns the new value and the caller rebinds.
# %%
def machine_bump(store, addr):            # returns (outcome, store')
    v = store[addr] + 1
    if v > 255: return ('panic', 'integerOverflow'), store          # store unchanged when the panic happens before the write
    s2 = dict(store); s2[addr] = v
    return ('ret', None), s2
def tr_bump(x): return u8_add(x, 1)
good = 0
for x in U8:
    out, st = machine_bump({0: x}, 0)
    r = tr_bump(x)
    if out[0] == 'ret': good += (r.tag == 'ok' and r.v == st[0])
    else: good += (r.tag == 'fail' and r.v == out[1])
check("3a. bump: machine (store, outcome) = translation (returned value, Result) for all u8", good == 256)
# Returned mutable borrow:  fn get_mut(v: &mut Vec<u8>, i: usize) -> &mut u8 ;  caller:  *get_mut(&mut v, i) += k
def machine_incr_at(vec, i, k):
    heap = {j: x for j, x in enumerate(vec)}            # elements are addressable places
    if i >= len(vec): return ('panic', 'arrayOutOfBounds'), vec
    v = heap[i] + k
    if v > 255: return ('panic', 'integerOverflow'), vec
    heap[i] = v
    return ('ret', None), [heap[j] for j in range(len(vec))]
def tr_get_mut(v, i):                      # forward value and BACKWARD function  back : u8 -> Vec
    if i >= len(v): return fail('arrayOutOfBounds')
    back = lambda new, v=v, i=i: tuple(new if j == i else x for j, x in enumerate(v))
    return ok((v[i], back))
def tr_incr_at(v, i, k):
    return bind(tr_get_mut(v, i), lambda p: bind(u8_add(p[0], k), lambda new: ok(p[1](new))))
cases = ties = 0
for n in range(0, 4):
    for vec in itertools.product([0, 1, 254, 255], repeat=n):
        for i in range(0, 5):
            for k in (0, 1, 2):
                out, st = machine_incr_at(list(vec), i, k)
                r = tr_incr_at(vec, i, k)
                cases += 1
                if out[0] == 'ret': ties += (r.tag == 'ok' and list(r.v) == st)
                else: ties += (r.tag == 'fail' and r.v == out[1])
check("3b. returned &mut: machine vs forward value + backward function, vec len<=3, i<=4, k<=2", ties == cases, f"{ties}/{cases} cases")
# %% [markdown]
# ## 4. Loops become recursion; non-termination is `div`
# `fn spin(x: u8) { while x != 0 { } }`.  The machine is run with a fuel bound F; `div` is what survives every F.
# `fn sum_to(n: u8) -> u8 { let mut s = 0; let mut i = 0; while i < n { i += 1; s += i } s }` : the accumulator overflows once n(n+1)/2 > 255, i.e. n >= 23.
# %%
def machine_spin(x, fuel):
    for _ in range(fuel):
        if x == 0: return ('ret', None)
    return ('fuel', None)
def tr_spin(x): return ok(None) if x == 0 else DIV
agree = all(((machine_spin(x, F)[0] == 'ret') == (tr_spin(x).tag == 'ok')) and ((machine_spin(x, F)[0] == 'fuel') == (tr_spin(x).tag == 'div')) for x in U8 for F in (1, 10, 1000))
check("4a. spin: 'div' corresponds to fuel exhaustion for every tested fuel", agree)
def machine_sum_to(n):
    s = i = 0
    while i < n:
        i += 1; s += i
        if s > 255: return ('panic', 'integerOverflow')
    return ('ret', s)
def tr_sum_loop(n, i=0, s=0):               # the recursive function Aeneas would emit
    if i < n:
        i2 = i + 1
        r = u8_add(s, i2)
        return bind(r, lambda s2: tr_sum_loop(n, i2, s2))
    return ok(s)
m = [machine_sum_to(n) for n in U8]
check("4b. sum_to: loop = recursive translation for all n in u8", all(obs(tr_sum_loop(n)) == m[n] for n in U8))
first_bad = min(n for n in U8 if m[n][0] == 'panic')
check("4b. first panicking n is 23 (n(n+1)/2 = 276 > 255, 22*23/2 = 253 <= 255)", first_bad == 23 and 22 * 23 // 2 == 253 and 23 * 24 // 2 == 276, f"first_bad={first_bad}")
# %% [markdown]
# ## 5. A proof of `not P`: closed proposition versus negation under binders
# Property of `avg`:  `P(a,b,r)`: `min(a,b) <= r <= max(a,b)`, evaluated on the translated function, with `Q(a,b) := exists r, avg a b = ok r /\ P(a,b,r)`.
# Closed statement  S := forall a b, H a b -> Q a b  (H trivially true here).   Its negation  not S  <=>  exists a b, H a b /\ not Q a b.
# Different, much stronger statement  S' := forall a b, H a b -> not Q a b.   Under binders, `not Q a b` for a free (a,b) is a claim about *every* instance.
# %%
def Q(a, b):
    r = tr_avg(a, b)
    return r.tag == 'ok' and min(a, b) <= r.v <= max(a, b)
H = lambda a, b: True
S      = all((not H(a, b)) or Q(a, b) for a, b in DOM)
notS   = any(H(a, b) and not Q(a, b) for a, b in DOM)
Sprime = all((not H(a, b)) or (not Q(a, b)) for a, b in DOM)
witnesses = [(a, b) for a, b in DOM if H(a, b) and not Q(a, b)]
check("5a. S is false and (not S) holds, certified by one witness", (not S) and notS and witnesses[0] == (1, 255), f"|W|={len(witnesses)}, first witness {witnesses[0]}")
check("5a. the 'under binders' strengthening S' (forall x, not Q x) is false: Q holds at (0,0)", (not Sprime) and Q(0, 0), f"|{{Q}}|={len(DOM)-len(witnesses)}")
check("5a. |W| = number of panicking inputs = 32640", len(witnesses) == 32640)
# With a precondition H a b := a + b <= 255 the closed statement is TRUE and its negation is not provable
H2 = lambda a, b: a + b <= 255
check("5b. with H := a+b<=255, S_H holds (so no certificate for not S_H exists)", all((not H2(a, b)) or Q(a, b) for a, b in DOM))
# Vacuity: if H is unsatisfiable, the under-binders form is provable and yet there is no witness at all
H3 = lambda a, b: a + b > 255 and a + b < 256
check("5c. unsatisfiable H: S'_H = forall x, H x -> not Q x holds vacuously, but the witness set is empty (no bug evidence)",
      all((not H3(a, b)) or (not Q(a, b)) for a, b in DOM) and not any(H3(a, b) for a, b in DOM))
# And with H satisfiable and S'_H true (every H-input fails): S'_H implies the closed negation
H4 = lambda a, b: a + b >= 256
check("5c. satisfiable H := a+b>=256: S'_H holds (every such input fails) and so implies not S_H", all((not H4(a, b)) or (not Q(a, b)) for a, b in DOM) and any(H4(a, b) for a, b in DOM))
# %% [markdown]
# ## 6. Reachability: the certificate is about `avg`, not about the program that calls it
# Public caller: `fn mean_pct(a: u8, b: u8) -> u8 { avg(min(a, c), min(b, c)) }` for a clamp `c` (a property of the *caller*).
# Reach_c = {(a', b'): a', b' <= c}.  Witnesses W from section 5 are real bugs of the program only if W intersect Reach_c is non-empty.
# %%
def count_formula(c):                       # pairs in [0,c]^2 with a+b >= 256
    return sum(max(0, min(c, 255) - max(0, 256 - a) + 1) for a in range(0, min(c, 255) + 1))
rows = []
for c in (50, 100, 127, 128, 150, 200, 255):
    brute = sum(1 for a, b in witnesses if a <= c and b <= c)
    rows.append((c, brute, count_formula(c)))
    assert brute == count_formula(c)
    print(f"  clamp c={c:3d}: |W ∩ Reach_c| = {brute:6d}   (closed form {count_formula(c)})")
check("6. |W ∩ Reach_c| equals the closed form for all tested c", all(b == f for _, b, f in rows))
check("6. c <= 127: certificate holds for avg yet the program mean_pct never panics (false alarm)", dict((c, b) for c, b, _ in rows)[127] == 0 and dict((c, b) for c, b, _ in rows)[128] == 1, "c=127 -> 0, c=128 -> 1 (128+128=256)")
# end-to-end: run the caller on its whole input domain
def mean_pct(a, b, c): return tr_avg(min(a, c), min(b, c))
for c in (100, 127, 128, 200):
    panics_c = sum(1 for a, b in DOM if mean_pct(a, b, c).tag == 'fail')
    print(f"  mean_pct with clamp {c}: panicking inputs over all 65536 raw (a,b): {panics_c}")
check("6. the caller with c=127 never fails on any of 65536 raw inputs; with c=128 it fails (a,b >= 128)", all(mean_pct(a, b, 127).tag == 'ok' for a, b in DOM) and sum(1 for a, b in DOM if mean_pct(a, b, 128).tag == 'fail') == 128 * 128)
# %% [markdown]
# ## 7. Translation soundness: wrong translations give wrong certificates in both directions
# * U1 (too lax): unbounded integers.  Then `S` is *provable*, the real program violates it: a missed bug.
# * U2 (wrapping, release-mode semantics): `Q` fails on the same inputs but as a wrong value, and panic-freedom becomes provable.  "The bug" depends on which Rust semantics is modelled.
# * U3 (too strict): checks overflow on the *widened* sum `(a as u16 + b as u16)/2`, which is fine in Rust, but translates the u16 add as a u8 add: spurious `fail`, a false certificate.
# %%
def Q_with(trf): return lambda a, b: (lambda r: r.tag == 'ok' and min(a, b) <= r.v <= max(a, b))(trf(a, b))
trU1 = tr_naive_avg
trU2 = lambda a, b: ok(((a + b) % 256) // 2)
def rust_wide(a, b): return ('ret', (a + b) // 2)                  # `((a as u16 + b as u16) / 2) as u8` never panics
trU3 = lambda a, b: bind(u8_add(a, b), lambda s: ok(s // 2))      # wrong: widened add checked at u8
Wset = lambda q: sum(1 for a, b in DOM if not q(a, b))
print("  failing sets |{not Q}|: truth(avg, checked) =", Wset(Q), ", U1 (unbounded) =", Wset(Q_with(trU1)), ", U2 (wrapping) =", Wset(Q_with(trU2)))
check("7a. U1: S is provable under the unbounded translation (no witness), but the program panics on 32640 inputs", Wset(Q_with(trU1)) == 0 and Wset(Q) == 32640)
pf_checked = all(tr_avg(a, b).tag == 'ok' for a, b in DOM); pf_wrap = all(trU2(a, b).tag == 'ok' for a, b in DOM)
same_set = {(a, b) for a, b in DOM if not Q_with(trU2)(a, b)} == set(witnesses)
check("7b. U2: wrapping gives the SAME 32640 witnesses for P but for a different reason (wrong value, not panic); panic-freedom is provable under U2 and false under checked semantics",
      same_set and pf_wrap and not pf_checked, f"|W_wrap|={Wset(Q_with(trU2))}, panic-free: wrap={pf_wrap}, checked={pf_checked}")
# wide version of the program: truth never panics, so S_wide holds; the strict translation certifies a bug
Qwide_true = lambda a, b: min(a, b) <= rust_wide(a, b)[1] <= max(a, b)
check("7c. U3: program `avg_wide` is correct on all inputs; the over-strict translation yields 32640 spurious witnesses", all(Qwide_true(a, b) for a, b in DOM) and Wset(Q_with(trU3)) == 32640)
# %% [markdown]
# ## Summary
# %%
for n, ok_, d in CHECKS: print(("PASS " if ok_ else "FAIL ") + n)
assert all(c[1] for c in CHECKS)
print(f"\n{len(CHECKS)} checks, all passed")