firestorm/MW4COMPARE/tools/decompile/datamap.py

#!/usr/bin/env python3
"""
datamap.py - field map for the mech `<chassis>.data{GameModel}` record (1636 B).

Built from the engine headers, not from value matching.

Why: value matching alone cannot separate keys that hold the same value in every
mech. `dampenWorldJoint` and `fallAdjustmentSeconds` are both 0.5 everywhere, so
34 of the 80 numeric keys came out ambiguous, and ordering them by their
appearance in the source file assigned several of them wrongly - it put
`tiltSpeed` at 672, which is really `slopeDecel2`, and swapped
`percentageOfTurnToStartTilt` with `percentageOfSpeedToStartTilt`.

So the layout is taken from the declaration order in

    Vehicle__GameModel   mw4/Code/MW4/Vehicle.hpp
    Mech__GameModel      mw4/Code/MW4/Mech.hpp

and each block is anchored using the fields that value matching *did* resolve
unambiguously. Every anchor agrees:

    Mech block    base 756: footReturnSeconds 764, dampenTorsoJoint 800,
                            undampenRootJoint 816, undampenHipJoint 824,
                            scaleInternalTiltDegree 832
    Vehicle block base 664: minSpeed 692, maxSpeed 696, acceleration 720,
                            decceleration 724, reverseAccelerationMultiplier 728

All fields in both blocks are 4-byte Scalar/Radian, so field i sits at
base + 4*i.
"""
import glob, math, os, re, struct, collections

REPO = "/home/rich/Repositories/firestorm/Gameleap"
SRC = REPO + "/mw4/Content/Mechs"
CODE = REPO + "/code"
REC = "/home/rich/Repositories/FS_Ours_extracted/core/mechs"
GAMEMODEL_SIZE = 1636

# The GameModel inheritance chain for a mech, base class first. Bases are NOT
# hardcoded: each block starts where the previous one ended. MWMover__GameModel
# is just `typedef Adept::Mover__GameModel` (MWMover.hpp:173), so it adds
# nothing. `Entity__GameModel` declares no base class and wraps its members in
# `#if NSWIZZLE`, which is never defined anywhere in the tree - the #else branch
# is live, and the two branches order their members DIFFERENTLY.
CHAIN = [
    ("Entity__GameModel", CODE + "/mw4/Libraries/Adept/Entity.hpp"),
    ("Mover__GameModel", CODE + "/mw4/Libraries/Adept/Mover.hpp"),
    ("MWObject__GameModel", CODE + "/mw4/Code/MW4/MWObject.hpp"),
    ("Vehicle__GameModel", CODE + "/mw4/Code/MW4/Vehicle.hpp"),
    ("Mech__GameModel", CODE + "/mw4/Code/MW4/Mech.hpp"),
]
# Independently measured block starts, used to check the computed chain.
ANCHORS = {"Vehicle__GameModel": 664, "Mech__GameModel": 756}
# The factories that convert authored degrees into stored radians.
FACTORIES = [CODE + "/mw4/Code/MW4/Mech_Tool.cpp", CODE + "/mw4/Code/MW4/Vehicle_Tool.cpp"]
SCALAR_TYPES = r'(?:Stuff::Scalar|Stuff::Radian|Stuff::Angle|float)'
DEG2RAD = math.pi / 180.0
MAX_STRING_LENGTH = 256          # Entity.hpp:199

# Member sizes. bool is ONE byte, not four -- getting this wrong shifts every
# field after the six m_canLoad* flags by exactly 16 bytes.
SIZES = {
    "Scalar": 4, "Radian": 4, "Angle": 4, "float": 4, "int": 4, "unsigned": 4,
    "DWORD": 4, "WORD": 2, "BYTE": 1, "bool": 1, "ResourceID": 4,
    "Point3D": 12, "Vector3D": 12, "UnitQuaternion": 16, "RGBAColor": 16,
    "Motion3D": 24, "LinearMatrix4D": 48, "char": 1,
    "ClassID": 4, "ReplicatorID": 4, "FactoryRequest": 4, "ObjectID": 4,
}
MEMBER_RE = re.compile(
    r'\b((?:Stuff::|Adept::)?(?:Scalar|Radian|Angle|Point3D|Vector3D|UnitQuaternion'
    r'|RGBAColor|Motion3D|LinearMatrix4D|ResourceID)|(?:Stuff::)?(?:RegisteredClass::)?ClassID'
    r'|ReplicatorID|ObjectID|(?:\w+::)?FactoryRequest|float|int|bool|BYTE|WORD|DWORD'
    r'|unsigned|char)\s+([A-Za-z_][\w\s,\[\]]*?);', re.S)


