feat: improvements to armature tracking (#665)

* opti(armature-lock): parallelize burst jobs for armature lock processing

* feat: continue armature tracking when the MAMA GameObject is disabled

Closes: #500

* feat: add global toggle for armature locking

Closes: #484

Editor/AvatarProcessor.cs

@@ -24,6 +24,7 @@
 
 using System;
 using System.Runtime.CompilerServices;
+using nadena.dev.modular_avatar.ui;
 using UnityEditor;
 using UnityEngine;
 
@@ -45,7 +46,7 @@ namespace nadena.dev.modular_avatar.core.editor
             ApplyToCurrentAvatar();
         }
 
-        [MenuItem("Tools/Modular Avatar/Manual bake avatar", true)]
+        [MenuItem(UnityMenuItems.TopMenu_ManualBakeAvatar, true, UnityMenuItems.TopMenu_ManualBakeAvatarOrder)]
         private static bool ValidateApplyToCurrentAvatar()
         {
             return ndmf.AvatarProcessor.CanProcessObject(Selection.activeGameObject);

Editor/Localization/Localization.cs

@@ -3,6 +3,7 @@ using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.IO;
 using System.Linq;
+using nadena.dev.modular_avatar.ui;
 using nadena.dev.ndmf.localization;
 using nadena.dev.ndmf.ui;
 using Newtonsoft.Json;
@@ -87,7 +88,7 @@ namespace nadena.dev.modular_avatar.core.editor
             }
         }
 
