"""The demo: a 256-byte arena, four 56-byte allocations, four frees, and a
96-byte request that fails on an arena with 256 free bytes in it.

Run:  python3 demo.py            # sections A-G, about a second
      python3 demo.py --census   # adds section H, a 986,580-trace sweep
                                 # over the fit policies (~1 minute)

Writes nothing to disk. No clock, no randomness: byte-identical every run.
"""

import itertools
import os
import struct
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from arena import Arena, OutOfMemory, FOOTER

ARENA, SIZE, N, REQ = 256, 56, 4, 96


def rule(title):
    print("\n" + title)
    print("-" * len(title))


def fill_and_free(order, req=None, **kw):
    """N allocations that exactly fill the arena, then N frees in `order`."""
    a = Arena(ARENA, **kw)
    pay = SIZE - (FOOTER if kw.get("boundary_tags") else 0)
    ptrs = [a.malloc(pay) for _ in range(N)]
    for i in order:
        a.free(ptrs[i])
    if req is None:
        return a, None
    try:
        a.malloc(req)
        return a, True
    except OutOfMemory:
        return a, False


def report(order, **kw):
    a, _ = fill_and_free(order, None, **kw)
    s, lay = a.stats(), a.layout()          # captured BEFORE the request
    print("  free order %s" % (list(order),))
    print("    free list : %s" % lay)
    print("    %d of %d bytes free (%d%%) in %d block(s); largest holds a "
          "%d-byte payload" % (s["free_bytes"], ARENA,
                               100 * s["free_bytes"] // ARENA, s["nfree"],
                               s["largest_payload"]))
    try:
        a.malloc(REQ)
        print("    malloc(%d) -> OK" % REQ)
    except OutOfMemory as e:
        print("    malloc(%d) -> FAILS: %s" % (REQ, e))
    return a


print("arena=%d bytes, header=8, %d x malloc(%d) fills it exactly "
      "(%d x %d = %d)." % (ARENA, N, SIZE, N, SIZE + 8, N * (SIZE + 8)))

rule("A. Same allocations, same frees, two orders")
report((0, 1, 2, 3))
report((3, 2, 1, 0))

rule("B. Counterfactual: turn coalescing off entirely")
a_off, _ = fill_and_free((0, 1, 2, 3), coalesce="none")
a_fwd, _ = fill_and_free((0, 1, 2, 3), coalesce="forward")
print("  coalesce='none'    : %s" % a_off.layout())
print("  coalesce='forward' : %s" % a_fwd.layout())
print("  byte-identical arenas? %s" % (a_off.mem == a_fwd.mem))

rule("C. All 24 free orders, forward-only coalescing")
seen = {}
for order in itertools.permutations(range(N)):
    a, _ = fill_and_free(order)
    seen.setdefault(a.layout(), []).append(order)
for lay in sorted(seen, key=lambda k: (-len(seen[k]), k)):
    print("  %-52s %2d order(s), e.g. %s"
          % (lay, len(seen[lay]), list(seen[lay][0])))
print()
print("  %-14s %s" % ("request", "orders that fail, of 24"))
for req in (48, 56, 88, 96, 120, 121, 152, 184, 200, 248):
    n = sum(1 for o in itertools.permutations(range(N))
            if fill_and_free(o, req)[1] is False)
    print("  %-14d %d" % (req, n))
print()
print("  Same shape, 6 x malloc(32) in a 240-byte arena, all 720 orders:")
for req in (56, 72, 96, 120):
    n = 0
    for order in itertools.permutations(range(6)):
        b = Arena(240)
        p = [b.malloc(32) for _ in range(6)]
        for i in order:
            b.free(p[i])
        try:
            b.malloc(req)
        except OutOfMemory:
            n += 1
    print("    request %3d -> %3d of 720 orders fail (%.1f%%)"
          % (req, n, 100 * n / 720))

rule("D. Does the fit policy rescue the failing arena?")
for fit in ("first", "best", "worst"):
    a, ok = fill_and_free((0, 1, 2, 3), REQ, fit=fit)
    print("  fit=%-6s malloc(%d) -> %-5s  %s"
          % (fit, REQ, "OK" if ok else "FAILS", a.layout()))
print()
print("  The knob is not inert -- give it three holes of different sizes and")
print("  malloc(24), which needs 32 bytes, lands somewhere different each time:")
for fit in ("first", "best", "worst"):
    b = Arena(ARENA, fit=fit, coalesce="both")
    p = [b.malloc(56), b.malloc(8), b.malloc(24), b.malloc(8)]
    b.free(p[0])
    b.free(p[2])
    if fit == "first":
        print("    holes : %s" % b.layout())
    b.malloc(24)
    print("    %-5s : %s" % (fit, b.layout()))

rule("E. What does rescue it, and what it costs")
print("  %-27s %-16s %-11s %s" % ("policy", "4 frees cost", "malloc(96)",
                                  "and costs"))
