"""Two green-thread runtimes over one program, differing in one thing:
*where* a task is allowed to suspend.

  STACKLESS  a task is a generator. Suspension is the `yield` keyword, so it
             can only happen in a frame that is itself a generator -- and
             every caller in between must be one too. That is "colouring".
  STACKFUL   a task owns a real stack. Suspension is a function *call*, so it
             happens at any depth, in any kind of frame, and no caller above
             it is edited or even recompiled.

Python cannot build a cheap stackful task, so `Stackful` uses OS threads and
hands a baton round-robin: exactly one runs at a time and it switches only at
an explicit `suspend()`. Both runtimes are therefore deterministic, and the
only difference left between the two columns is the suspension point.

The workload is a bank. T tellers each deposit 1 into one account D times,
through five frames: teller -> deposit -> repo_get -> db_fetch -> io_read.
Nothing in it looks concurrent.
"""
import collections
import threading

# ---------------------------------------------------------------- runtime A


class Green:
    """STACKLESS. A task is a generator, driven with send(); a task suspends
    by returning from send() and resumes on the next one. Plain FIFO."""

    def __init__(self):
        self.q = collections.deque()
        self.switches = 0

    def spawn(self, gen):
        self.q.append(gen)

    def run(self):
        while self.q:
            g = self.q.popleft()
            try:
                g.send(None)
            except StopIteration:
                continue
            self.switches += 1
            self.q.append(g)


# ---------------------------------------------------------------- runtime B


class Stackful:
    """STACKFUL. A task is an OS thread holding a real C+Python stack. The
    scheduler passes a baton round-robin so exactly one thread runs at a
    time, which buys the same determinism `Green` gets for free."""

    def __init__(self):
        self.slots = []
        self.back = threading.Semaphore(0)
        self.cur = threading.local()
        self.switches = 0

    def spawn(self, fn, *args):
        slot = {"sem": threading.Semaphore(0), "alive": True}

        def body():
            slot["sem"].acquire()           # wait to be given the baton
            self.cur.slot = slot
            try:
                fn(*args)
            finally:
                slot["alive"] = False
                self.back.release()         # hand it back for the last time

        slot["thread"] = threading.Thread(target=body, daemon=True)
        self.slots.append(slot)
        slot["thread"].start()

    def suspend(self):
        """Hand the baton back and wait to be given it again.

        This is an ordinary method call, so it works at ANY stack depth, in
        any kind of frame. That single fact is the whole toy.
        """
        slot = self.cur.slot
        self.back.release()
        slot["sem"].acquire()

    def run(self):
        while any(s["alive"] for s in self.slots):
            for s in list(self.slots):
                if not s["alive"]:
                    continue
                s["sem"].release()
                self.back.acquire()
                self.switches += 1
        for s in self.slots:
            s["thread"].join()


# ------------------------------------------------------------- the workload

START, TELLERS, DEPOSITS = 100, 2, 50

SCHED = None            # the Stackful runtime currently running
WHERE = "after_read"    # which edition of the library's io_read is installed


def suspend():
    """The stackful suspension primitive, reachable from anywhere -- exactly
    as `runtime.Gosched()` is in Go and a blocking call is on a Loom virtual
    thread. Nothing in the signature of a caller records that it can happen.
    """
    SCHED.suspend()


# --- the app, STACKFUL. Read the four frames above io_read: not one of them
# --- carries a marker of any kind, in any edition of the library.

def sf_io_read(bank, acct):
    if WHERE == "before_read":
        suspend()
    v = bank[acct]                  # the row is read here...
    if WHERE == "after_read":
        suspend()                   # <-- THE ONE LINE THE LIBRARY ADDS
    return v                        # ...and the reply reaches us here

def sf_db_fetch(bank, acct):
    return sf_io_read(bank, acct)

def sf_repo_get(bank, acct):
    return sf_db_fetch(bank, acct)

def sf_deposit(bank, acct, amt):
    b = sf_repo_get(bank, acct)     # critical section starts
    b = b + amt
    bank[acct] = b                  # critical section ends

def sf_teller(bank, acct, n):
    for _ in range(n):
        sf_deposit(bank, acct, 1)

# --- the app, STACKLESS, library v1: io_read does not suspend, so it is an
# --- ordinary function and so is every frame above it.

def g1_io_read(bank, acct):
    v = bank[acct]
    return v

def g1_db_fetch(bank, acct):
    return g1_io_read(bank, acct)

