"""A malloc/free allocator small enough to read in one sitting: one fixed
byte array, an 8-byte header in front of every block, split on allocate and
merge on free.

    +--------+-----------------+--------+------------------+
    | hdr(8) | payload         | hdr(8) | payload          |  ...
    +--------+-----------------+--------+------------------+
    ^        ^
    block    the "pointer" malloc hands back (an int offset, not an address)

The header is `<II`: the block's TOTAL size including the header, and a flag
word. There is no free list threaded through the blocks -- `_find` walks the
heap, because `off + total` is the address of the next block. Which is the
whole point of the toy:

1. A header-only layout can walk FORWARD (`off + total`) and cannot walk
   BACKWARD. So `free` can merge a block with its successor and not with its
   predecessor.
2. Merging is local, and it happens at free time. Freeing blocks in
   ASCENDING address order therefore merges nothing at all: every successor
   is still live at the moment its predecessor is freed.
3. So the arena can be 100% free by bytes and still fail a request smaller
   than its total free space. `coalesce="both"` fixes it with an O(n) heap
   walk; `boundary_tags=True` fixes it in O(1) with a 4-byte footer -- and
   the footer is only readable because of the PREV_FREE bit, which is what
   glibc calls PREV_INUSE.

Deterministic: no clock, no randomness, offsets rather than addresses, and
`_find` scans in address order. Every run prints the same bytes.
"""

import struct

HEADER = 8       # <II: total size (header included), flags
FOOTER = 4       # <I:  total size again, at the tail of a FREE block
FREE = 1         # flags bit 0: this block is free
PREV_FREE = 2    # flags bit 1: the physically PREVIOUS block is free


class OutOfMemory(Exception):
    pass


class Arena:
    """`size` bytes of memory and five policy knobs.

    fit                 "first" | "best" | "worst" -- which free block wins
    coalesce            "none" | "forward" | "both" -- who gets merged, when
    split_min           smallest payload worth splitting a remainder off for
    boundary_tags       spend FOOTER bytes per block to make "both" O(1)
    consolidate_on_fail sweep and merge everything when a malloc fails
    """

    def __init__(self, size, fit="first", coalesce="forward", split_min=0,
                 boundary_tags=False, consolidate_on_fail=False):
        self.mem = bytearray(size)
        self.size = size
        self.fit = fit
        self.coalesce = coalesce
        self.split_min = split_min
        self.boundary_tags = boundary_tags
        self.consolidate_on_fail = consolidate_on_fail
        # What a block costs before any payload. The footer has to be
        # reserved in every block, free or not: nothing knows at allocation
        # time whether this block will be freed later.
        self.overhead = HEADER + (FOOTER if boundary_tags else 0)
        self.reads = 0           # headers read -- the toy's cost unit
        self._mark(0, size, True, prev_free=False)

    # ---------------------------------------------------------------- layout

    def _hdr(self, off):
        """Read one header. Every read goes through here so `self.reads` is
        an honest count of the work a policy does."""
        self.reads += 1
        return struct.unpack_from("<II", self.mem, off)

    def _mark(self, off, total, free, prev_free=None):
        """Write a block's state. With boundary tags that state lives in
        THREE places -- header, footer, and the successor's PREV_FREE bit --
        so every write goes through one function and they cannot drift.
        """
        if prev_free is None:
            prev_free = bool(self._hdr(off)[1] & PREV_FREE)
        flags = (FREE if free else 0) | (PREV_FREE if prev_free else 0)
        struct.pack_into("<II", self.mem, off, total, flags)
        if not self.boundary_tags:
            return
        if free:
            struct.pack_into("<I", self.mem, off + total - FOOTER, total)
        nxt = off + total
        if nxt < self.size:
            ntotal, nflags = self._hdr(nxt)
            nflags = (nflags & FREE) | (PREV_FREE if free else 0)
            struct.pack_into("<II", self.mem, nxt, ntotal, nflags)

    def blocks(self):
        """Every block, front to back. `off + total` is the next header --
        the only navigation an 8-byte header can offer."""
        off = 0
        while off < self.size:
            total, flags = self._hdr(off)
            yield off, total, bool(flags & FREE)
            off += total

    # -------------------------------------------------------------- allocate

    def malloc(self, nbytes):
        need = nbytes + self.overhead
        off = self._find(need)
        if off is None and self.consolidate_on_fail:
            self.consolidate()
            off = self._find(need)
        if off is None:
            raise OutOfMemory("no block for %d bytes" % nbytes)
        total, _ = self._hdr(off)
        rest = total - need
        if rest >= self.overhead + self.split_min:
            # Split. The remainder is written FIRST: marking the allocated
            # block clears the remainder's PREV_FREE, which is only correct
            # if the remainder's header already exists.
            self._mark(off + need, rest, True, prev_free=False)
            self._mark(off, need, False)
        else:
            # Too small to be a block of its own: the request keeps the
            # slack. Internal fragmentation, deliberately.
            self._mark(off, total, False)
        return off + HEADER

    def _find(self, need):
        """The knob everyone argues about. It can only choose among blocks
        that already fit -- it cannot make one bigger."""
        best = None
        for off, total, free in self.blocks():
            if not free or total < need:
                continue
            if self.fit == "first":
                return off
            if (best is None
                    or (self.fit == "best" and total < best[1])
                    or (self.fit == "worst" and total > best[1])):
                best = (off, total)
        return None if best is None else best[0]

    # ------------------------------------------------------------------ free

    def free(self, ptr):
        off = ptr - HEADER
        total, _ = self._hdr(off)
        self._mark(off, total, True)
        if self.coalesce == "none":
            return
        nxt = off + total                      # forward: one addition
        if nxt < self.size:
            ntotal, nflags = self._hdr(nxt)
            if nflags & FREE:
                total += ntotal
                self._mark(off, total, True)
        if self.coalesce == "both":
            prev = self._prev(off)             # backward: the hard direction
            if prev is not None:
                ptotal, _ = self._hdr(prev)
                self._mark(prev, ptotal + total, True)

    def _prev(self, off):
        """The physically previous block, if it is free. A header-only heap
        cannot answer this in less than a walk from byte 0; a footer answers
        it in one read, but ONLY if PREV_FREE says those 4 bytes are a
        footer and not the tail of somebody's payload.
        """
        if off == 0:
            return None
        if self.boundary_tags:
            if not self._hdr(off)[1] & PREV_FREE:
                return None
            ptotal, = struct.unpack_from("<I", self.mem, off - FOOTER)
            return off - ptotal
        prev = None
        for o, total, free in self.blocks():
            if o == off:
                return prev
            prev = o if free else None
        return None

    def consolidate(self):
        """Deferred coalescing: one forward sweep merging every run of free
        blocks. Needs no footer, and no free order can defeat it -- but it
        costs a full heap walk, so it is paid on failure, not on every free.
        """
        off = 0
        while off < self.size:
            total, flags = self._hdr(off)
            if flags & FREE:
                nxt = off + total
                while nxt < self.size:
                    ntotal, nflags = self._hdr(nxt)
                    if not nflags & FREE:
                        break
                    total += ntotal
                    nxt += ntotal
                self._mark(off, total, True)
            off += total

    # ------------------------------------------------------------- reporting

    def stats(self):
        free = [t for _, t, f in self.blocks() if f]
        largest = max(free, default=0)
        return dict(free_bytes=sum(free), nfree=len(free), largest=largest,
                    largest_payload=max(0, largest - self.overhead))

    def layout(self):
        return " ".join("[%s %d@%d]" % ("free" if f else "USED", t, o)
                        for o, t, f in self.blocks())