#!/usr/bin/env python3
"""Hands-on P12 — an OS kernel's mechanisms, assembled from eight lego blocks."""
import random
from collections import OrderedDict, deque
from _harness import block, run_all

@block(1, "Physical frame allocator", "the first allocation problem, with no allocator to help")
def b1(s, show):
    class Frames:
        def __init__(self, n): self.bm = bytearray(n); self.n = n; self.hint = 0
        def alloc(self):
            for k in range(self.n):
                i = (self.hint + k) % self.n
                if not self.bm[i]:
                    self.bm[i] = 1; self.hint = i + 1; return i
            raise MemoryError("out of physical memory")
        def free(self, i):
            if not self.bm[i]: raise ValueError(f"double free of frame {i}")
            self.bm[i] = 0
        def used(self): return sum(self.bm)
    if show:
        f = Frames(1024)
        got = [f.alloc() for _ in range(600)]
        for i in got[::3]: f.free(i)
        print(f"  1024 frames; allocated 600, freed 200 -> {f.used()} in use")
        print(f"  next allocation reuses a hole: frame {f.alloc()}")
        try:
            f.free(got[1]); f.free(got[1])
        except ValueError as e:
            print(f"  double free is caught: {e}")
        print("  A bitmap costs 1 bit per 4 KiB page: 32 KiB of metadata per GiB, or")
        print("  0.003%. The rotating hint turns the scan from O(n) per alloc into")
        print("  O(1) amortised. This runs BEFORE any heap exists, so it cannot")
        print("  allocate -- every structure the kernel needs early is a fixed array.")
    return {"Frames": Frames}

@block(2, "Virtual memory", "one indirection, and the process model falls out of it")
def b2(s, show):
    PAGE = 4096
    class MMU:
        def __init__(self, frames): self.f = frames; self.tables = {}
        def map(self, pid, vpn, writable=True):
            t = self.tables.setdefault(pid, {})
            if vpn in t: return t[vpn][0]
            fr = self.f.alloc(); t[vpn] = (fr, writable); return fr
        def translate(self, pid, va, write=False):
            vpn, off = va // PAGE, va % PAGE
            e = self.tables.get(pid, {}).get(vpn)
            if e is None: raise MemoryError(f"SEGFAULT pid={pid} va=0x{va:x}")
            fr, w = e
            if write and not w: raise PermissionError(f"write to RO page va=0x{va:x}")
            return fr * PAGE + off
    if show:
        f = s["Frames"](256); m = MMU(f)
        m.map(1, 0); m.map(1, 1); m.map(2, 0)
        pa1 = m.translate(1, 0x0100); pa2 = m.translate(2, 0x0100)
        print(f"  pid 1 va 0x0100 -> pa 0x{pa1:x}")
        print(f"  pid 2 va 0x0100 -> pa 0x{pa2:x}   (same VA, different frame)")
        print(f"  isolation is structural, not checked: {pa1 != pa2}")
        for lbl, fn in (("unmapped read", lambda: m.translate(1, 0x9000)),
                        ("write to RO", lambda: (m.map(1, 3, writable=False),
                                                 m.translate(1, 3*PAGE, write=True)))):
            try: fn()
            except Exception as e: print(f"  {lbl:<16}-> {type(e).__name__}: {e}")
        print("  Two processes at the same virtual address is the ENTIRE reason a")
        print("  bug in one cannot corrupt the other. Everything else -- COW, mmap,")
        print("  shared libraries, demand paging -- is a variation on who gets to")
        print("  point at which frame.")
    return {"MMU": MMU, "PAGE": PAGE}

