firestorm/MW4COMPARE/tools/decompile/data.py

#!/usr/bin/env python3
"""
data.py - decompiler for a mech `<chassis>.data` file.

Rebuilds the `[GameData]` page from the compiled records:

    <chassis>.data{GameModel}   1636-byte flat struct  (datamap.py)
    <chassis>.data{FootSteps}   foot-step texture list
    <chassis>.data[shadow]      inline Shadow notation block

Value sources, in order of preference:

  * a struct member          - typed read through datamap.chain_layout()
  * a ResourceID member      - record id (HIGH word) resolved via the manifest
  * a symbolic constant      - int reversed through constants.py
  * an explicit factory key  - handled below, because SaveGameModel writes these
                               outside the attribute table

Four keys cannot be recovered and are deliberately omitted: BattleDamageRatio,
BattleKillBonus, DragoonValue and VehicleTradeValue. They appear in every source
.data but are read by NOTHING in the engine - no factory, no attribute
registration, no runtime reference - so they never enter the package. They are
authoring metadata; omitting them changes no behaviour.

    python3 data.py <chassis-record-dir> [-o out.data]
    python3 data.py --verify
"""
import argparse, collections, os, glob, re, struct, subprocess, sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import datamap, constants, subsystems

MANIFEST = "/home/rich/Repositories/FS_Ours_extracted/_manifest.tsv"

# Adept.hpp:212. NoMaterial is 0.
MATERIALS = [
    "NoMaterial", "Grass", "Water", "Concrete", "GreyDirt", "BrownDirt", "Rock",
    "DarkConcrete", "DarkGreyDirt", "DarkBrownDirt", "DarkRock", "Blacktop",
    "Snow", "Wood", "Lava", "Glass", "Steel", "Us", "Them", "LightMineral",
    "DarkMineral", "Ash", "CrackedLava", "OpenLava",
]

# Identical in all 64 chassis; SaveGameModel writes them outside the attribute table.
CONSTANTS = {
    "Class": "MechWarrior4::Mech",
    "FaceLighting": "yes",
    "LookupLighting": "yes",
    "VertexLighting": "yes",
    "LightMapLighting": "no",
    # byte-identical in all 89 mech .data files, destroyed variants included
    "Shadow": "{\n[shadow]\nLightType=Shadow\nInnerRadius=4.0\nOuterRadius=10.0\n"
              "BlobDistance=200.0\nShadowMap=ShadowMask\nIntensity=0.4\n}",
}

# Keys whose source name differs from the struct member name.
ALIASES = {
    "AnimationScript": "animScriptName",
    "HeatManager": "heatManagerResource",
    "FootEffectsFile": "footFallEffectsTable",
}

# CraterName is stored as MString::GetHashValue (DeathEntity_Tool.cpp:34), which
# is one-way. All 64 chassis hold the same hash, so the single authored value is
# recoverable by constancy rather than by inversion.
CRATER_HASH = {217688981: "crater01"}

# Read by nothing in the engine - see module docstring.
UNRECOVERABLE = ["BattleDamageRatio", "BattleKillBonus", "DragoonValue", "VehicleTradeValue"]

SYMBOLIC = ("MechID", "TechType", "NameIndex", "MoveTypeFlag")


def stem(mech_dir):
    """File stem used inside a mech folder; it need not match the folder name.

    jenner2c/ holds jenner_2c.*, the same trap Black Hawk/nova sprang on
    .subsystems -- never assume the folder name.
    """
    names = [f for f in os.listdir(mech_dir) if f.lower().endswith(".data")]
    if names:
        return names[0][:-len(".data")]
    for f in os.listdir(mech_dir):
        m = re.match(r'(.+)\.data[\{\[]', f, re.I)
        if m:
            return m.group(1)
    return os.path.basename(mech_dir.rstrip("/"))


def record(mech_dir, suffix):
    """Read one qualified record. Not glob -- '[shadow]' is a character class."""
    want = (stem(mech_dir) + suffix).lower()
    for fn in os.listdir(mech_dir):
        if fn.lower() == want:
            return open(os.path.join(mech_dir, fn), "rb").read()
    return None


def bool_words():
    """-> {key: (trueWord, falseWord)}; sources spell these inconsistently.

    Collider and CanBeShot are written true/false, the CanLoad* flags Yes/No.
    """
    seen = collections.defaultdict(collections.Counter)
    for _ch, kv, _b in datamap.corpus():
        for k, v in kv.items():
            w = v.strip().lower()
            if w in ("true", "false", "yes", "no"):
                seen[k][w] += 1
    out = {}
    for k, c in seen.items():
        plain = c["true"] + c["false"] >= c["yes"] + c["no"]
        out[k] = ("true", "false") if plain else ("Yes", "No")
    return out


def foot_steps(blob):
    """-> (defaultTexture, [(texture, materialName)]).

    Stream written by Mech_Tool.cpp:196: int material (-1 = default), int length
    NOT counting the terminator, the characters, a NUL, then a one-byte
    isDefault flag.
    """
    default, rows, off = None, [], 0
    while off + 8 <= len(blob):
        material, length = struct.unpack_from("<iI", blob, off)
        off += 8
        text = blob[off:off + length].decode("latin-1")
        off += length + 1                       # skip the terminator
        is_default = blob[off] if off < len(blob) else 0
        off += 1
        if is_default or material < 0:
            default = text
        else:
            name = MATERIALS[material] if 0 <= material < len(MATERIALS) else str(material)
            rows.append((text, name))
    return default, rows


