"""Every number the commentary claims, pinned. Plain asserts, no pytest.

Run:  python3 test_arena.py

Writes nothing. The 24- and 720-permutation sweeps run in well under a
second; the 986,580-trace census does not live here (it is
`python3 demo.py --census`, about a minute).
"""

import itertools
import os
import struct
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from arena import Arena, OutOfMemory, HEADER, FOOTER, FREE, PREV_FREE

TESTS = []


def test(fn):
    TESTS.append(fn)
    return fn


def filled(order, **kw):
    """4 x malloc(56) fills a 256-byte arena exactly; then free in `order`."""
    a = Arena(256, **kw)
    pay = 56 - (FOOTER if kw.get("boundary_tags") else 0)
    ptrs = [a.malloc(pay) for _ in range(4)]
    for i in order:
        a.free(ptrs[i])
    return a


def fits(a, n):
    try:
        a.malloc(n)
        return True
    except OutOfMemory:
        return False


# --------------------------------------------------------------- the layout

@test
def test_a_fresh_arena_is_one_free_block():
    a = Arena(256)
    assert a.layout() == "[free 256@0]"
    assert a.stats() == dict(free_bytes=256, nfree=1, largest=256,
                             largest_payload=248)


@test
def test_malloc_returns_the_offset_just_past_the_header():
    a = Arena(256)
    assert a.malloc(56) == HEADER


@test
def test_malloc_splits_and_the_remainder_is_free():
    a = Arena(256)
    a.malloc(56)
    assert a.layout() == "[USED 64@0] [free 192@64]"


@test
def test_four_56_byte_allocations_fill_256_bytes_exactly():
    a = Arena(256)
    for _ in range(4):
        a.malloc(56)
    assert a.layout() == ("[USED 64@0] [USED 64@64] [USED 64@128] "
                          "[USED 64@192]")
    assert a.stats()["free_bytes"] == 0
    assert not fits(a, 1)


@test
def test_a_remainder_too_small_for_a_header_is_not_split_off():
    a = Arena(256)
    a.malloc(244)                      # 244 + 8 = 252, leaving 4 < HEADER
    assert a.layout() == "[USED 256@0]"


@test
def test_split_min_keeps_small_remainders_inside_the_allocation():
    a = Arena(256, split_min=32)
    a.malloc(210)                      # remainder 38 >= 8 but < 8 + 32
    assert a.layout() == "[USED 256@0]"
    b = Arena(256, split_min=0)
    b.malloc(210)
    assert b.layout() == "[USED 218@0] [free 38@218]"


# ------------------------------------------------------------- the headline

@test
def test_ascending_frees_leave_a_100_percent_free_arena_in_four_pieces():
    a = filled((0, 1, 2, 3))
    assert a.layout() == ("[free 64@0] [free 64@64] [free 64@128] "
                          "[free 64@192]")
    assert a.stats() == dict(free_bytes=256, nfree=4, largest=64,
                             largest_payload=56)


@test
def test_the_96_byte_request_fails_on_a_100_percent_free_arena():
    a = filled((0, 1, 2, 3))
    assert a.stats()["free_bytes"] == 256          # every byte is free
    assert 96 < 256                                # and the request is smaller
    assert not fits(a, 96)                         # and it fails anyway
    assert 96 + HEADER > 64                        # 104 > 64: the arithmetic


@test
def test_descending_frees_leave_one_block_and_serve_the_same_request():
    a = filled((3, 2, 1, 0))
    assert a.layout() == "[free 256@0]"
    assert fits(a, 96)
    assert fits(Arena(256), 248) and not fits(Arena(256), 249)


@test
def test_forward_coalescing_merges_nothing_on_an_ascending_free_order():
    """The counterfactual: turning coalescing off changes not one byte."""
    off = filled((0, 1, 2, 3), coalesce="none")
    fwd = filled((0, 1, 2, 3), coalesce="forward")
    assert off.mem == fwd.mem


