reimplement animator merging in yaml parser

Object IDs are allocated optimally now, but it's a bit slower due to
long parse times.

Also fix minor bug in generate_fx.py.

1 files changed, 294 insertions, 307 deletions

diff --git a/libunity.py b/libunity.py
index 5559c5f..097e8b8 100644
--- a/libunity.py
+++ b/libunity.py
@@ -1,190 +1,11 @@
 #!/usr/bin/env python3
 
+import copy
 import enum
+import sys
 # python3 -m pip install pyyaml
 import yaml
 
-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):
-    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")
-
-# Import a Unity YAML, modifying it to make it legal YAML 1.1.
-def importUnityYaml(fx_old, fx_new):
-    lines = []
-    with open(fx_old, "r") as f:
-        first_document = True
-        for line in f:
-            # Add end-of-document indicators.
-            if line.startswith("---"):
-                if not first_document:
-                    lines.append("...\n")
-                first_document = False
-
-            # Remove class ID tag from each block.
-            if line.startswith("---"):
-                parts = line.split()
-                lines.append(parts[0] + " " + parts[2] + "\n")
-                continue
-            lines.append(line)
-
-    lines.append("...\n")
-    with open(fx_new, "w") as f:
-        for line in lines:
-            f.write(f"{line}")
-
-def AnimatorControllerConstructor(loader, tag_suffix, node):
-    print("Got animator node: {}".format(node))
-    return None
-
 class Node:
     def __init__(self):
         # Optional. In Unity, this is the fileID of an object. Not all YAML
@@ -248,6 +69,10 @@ class Sequence(Node):
 
         return child
 
+    def forEach(self, cb):
+        for k in self.sequence:
+            cb(k)
+
 class Mapping(Node):
     def __init__(self):
         super().__init__()
@@ -307,33 +132,178 @@ class Mapping(Node):
 
         return child
 
+    def forEach(self, cb):
+        for k, v in self.mapping.items():
+            cb(v)
+
+def classId(anchor):
+    # AnimatorController
+    if anchor.startswith("91"):
+        return "91"
+    # MonoBehaviour
+    if anchor.startswith("114"):
+        return "114"
+    # BlendTree
+    if anchor.startswith("206"):
+        return "206"
+    # AnimatorStateTransition
+    if anchor.startswith("1101"):
+        return "1101"
+    # AnimatorState
+    if anchor.startswith("1102"):
+        return "1102"
+    # AnimatorStateMachine
+    if anchor.startswith("1107"):
+        return "1107"
+    # AnimatorTransition
+    if anchor.startswith("1109"):
+        return "1109"
+
+    # IDK what this is lmao
+    if anchor.startswith("74"):
+        return "74"
+    if anchor.startswith("115"):
+        return "115"
+
+    raise Exception("Unrecognized object: {}".format(anchor))
+
 class UnityDocument(Mapping):
     def __str__(self):
         return super().__str__()
 
     def classId(self):