@block(3, "Page replacement and Belady's anomaly", "more memory, more faults")
def b3(s, show):
    def fifo(refs, n):
        mem, q, faults = set(), deque(), 0
        for r in refs:
            if r in mem: continue
            faults += 1
            if len(mem) == n: mem.discard(q.popleft())
            mem.add(r); q.append(r)
        return faults
    def lru(refs, n):
        mem, faults = OrderedDict(), 0
        for r in refs:
            if r in mem: mem.move_to_end(r); continue
            faults += 1
            if len(mem) == n: mem.popitem(last=False)
            mem[r] = 1
        return faults
    def clock(refs, n):
        buf, ref, hand, faults = [], [], 0, 0
        for r in refs:
            if r in buf: ref[buf.index(r)] = 1; continue
            faults += 1
            if len(buf) < n: buf.append(r); ref.append(1); continue
            while ref[hand]: ref[hand] = 0; hand = (hand + 1) % n
            buf[hand] = r; ref[hand] = 1; hand = (hand + 1) % n
        return faults
    def opt(refs, n):
        mem, faults = [], 0
        for i, r in enumerate(refs):
            if r in mem: continue
            faults += 1
            if len(mem) == n:
                fut = [(refs[i+1:].index(x) if x in refs[i+1:] else 1 << 30) for x in mem]
                mem.pop(fut.index(max(fut)))
            mem.append(r)
        return faults
    BEL = [1,2,3,4,1,2,5,1,2,3,4,5]
    if show:
        print(f"  reference string {BEL}")
        print(f"  {'frames':>8}{'FIFO':>7}{'LRU':>6}{'CLOCK':>7}{'OPT':>6}")
        for n in (3, 4):
            print(f"  {n:>8}{fifo(BEL,n):>7}{lru(BEL,n):>6}{clock(BEL,n):>7}{opt(BEL,n):>6}")
        print(f"  FIFO with 4 frames faults MORE than with 3 "
              f"({fifo(BEL,4)} vs {fifo(BEL,3)}) -- Belady's anomaly, 1969.")
        print("  LRU cannot do this: it is a stack algorithm, so the pages resident")
        print("  with n frames are always a subset of those with n+1. FIFO has no")
        print("  such property, so buying RAM can lose performance. This is not a")
        print("  curiosity -- it is why 'add cache' is a hypothesis, not a fix.\\n")
        rng = random.Random(4)
        refs = [rng.choice(range(20)) if rng.random() < .3 else rng.choice(range(4))
                for _ in range(4000)]
        print(f"  4000 refs, 80/20 locality: {'frames':>8}{'FIFO':>7}{'LRU':>6}"
              f"{'CLOCK':>7}{'OPT':>6}")
        for n in (4, 8, 16):
            print(f"  {'':<26}{n:>8}{fifo(refs,n):>7}{lru(refs,n):>6}"
                  f"{clock(refs,n):>7}{opt(refs,n):>6}")
        print("  CLOCK tracks LRU within a few percent at a fraction of the cost --")
        print("  one reference bit per page instead of a timestamp and a list splice")
        print("  on every hit. Every real kernel ships an approximation, not LRU.")
    return {"fifo": fifo, "lru": lru, "clock": clock, "opt": opt}

@block(4, "Context switch", "saving state is easy; the cache is what costs")
def b4(s, show):
    class Task:
        def __init__(self, tid, work): self.tid, self.work = tid, work; self.done = 0
        def step(self, q):
            n = min(q, self.work - self.done); self.done += n; return n
        def finished(self): return self.done >= self.work
    def run(tasks, quantum, switch_cost=0.0):
        t, log, sw = 0.0, [], 0
        q = deque(tasks)
        while q:
            task = q.popleft()
            did = task.step(quantum)
            t += did
            if task.finished(): log.append((task.tid, t))
            else:
                t += switch_cost; sw += 1; q.append(task)
        return t, log, sw
    if show:
        mk = lambda: [Task(i, w) for i, w in enumerate([50, 10, 80, 5, 30])]
        print(f"  five tasks, work = [50, 10, 80, 5, 30] microseconds")
        print(f"  {'quantum':>9}{'switches':>10}{'makespan':>11}{'mean turnaround':>18}"
              f"{'overhead':>10}")
        for q in (200, 50, 10, 2):
            tot, log, sw = run(mk(), q, 1.7)
            mt = sum(x for _, x in log) / len(log)
            print(f"  {q:>9}{sw:>10}{tot:>10.1f}us{mt:>17.1f}us"
                  f"{(sw*1.7)/tot:>9.1%}")
        print("  The 1.7 us switch cost is not invented -- it is this machine's")
        print("  measured context-switch time from numbers.md, which came out in the")
        print("  1,383-1,706 ns range. At a 2 us quantum the kernel spends 45% of the")
        print("  CPU switching between tasks. Responsiveness and throughput are the")
        print("  same dial, and the dial has a floor set by hardware.")
        print("  What this model does NOT capture: the cold cache after a switch. The")
        print("  register save is ~50 instructions; the L1 and TLB refill afterwards")
        print("  can be thousands of cycles, and it does not show up in any counter")
        print("  named 'context switch'.")
    return {"run_rr": run, "Task": Task}