def decompile(mech_dir, manifest=None, retarget_ids=False, obb_dir=None):
    """-> OrderedDict key -> value or [values]."""
    manifest = manifest if manifest is not None else subsystems.load_manifest(MANIFEST)
    ch = os.path.basename(mech_dir.rstrip("/"))
    gm = record(mech_dir, ".data{GameModel}")
    if gm is None:
        raise SystemExit(f"no GameModel record in {mech_dir}")

    layout = datamap.chain_layout()[0]
    angles = datamap.angle_fields()
    tables = constants.tables()
    words = bool_words()
    by_member = {datamap.norm(n): n for n in layout}
    out = collections.OrderedDict()

    def member(key):
        n = ALIASES.get(key) or by_member.get(datamap.norm(key))
        return (n, *layout[n]) if n in layout else None

    # every key the corpus knows about, so output matches the authored shape.
    # Keys only a handful of mechs author (VehicleBattleValue: 1 of 64) are left
    # out rather than emitted at their default.
    corpus = datamap.corpus()
    common = collections.Counter(k for _c, kv, _b in corpus for k in kv)
    for key in sorted(common):
        if key in UNRECOVERABLE or common[key] * 2 < len(corpus):
            continue
        if key in CONSTANTS:
            out[key] = CONSTANTS[key]
            continue
        if key == "CraterName":
            name = CRATER_HASH.get(datamap.read(gm, *layout["m_craterID"]))
            if name:
                out[key] = name
            continue
        m = member(key)
        if key in SYMBOLIC and m:
            _n, off, typ, size = m
            raw = datamap.read(gm, off, typ, size)
            sym = constants.symbol(tables, key, raw, hint=ch)
            # MechID/NameIndex are authored as $(M_Chassis)/$(IDS_Chassis). A
            # foreign package may store an int from a different roster -- V4H's
            # are shifted by one against ours -- so --retarget-ids emits the
            # chassis's own symbol and lets the build resolve it.
            if retarget_ids and key in ("MechID", "NameIndex"):
                prefix = "M_" if key == "MechID" else "IDS_"
                if not sym or constants.norm(ch) not in constants.norm(sym):
                    out[key] = f"$({prefix}{ch[:1].upper() + ch[1:]})"
                    continue
            if sym:
                out[key] = sym
            continue
        if m:
            name, off, typ, size = m
            val = datamap.read(gm, off, typ, size, datamap.norm(name) in angles
                               or typ == "Radian")
            if typ == "ResourceID":
                path = manifest.get(val >> 16)
                if path:
                    out[key] = path
            elif typ in datamap.VECTORS:
                out[key] = " ".join(f"{v:g}" for v in val)
            elif typ in ("bool", "BYTE"):
                yes, no = words.get(key, ("true", "false"))
                out[key] = yes if val else no
            elif typ in datamap.FLOATS:
                out[key] = f"{val:g}"
            elif typ == "char":
                if val:
                    out[key] = val
            else:
                out[key] = str(val)

    # The OBB filenames are source-side names the package never stores. Prefer the
    # .obb files actually shipped next to the output so the keys cannot disagree
    # with them; fall back to the usual convention when none are present.
    solid = hier = None
    if obb_dir and os.path.isdir(obb_dir):
        for fn in sorted(os.listdir(obb_dir)):
            if not fn.lower().endswith(".obb"):
                continue
            if fn.lower().endswith("_solid.obb"):
                solid = fn
            else:
                hier = fn
    out["SolidOBB"] = solid or f"{ch}_Skeleton_SOLID.obb"
    out["HierarchicalOBB"] = hier or f"{ch}_Skeleton.obb"

    fs = record(mech_dir, ".data{FootSteps}")
    if fs:
        default, rows = foot_steps(fs)
        if default:
            out["DefaultFootStepTexture"] = default
        if rows:
            out["FootStepTexture"] = [f"{t},{m}" for t, m in rows]

    sh = record(mech_dir, ".data[shadow]")
    if sh is None:
        out.pop("Shadow", None)
    return out


def emit(kv):
    lines = ["[GameData]"]
    for key, val in kv.items():
        for v in (val if isinstance(val, list) else [val]):
            lines.append(f"{key}={v}")
    return "\r\n".join(lines) + "\r\n"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("mech_dir", nargs="?")
    ap.add_argument("-o", "--output")
    ap.add_argument("-m", "--manifest", default=MANIFEST,
                    help="package manifest for resolving ResourceIDs")
    ap.add_argument("--retarget-ids", action="store_true",
                    help="emit the chassis's own MechID/NameIndex symbol when the "
                         "stored int belongs to a different roster")
    ap.add_argument("--verify", action="store_true")
    args = ap.parse_args()
    if args.verify:
        here = os.path.dirname(os.path.abspath(__file__))
        raise SystemExit(subprocess.call([sys.executable,
                                          os.path.join(here, "verify_roundtrip.py")]))
    if not args.mech_dir:
        ap.error("mech_dir required")
    text = emit(decompile(args.mech_dir, subsystems.load_manifest(args.manifest),
                          retarget_ids=args.retarget_ids,
                          obb_dir=os.path.dirname(args.output) if args.output else None))
    if args.output:
        with open(args.output, "wb") as fh:
            fh.write(text.encode("latin-1"))
        print(f"wrote {args.output}")
    else:
        sys.stdout.write(text)


if __name__ == "__main__":
    main()