-        # AnimatorController
-        if self.anchor.startswith("91"):
-            return "91"
-        # MonoBehaviour
-        if self.anchor.startswith("114"):
-            return "114"
-        # BlendTree
-        if self.anchor.startswith("206"):
-            return "206"
-        # AnimatorStateTransition
-        if self.anchor.startswith("1101"):
-            return "1101"
-        # AnimatorState
-        if self.anchor.startswith("1102"):
-            return "1102"
-        # AnimatorStateMachine
-        if self.anchor.startswith("1107"):
-            return "1107"
-        # AnimatorTransition
-        if self.anchor.startswith("1109"):
-            return "1109"
-        raise Exception("Unrecognized object: {}".format(self.anchor))
+        return classId(self.anchor)
+
+class UnityAnimator():
+    def __init__(self):
+        self.nodes = []
+        self.id_to_node = {}
+
+    def addNodes(self, nodes):
+        for node in nodes:
+            self.nodes.append(node)
+            if node.anchor == None:
+                raise Exception("Node is missing anchor: {}".format(str(node)))
+            if node.anchor in self.id_to_node:
+                raise Exception("Duplicate anchor: {}, node 1: {}, node 2: {}".format(node.anchor, str(node), str(self.id_to_node[node.anchor])))
+            self.id_to_node[node.anchor] = node
+
+    def getUniqueId(self, anchor):
+        if anchor in self.id_mapping.keys():
+            return self.id_mapping[anchor]
+
+        if classId(anchor) in self.class_to_next_id:
+            new_id = self.class_to_next_id[classId(anchor)]
+            self.class_to_next_id[classId(anchor)] += 1
+            self.id_mapping[anchor] = new_id
+        else:
+            new_id = int("%s%05d" % (classId(anchor), 0))
+            next_id = new_id + 1
+            self.class_to_next_id[classId(anchor)] = next_id
+            self.id_mapping[anchor] = new_id
+
+        #print("Map {} to {}".format(anchor, self.id_mapping[anchor]), file=sys.stderr)
+        return self.id_mapping[anchor]
+
+    def mergeIterator(self, v):
+        if hasattr(v, "mapping"):
+            # Don't relabel anything that's defined in an external file.
+            # TODO(yum) do this.
+            if 'fileID' in v.mapping and not 'guid' in v.mapping:
+                if v.mapping['fileID'] != '0':
+                    old_id = v.mapping['fileID']
+                    new_id = self.getUniqueId(old_id)
+                    v.mapping['fileID'] = str(new_id)
+        if hasattr(v, "forEach"):
+            v.forEach(self.mergeIterator)
+
+    def popNodeOfClass(self, classId):
+        result = None
+        for node in self.nodes:
+            if node.classId() == classId:
+                result = node
+                self.nodes.remove(result)
+                break
+        del self.id_to_node[result.anchor]
+        return result
+
+    def pushNode(self, node):
+        self.nodes.append(node)
+        self.id_to_node[node.anchor] = node
+
+    def mergeAnimatorControllers(self, ctrl0, ctrl1):
+        ctrl0 = copy.deepcopy(ctrl0)
+        ctrl1 = copy.deepcopy(ctrl1)
+
+        self.id_mapping0 = {}
+        self.id_mapping1 = {}
+
+        p0 = ctrl0.mapping['AnimatorController'].mapping['m_AnimatorParameters']
+        p1 = ctrl1.mapping['AnimatorController'].mapping['m_AnimatorParameters']
+
+        a0 = ctrl0.mapping['AnimatorController'].mapping['m_AnimatorLayers']
+        a1 = ctrl1.mapping['AnimatorController'].mapping['m_AnimatorLayers']
+
+        self.id_mapping = self.id_mapping0
+        p0.forEach(self.mergeIterator)
+        a0.forEach(self.mergeIterator)
+
+        # Hack to prevent ctrl1 from getting a new ID for the animator.
+        del self.class_to_next_id['91']
+
+        self.id_mapping = self.id_mapping1
+        p1.forEach(self.mergeIterator)
+        a1.forEach(self.mergeIterator)
+
+        p0.sequence += p1.sequence
+        a0.sequence += a1.sequence
+
+        return ctrl0
+
+    def merge(self, other):
+        # Mapping from class ID (string) to next ID (int)
+        self.class_to_next_id = {}
+
+        ctrl0 = self.popNodeOfClass('91')
+        ctrl1 = other.popNodeOfClass('91')
+        merged_anim = self.mergeAnimatorControllers(ctrl0, ctrl1)
+
+        # Mapping from class ID (string) to new class ID (int)
+        self.id_mapping = self.id_mapping0
+        for node in self.nodes:
+            new_id = self.getUniqueId(node.anchor)
+            node.anchor = str(new_id)
+            node.forEach(self.mergeIterator)
+
+        self.id_mapping = self.id_mapping1
+        for node in other.nodes:
+            new_id = self.getUniqueId(node.anchor)
+            node.anchor = str(new_id)
+            node.forEach(self.mergeIterator)
+
+        nodes = self.nodes
+        self.nodes = []
+        self.id_to_node = {}
+        self.pushNode(merged_anim)
+        self.addNodes(nodes)
+        self.addNodes(other.nodes)
+
+def unityAnimatorToString(nodes):
+    lines = []
+    preamble = """
+%YAML 1.1
+%TAG !u! tag:unity3d.com,2011:
+"""[1:][:-1]
+    lines.append(preamble)
+    for doc in nodes:
+        lines.append("--- !u!" + doc.classId() + " &" + doc.anchor)
+        lines.append(str(doc))
+    result = '\n'.join(lines)
+
+    for i in range(0,10):
+        result = result.replace("\n\n", "\n")
+
+    return result
 
 class UnityParser:
     STREAM_START = 100
@@ -347,137 +317,154 @@ class UnityParser:
 
     SEQUENCE_VALUE = 400
 
-    state = STREAM_START
-    cur_scalar = None
-    cur_node = None
+    def __init__(self):
+        self.state = self.STREAM_START
+        self.cur_scalar = None
+        self.cur_node = None
 
-    # Simple list of parsed documents. Populated by parse().
-    nodes = []
-    prev_states = []
+        # Simple list of parsed documents. Populated by parse().
+        self.nodes = []
+        self.prev_states = []
 
     def __str__(self):