@test
def test_forward_coalescing_does_work_when_the_successor_is_already_free():
    a = filled((3, 2))                             # 2 merges with 3
    assert a.layout() == "[USED 64@0] [USED 64@64] [free 128@128]"


# ------------------------------------------------------- the 24-order census

@test
def test_exactly_one_of_the_24_free_orders_fails_at_96_bytes():
    bad = [o for o in itertools.permutations(range(4))
           if not fits(filled(o), 96)]
    assert bad == [(0, 1, 2, 3)]


@test
def test_the_failing_orders_per_request_size():
    counts = {}
    for req in (48, 56, 88, 96, 120, 121, 152, 184, 200, 248):
        counts[req] = sum(1 for o in itertools.permutations(range(4))
                          if not fits(filled(o), req))
    assert counts == {48: 0, 56: 0, 88: 1, 96: 1, 120: 1, 121: 17, 152: 17,
                      184: 17, 200: 23, 248: 23}


@test
def test_only_the_fully_descending_order_coalesces_the_whole_arena():
    whole = [o for o in itertools.permutations(range(4))
             if filled(o).stats()["nfree"] == 1]
    assert whole == [(3, 2, 1, 0)]


@test
def test_six_blocks_720_orders():
    def one(order, req):
        a = Arena(240)
        p = [a.malloc(32) for _ in range(6)]
        for i in order:
            a.free(p[i])
        return fits(a, req)
    counts = {}
    for req in (56, 72, 96, 120):
        counts[req] = sum(1 for o in itertools.permutations(range(6))
                          if not one(o, req))
    assert counts == {56: 1, 72: 1, 96: 349, 120: 642}


# ---------------------------------------------------------------- the knobs

@test
def test_no_fit_policy_rescues_the_failing_request():
    for fit in ("first", "best", "worst"):
        assert not fits(filled((0, 1, 2, 3), fit=fit), 96)


@test
def test_the_fit_policies_do_differ_when_a_choice_exists():
    """Same two holes, three answers -- so the knob is not inert, it just
    cannot help when nothing fits."""
    def three_holes(fit, place=True):
        a = Arena(256, fit=fit, coalesce="both")
        p = [a.malloc(56), a.malloc(8), a.malloc(24), a.malloc(8)]
        a.free(p[0])                      # 64-byte hole at 0
        a.free(p[2])                      # 32-byte hole at 80; 128 free at end
        if place:
            a.malloc(24)                  # needs 32 -- all three holes fit
        return a.layout()
    assert three_holes("first", place=False) == (
        "[free 64@0] [USED 16@64] [free 32@80] [USED 16@112] [free 128@128]")
    assert three_holes("first") == (
        "[USED 32@0] [free 32@32] [USED 16@64] [free 32@80] [USED 16@112] "
        "[free 128@128]")
    assert three_holes("best") == (
        "[free 64@0] [USED 16@64] [USED 32@80] [USED 16@112] [free 128@128]")
    assert three_holes("worst") == (
        "[free 64@0] [USED 16@64] [free 32@80] [USED 16@112] [USED 32@128] "
        "[free 96@160]")


@test
def test_backward_coalescing_rescues_every_one_of_the_24_orders():
    for o in itertools.permutations(range(4)):
        assert filled(o, coalesce="both").stats() == dict(
            free_bytes=256, nfree=1, largest=256, largest_payload=248)


@test
def test_deferred_consolidation_rescues_it_too():
    a = filled((0, 1, 2, 3), coalesce="none", consolidate_on_fail=True)
    assert a.stats()["nfree"] == 4                 # nothing merged at free
    assert fits(a, 96)                             # merged on the failure
    assert a.layout() == "[USED 104@0] [free 152@104]"


@test
def test_consolidate_merges_every_run_but_not_across_a_live_block():
    a = Arena(256, coalesce="none")
    p = [a.malloc(56) for _ in range(4)]
    a.free(p[0])
    a.free(p[1])
    a.free(p[3])
    a.consolidate()
    assert a.layout() == "[free 128@0] [USED 64@128] [free 64@192]"