def members(cls, path):
    """-> [(typeName, memberName, sizeInBytes)] in declaration order."""
    txt = open(path, encoding="latin-1").read()
    m = re.search(rf'class\s+{cls}\s*(?::[^{{;]*)?\{{', txt, re.S)
    if not m:
        return []
    depth, i = 1, m.end()          # brace-match; indentation is inconsistent between headers
    while i < len(txt) and depth:
        depth += (txt[i] == "{") - (txt[i] == "}")
        i += 1
    body = txt[m.end():i - 1]
    ctor = body.find(f"{cls}(")
    if ctor > 0:
        body = body[:ctor]
    body = re.sub(r'//[^\n]*', '', body)
    # `typedef int AttributeID;` is not a member. This masked a missing ClassID
    # for a while: two 4-byte errors that happened to cancel.
    body = re.sub(r'\btypedef\b[^;]*;', '', body, flags=re.S)
    body = re.sub(r'#\s*if\s+NSWIZZLE\b.*?#\s*else', '', body, flags=re.S)
    body = re.sub(r'#\s*(endif|else|if\w*|ifdef|ifndef)[^\n]*', '', body)
    out = []
    for d in MEMBER_RE.finditer(body):
        base = (d.group(1).replace("Stuff::", "").replace("Adept::", "")
                        .replace("RegisteredClass::", ""))
        base = base.rsplit("::", 1)[-1]
        for name in d.group(2).split(","):
            name = name.strip()
            arr = re.fullmatch(r'([A-Za-z_]\w*)\s*\[\s*([A-Za-z_]\w*|\d+)\s*\]', name)
            if arr:
                n = arr.group(2)
                count = MAX_STRING_LENGTH if n == "MaxStringLength" else int(n)
                out.append((base, arr.group(1), SIZES[base] * count))
            elif re.fullmatch(r'[A-Za-z_]\w*', name):
                out.append((base, name, SIZES[base]))
    return out


def chain_layout(chain=None, anchors=None, start=0):
    """-> ({memberName: (offset, typeName, size)}, {className: baseOffset}).

    Each block starts where the previous ended; /Zp4 means align = min(4, size).
    Raises if a computed base contradicts an independently measured anchor.
    Defaults to the mech chain; pass another for Torso/Engine and friends.
    """
    chain = CHAIN if chain is None else chain
    anchors = ANCHORS if anchors is None else anchors
    fields, bases, off = {}, {}, start
    for cls, path in chain:
        off += (-off) % 4
        bases[cls] = off
        want = anchors.get(cls)
        if want is not None and off != want:
            raise AssertionError(f"{cls} computed at {off}, measured {want}")
        for typ, name, size in members(cls, path):
            off += (-off) % min(4, size)
            fields.setdefault(name, (off, typ, size))
            off += size
    return fields, bases


def angle_fields():
    """Fields the tool factory scales by Radians_Per_Degree.

    This cannot be inferred from the header: torsoHitSpringMotionLimit and its
    siblings are declared plain Stuff::Scalar, yet Mech_Tool.cpp:889 stores
    `model->torsoHitSpringMotionLimit * Radians_Per_Degree`. Only the writer
    knows. Stuff::Radian fields (tiltSpeed, tiltDegree, topSpeedTurnRate,
    fullStopTurnRate) are handled by their declared type as well.
    """
    out = set()
    for path in FACTORIES:
        if not os.path.exists(path):
            continue
        txt = open(path, encoding="latin-1").read()
        for m in re.finditer(r'model->(\w+)\s*=\s*model->\w+\s*\*\s*Radians_Per_Degree', txt):
            out.add(norm(m.group(1)))
    return out