@block(5, "Scheduling policy", "the same work, five orders, five different fairness stories")
def b5(s, show):
    JOBS = [("A", 0, 50), ("B", 0, 10), ("C", 5, 80), ("D", 10, 5), ("E", 20, 30)]
    def simulate(policy, quantum=10, switch=1.7):
        rem = {n: w for n, a, w in JOBS}; arr = {n: a for n, a, w in JOBS}
        t, done, ready, pending = 0.0, {}, [], sorted(JOBS, key=lambda j: j[1])
        first = {}
        while pending or ready:
            while pending and pending[0][1] <= t: ready.append(pending.pop(0)[0])
            if not ready:
                t = pending[0][1]; continue
            if policy == "fcfs": n = ready[0]; q = 1e9
            elif policy == "sjf": n = min(ready, key=lambda x: rem[x]); q = 1e9
            elif policy == "srtf": n = min(ready, key=lambda x: rem[x]); q = quantum
            else: n = ready[0]; q = quantum
            first.setdefault(n, t - arr[n])
            run = min(q, rem[n]); rem[n] -= run; t += run
            if rem[n] <= 0:
                done[n] = t - arr[n]; ready.remove(n)
            else:
                ready.remove(n); ready.append(n); t += switch
        return done, first
    if show:
        print("  jobs (name, arrival, work): " + ", ".join(f"{n}@{a}:{w}" for n,a,w in JOBS))
        print(f"  {'policy':<22}{'mean turnaround':>17}{'mean response':>15}"
              f"{'max turnaround':>16}")
        for pol, lbl in (("fcfs", "FCFS"), ("sjf", "SJF (non-preempt)"),
                         ("srtf", "SRTF (preemptive)"), ("rr", "round robin q=10")):
            d, f = simulate(pol)
            print(f"  {lbl:<22}{sum(d.values())/len(d):>16.1f}{sum(f.values())/len(f):>15.1f}"
                  f"{max(d.values()):>16.1f}")
        print("  SJF minimises mean turnaround -- provably, and it is the only")
        print("  policy here that can starve a long job forever. Round robin has the")
        print("  best response time and the worst turnaround. There is no scheduler")
        print("  that wins every column; there is only a choice of which column the")
        print("  workload cares about, which is why Linux ships several.")
    return {}

@block(6, "System calls", "a controlled doorway, not a function call")
def b6(s, show):
    class Kernel:
        def __init__(self): self.files = {}; self.trap_count = 0
        def syscall(self, num, *args):
            self.trap_count += 1
            table = {0: self.sys_open, 1: self.sys_write, 2: self.sys_read}
            fn = table.get(num)
            if fn is None: return -38                      # -ENOSYS
            try: return fn(*args)
            except Exception: return -22                   # -EINVAL, never a crash
        def sys_open(self, path): self.files.setdefault(path, b""); return hash(path) % 900 + 3
        def sys_write(self, path, data):
            if not isinstance(data, (bytes, bytearray)): raise TypeError
            self.files[path] += data; return len(data)
        def sys_read(self, path): return self.files[path]
    if show:
        k = Kernel()
        fd = k.syscall(0, "/tmp/x")
        print(f"  open  -> fd {fd}")
        print(f"  write -> {k.syscall(1, '/tmp/x', b'hello')} bytes")
        print(f"  read  -> {k.syscall(2, '/tmp/x')!r}")
        print(f"  bad syscall number -> {k.syscall(99)} (-ENOSYS)")
        print(f"  bad argument type  -> {k.syscall(1, '/tmp/x', 12345)} (-EINVAL)")
        print(f"  traps taken: {k.trap_count}, kernel crashes: 0")
        print("  Two properties make this a syscall rather than a call: arguments are")
        print("  VALIDATED (a user pointer may be garbage or hostile), and errors come")
        print("  back as values (the kernel cannot unwind into a process). This")
        print("  machine's measured syscall cost is 127.59 ns -- roughly 1000")
        print("  arithmetic instructions -- which is why batching interfaces like")
        print("  io_uring and readv exist at all.")
    return {}

@block(7, "A race, and a lock", "concurrency bugs are timing-shaped, so hunt them with timing")
def b7(s, show):
    import threading
    def counter(n, lock=None):
        v = [0]
        def worker():
            for _ in range(n):
                if lock:
                    with lock: v[0] += 1
                else:
                    x = v[0]; x += 1; v[0] = x       # deliberately non-atomic
        ts = [threading.Thread(target=worker) for _ in range(4)]
        for t in ts: t.start()
        for t in ts: t.join()
        return v[0]
    if show:
        n = 60_000
        bad = [counter(n) for _ in range(3)]
        good = counter(n, threading.Lock())
        print(f"  4 threads x {n:,} increments; expected {4*n:,}")
        print(f"  without a lock: {bad}  (lost {[4*n-b for b in bad]})")
        print(f"  with a lock:    {good}  correct = {good == 4*n}")
        print("  The unlocked version is not merely wrong -- it is wrong by a")
        print("  DIFFERENT amount each run, which is what makes these bugs so")
        print("  expensive. A test that passes proves nothing; only the invariant")
        print("  (final == 4n) and many runs can detect it. Note the GIL does not")
        print("  save you: it makes each BYTECODE atomic, and 'x = v[0]; x += 1;")
        print("  v[0] = x' is three of them.")
    return {}