# ------------------------------------------------------------ boundary tags

@test
def test_boundary_tags_also_rescue_every_order_at_12_bytes_a_block():
    for o in itertools.permutations(range(4)):
        a = filled(o, coalesce="both", boundary_tags=True)
        assert a.overhead == 12
        assert a.stats() == dict(free_bytes=256, nfree=1, largest=256,
                                 largest_payload=244)


@test
def test_the_footer_costs_one_of_the_four_allocations():
    served = {}
    for bt in (False, True):
        a = Arena(256, coalesce="both", boundary_tags=bt)
        n = 0
        try:
            for _ in range(4):
                a.malloc(56)
                n += 1
        except OutOfMemory:
            pass
        served[bt] = n
    assert served == {False: 4, True: 3}


@test
def test_prev_free_is_maintained_on_both_sides_of_an_allocation():
    a = Arena(256, coalesce="both", boundary_tags=True)
    p = a.malloc(52)
    assert struct.unpack_from("<II", a.mem, 64)[1] == FREE        # not PREV_FREE
    a.free(p)
    assert struct.unpack_from("<II", a.mem, 0)[1] & FREE
    assert a.layout() == "[free 256@0]"


@test
def test_the_footer_is_only_readable_because_of_the_prev_free_bit():
    """Without the bit, the 4 bytes below a header are somebody's payload.
    Read them anyway and a block merges with itself."""
    class NaiveFooter(Arena):
        def _prev(self, off):
            if off == 0:
                return None
            ptotal, = struct.unpack_from("<I", self.mem, off - FOOTER)
            return off - ptotal

    a = NaiveFooter(512, coalesce="both", boundary_tags=True)
    p = [a.malloc(20) for _ in range(16)]
    assert struct.unpack_from("<I", a.mem, 480 - FOOTER)[0] == 0
    try:
        for i in range(15, -1, -1):
            a.free(p[i])
        raise AssertionError("expected the naive footer to blow up")
    except struct.error as e:
        assert "540" in str(e) and "512" in str(e)

    good = Arena(512, coalesce="both", boundary_tags=True)
    q = [good.malloc(20) for _ in range(16)]
    for i in range(15, -1, -1):
        good.free(q[i])
    assert good.layout() == "[free 512@0]"


@test
def test_the_footer_turns_a_quadratic_walk_into_a_linear_one():
    cost = {}
    for n in (4, 8, 16, 32, 64):
        for bt in (False, True):
            a = Arena(n * 32, coalesce="both", boundary_tags=bt)
            p = [a.malloc(24 - (FOOTER if bt else 0)) for _ in range(n)]
            a.reads = 0
            for i in range(n - 1, -1, -1):
                a.free(p[i])
            cost[(n, bt)] = a.reads
            assert a.stats()["nfree"] == 1
    assert [cost[(n, False)] for n in (4, 8, 16, 32, 64)] == [23, 65, 197,
                                                              653, 2333]
    assert [cost[(n, True)] for n in (4, 8, 16, 32, 64)] == [20, 44, 92,
                                                             188, 380]
    for n in (4, 8, 16, 32, 64):                   # the closed forms
        assert cost[(n, False)] == n * n // 2 + 9 * n // 2 - 3
        assert cost[(n, True)] == 6 * n - 4


@test
def test_what_each_policy_costs_on_the_headline_trace():
    cost = {}
    for label, kw in (("forward", dict(coalesce="forward")),
                      ("both", dict(coalesce="both")),
                      ("footers", dict(coalesce="both", boundary_tags=True)),
                      ("deferred", dict(coalesce="none",
                                        consolidate_on_fail=True))):
        a = Arena(256, **kw)
        pay = 56 - (FOOTER if kw.get("boundary_tags") else 0)
        p = [a.malloc(pay) for _ in range(4)]
        a.reads = 0
        for i in range(4):
            a.free(p[i])
        frees = a.reads
        a.reads = 0
        fits(a, 96)
        cost[label] = (frees, a.reads)
    assert cost == {"forward": (11, 4), "both": (23, 3), "footers": (25, 4),
                    "deferred": (8, 12)}