def declared_fields(cls, path):
    """-> [(fieldName, isAngle)] in declaration order, scalars only.

    isAngle marks Stuff::Radian / Stuff::Angle. Those are authored in DEGREES in
    the .data but stored in RADIANS in the record, so they need a pi/180 factor
    on the way in and 180/pi on the way back out.
    """
    txt = open(path, encoding="latin-1").read()
    m = re.search(rf'class\s+{cls}\s*:(.*?)^\t\t\}};', txt, re.S | re.M)
    if not m:
        return []
    body = m.group(1)
    ctor = body.find(f"{cls}(")
    if ctor > 0:
        body = body[:ctor]
    body = re.sub(r'//[^\n]*', '', body)
    out = []
    for decl in re.finditer(rf'\b({SCALAR_TYPES})\s+([A-Za-z_][\w\s,]*?);', body, re.S):
        angle = decl.group(1) in ("Stuff::Radian", "Stuff::Angle")
        for name in decl.group(2).split(","):
            n = name.strip()
            if re.fullmatch(r'[A-Za-z_]\w*', n):
                out.append((n, angle))
    return out


def norm(name):
    return re.sub(r'^m_', '', name).replace("_", "").lower()


def header_field_map():
    """normalised member name -> (offset, typeName, size, isAngle)."""
    angles = angle_fields()
    out = {}
    for name, (off, typ, size) in chain_layout()[0].items():
        key = norm(name)
        out.setdefault(key, (off, typ, size, typ == "Radian" or key in angles))
    return out


def gamedata(path):
    txt = open(path, "rb").read().decode("latin-1")
    # Shadow={...} contains a line reading "[shadow]"; without hiding braced
    # blocks the page scan stops there and every later key is lost silently.
    txt = re.sub(r'\{.*?\}',
                 lambda x: x.group(0).replace("\r", "").replace("\n", "\x01"),
                 txt, flags=re.S)
    m = re.search(r'^\[GameData\]\r?\n(.*?)(?=^\[[A-Za-z]|\Z)', txt, re.M | re.S)
    if not m:
        return collections.OrderedDict()
    kv = collections.OrderedDict()
    for k, v in re.findall(r'^([A-Za-z0-9_]+)=([^\r\n]*)', m.group(1), re.M):
        kv.setdefault(k, v.replace("\x01", "\n"))
    return kv


def corpus():
    """-> [(chassis, {key: value}, gameModelBytes)] for our 64 mech chassis."""
    out = []
    for d in sorted(glob.glob(REC + "/*")):
        ch = os.path.basename(d)
        gm = [g for g in glob.glob(d + "/*.data{GameModel}")
              if os.path.basename(g).lower().startswith(ch.lower() + ".data")]
        if not gm:
            continue
        blob = open(gm[0], "rb").read()
        if len(blob) != GAMEMODEL_SIZE:
            continue
        s = [p for p in glob.glob(SRC + "/*/*.data")
             if os.path.basename(p).lower() == ch.lower() + ".data"]
        if s:
            out.append((ch, gamedata(s[0]), blob))
    return out


def build():
    """-> (corpus, {sourceKey: (offset, typeName, size, isAngle)})"""
    pairs = corpus()
    hmap = header_field_map()
    keys = collections.Counter()
    for _c, kv, _b in pairs:
        keys.update(kv.keys())
    return pairs, {k: hmap[norm(k)] for k in keys if norm(k) in hmap}


FLOATS = {"Scalar", "Radian", "Angle", "float"}
VECTORS = {"Point3D": 3, "Vector3D": 3, "RGBAColor": 4, "UnitQuaternion": 4}


def read(blob, off, typ="Scalar", size=4, angle=False):
    """Decoded value in the units and form the .data source uses."""
    if typ in FLOATS:
        v = struct.unpack_from("<f", blob, off)[0]
        return v / DEG2RAD if angle else v
    if typ in VECTORS:
        return list(struct.unpack_from(f"<{VECTORS[typ]}f", blob, off))
    if typ == "char":
        return blob[off:off + size].split(b"\0")[0].decode("latin-1")
    if typ in ("bool", "BYTE"):
        return blob[off]
    if typ == "WORD":
        return struct.unpack_from("<H", blob, off)[0]
    if typ == "ResourceID":
        return struct.unpack_from("<I", blob, off)[0]
    return struct.unpack_from("<i", blob, off)[0]


if __name__ == "__main__":
    pairs, field_map = build()
    print(f"chassis: {len(pairs)}   mapped keys: {len(field_map)}")
    for k, (off, typ, size, angle) in sorted(field_map.items(), key=lambda kv: kv[1][0]):
        note = "  [degrees]" if angle else ""
        print(f"  {off:5d}  {typ:14s} {k}{note}")