def g1_repo_get(bank, acct):
    return g1_db_fetch(bank, acct)

def g1_deposit(bank, acct, amt):
    b = g1_repo_get(bank, acct)
    b = b + amt
    bank[acct] = b

def g1_teller(bank, acct, n):
    for _ in range(n):
        g1_deposit(bank, acct, 1)
    yield                           # a task must be a generator at all

# --- the app, STACKLESS, library v2. Only io_read changed: it gained a
# --- `yield` and is now a generator function. Every frame above it is
# --- byte-identical to v1.

def g2_io_read(bank, acct):
    v = bank[acct]
    yield                           # <-- THE ONE LINE THE LIBRARY ADDS
    return v

def g2_db_fetch(bank, acct):
    return g2_io_read(bank, acct)

def g2_repo_get(bank, acct):
    return g2_db_fetch(bank, acct)

def g2_deposit(bank, acct, amt):
    b = g2_repo_get(bank, acct)
    b = b + amt
    bank[acct] = b

def g2_teller(bank, acct, n):
    for _ in range(n):
        g2_deposit(bank, acct, 1)
    yield

# --- the app, STACKLESS, library v2, after colouring all four callers.

def g3_io_read(bank, acct):
    v = bank[acct]
    yield
    return v

def g3_db_fetch(bank, acct):
    return (yield from g3_io_read(bank, acct))

def g3_repo_get(bank, acct):
    return (yield from g3_db_fetch(bank, acct))

def g3_deposit(bank, acct, amt):
    b = yield from g3_repo_get(bank, acct)
    b = b + amt
    bank[acct] = b

def g3_teller(bank, acct, n):
    for _ in range(n):
        yield from g3_deposit(bank, acct, 1)

# ------------------------------------------------------------------ drivers


def run_stackful(where=WHERE, tellers=TELLERS, deposits=DEPOSITS):
    global SCHED, WHERE
    SCHED, WHERE = Stackful(), where
    bank = {"acct": START}
    for _ in range(tellers):
        SCHED.spawn(sf_teller, bank, "acct", deposits)
    SCHED.run()
    return bank["acct"], SCHED.switches


def run_green(teller, tellers=TELLERS, deposits=DEPOSITS):
    sched = Green()
    bank = {"acct": START}
    for _ in range(tellers):
        sched.spawn(teller(bank, "acct", deposits))
    sched.run()
    return bank["acct"], sched.switches


# -------------------------------------------------- call depth, generated
# D = how many frames sit above the suspending call inside the task, the task
# body included. The bank workload is D = 4 (db_fetch, repo_get, deposit,
# teller); D = 0 means the task body suspends directly. Colouring costs D
# edits and a stackful runtime costs 0 at every D -- built and run, not
# asserted.


def chain(depth, style):
    """Source for a task with `depth` frames above its suspending call.

    style: 'stackful' | 'coloured' | 'uncoloured' (the last two stackless).
    Returns (source, edits), `edits` counting the frames that had to gain a
    suspension marker they did not previously carry.
    """
    mark = "    suspend()" if style == "stackful" else "    yield"
    if depth == 0:                          # the task body suspends directly
        return "\n".join(["def task():", "    v = 7", mark,
                          "    out.append(v + 1)"]), 0
    src = ["def leaf():", "    v = 7", mark, "    return v"]
    call, edits = "leaf()", 0
    for i in range(depth - 1):
        if style == "coloured":
            src += [f"def f{i}():", f"    return (yield from {call})"]
            edits += 1
        else:
            src += [f"def f{i}():", f"    return {call}"]
        call = f"f{i}()"
    if style == "coloured":
        src += ["def task():", f"    v = yield from {call}",
                "    out.append(v + 1)"]
        edits += 1
    else:
        src += ["def task():", f"    out.append({call} + 1)"]
        if style == "uncoloured":
            src += ["    yield"]           # a Green task must be a generator
    return "\n".join(src), edits


def run_chain(depth, style):
    """Build and run one chain. Returns (out, edits, error-or-None)."""
    global SCHED
    src, edits = chain(depth, style)
    ns = {"out": [], "suspend": suspend}
    exec(src, ns)
    try:
        if style == "stackful":
            SCHED = Stackful()
            SCHED.spawn(ns["task"])
            SCHED.run()
        else:
            g = Green()
            g.spawn(ns["task"]())
            g.run()
    except Exception as e:                      # an uncoloured caller
        return ns["out"], edits, f"{type(e).__name__}: {e}"
    return ns["out"], edits, None