"""A second model checker for the Futex page, written from the C, not the engine.

The engine (site/js/futex.js) describes each mutex as a table of step objects
and walks it with one generic interpreter, depth of the queue first. This file
writes each mutex as a Python function from (word, program point, register) to
what happens next, keeps a state as a flat tuple, and finds the longest run of
failed waits with Kosaraju's two passes rather than Tarjan's one. The two share
no data structure, no traversal and no language, so when they agree on a state
count it is evidence rather than a tautology.

What they must share is the definition of a state, because a count depends on
it: the futex word, each thread's program point, whether it is asleep, and its
register `c` only where something will still read it -- in the first mutex,
between the increment and the futex_wait that uses it. Everywhere else the C
branches on `c` at once and never looks at it again.

    python3 tests/futex_oracle.py          prints one JSON object
"""
import json
import sys
from collections import deque

# Program points, named after the paper's code rather than numbered.
INC, WAIT1 = 'inc', 'wait'
C01, C12, W2, C02 = 'cas01', 'cas12', 'wait2', 'cas02'
X2A, X2B = 'xchg-first', 'xchg-loop'
DEC, STORE, WAKE = 'dec', 'store', 'wake'
CS, REM = 'cs', 'rem'


def lock_entry(mutex):
    return INC if mutex == 'mutex' else C01


def act(mutex, word, pt, c, width):
    """One thread's next instruction from this point.

    Returns ('rmw', new word, next point, kept register, old value),
    ('store', next point), ('wait', expected value, point after),
    ('wake', next point), or ('move', next point) for leaving the critical
    section or the remainder, which nobody else can observe.
    """
    if pt == CS:
        return ('move', STORE if mutex == 'mutex' else DEC)
    if pt == REM:
        return ('move', lock_entry(mutex))

    if mutex == 'mutex':
        if pt == INC:
            nxt = CS if word == 0 else WAIT1
            return ('rmw', (word + 1) % width, nxt, word, word)
        if pt == WAIT1:
            return ('wait', (c + 1) % width, INC)
        if pt == STORE:
            return ('store', WAKE)
        if pt == WAKE:
            return ('wake', REM)

    retake = 1 if mutex == 'mutex2one' else 2
    if mutex in ('mutex2', 'mutex2one'):
        if pt in (C01, C02):
            want = 1 if pt == C01 else retake
            new = want if word == 0 else word
            nxt = CS if word == 0 else (W2 if word == 2 else C12)
            return ('rmw', new, nxt, 0, word)
        if pt == C12:
            new = 2 if word == 1 else word
            return ('rmw', new, W2 if word != 0 else C02, 0, word)
        if pt == W2:
            return ('wait', 2, C02)

    if mutex == 'mutex3':
        if pt == C01:
            new = 1 if word == 0 else word
            nxt = CS if word == 0 else (X2A if word != 2 else W2)
            return ('rmw', new, nxt, 0, word)
        if pt in (X2A, X2B):
            return ('rmw', 2, CS if word == 0 else W2, 0, word)
        if pt == W2:
            return ('wait', 2, X2B)

    # The unlock the second and third share.
    if pt == DEC:
        return ('rmw', word - 1, STORE if word != 1 else REM, 0, word)
    if pt == STORE:
        return ('store', WAKE)
    if pt == WAKE:
        return ('wake', REM)
    raise ValueError((mutex, pt))


def moves(mutex, state, width, frozen=None, no_wrap=False):
    """(next state, thread, what happened, old value or None) for every move."""
    word, threads = state[0], state[1:]
    out = []
    for i, (pt, c, asleep) in enumerate(threads):
        if asleep or i == frozen:
            continue
        a = act(mutex, word, pt, c, width)
        t = list(threads)
        if a[0] == 'move':
            t[i] = (a[1], 0, False)
            out.append(((word,) + tuple(t), i, 'move', None))
        elif a[0] == 'rmw':
            _, new, nxt, keep, old = a
            if no_wrap and mutex == 'mutex' and pt == INC and word == width - 1:
                continue
            t[i] = (nxt, keep, False)
            out.append(((new,) + tuple(t), i, 'atomic', old))
        elif a[0] == 'store':
            t[i] = (a[1], 0, False)
            out.append(((0,) + tuple(t), i, 'store', None))
        elif a[0] == 'wait':
            _, expected, after = a
            sleeps = word == expected
            t[i] = (after, 0, sleeps)
            out.append(((word,) + tuple(t), i, 'sleep' if sleeps else 'ewouldblock', None))
        else:
            sleepers = [j for j, th in enumerate(threads) if th[2]]
            t[i] = (a[1], 0, False)
            if not sleepers:
                out.append(((word,) + tuple(t), i, 'wake-nobody', None))
            for j in sleepers:
                u = list(t)
                u[j] = (u[j][0], 0, False)
                out.append(((word,) + tuple(u), i, 'wake', None))
    return out


