//! Bounded proofs over every reachable state of a small frame table.
//! Each harness starts from an arbitrary *well-formed* table (free count
//! matches the number of empty slots) and checks one operation.

use super::*;
use types::address::{PhysicalAddress, VirtualAddress};
use types::id::{ProcessId, RegionId};
use types::mapping::PagePerms;

const N: usize = 4;
const BASE: usize = 0x8400_0000;

fn any_perms() -> PagePerms {
    match kani::any::<u8>() % 3 {
        0 => PagePerms::ReadOnly,
        1 => PagePerms::ReadWrite,
        _ => PagePerms::ReadExecute,
    }
}

fn any_use() -> FrameUse {
    // Small id ranges keep owner collisions likely, which is where the
    // interesting cases are.
    let owner = ProcessId::from_raw(kani::any::<u16>() % 3);
    match kani::any::<u8>() % 3 {
        0 => FrameUse::User {
            owner,
            page: kani::any::<u32>() % 4,
            perms: any_perms(),
        },
        1 => FrameUse::PageTable { owner },
        _ => FrameUse::Region {
            region: RegionId::from_raw(kani::any::<u16>() % 2),
        },
    }
}

fn any_slot() -> Option<FrameUse> {
    if kani::any() { Some(any_use()) } else { None }
}

fn any_table() -> FrameTable<N> {
    let frames = [any_slot(), any_slot(), any_slot(), any_slot()];
    let free = frames.iter().filter(|f| f.is_none()).count();
    FrameTable {
        base: PhysicalAddress::from_raw(BASE),
        total: N,
        frames,
        free,
        // Next-fit may resume anywhere, so the proofs start it anywhere.
        cursor: kani::any::<usize>() % (N + 1),
    }
}

fn well_formed(t: &FrameTable<N>) -> bool {
    t.free == t.frames[..t.total].iter().filter(|f| f.is_none()).count()
}

fn any_pa() -> PhysicalAddress {
    // Anything from below the pool to past its end, aligned or not.
    PhysicalAddress::from_raw(BASE - 4096 + kani::any::<usize>() % ((N + 2) * 4096))
}

#[kani::proof]
#[kani::unwind(6)]
fn alloc_only_hands_out_free_frames() {
    let mut t = any_table();
    let before = t.frames;
    let usage = any_use();
    match t.alloc(usage) {
        Some(pa) => {
            let idx = (pa.raw() - BASE) / PAGE_SIZE;
            assert!(idx < N);
            assert!(before[idx].is_none());
            assert!(t.frames[idx] == Some(usage));
            for i in 0..N {
                if i != idx {
                    assert!(t.frames[i] == before[i]);
                }
            }
        }
        None => {
            assert!(before.iter().all(Option::is_some));
            assert!(t.frames == before);
        }
    }
    assert!(well_formed(&t));
}

#[kani::proof]
#[kani::unwind(6)]
fn free_releases_exactly_one_allocated_frame_or_nothing() {
    let mut t = any_table();
    let before = t.frames;
    let pa = any_pa();
    match t.free(pa) {
        Some(usage) => {
            let idx = (pa.raw() - BASE) / PAGE_SIZE;
            assert!(pa.raw() % PAGE_SIZE == 0);
            assert!(before[idx] == Some(usage));
            assert!(t.frames[idx].is_none());
            for i in 0..N {
                if i != idx {
                    assert!(t.frames[i] == before[i]);
                }
            }
        }
        None => assert!(t.frames == before),
    }
    assert!(well_formed(&t));
}

#[kani::proof]
#[kani::unwind(6)]
fn retag_never_allocates() {
    let mut t = any_table();
    let before = t.frames;
    let pa = any_pa();
    let usage = any_use();
    let changed = t.retag(pa, usage);
    for i in 0..N {
        if before[i].is_none() {
            assert!(t.frames[i].is_none());
        }
    }
    if !changed {
        assert!(t.frames == before);
    }
    assert!(well_formed(&t));
}

#[kani::proof]
#[kani::unwind(6)]
fn free_owned_by_spares_everyone_else() {
    let mut t = any_table();
    let before = t.frames;
    let owner = ProcessId::from_raw(kani::any::<u16>() % 3);
    t.free_owned_by(owner);
    for i in 0..N {
        let was_owned = matches!(
            before[i],
            Some(FrameUse::User { owner: o, .. } | FrameUse::PageTable { owner: o }) if o == owner
        );
        if was_owned {
            assert!(t.frames[i].is_none());
        } else {
            assert!(t.frames[i] == before[i]);
        }
    }
    assert!(well_formed(&t));
}

#[kani::proof]
#[kani::unwind(6)]
fn find_user_only_reports_the_owners_mapping() {
    let t = any_table();
    let owner = ProcessId::from_raw(kani::any::<u16>() % 3);
    let va = VirtualAddress::from_raw(kani::any::<usize>() % (5 * 4096));
    if let Some((pa, perms)) = t.find_user(owner, va) {
        let idx = (pa.raw() - BASE) / PAGE_SIZE;
        assert!(idx < N);
        assert!(va.raw() % PAGE_SIZE == 0);
        assert!(
            t.frames[idx]
                == Some(FrameUse::User {
                    owner,
                    page: (va.raw() / PAGE_SIZE) as u32,
                    perms,
                })
        );
    }
}

#[kani::proof]
#[kani::unwind(6)]
fn alloc_contiguous_only_claims_a_free_run() {
    let mut t = any_table();
    let before = t.frames;
    let count = kani::any::<usize>() % (N + 2);
    let region = RegionId::from_raw(1);
    match t.alloc_contiguous(count, |_| FrameUse::Region { region }) {
        Some(pa) => {
            let start = (pa.raw() - BASE) / PAGE_SIZE;
            assert!(count > 0 && start + count <= N);
            for i in 0..N {
                if i >= start && i < start + count {
                    assert!(before[i].is_none());
                    assert!(t.frames[i] == Some(FrameUse::Region { region }));
                } else {
                    assert!(t.frames[i] == before[i]);
                }
            }
        }
        None => assert!(t.frames == before),
    }
    assert!(well_formed(&t));
}