#!/usr/bin/env python3
"""Find the structural link between the 116 roster records and the live craft.

Two structures coexist in a mission (mission-liveness-probe.md):
  * 116 objects with vtable 0x820AF030 -- one per UnitGroup roster member,
    confirmed by an exact unit-composition match;
  * 298 heap objects, one per actual craft, located as (def-pointer site - 0x130).

Counting has failed three times to explain 116 -> 298 (Count, n, formation
slots). This looks for a POINTER instead: does a roster record reference its
craft, or a craft reference its roster record?
"""
import os, sys, struct, collections
sys.path.insert(0, __file__.rsplit('/', 1)[0])
import gmem, gworld, entities2

ROSTER_VT = struct.pack('>I', 0x820AF030)
DELTA = 0x130
WIN = 0x400

def scan_vt(fd, size, vt):
    out = []
    for a, b in gmem.extents(fd, size):
        pos = a
        while pos < b:
            n = min(1 << 24, b - pos)
            blob = os.pread(fd, n, pos)
            i = blob.find(vt)
            while i != -1:
                if (pos + i) % 4 == 0: out.append(pos + i)
                i = blob.find(vt, i + 1)
            pos += n
    return sorted(out)

def craft(fd, defs):
    lo, hi = gmem.va_to_off(entities2.ENT_VA_LO), gmem.va_to_off(entities2.ENT_VA_HI)
    out, pos = [], lo
    while pos < hi:
        n = min(1 << 24, hi - pos)
        blob = os.pread(fd, n, pos)
        for k in range(0, len(blob) - 3, 4):
            nm = defs.get(blob[k:k+4])
            if nm: out.append((pos + k - DELTA, nm))
        pos += n
    return out

def main():
    w = gworld.World(); fd = w.fd
    defs = entities2.definitions(w)
    if not defs: print('NOT IN A MISSION'); return 2
    roster = scan_vt(fd, w.size, ROSTER_VT)
    cr = craft(fd, defs)
    print('roster records: %d   live craft: %d   definitions: %d'
          % (len(roster), len(cr), len(defs)))

    craft_va = {}
    for off, nm in cr:
        va = gmem.primary_va(off)
        if va is not None: craft_va[va] = (off, nm)
    cset = set(craft_va)
    print('craft with resolvable VA: %d' % len(cset))

    # forward: roster record -> craft pointer
    fwd = collections.Counter(); hits_at = collections.Counter(); linked = 0
    for ro in roster:
        blob = os.pread(fd, WIN, ro)
        n = 0
        for k in range(0, len(blob) - 3, 4):
            (p,) = struct.unpack_from('>I', blob, k)
            if p in cset:
                n += 1; hits_at[k] += 1
        fwd[n] += 1
        if n: linked += 1
    print('\n--- forward: roster record -> craft VA (window %#x) ---' % WIN)
    print('  records with >=1 craft pointer: %d / %d' % (linked, len(roster)))
    print('  fan-out histogram:', sorted(fwd.items())[:10])
    print('  hit offsets in record:', [('%#x' % k, c) for k, c in hits_at.most_common(8)])

    # reverse: craft -> roster record pointer
    rset = {}
    for ro in roster:
        va = gmem.primary_va(ro)
        if va is not None: rset[va] = ro
    rev = collections.Counter(); roff = collections.Counter(); rlinked = 0
    for off, nm in cr:
        blob = os.pread(fd, WIN, off)
        n = 0
        for k in range(0, len(blob) - 3, 4):
            (p,) = struct.unpack_from('>I', blob, k)
            if p in rset:
                n += 1; roff[k] += 1
        rev[n] += 1
        if n: rlinked += 1
    # Is this a real "no link", or is my VA convention simply wrong? Count
    # pointers into the RANGES rather than at exact computed bases.
    ENT_LO, ENT_HI = entities2.ENT_VA_LO, entities2.ENT_VA_HI
    rvas = [gmem.primary_va(o) for o in roster]
    rvas = [v for v in rvas if v is not None]
    print('\n--- sanity: address ranges ---')
    print('  roster VAs: %#x .. %#x (%d resolvable)'
          % (min(rvas), max(rvas), len(rvas)))
    cv = sorted(cset)
    print('  craft  VAs: %#x .. %#x' % (min(cv), max(cv)))
    into_ent = near = 0
    for ro in roster:
        blob = os.pread(fd, WIN, ro)
        for k in range(0, len(blob) - 3, 4):
            (p,) = struct.unpack_from('>I', blob, k)
            if ENT_LO <= p < ENT_HI:
                into_ent += 1
                if any(abs(p - c) <= 0x400 for c in cv[:4000]): near += 1
    print('  roster words pointing into the entity heap: %d (within 0x400 of a '
          'known craft base: %d)' % (into_ent, near))
    rlo, rhi = min(rvas), max(rvas)
    into_ros = 0
    for off, nm in cr:
        blob = os.pread(fd, WIN, off)
        for k in range(0, len(blob) - 3, 4):
            (p,) = struct.unpack_from('>I', blob, k)
            if rlo <= p <= rhi + 0x400: into_ros += 1
    print('  craft words pointing into the roster VA range: %d' % into_ros)

    # Settle the 5101: if craft point at roster_base + X for a fixed X, the
    # distance to the nearest roster base below spikes at one delta. Flat means
    # they are unrelated allocations sharing a heap.
    import bisect
    rb = sorted(rvas)
    dh = collections.Counter(); far = 0
    for off, nm in cr:
        blob = os.pread(fd, WIN, off)
        for k in range(0, len(blob) - 3, 4):
            (p,) = struct.unpack_from('>I', blob, k)
            if not (rlo <= p <= rhi + 0x400): continue
            i = bisect.bisect_right(rb, p) - 1
            if i < 0: continue
            d = p - rb[i]
            if d <= 0x4000: dh[d] += 1
            else: far += 1
    print('\n--- craft pointer -> distance to nearest roster base ---')
    tot = sum(dh.values()) + far
    print('  in range: %d  (>0x4000 from any base: %d)' % (tot, far))
    for d, c in dh.most_common(12):
        print('    +%#07x  x%-5d  %5.1f%%' % (d, c, 100.0 * c / max(1, tot)))
    if dh:
        top = dh.most_common(1)[0]
        print('  top delta accounts for %.1f%% of in-range pointers'
              % (100.0 * top[1] / max(1, tot)))

    print('\n--- reverse: craft -> roster record VA ---')
    print('  craft with >=1 roster pointer: %d / %d' % (rlinked, len(cr)))
    print('  fan-in histogram:', sorted(rev.items())[:10])
    print('  hit offsets in craft:', [('%#x' % k, c) for k, c in roff.most_common(8)])
    return 0

if __name__ == '__main__':
    sys.exit(main())