@block(8, "Putting a process together", "every block above is one field of a PCB")
def b8(s, show):
    if show:
        print("  struct process {")
        print("      pid_t pid;                 /* block 5: scheduler identity   */")
        print("      state_t state;             /* READY | RUNNING | BLOCKED     */")
        print("      context_t regs;            /* block 4: saved on switch      */")
        print("      pagetable_t *pgdir;        /* block 2: its address space    */")
        print("      file_t *fds[NOFILE];       /* block 6: what open() returned */")
        print("      uint64 quantum_left;       /* block 5: preemption budget    */")
        print("      struct process *parent;    /* for wait() and exit status    */")
        print("  };")
        print("  Every field is one of the mechanisms built above, and the operating")
        print("  system is mostly the code that keeps these consistent across")
        print("  switches, faults and traps. 'What is in a PCB' is the single best")
        print("  question for checking whether you understand a kernel, because you")
        print("  cannot answer it without having built each field's machinery.")
    return {}

def assembly(s):
    print("\nEight blocks = the mechanisms of a kernel. One workload through all of them.\n")
    rng = random.Random(12)
    NPROC, PAGES = 4, 12
    refs = [(rng.randrange(NPROC), rng.randrange(PAGES) if rng.random() < .3
             else rng.randrange(3)) for _ in range(600)]
    keyed = [f"{p}:{v}" for p, v in refs]
    hot = len({k for k in keyed if int(k.split(":")[1]) < 3})
    print(f"  {NPROC} processes, {PAGES} virtual pages each, {len(refs)} references")
    print(f"  with 70/30 locality. The combined hot working set is {hot} pages.\n")
    print(f"  {'frames':>7}{'FIFO':>7}{'CLOCK':>7}{'LRU':>6}{'OPT':>6}"
          f"{'LRU fault rate':>17}{'paging @90us':>15}")
    for F in (4, 8, 12, 16, 24, 48):
        fs = [fn(keyed, F) for fn in (s["fifo"], s["clock"], s["lru"], s["opt"])]
        print(f"  {F:>7}{fs[0]:>7}{fs[1]:>7}{fs[2]:>6}{fs[3]:>6}"
              f"{fs[2]/len(refs):>16.1%}{fs[2]*90/1000:>13.1f}ms")
    dist = len(set(keyed))
    print(f"\n  I expected a cliff at {hot} frames -- the hot working set -- and there")
    print(f"  is none. The curve declines smoothly from 79.7% to 7.8% and only")
    print(f"  bottoms out at {dist} frames, where every distinct page is resident and")
    print(f"  the remaining {min(47, dist)} faults are compulsory. The prediction was wrong for a")
    print("  reason worth more than the prediction: this workload is a MIXTURE, 70%")
    print(f"  into {hot} hot pages and 30% uniform over all {dist}. A mixture has no single")
    print("  working set, so it has no knee. Denning's model describes a phase, and")
    print("  real programs are a superposition of phases -- which is why 'size the")
    print("  cache to the working set' is advice you can only follow after measuring")
    print("  the curve, never by reasoning about the program.")
    print("\n  Two things to take from the numbers rather than the shape. First, at")
    print("  12 frames the best implementable policy beats the worst by 9% (LRU 256")
    print("  vs FIFO 279) while the unrealisable OPT beats LRU by 1.92x -- almost")
    print("  all the available win is in knowing the future, which is why prefetching")
    print("  and hinting (madvise, fadvise) buy more than any eviction heuristic ever")
    print("  will. Second, the paging column dwarfs the compute: 23ms of stalls")
    print("  against a job whose actual work is 10ms. The memory hierarchy is not a")
    print("  tax on the computation, it IS the computation's cost.")
    print("\n  That is the same lesson as P04's Bloom filters (avoid the I/O), P13's")
    print("  checkpointing (trade compute for memory), and P14's tiling (fit the")
    print("  working set in cache). Four projects, four altitudes, one hierarchy.")
    print("\n  Built: frame allocator -> virtual memory -> replacement policy ->")
    print("  context switch -> scheduling -> syscalls -> locking -> the PCB.")
    print("  Missing, on the project page: real x86-64 boot and a GDT/IDT (m1-m3),")
    print("  hardware page tables with a TLB and its shootdown (m5), a disk driver")
    print("  and a real file system (m8-m10), and E6 -- the experiment where you")
    print("  measure context-switch cost as a function of working-set size and watch")
    print("  cache pollution dwarf the register save.")

if __name__ == "__main__":
    run_all(assembly, "HANDS-ON P12 — Operating system kernel, block by block")