for label, kw in (("coalesce='forward'", dict(coalesce="forward")),
                  ("coalesce='both' (heap walk)", dict(coalesce="both")),
                  ("coalesce='both' + footers",
                   dict(coalesce="both", boundary_tags=True)),
                  ("'none' + consolidate", dict(coalesce="none",
                                                consolidate_on_fail=True))):
    a = Arena(ARENA, **kw)
    pay = SIZE - (FOOTER if kw.get("boundary_tags") else 0)
    ptrs = [a.malloc(pay) for _ in range(N)]
    a.reads = 0
    for i in range(N):
        a.free(ptrs[i])
    frees = a.reads
    a.reads = 0
    try:
        a.malloc(REQ)
        got = "OK"
    except OutOfMemory:
        got = "FAILS"
    print("  %-27s %6d reads   %-11s %6d reads"
          % (label, frees, got, a.reads))

rule("F. Boundary tags: the cost, the price, and the trap")
print("  Backward merge over n blocks, freed newest-first (every free merges):")
print("  %4s %18s %18s" % ("n", "heap walk", "footer + PREV_FREE"))
for n in (4, 8, 16, 32, 64):
    row = []
    for bt in (False, True):
        b = Arena(n * 32, coalesce="both", boundary_tags=bt)
        p = [b.malloc(24 - (FOOTER if bt else 0)) for _ in range(n)]
        b.reads = 0
        for i in range(n - 1, -1, -1):
            b.free(p[i])
        row.append(b.reads)
        assert b.stats()["nfree"] == 1
    print("  %4d %18d %18d" % (n, row[0], row[1]))

print()
print("  The price, on this toy's own trace (4 x malloc(56) in 256 bytes):")
for bt in (False, True):
    b = Arena(ARENA, coalesce="both", boundary_tags=bt)
    served = 0
    try:
        for _ in range(N):
            b.malloc(SIZE)
            served += 1
    except OutOfMemory:
        pass
    print("    boundary_tags=%-5s overhead=%2d bytes/block -> %d of %d "
          "allocations served" % (bt, b.overhead, served, N))


class NaiveFooter(Arena):
    """Boundary tags as you would write them the first time: the 4 bytes
    below a header are the previous block's size, so read them and believe
    them. No PREV_FREE bit anywhere.
    """

    def _prev(self, off):
        if off == 0:
            return None
        ptotal, = struct.unpack_from("<I", self.mem, off - FOOTER)
        return off - ptotal


print()
print("  And the trap -- the same trace, with the PREV_FREE check removed:")
b = NaiveFooter(512, coalesce="both", boundary_tags=True)
p = [b.malloc(20) for _ in range(16)]
try:
    for i in range(15, -1, -1):
        b.free(p[i])
    print("    no error")
except struct.error as e:
    print("    struct.error: %s" % e)
print("    the block below is USED, so those 4 bytes are payload: %d."
      % struct.unpack_from("<I", b.mem, 480 - FOOTER)[0])
print("    prev = off - 0 = off, so the block merges with ITSELF: 32 -> 64,")
print("    and a footer for a 64-byte block at offset 480 lands at 540.")

rule("G. Where it vanishes: one size, no mixing")
b = Arena(ARENA)
p = [b.malloc(SIZE) for _ in range(N)]
for i in range(N):
    b.free(p[i])
print("  ascending frees, forward-only: %s" % b.layout())
served = 0
try:
    for _ in range(N):
        b.malloc(SIZE)
        served += 1
except OutOfMemory:
    pass
print("  %d of %d further malloc(%d) served from that same arena."
      % (served, N, SIZE))

if "--census" in sys.argv:
    rule("H. Fit policy over 986,580 traces (this takes about a minute)")
    SIZES = [16, 24, 32, 40, 56, 64, 88, 104, 120]

    def traces():
        for a1, a2, a3, a4 in itertools.product(SIZES, repeat=4):
            if a1 + a2 + a3 + a4 > 200:
                continue
            for frees in itertools.combinations(range(4), 2):
                for r1 in SIZES:
                    for r2 in SIZES:
                        yield ([("m", a1), ("m", a2), ("m", a3), ("m", a4)]
                               + [("f", i) for i in frees]
                               + [("m", r1), ("m", r2)])

    def play(trace, **kw):
        a = Arena(ARENA, coalesce="both", **kw)
        ptrs = []
        for k, op in enumerate(trace):
            if op[0] == "m":
                try:
                    ptrs.append(a.malloc(op[1]))
                except OutOfMemory:
                    return k
            else:
                a.free(ptrs[op[1]])
        return None

    total = 0
    served = {"first": 0, "best": 0, "worst": 0}
    sole_loser = {"first": 0, "best": 0, "worst": 0}
    for t in traces():
        total += 1
        res = {f: play(t, fit=f) for f in served}
        for f in served:
            if res[f] is None:
                served[f] += 1
        losers = [f for f in res if res[f] is not None]
        if len(losers) == 1:
            sole_loser[losers[0]] += 1
    print("  %d traces of 4 allocations, 2 frees, 2 more requests" % total)
    for f in ("first", "best", "worst"):
        print("  fit=%-6s served every request in %6d traces (%.2f%%); "
              "sole loser in %5d" % (f, served[f], 100 * served[f] / total,
                                     sole_loser[f]))
    print()
    print("  Splitting threshold, same traces, fit=first:")
    for sm in (0, 8, 16, 32, 64):
        ok = sum(1 for t in traces() if play(t, fit="first", split_min=sm) is None)
        print("    split_min=%2d -> %6d traces fully served" % (sm, ok))