def start(mutex, n):
    return (0,) + tuple((lock_entry(mutex), 0, False) for _ in range(n))


def graph(mutex, n, width, first=None, frozen=None, no_wrap=False):
    first = first or start(mutex, n)
    succ = {}
    todo = [first]
    succ[first] = None
    while todo:
        s = todo.pop()
        ms = moves(mutex, s, width, frozen, no_wrap)
        succ[s] = ms
        for (t, _, _, _) in ms:
            if t not in succ:
                succ[t] = None
                todo.append(t)
    return first, succ


def count_inside(s):
    return sum(1 for (pt, _, _) in s[1:] if pt == CS)


def census(mutex, n, width, no_wrap=False):
    _, g = graph(mutex, n, width, no_wrap=no_wrap)
    two = sum(1 for s in g if count_inside(s) > 1)
    lost = 0
    for s in g:
        th = s[1:]
        if any(a for (_, _, a) in th) and all(a or pt == REM for (pt, _, a) in th):
            lost += 1
    # futex_wake calls: each branch, the branches that woke nobody, and the
    # places -- a state and the thread calling -- which a two-sleeper wake
    # reaches by two branches.
    wake = sum(1 for s in g for (_, _, k, _) in g[s] if k == 'wake')
    nobody = sum(1 for s in g for (_, _, k, _) in g[s] if k == 'wake-nobody')
    places = len({(s, who) for s in g for (_, who, k, _) in g[s] if k in ('wake', 'wake-nobody')})
    return {'states': len(g), 'two': two, 'forgotten': lost,
            'maxword': max(s[0] for s in g), 'minword': min(s[0] for s in g),
            'wake': wake, 'wakenobody': nobody, 'wakeplaces': places}


def failed_run(mutex, n, width):
    """Hold thread 1 inside and find the longest run of failed waits."""
    first = (1, (CS, 0, False)) + tuple((lock_entry(mutex), 0, False) for _ in range(n - 1))
    first, g = graph(mutex, n, width, first=first, frozen=0)
    region = {s for s in g if count_inside(s) == 1}
    fwd = {s: [t for (t, _, _, _) in g[s] if t in region] for s in region}
    rev = {s: [] for s in region}
    for s in region:
        for t in fwd[s]:
            rev[t].append(s)
    # Kosaraju: finish order on the graph, then components on its reverse.
    order, done = [], set()
    for root in region:
        if root in done:
            continue
        done.add(root)
        stack = [(root, iter(fwd[root]))]
        while stack:
            v, it = stack[-1]
            for w in it:
                if w not in done:
                    done.add(w)
                    stack.append((w, iter(fwd[w])))
                    break
            else:
                order.append(v)
                stack.pop()
    comp = {}
    for root in reversed(order):
        if root in comp:
            continue
        comp[root] = root
        stack = [root]
        while stack:
            v = stack.pop()
            for w in rev[v]:
                if w not in comp:
                    comp[w] = root
                    stack.append(w)
    for s in region:
        for (t, _, kind, _) in g[s]:
            if kind == 'ewouldblock' and t in region and comp[t] == comp[s]:
                return {'longest': None, 'two': len(g) - len(region)}
    # Longest weighted path by memoised search over the component DAG.
    members = {}
    for s, r in comp.items():
        members.setdefault(r, []).append(s)
    memo = {}

    def longest(r):
        if r in memo:
            return memo[r]
        stack = [r]
        while stack:
            x = stack[-1]
            pending = [comp[t] for s in members[x] for (t, _, _, _) in g[s]
                       if t in region and comp[t] != x and comp[t] not in memo]
            if pending:
                stack.extend(pending)
                continue
            best = 0
            for s in members[x]:
                for (t, _, kind, _) in g[s]:
                    if t in region and comp[t] != x:
                        best = max(best, memo[comp[t]] + (kind == 'ewouldblock'))
            memo[x] = best
            stack.pop()
        return memo[r]

    return {'longest': longest(comp[first]), 'two': len(g) - len(region),
            'states': len(g)}