# ------------------------------------------------------------- the boundary

@test
def test_one_size_fits_all_makes_the_free_order_irrelevant():
    for o in itertools.permutations(range(4)):
        a = filled(o)
        assert all(fits(a, 56) for _ in range(4))


@test
def test_merging_the_nearest_earlier_block_hands_out_live_memory():
    """§7.1: `prev = o if free else None` without the `else None`."""
    class MergeNearestEarlier(Arena):
        def _prev(self, off):
            if off == 0:
                return None
            prev = None
            for o, total, free in self.blocks():
                if o == off:
                    return prev
                if free:
                    prev = o
            return None

    a = MergeNearestEarlier(256, coalesce="both")
    p = [a.malloc(56) for _ in range(4)]
    a.free(p[0])
    a.free(p[2])
    assert a.layout() == "[free 128@0] [free 64@128] [USED 64@192]"
    a.mem[p[1]:p[1] + 5] = b"MINE!"
    new = a.malloc(120)
    a.mem[new:new + 120] = b"X" * 120
    assert new == 8 and p[1] == 72
    assert bytes(a.mem[p[1]:p[1] + 5]) == b"XXXXX"   # block 1 overwritten
    good = Arena(256, coalesce="both")
    q = [good.malloc(56) for _ in range(4)]
    good.free(q[0])
    good.free(q[2])
    assert good.layout() == ("[free 64@0] [USED 64@64] [free 64@128] "
                             "[USED 64@192]")


@test
def test_freeing_only_three_of_the_four_blocks_fails_the_same_way():
    """§7.3: the failure needs fragmentation, not an empty arena."""
    a = Arena(256)
    p = [a.malloc(56) for _ in range(4)]
    for i in (0, 1, 3):
        a.free(p[i])
    assert a.layout() == ("[free 64@0] [free 64@64] [USED 64@128] "
                          "[free 64@192]")
    assert a.stats()["free_bytes"] == 192
    assert not fits(a, 96)


@test
def test_the_largest_request_the_failing_arena_can_serve_is_56():
    """§9 Q1."""
    assert max(n for n in range(1, 257) if fits(filled((0, 1, 2, 3)), n)) == 56


@test
def test_three_per_thread_arenas_hold_768_free_bytes_and_serve_none_of_them():
    """§9 Q4: sharding multiplies free bytes, not contiguity."""
    arenas = [filled((0, 1, 2, 3)) for _ in range(3)]
    assert sum(a.stats()["free_bytes"] for a in arenas) == 768
    assert sum(1 for a in arenas if fits(a, 96)) == 0


@test
def test_with_footers_the_fourth_allocation_has_to_shrink_to_40():
    """§9 Q5."""
    def head_room(n):
        a = Arena(256, coalesce="both", boundary_tags=True)
        for _ in range(3):
            a.malloc(56)
        assert a.layout() == ("[USED 68@0] [USED 68@68] [USED 68@136] "
                              "[free 52@204]")
        return fits(a, n)
    assert head_room(40) and not head_room(41)


@test
def test_payloads_survive_their_neighbours_being_freed_and_reused():
    a = Arena(256, coalesce="both")
    keep = a.malloc(56)
    a.mem[keep:keep + 5] = b"MINE!"
    doomed = a.malloc(56)
    a.mem[doomed:doomed + 5] = b"THEIR"
    a.free(doomed)
    a.malloc(24)
    assert bytes(a.mem[keep:keep + 5]) == b"MINE!"


if __name__ == "__main__":
    for fn in TESTS:
        fn()
        print("ok   %s" % fn.__name__[5:].replace("_", " "))
    print("\n%d tests passed" % len(TESTS))