-        lines = []
-        preamble = """
-%YAML 1.1
-%TAG !u! tag:unity3d.com,2011:
-"""[1:][:-1]
-        lines.append(preamble)
-        for doc in self.nodes:
-            lines.append("--- !u!" + doc.classId() + " &" + doc.anchor)
-            lines.append(str(doc))
-        result = '\n'.join(lines)
-
-        for i in range(0,10):
-            result = result.replace("\n\n", "\n")
-
-        return result
+        return unityAnimatorToString(self.nodes)
 
     def pushState(self, state):
         self.prev_states.append(self.state)
         self.state = state
         #print("state {} ({})".format(self.state, len(self.prev_states)))
 
-    def peekState(self):
-        return self.state
-
     def popState(self):
         self.state = self.prev_states[-1]
         self.prev_states = self.prev_states[0:len(self.prev_states) - 1]
         #print("state {} ({})".format(self.state, len(self.prev_states)))
         return self.state
 
+    def cleanYaml(self, yaml_str):
+        lines = []
+        first_document = True
+        for line in yaml_str.split("\n"):
+            # Add end-of-document indicators.
+            if line.startswith("---"):
+                if not first_document:
+                    lines.append("...\n")
+                first_document = False
+
+            # Remove class ID tag from each block.
+            if line.startswith("---"):
+                parts = line.split()
+                lines.append(parts[0] + " " + parts[2] + "\n")
+                continue
+            lines.append(line)
+
+        lines.append("...\n")
+        return '\n'.join(lines)
+
     def parse(self, yaml_file):
+        yaml_str = ""
         with open(yaml_file, "r") as f:
             yaml_str = f.read()
-            for event in yaml.parse(yaml_str):
-                #print("event {}".format(event))
-
-                if isinstance(event, yaml.StreamStartEvent):
-                    if len(self.prev_states) > 0:
-                        raise Exception("Multiple StreamStartEvents received")
-                    self.pushState(self.STREAM_START)
-                    continue
-
-                if isinstance(event, yaml.StreamEndEvent):
-                    if self.peekState() != self.STREAM_START:
-                        raise Exception("Document end received after state {}".format(self.peekState()))
-                    self.popState()
-                    if len(self.prev_states) > 0:
-                        raise Exception("Extra states at stream end")
-                    continue
-
-                if isinstance(event, yaml.DocumentStartEvent):
-                    if self.peekState() != self.STREAM_START and self.peekState() != self.DOCUMENT_END:
-                        raise Exception("Document start received after state {}".format(self.peekState()))
-                    self.pushState(self.DOCUMENT_START)
-                    continue
-
-                if isinstance(event, yaml.DocumentEndEvent):
-                    if self.peekState() != self.DOCUMENT_START:
-                        raise Exception("Document end received after state {}".format(self.peekState()))
-                    self.popState()
-                    self.nodes.append(self.cur_node)
-                    self.cur_node = None
-                    continue
-
-                if isinstance(event, yaml.MappingStartEvent):
-                    if self.cur_node == None:
-                        self.cur_node = UnityDocument()
-                        self.cur_node.anchor = event.anchor
-                    else:
-                        self.cur_node = self.cur_node.addChildMapping(self.cur_scalar)
-                    self.pushState(self.MAPPING_START)
-                    continue
-
-                if isinstance(event, yaml.MappingEndEvent):
-                    if self.peekState() != self.MAPPING_START:
-                        raise Exception("Mapping end received after state {}".format(self.peekState()))
+        yaml_str = self.cleanYaml(yaml_str)
+
+        for event in yaml.parse(yaml_str):
+            if isinstance(event, yaml.StreamStartEvent):
+                if len(self.prev_states) > 0:
+                    raise Exception("Multiple StreamStartEvents received")
+                self.pushState(self.STREAM_START)
+
+            elif isinstance(event, yaml.StreamEndEvent):
+                if self.state != self.STREAM_START:
+                    raise Exception("Document end received after state {}".format(self.state))
+                self.popState()
+                if len(self.prev_states) > 0:
+                    raise Exception("Extra states at stream end")
+
+            elif isinstance(event, yaml.DocumentStartEvent):
+                if self.state != self.STREAM_START and self.state != self.DOCUMENT_END:
+                    raise Exception("Document start received after state {}".format(self.state))
+                self.pushState(self.DOCUMENT_START)
+
+            elif isinstance(event, yaml.DocumentEndEvent):
+                if self.state != self.DOCUMENT_START:
+                    raise Exception("Document end received after state {}".format(self.state))
+                self.popState()
+                self.nodes.append(self.cur_node)
+                self.cur_node = None
+
+            elif isinstance(event, yaml.MappingStartEvent):
+                if self.cur_node == None:
+                    self.cur_node = UnityDocument()
+                    self.cur_node.anchor = event.anchor
+                else:
+                    self.cur_node = self.cur_node.addChildMapping(self.cur_scalar)
+                self.pushState(self.MAPPING_START)
+
+            elif isinstance(event, yaml.MappingEndEvent):
+                if self.state != self.MAPPING_START:
+                    raise Exception("Mapping end received after state {}".format(self.state))
+                self.popState()
+                if self.state == self.MAPPING_KEY:
                     self.popState()