-        [MenuItem("Tools/Modular Avatar/Reload localizations")]
+        [MenuItem(UnityMenuItems.TopMenu_ReloadLocalizations, false, UnityMenuItems.TopMenu_ReloadLocalizationsOrder)]
         public static void Reload()
         {
             Localizer.ReloadLocalizations();

Runtime/ArmatureAwase/ArmatureLockController.cs

@@ -1,13 +1,65 @@
 using System;
 using System.Collections.Generic;
+using nadena.dev.modular_avatar.ui;
 using UnityEngine;
-
 #if UNITY_EDITOR
 using UnityEditor;
 #endif
 
+
 namespace nadena.dev.modular_avatar.core.armature_lock
 {
+    internal class ArmatureLockConfig
+#if UNITY_EDITOR
+        : UnityEditor.ScriptableSingleton<ArmatureLockConfig>
+#endif
+    {
+#if !UNITY_EDITOR
+        internal static ArmatureLockConfig instance { get; } = new ArmatureLockConfig();
+#endif
+
+        [SerializeField]
+        private bool _globalEnable = true;
+        
+        internal bool GlobalEnable
+        {
+            get => _globalEnable;
+            set
+            {
+                if (value == _globalEnable) return;
+                
+                #if UNITY_EDITOR
+                Undo.RecordObject(this, "Toggle Edit Mode Bone Sync");
+                Menu.SetChecked(UnityMenuItems.TopMenu_EditModeBoneSync, value);
+                #endif
+                
+                _globalEnable = value;
+
+                if (!value)
+                {
+                    // Run prepare one last time to dispose of lock structures
+                    UpdateLoopController.InvokeArmatureLockPrepare();
+                }
+            }
+        }
+
+#if UNITY_EDITOR
+        [InitializeOnLoadMethod]
+        static void Init()
+        {
+            EditorApplication.delayCall += () => {
+                Menu.SetChecked(UnityMenuItems.TopMenu_EditModeBoneSync, instance._globalEnable);
+            };
+        }
+        
+        [MenuItem(UnityMenuItems.TopMenu_EditModeBoneSync, false, UnityMenuItems.TopMenu_EditModeBoneSyncOrder)]
+        static void ToggleBoneSync()
+        {
+            instance.GlobalEnable = !instance.GlobalEnable;
+        }
+#endif
+    }
+
     internal class ArmatureLockController : IDisposable
     {
         private static long lastMovedFrame = 0;
@@ -24,6 +76,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         private readonly GetTransformsDelegate _getTransforms;
         private IArmatureLock _lock;
 
+        private bool GlobalEnable => ArmatureLockConfig.instance.GlobalEnable;
         private bool _updateActive;
 
         private bool UpdateActive
@@ -35,11 +88,13 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 #if UNITY_EDITOR
                 if (value)
                 {
-                    UpdateLoopController.OnArmatureLockUpdate += VoidUpdate;
+                    UpdateLoopController.OnArmatureLockPrepare += UpdateLoopPrepare;
+                    UpdateLoopController.OnArmatureLockUpdate += UpdateLoopFinish;
                 }
                 else
                 {
-                    UpdateLoopController.OnArmatureLockUpdate -= VoidUpdate;
+                    UpdateLoopController.OnArmatureLockPrepare -= UpdateLoopPrepare;
+                    UpdateLoopController.OnArmatureLockUpdate -= UpdateLoopFinish;
                 }
 
                 _updateActive = value;
@@ -70,6 +125,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
             set
             {
                 if (Enabled == value) return;
+
                 _enabled = value;
                 if (_enabled) UpdateActive = true;
             }
@@ -105,24 +161,73 @@ namespace nadena.dev.modular_avatar.core.armature_lock
             return RebuildLock() && (_lock?.IsStable() ?? false);
         }
 
-        private void VoidUpdate()
+        private void VoidPrepare()
         {
-            Update();
+            UpdateLoopPrepare();
+        }
+        
+        private void UpdateLoopFinish()
+        {
+            DoFinish();
         }
 
         internal bool Update()
         {
-            LockResult result;
+            UpdateLoopPrepare();
+            return DoFinish();
+        }
+
+        private bool IsPrepared = false;
+        
+        private void UpdateLoopPrepare()
+        {
+            if (_mama == null || !_mama.gameObject.scene.IsValid())
+            {
+                UpdateActive = false;
+                return;
+            }
+            
             if (!Enabled)
             {
                 UpdateActive = false;
+                _lock?.Dispose();
+                _lock = null;
+                return;
+            }
+
+            if (!GlobalEnable)
+            {
+                _lock?.Dispose();
+                _lock = null;
+                return;
+            }
+            
+            if (_curMode == _mode)
+            {
+                _lock?.Prepare();
+                IsPrepared = _lock != null;
+            }
+        }
+
+        private bool DoFinish()
+        {
+            LockResult result;
+            
+            if (!GlobalEnable)
+            {
                 _lock?.Dispose();
                 _lock = null;
                 return true;
             }
 
+            var wasPrepared = IsPrepared;
+            IsPrepared = false;
+            
+            if (!Enabled) return true;
+            
             if (_curMode == _mode)
             {
+                if (!wasPrepared) _lock?.Prepare();
                 result = _lock?.Execute() ?? LockResult.Failed;
                 if (result == LockResult.Success)
                 {
@@ -134,6 +239,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
             if (!RebuildLock()) return false;
 
+            _lock?.Prepare();
             result = (_lock?.Execute() ?? LockResult.Failed);
 
             return result != LockResult.Failed;

Runtime/ArmatureAwase/BidirectionalArmatureLock.cs

@@ -20,6 +20,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         private NativeIntPtr WroteAny;
 
         private JobHandle LastOp;
+        private JobHandle LastPrepare;
 
         public BidirectionalArmatureLock(IReadOnlyList<(Transform, Transform)> bones)
         {
@@ -142,21 +143,20 @@ namespace nadena.dev.modular_avatar.core.armature_lock
             _disposed = true;
         }
 
-        private bool DoCompute(out JobHandle handle)
+        public void Prepare()
         {
-            handle = default;
-
-            if (_disposed) return false;
-
-            WroteAny.Value = 0;
-
+            if (_disposed) return;
+            
             LastOp.Complete();
+            
+            WroteAny.Value = 0;
 
             var readBase = new ReadBone()
             {
                 _state = BaseBones,
             }.Schedule(_baseBoneAccess);
-            LastOp = handle = new Compute()
+
+            LastOp = LastPrepare = new Compute()
             {
                 BaseBones = BaseBones,
                 MergeBones = MergeBones,
@@ -165,7 +165,12 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                 ShouldWriteMerge = ShouldWriteMerge,
                 WroteAny = WroteAny.GetParallel(),
             }.Schedule(_mergeBoneAccess, readBase);
+        }
 
+        private bool CheckConsistency()
+        {
+            if (_disposed) return false;
+            
             // Check parents haven't changed
             for (int i = 0; i < _baseBones.Length; i++)
             {
@@ -186,27 +191,27 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
         public bool IsStable()
         {
-            if (!DoCompute(out var compute)) return false;
-
-            compute.Complete();
-
+            Prepare();
+            if (!CheckConsistency()) return false;
+            LastPrepare.Complete();
+            
             return WroteAny.Value == 0;
         }
 
         public LockResult Execute()
         {
-            if (!DoCompute(out var compute)) return LockResult.Failed;
+            if (!CheckConsistency()) return LockResult.Failed;
 
             var commitBase = new Commit()
             {
                 BoneState = BaseBones,
                 ShouldWrite = ShouldWriteBase,
-            }.Schedule(_baseBoneAccess, compute);
+            }.Schedule(_baseBoneAccess, LastPrepare);
             var commitMerge = new Commit()
             {
                 BoneState = MergeBones,
                 ShouldWrite = ShouldWriteMerge,
-            }.Schedule(_mergeBoneAccess, compute);
+            }.Schedule(_mergeBoneAccess, LastPrepare);
 
             commitBase.Complete();
             commitMerge.Complete();

Runtime/ArmatureAwase/IArmatureLock.cs

@@ -4,6 +4,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 {
     internal interface IArmatureLock : IDisposable
     {
+        void Prepare();
         LockResult Execute();
         bool IsStable();
     }

Runtime/ArmatureAwase/OnewayArmatureLock.cs

@@ -4,11 +4,11 @@ using nadena.dev.modular_avatar.JacksonDunstan.NativeCollections;
 using Unity.Burst;
 using Unity.Collections;
 using Unity.Jobs;
+using UnityEngine;
+using UnityEngine.Jobs;
 #if UNITY_EDITOR
 using UnityEditor;
 #endif
-using UnityEngine;
-using UnityEngine.Jobs;
 
 namespace nadena.dev.modular_avatar.core.armature_lock
 {
@@ -29,7 +29,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         private TransformAccessArray _baseBonesAccessor, _mergeBonesAccessor;
 
         private bool _disposed;
-        private JobHandle LastOp;
+        private JobHandle LastOp, LastPrepare;
 
         [BurstCompile]
         struct WriteBone : IJobParallelForTransform
@@ -197,13 +197,13 @@ namespace nadena.dev.modular_avatar.core.armature_lock
             return (scale.x < epsilon || scale.y < epsilon || scale.z < epsilon);
         }
 
-        private bool DoCompute(out JobHandle handle)
+        public void Prepare()
         {
-            handle = default;
-            if (_disposed) return false;
-
+            if (_disposed) return;
+            
             LastOp.Complete();
 
+            
             _fault.Value = 0;
             _wroteAny.Value = 0;
 
@@ -216,7 +216,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                 _state = _mergeState
             }.Schedule(_mergeBonesAccessor);
             var readAll = JobHandle.CombineDependencies(jobReadBase, jobReadMerged);
-            LastOp = handle = new ComputePosition
+            LastOp = LastPrepare = new ComputePosition
             {
                 _boneStatic = _boneStaticData,
                 _mergeState = _mergeState,
@@ -225,6 +225,11 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                 _fault = _fault.GetParallel(),
                 _wroteAny = _wroteAny.GetParallel(),
             }.Schedule(_baseBones.Length, 32, readAll);
+        }
+
+        private bool CheckConsistency()
+        {
+            if (_disposed) return false;
 
             // Validate parents while that job is running
             for (int i = 0; i < _baseBones.Length; i++)
@@ -246,9 +251,10 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
         public bool IsStable()
         {
-            if (!DoCompute(out var compute)) return false;
+            Prepare();
+            if (!CheckConsistency()) return false;
 
-            compute.Complete();
+            LastPrepare.Complete();
 
             return _fault.Value == 0 && _wroteAny.Value == 0;
         }
@@ -259,14 +265,14 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         /// <returns>True if successful, false if cached data was invalidated and needs recreating</returns>
         public LockResult Execute()
         {
-            if (!DoCompute(out var compute)) return LockResult.Failed;
-
+            if (!CheckConsistency()) return LockResult.Failed;
+            
             var commit = new WriteBone()
             {
                 _fault = _fault,
                 _values = _mergeSavedState,
                 _shouldWrite = _wroteAny
-            }.Schedule(_mergeBonesAccessor, compute);
+            }.Schedule(_mergeBonesAccessor, LastPrepare);
 
             commit.Complete();
 

Runtime/ArmatureAwase/UpdateLoopController.cs

@@ -4,6 +4,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 {
     internal static class UpdateLoopController
     {
+        internal static event Action OnArmatureLockPrepare;
         internal static event Action OnArmatureLockUpdate;
         internal static event Action OnMoveIndependentlyUpdate;
 
@@ -11,14 +12,22 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         [UnityEditor.InitializeOnLoadMethod]
         private static void Init()
         {
-#if UNITY_EDITOR
             UnityEditor.EditorApplication.update += () =>
             {
-                OnArmatureLockUpdate?.Invoke();
+                if (ArmatureLockConfig.instance.GlobalEnable)
+                {
+                    OnArmatureLockPrepare?.Invoke();
+                    OnArmatureLockUpdate?.Invoke();
+                }
+
                 OnMoveIndependentlyUpdate?.Invoke();
             };
-#endif
         }
 #endif
+
+        internal static void InvokeArmatureLockPrepare()
+        {
+            OnArmatureLockPrepare?.Invoke();
+        }
     }
 }

Runtime/ModularAvatarMergeArmature.cs

@@ -74,6 +74,7 @@ namespace nadena.dev.modular_avatar.core
             base.OnValidate();
             MigrateLockConfig();
             RuntimeUtil.delayCall(SetLockMode);
+            Debug.Log("$$$ OnValidate");
         }
 
         internal void ResetArmatureLock()
@@ -106,7 +107,7 @@ namespace nadena.dev.modular_avatar.core
                 }
             }
 
-            _lockController.Enabled = isActiveAndEnabled;
+            _lockController.Enabled = enabled;
         }
 
         private void MigrateLockConfig()
@@ -126,7 +127,8 @@ namespace nadena.dev.modular_avatar.core
 
         private void OnDisable()
         {
-            _lockController.Enabled = false;
+            // we use enabled instead of activeAndEnabled to ensure we track even when the GameObject is disabled
+            _lockController.Enabled = enabled;
         }
 
         protected override void OnDestroy()

Runtime/UI.meta

@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: e36cbbfd8a704421becea83cb2d1e72f
+timeCreated: 1707812798

Runtime/UI/UnityMenuItems.cs

@@ -0,0 +1,14 @@
+namespace nadena.dev.modular_avatar.ui
+{
+    internal class UnityMenuItems
+    {
+        internal const string TopMenu_EditModeBoneSync = "Tools/Modular Avatar/Sync Bones in Edit Mode";
+        internal const int TopMenu_EditModeBoneSyncOrder = 100;
+        
+        internal const string TopMenu_ManualBakeAvatar = "Tools/Modular Avatar/Manual Bake Avatar";
+        internal const int TopMenu_ManualBakeAvatarOrder = 1000;
+        
+        internal const string TopMenu_ReloadLocalizations = "Tools/Modular Avatar/Reload Localizations";
+        internal const int TopMenu_ReloadLocalizationsOrder = 1001;
+    }
+}

Runtime/UI/UnityMenuItems.cs.meta

@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: c5b3fcacdeb74934b0fb8554490d5d9c
+timeCreated: 1707812807