From bb0db51131f2c00d6aa08093ea4514ef98c15a57 Mon Sep 17 00:00:00 2001
From: yum <yum.food.vr@gmail.com>
Date: Wed, 19 Feb 2025 03:50:09 -0800
Subject: Append YOTS animator to original one

---
 Scripts/YOTSNDMFGenerator.cs | 72 ++++++++++++++++++++++++++++++++++++--------
 1 file changed, 59 insertions(+), 13 deletions(-)

(limited to 'Scripts/YOTSNDMFGenerator.cs')

diff --git a/Scripts/YOTSNDMFGenerator.cs b/Scripts/YOTSNDMFGenerator.cs
index e6e0585..864edf8 100644
--- a/Scripts/YOTSNDMFGenerator.cs
+++ b/Scripts/YOTSNDMFGenerator.cs
@@ -9,6 +9,8 @@ using VRC.SDK3.Avatars.ScriptableObjects;
 using UnityEditor;
 using System;
 using System.Collections.Generic;
+using System.Linq;
+using UnityEditor.Animations;
 
 [assembly: ExportsPlugin(typeof(YOTS.YOTSNDMFGenerator))]
 
@@ -39,7 +41,7 @@ namespace YOTS
         // Shoutsout anatawa12/AvatarOptimizer
         .BeforePass(RemoveEditorOnlyPass.Instance);
 
-      // Second pass: Generate animator
+      // Second pass: Generate animator and merge with the original
       InPhase(BuildPhase.Transforming)
         .Run("Generate YOTS Animator", ctx => {
           var state = ctx.GetState<YOTSBuildState>();
@@ -48,7 +50,7 @@ namespace YOTS
                 "No YOTS config found on the avatar.");
             return;
           }
-          // Get config
+
           var config = ctx.GetState<YOTSBuildState>();
           if (config == null) {
             ErrorReport.ReportError(localizer, ErrorSeverity.Error, "yots.error.no_config", 
@@ -60,19 +62,15 @@ namespace YOTS
                 "YOTS config component is missing JSON config file.");
             return;
           }
-          // Get descriptor
+
+          // Get menu and parameters
           var descriptor = ctx.AvatarDescriptor;
           if (descriptor == null) {
             ErrorReport.ReportError(localizer, ErrorSeverity.Error, "yots.error.no_descriptor", 
                 "Avatar descriptor is missing.");
             return;
           }
-          RuntimeAnimatorController animator = descriptor.baseAnimationLayers[4].animatorController;
-          if (animator == null) {
-            ErrorReport.ReportError(localizer, ErrorSeverity.Error, "yots.error.no_animator", 
-                "FX layer is missing.");
-            return;
-          }
+          RuntimeAnimatorController originalAnimator = descriptor.baseAnimationLayers[4].animatorController;
           var menu = descriptor.expressionsMenu;
           var parameters = descriptor.expressionParameters;
           if (parameters == null || menu == null)
@@ -81,10 +79,13 @@ namespace YOTS
                 "Avatar parameters or menu is missing.");
             return;
           }
+          // TODO do we need to make copies?
           menu = UnityEngine.Object.Instantiate(menu);
           parameters = UnityEngine.Object.Instantiate(parameters);
           descriptor.expressionsMenu = menu;
           descriptor.expressionParameters = parameters;
+
+          // Generate the YOTS animator.
           RuntimeAnimatorController generatedAnimator = YOTSCore.GenerateAnimator(
               state.jsonConfig,
               parameters,
@@ -92,13 +93,58 @@ namespace YOTS
           );
           if (generatedAnimator == null) {
               ErrorReport.ReportError(localizer, ErrorSeverity.Error, "yots.error.generation_failed", 
-                  "Failed to generate animator controller.");
+                  "Failed to generate animator.");
               return;
           }
-          // TODO merge animators
-          descriptor.baseAnimationLayers[4].animatorController = generatedAnimator;
+
+          if (originalAnimator == null)
+          {
+            descriptor.baseAnimationLayers[4].animatorController = generatedAnimator;
+            return;
+          }
+
+          AnimatorController originalController = originalAnimator as AnimatorController;
+          AnimatorController generatedController = generatedAnimator as AnimatorController;
+          MergeAnimatorControllers(localizer, originalController, generatedController);
+          descriptor.baseAnimationLayers[4].animatorController = originalController;
+        });
+    }
+
+    // Simply append generated params and layers to the original animator.
+    private static void MergeAnimatorControllers(Localizer localizer, AnimatorController original, AnimatorController generated)
+    {
+      // Merge parameters from generated into original.
+      foreach (var genParam in generated.parameters)
+      {
+        // This is an O(m*n) check but m and n should be small enough to not matter.
+        if (original.parameters.Any(p => p.name == genParam.name))
+        {
+          ErrorReport.ReportError(localizer, ErrorSeverity.Error, "yots.error.parameter_conflict",
+              $"Parameter '{genParam.name}' already exists in the original animator.");
+          return;
         }
-      );
+        original.AddParameter(genParam);
+      }
+
+      // Append each YOTS layer after the original layers.
+      foreach (var genLayer in generated.layers)
+      {
+        // This isn't strictly an error but if someone already has layers named
+        // YOTS_* that's probably not on purpose.
+        if (original.layers.Any(l => l.name == genLayer.name))
+        {
+          ErrorReport.ReportError(localizer, ErrorSeverity.Error, "yots.error.layer_conflict",
+              $"Layer with name '{genLayer.name}' already exists in the original animator.");
+          return;
+        }
+        var newLayer = new AnimatorControllerLayer
+        {
+          name = genLayer.name,
+          defaultWeight = genLayer.defaultWeight,
+          stateMachine = genLayer.stateMachine
+        };
+        original.AddLayer(newLayer);
+      }
     }
 
     private class YOTSBuildState
-- 
cgit v1.2.3