Add ability to merge FX controllers

TODO
* write default detection/correction
* real cmdline interface

1 files changed, 157 insertions, 1 deletions

diff --git a/libunity.py b/libunity.py
index 897518f..3470ea0 100644
--- a/libunity.py
+++ b/libunity.py
@@ -1,4 +1,160 @@
+#!/usr/bin/env python3
+
+def getObjectIds(fx):
+    fx_ids = set()
+    with open(fx, "r") as f:
+        for line in f:
+            parts = line.split()
+            if line.startswith("---") and len(parts) == 3:
+                obj_id = parts[2]
+                # Remove & from beginning of object ID.
+                obj_id = obj_id[1:]
+                fx_ids.add(obj_id)
+    return fx_ids
+
+def getOldIds(old_fx_ids, unified_fx_ids):
+    id_mapping = {}
+
+    for old_id in old_fx_ids:
+        id_mapping[old_id] = str(old_id)
+        unified_fx_ids.add(int(old_id))
+
+    return id_mapping
+
+# Map an object ID from old namespace to unified.
+def getNewIds(old_fx_ids, unified_fx_ids):
+    id_mapping = {}
+
+    for old_id in old_fx_ids:
+        new_id = int(old_id)
+
+        # 9100000 is the ID of the animator. There's only one, and we take care
+        # to merge it. So no need to create a new identifier for it.
+        if new_id != 9100000:
+            while (new_id in unified_fx_ids or new_id in old_fx_ids):
+                new_id += 1
+
+        unified_fx_ids.add(new_id)
+        id_mapping[old_id] = str(new_id)
+
+    return id_mapping
+
+def replaceIds(old_str, id_map):
+    result = old_str
+    lines = []
+    for line in result.split("\n"):
+        for old_id, new_id in id_map.items():
+            line = line.replace(old_id, new_id)
+        lines.append(line)
+    return '\n'.join(lines)
+
+def cat(path):
+    lines = []
+    with open(path, "r") as f:
+        for line in f:
+            lines.append(line)
+    return ''.join(lines)
+
+# Extract the first block whose '---' line matches `prefix`.
+def extractFirstBlock(fx, begin_prefix, end_prefix):
+    lines = []
+    in_block = False
+    for line in fx.split("\n"):
+        if in_block and line.startswith(end_prefix):
+            break
+        if line.startswith(begin_prefix):
+            in_block = True
+        if in_block:
+            lines.append(line)
+    return '\n'.join(lines) + "\n"
+
+def mergeAnimatorControllers(fx0, fx1, fx0_ids, fx1_ids):
+    fx0_anim = extractFirstBlock(cat(fx0), "--- !u!91", "---")
+    fx1_anim = extractFirstBlock(cat(fx1), "--- !u!91", "---")
+
+    fx0_anim_params = extractFirstBlock(fx0_anim, "  m_AnimatorParameters:", "  m_AnimatorLayers:")
+    fx0_anim_params = "\n".join(fx0_anim_params.split("\n")[1:])
+    fx0_anim_params = replaceIds(fx0_anim_params, fx0_ids)
+
+    fx0_anim_layers = extractFirstBlock(fx0_anim, "  m_AnimatorLayers:", "---")
+    fx0_anim_layers = "\n".join(fx0_anim_layers.split("\n")[1:])
+    fx0_anim_layers = replaceIds(fx0_anim_layers, fx0_ids)
+
+    fx1_anim_params = extractFirstBlock(fx1_anim, "  m_AnimatorParameters:", "  m_AnimatorLayers:")
+    fx1_anim_params = "\n".join(fx1_anim_params.split("\n")[1:])
+    fx1_anim_params = replaceIds(fx1_anim_params, fx1_ids)
+
+    fx1_anim_layers = extractFirstBlock(fx1_anim, "  m_AnimatorLayers:", "---")
+    fx1_anim_layers = "\n".join(fx1_anim_layers.split("\n")[1:])
+    fx1_anim_layers = replaceIds(fx1_anim_layers, fx1_ids)
+
+    result = """
+--- !u!91 &9100000
+AnimatorController:
+  m_ObjectHideFlags: 0
+  m_CorrespondingSourceObject: {fileID: 0}
+  m_PrefabInstance: {fileID: 0}
+  m_PrefabAsset: {fileID: 0}
+  m_Name: FX
+  serializedVersion: 5
+  m_AnimatorParameters:
+"""[1:]
+
+    result += fx0_anim_params
+    result += fx1_anim_params
+
+    result += "  m_AnimatorLayers:\n"
+
+    result += fx0_anim_layers
+    result += fx1_anim_layers
+
+    return result
+
 # Merge two FX layers.
 # fx0, fx1 are both paths to animators.
+# The object IDs for fx1 are reassigned to not collide with those used by fx0.
 def mergeFX(fx0, fx1):
-    pass
+    fx0_ids = getObjectIds(fx0)
+    fx1_ids = getObjectIds(fx1)
+
+    # Get unique identifiers for all objects in both layers.
+    new_ids = set()
+    fx0_id_mapping = getOldIds(fx0_ids, new_ids)
+    fx1_id_mapping = getNewIds(fx1_ids, new_ids)
+
+    # Merge animators
+    anim_ctrl = mergeAnimatorControllers(fx0, fx1, fx0_id_mapping, fx1_id_mapping)
+
+    # Remove animators and prefix
+    fx0_str = cat(fx0)
+    fx1_str = cat(fx1)
+
+    fx0_anim = extractFirstBlock(fx0_str, "--- !u!91", "---")
+    fx1_anim = extractFirstBlock(fx1_str, "--- !u!91", "---")
+
+    fx0_str = fx0_str.replace(fx0_anim, "")
+    fx1_str = fx1_str.replace(fx1_anim, "")
+
+    prefix = """
+%YAML 1.1
+%TAG !u! tag:unity3d.com,2011:
+"""[1:]
+
+    fx0_str = fx0_str.replace(prefix, "")
+    fx1_str = fx1_str.replace(prefix, "")
+
+    # Replace old IDs
+    #fx0_str = replaceIds(fx0_str, fx0_id_mapping)
+    fx1_str = replaceIds(fx1_str, fx1_id_mapping)
+
+    # Output
+    return deleteEmptyLines(prefix + anim_ctrl + fx0_str + fx1_str)
+
+def deleteEmptyLines(fx):
+    return fx.replace("\n\n", "\n")
+
+if __name__ == "__main__":
+    #mergeFX("FX.controller", "TaSTT/generated_fx.controller")
+    #print(extractFirstBlock("../FX.controller", "--- !u!91", "---"))
+    print(mergeFX("TaSTT_fx.controller", "../FX.controller"))
+