def costs(mutex, n, width, waits_cap=3, spin_cap=4):
    """Follow thread 1 through every lock() and unlock() call it makes.

    The record of a lock() is what its first atomic found (capped at 2), the
    atomics in each stretch between futex_waits, and the system calls. A call
    past the caps is set aside, as the engine does, so the two can be compared
    exactly rather than approximately.
    """
    fresh_lock = ('L', None, (0,), 0, False, None)
    first = (start(mutex, n), fresh_lock)
    seen = {first}
    queue = deque([first])
    locks, unlocks = set(), set()
    while queue:
        s, acc = queue.popleft()
        for (t, who, kind, old) in moves(mutex, s, width, no_wrap=True):
            na = acc
            if who == 0:
                mode, found, segs, sysc, gone, woke = acc
                prev_pt, now_pt = s[1][0], t[1][0]
                if prev_pt == REM:
                    na = ('L', None, (0,), 0, False, None)
                elif prev_pt == CS:
                    na = ('U', None, (0,), 0, False, None)
                else:
                    if not gone:
                        if kind == 'atomic':
                            if found is None:
                                found = min(old, 2) if mode == 'L' else old
                            segs = segs[:-1] + (segs[-1] + 1,)
                            gone = segs[-1] > spin_cap
                        elif kind in ('sleep', 'ewouldblock'):
                            sysc += 1
                            segs = segs + (0,)
                            gone = len(segs) - 1 > waits_cap
                        elif kind in ('wake', 'wake-nobody'):
                            sysc += 1
                            woke = kind == 'wake'
                    na = (mode, found, segs, sysc, gone, woke)
                    if mode == 'L' and now_pt == CS and not gone:
                        locks.add((found, segs, sysc))
                    if mode == 'U' and now_pt == REM:
                        unlocks.add((found, segs[0], sysc,
                                     ('woke' if woke else 'nobody') if sysc else 'none'))
            item = (t, na)
            if item not in seen:
                seen.add(item)
                queue.append(item)
    return {'locks': sorted([f, list(sg), sy] for (f, sg, sy) in locks),
            'unlocks': sorted([f if f is not None else -1, a, sy, w] for (f, a, sy, w) in unlocks)}


def main():
    out = {'census': {}, 'nowrap': {}, 'failed': {}, 'costs': {}}
    for mutex in ('mutex', 'mutex2', 'mutex3', 'mutex2one'):
        for n in (2, 3, 4):
            for bits in ((2, 3, 4) if mutex == 'mutex' else (4,)):
                if mutex == 'mutex' and n == 4 and bits > 3:
                    continue
                out['census'][f'{mutex}/{n}/{bits}'] = census(mutex, n, 1 << bits)
    # Step 1's own counts: the first mutex with the counter held below 16
    # rather than wrapping, at each thread count its select offers.
    out['nowrap'] = {f'mutex/{n}/4': census('mutex', n, 16, no_wrap=True) for n in (1, 2, 3)}
    for bits in range(2, 7):
        out['failed'][f'mutex/3/{bits}'] = failed_run('mutex', 3, 1 << bits)
    for mutex in ('mutex2', 'mutex3'):
        out['failed'][f'{mutex}/4/4'] = failed_run(mutex, 4, 16)
    for mutex in ('mutex', 'mutex2', 'mutex3'):
        for n in (1, 3):
            out['costs'][f'{mutex}/{n}'] = costs(mutex, n, 16)
    json.dump(out, sys.stdout)
    print()


if __name__ == '__main__':
    main()