-                    if self.peekState() == self.MAPPING_KEY:
-                        self.popState()
-                    if self.cur_node.parent != None:
-                        self.cur_node = self.cur_node.parent
-                    continue
-
-                if isinstance(event, yaml.SequenceStartEvent):
-                    self.cur_node = self.cur_node.addChildSequence(self.cur_scalar)
-                    self.pushState(self.SEQUENCE_VALUE)
-                    continue
-
-                if isinstance(event, yaml.SequenceEndEvent):
-                    if self.peekState() != self.SEQUENCE_VALUE:
-                        raise Exception("Sequence end received after state {}".format(self.peekState()))
-                    self.popState()
-                    if self.peekState() == self.MAPPING_KEY:
-                        self.popState()
+                if self.cur_node.parent != None:
                     self.cur_node = self.cur_node.parent
-                    continue
-
-                if isinstance(event, yaml.ScalarEvent):
-                    if self.peekState() == self.MAPPING_START:
-                        self.cur_scalar = event.value
-                        self.pushState(self.MAPPING_KEY)
-                    elif self.peekState() == self.MAPPING_KEY:
-                        self.cur_node.mapping[self.cur_scalar] = event.value
-                        self.popState()
-                    elif self.peekState() == self.SEQUENCE_VALUE:
-                        self.cur_node.sequence.append(event.value)
-                    else:
-                        raise Exception("Scalar event received after state {}".format(self.peekState()))
-                    continue
 
+            elif isinstance(event, yaml.SequenceStartEvent):
+                self.cur_node = self.cur_node.addChildSequence(self.cur_scalar)
+                self.pushState(self.SEQUENCE_VALUE)
+
+            elif isinstance(event, yaml.SequenceEndEvent):
+                if self.state != self.SEQUENCE_VALUE:
+                    raise Exception("Sequence end received after state {}".format(self.state))
+                self.popState()
+                if self.state == self.MAPPING_KEY:
+                    self.popState()
+                self.cur_node = self.cur_node.parent
+
+            elif isinstance(event, yaml.ScalarEvent):
+                if self.state == self.MAPPING_START:
+                    self.cur_scalar = event.value
+                    self.pushState(self.MAPPING_KEY)
+                elif self.state == self.MAPPING_KEY:
+                    self.cur_node.mapping[self.cur_scalar] = event.value
+                    self.popState()
+                elif self.state == self.SEQUENCE_VALUE:
+                    self.cur_node.sequence.append(event.value)
+                else:
+                    raise Exception("Scalar event received after state {}".format(self.state))
+            else:
                 raise Exception("Unhandled event {}".format(event))
+            continue
 
 if __name__ == "__main__":
-    arg0 = "tst.fx"
 
-    tmp0 = "fx0.controller"
-    importUnityYaml(arg0, tmp0)
+    arg0 = "../FX.controller"
+    print("Parsing {}".format(arg0), file=sys.stderr)
+    parser0 = UnityParser()
+    try:
+        parser0.parse(arg0)
+    except Exception as e:
+        print("exception: {}".format(e))
+
+    anim0 = UnityAnimator()
+    anim0.addNodes(parser0.nodes)
 
-    parser = UnityParser()
+    arg1 = "TaSTT_fx.controller"
+    print("Parsing {}".format(arg1), file=sys.stderr)
+    parser1 = UnityParser()
     try:
-        parser.parse(tmp0)
+        parser1.parse(arg1)
     except Exception as e:
         print("exception: {}".format(e))
-    print(parser)
+
+    anim1 = UnityAnimator()
+    anim1.addNodes(parser1.nodes)
+
+    #anim1.nodes = []
+    #anim1.id_to_node = {}
+    print("Merging animators", file=sys.stderr)
+    anim0.merge(anim1)
+
+    #print(unityAnimatorToString(parser0.nodes))
+    print("Serializing", file=sys.stderr)
+    print(unityAnimatorToString(anim0.nodes))