opti: perf improvements for armature lock (#714)

* opti: perf improvements for armature lock

* chore: unity 2019 compatibility

* chore: update comments

Runtime/ArmatureAwase/ArmatureLock.cs

@@ -1,41 +0,0 @@
-using System;
-
-namespace nadena.dev.modular_avatar.core.armature_lock
-{
-    internal abstract class ArmatureLock : IDisposable
-    {
-        private bool _enableAssemblyReloadCallback;
-
-        protected bool EnableAssemblyReloadCallback
-        {
-            get => _enableAssemblyReloadCallback;
-            set
-            {
-                if (_enableAssemblyReloadCallback == value) return;
-                _enableAssemblyReloadCallback = value;
-#if UNITY_EDITOR
-                if (value)
-                {
-                    UnityEditor.AssemblyReloadEvents.beforeAssemblyReload += OnDomainUnload;
-                }
-                else
-                {
-                    UnityEditor.AssemblyReloadEvents.beforeAssemblyReload -= OnDomainUnload;
-                }
-#endif
-            }
-        }
-
-        public abstract void Prepare();
-        public abstract LockResult Execute();
-        public abstract bool IsStable();
-        public abstract void Dispose();
-
-        private void OnDomainUnload()
-        {
-            // Unity 2019 does not call deferred callbacks before domain unload completes,
-            // so we need to make sure to immediately destroy all our TransformAccessArrays.
-            DeferDestroy.DestroyImmediate(this);
-        }
-    }
-}

Runtime/ArmatureAwase/ArmatureLock.cs.meta

@@ -1,3 +0,0 @@
-fileFormatVersion: 2
-guid: 7b4b88c94c2144128ffbe7f271b28ba2
-timeCreated: 1693712261

Runtime/ArmatureAwase/ArmatureLockController.cs

@@ -1,4 +1,6 @@
-using System;
+#region
+
+using System;
 using System.Collections.Generic;
 using nadena.dev.modular_avatar.ui;
 using UnityEngine;
@@ -6,12 +8,14 @@ using UnityEngine;
 using UnityEditor;
 #endif
 
+#endregion
+
 
 namespace nadena.dev.modular_avatar.core.armature_lock
 {
     internal class ArmatureLockConfig
 #if UNITY_EDITOR
-        : UnityEditor.ScriptableSingleton<ArmatureLockConfig>
+        : ScriptableSingleton<ArmatureLockConfig>
 #endif
     {
 #if !UNITY_EDITOR
@@ -19,6 +23,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 #endif
 
         [SerializeField] private bool _globalEnable = true;
+        internal event Action OnGlobalEnableChange;
 
         internal bool GlobalEnable
         {
@@ -34,11 +39,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
                 _globalEnable = value;
 
-                if (!value)
+                OnGlobalEnableChange?.Invoke();
-                {
-                    // Run prepare one last time to dispose of lock structures
-                    UpdateLoopController.InvokeArmatureLockPrepare();
-                }
             }
         }
 
@@ -60,65 +61,54 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 #endif
     }
 
+
     internal class ArmatureLockController : IDisposable
     {
+        public delegate IReadOnlyList<(Transform, Transform)> GetTransformsDelegate();
+
         private static long lastMovedFrame = 0;
-        public static bool MovedThisFrame => Time.frameCount == lastMovedFrame;
 
         // Undo operations can reinitialize the MAMA component, which destroys critical lock controller state.
         // Avoid this issue by keeping a static reference to the controller for each MAMA component.
         private static Dictionary<ModularAvatarMergeArmature, ArmatureLockController>
             _controllers = new Dictionary<ModularAvatarMergeArmature, ArmatureLockController>();
 
-        public delegate IReadOnlyList<(Transform, Transform)> GetTransformsDelegate();
+        private readonly GetTransformsDelegate _getTransforms;
 
         private readonly ModularAvatarMergeArmature _mama;
-        private readonly GetTransformsDelegate _getTransforms;
-        private ArmatureLock _lock;
-
-        private bool GlobalEnable => ArmatureLockConfig.instance.GlobalEnable;
-        private bool _updateActive;
-
-        private bool UpdateActive
-        {
-            get => _updateActive;
-            set
-            {
-                if (UpdateActive == value) return;
-#if UNITY_EDITOR
-                if (value)
-                {
-                    UpdateLoopController.OnArmatureLockPrepare += UpdateLoopPrepare;
-                    UpdateLoopController.OnArmatureLockUpdate += UpdateLoopFinish;
-                }
-                else
-                {
-                    UpdateLoopController.OnArmatureLockPrepare -= UpdateLoopPrepare;
-                    UpdateLoopController.OnArmatureLockUpdate -= UpdateLoopFinish;
-                }
-
-                _updateActive = value;
-#endif
-            }
-        }
 
         private ArmatureLockMode _curMode, _mode;
 
+        private bool _enabled;
+        private ArmatureLockJob _job;
+
+        public ArmatureLockController(ModularAvatarMergeArmature mama, GetTransformsDelegate getTransforms)
+        {
+#if UNITY_EDITOR
+            AssemblyReloadEvents.beforeAssemblyReload += OnDomainUnload;
+#endif
+
+            _mama = mama;
+            _getTransforms = getTransforms;
+        }
+
+        public static bool MovedThisFrame => Time.frameCount == lastMovedFrame;
+
+        private bool GlobalEnable => ArmatureLockConfig.instance.GlobalEnable;
+
         public ArmatureLockMode Mode
         {
             get => _mode;
             set
             {
-                if (value == _mode) return;
+                if (value == _mode && _job != null) return;
 
                 _mode = value;
 
-                UpdateActive = true;
+                RebuildLock();
             }
         }
 
-        private bool _enabled;
-
         public bool Enabled
         {
             get => _enabled;
@@ -127,20 +117,25 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                 if (Enabled == value) return;
 
                 _enabled = value;
-                if (_enabled) UpdateActive = true;
+
+                RebuildLock();
             }
         }
 
-        public ArmatureLockController(ModularAvatarMergeArmature mama, GetTransformsDelegate getTransforms)
+        public void Dispose()
         {
+            _job?.Dispose();
+            _job = null;
+
 #if UNITY_EDITOR
-            AssemblyReloadEvents.beforeAssemblyReload += OnDomainUnload;
+            AssemblyReloadEvents.beforeAssemblyReload -= OnDomainUnload;
 #endif
 
-            this._mama = mama;
+            _controllers.Remove(_mama);
-            this._getTransforms = getTransforms;
         }
 
+        internal event Action WhenUnstable;
+
         public static ArmatureLockController ForMerge(ModularAvatarMergeArmature mama,
             GetTransformsDelegate getTransforms)
         {
@@ -153,102 +148,32 @@ namespace nadena.dev.modular_avatar.core.armature_lock
             return controller;
         }
 
-        public bool IsStable()
+        internal void CheckLockJob()
         {
-            if (Mode == ArmatureLockMode.NotLocked) return true;
+            if (_mama == null || !_mama.gameObject.scene.IsValid() || !Enabled)
-
-            if (_curMode == _mode && _lock?.IsStable() == true) return true;
-            return RebuildLock() && (_lock?.IsStable() ?? false);
-        }
-
-        private void VoidPrepare()
-        {
-            UpdateLoopPrepare();
-        }
-
-        private void UpdateLoopFinish()
-        {
-            DoFinish();
-        }
-
-        internal bool Update()
-        {
-            UpdateLoopPrepare();
-            return DoFinish();
-        }
-
-        private bool IsPrepared = false;
-
-        private void UpdateLoopPrepare()
-        {
-            if (_mama == null || !_mama.gameObject.scene.IsValid())
             {
-                UpdateActive = false;
+                _job?.Dispose();
                 return;
             }
 
-            if (!Enabled)
+            if (_curMode != _mode || _job == null || !_job.IsValid)
             {
-                UpdateActive = false;
+                if (_job != null && _job.FailedOnStartup)
-                _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)
                 {
-                    lastMovedFrame = Time.frameCount;
+                    WhenUnstable?.Invoke();
+                    Enabled = false;
+                    _job?.Dispose();
+                    return;
                 }
 
-                if (result != LockResult.Failed) return true;
+                RebuildLock();
             }
-
-            if (!RebuildLock()) return false;
-
-            _lock?.Prepare();
-            result = (_lock?.Execute() ?? LockResult.Failed);
-
-            return result != LockResult.Failed;
         }
 
         private bool RebuildLock()
         {
-            _lock?.Dispose();
+            _job?.Dispose();
-            _lock = null;
+            _job = null;
 
             var xforms = _getTransforms();
             if (xforms == null)
@@ -261,40 +186,34 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                 switch (Mode)
                 {
                     case ArmatureLockMode.BidirectionalExact:
-                        _lock = new BidirectionalArmatureLock(_getTransforms());
+                        _job = BidirectionalArmatureLockOperator.Instance.RegisterLock(xforms);
                         break;
                     case ArmatureLockMode.BaseToMerge:
-                        _lock = new OnewayArmatureLock(_getTransforms());
+                        _job = OnewayArmatureLockOperator.Instance.RegisterLock(xforms);
                         break;
                     default:
-                        UpdateActive = false;
+                        Enabled = false;
                         break;
                 }
             }
             catch (Exception)
             {
-                _lock = null;
+                _job = null;
                 return false;
             }
 
+            if (_job != null)
+            {
+#if UNITY_EDITOR
+                _job.OnInvalidation += () => { EditorApplication.delayCall += CheckLockJob; };
+#endif
+            }
+
             _curMode = _mode;
 
             return true;
         }
 
-        public void Dispose()
-        {
-            _lock?.Dispose();
-            _lock = null;
-
-            #if UNITY_EDITOR
-            AssemblyReloadEvents.beforeAssemblyReload -= OnDomainUnload;
-            #endif
-
-            _controllers.Remove(_mama);
-            UpdateActive = false;
-        }
-
         private void OnDomainUnload()
         {
             // Unity 2019 does not call deferred callbacks before domain unload completes,

Runtime/ArmatureAwase/ArmatureLockJob.cs

@@ -0,0 +1,100 @@
+#region
+
+using System;
+using System.Collections.Immutable;
+using System.Linq;
+using UnityEditor;
+using UnityEngine;
+
+#endregion
+
+namespace nadena.dev.modular_avatar.core.armature_lock
+{
+    internal sealed class ArmatureLockJob : IDisposable
+    {
+        private bool _didLoop = false;
+
+        private Action _dispose;
+
+        private bool _isValid = true;
+        private long _lastHierarchyCheck = -1;
+        private Action _update;
+
+        internal ImmutableList<(Transform, Transform)> RecordedParents;
+        internal ImmutableList<(Transform, Transform)> Transforms;
+
+        internal ArmatureLockJob(ImmutableList<(Transform, Transform)> transforms, Action dispose, Action update)
+        {
+            Transforms = transforms;
+            RecordedParents = transforms.Select(((tuple, _) => (tuple.Item1.parent, tuple.Item2.parent)))
+                .ToImmutableList();
+            _dispose = dispose;
+            _update = update;
+        }
+
+        internal bool FailedOnStartup => !_isValid && !_didLoop;
+
+        internal bool HierarchyChanged
+        {
+            get
+            {
+                var unchanged = RecordedParents.Zip(Transforms,
+                    (p, t) =>
+                    {
+                        return t.Item1 != null && t.Item2 != null && t.Item1.parent == p.Item1 &&
+                               t.Item2.parent == p.Item2;
+                    }).All(b => b);
+
+                return !unchanged;
+            }
+        }
+
+        internal bool IsValid
+        {
+            get => _isValid;
+            set
+            {
+                var transitioned = (_isValid && !value);
+                _isValid = value;
+
+                if (transitioned)
+                {
+                    Debug.Log("Invalidated job!");
+#if UNITY_EDITOR
+                    EditorApplication.delayCall += () => OnInvalidation?.Invoke();
+#endif
+                }
+            }
+        }
+
+        internal bool WroteAny { get; set; }
+
+        public void Dispose()
+        {
+            _dispose?.Invoke();
+            _dispose = null;
+            _update = null;
+        }
+
+        internal event Action OnInvalidation;
+
+        internal void MarkLoop()
+        {
+            _didLoop = _didLoop || _isValid;
+        }
+
+        internal bool BoneChanged(int boneIndex)
+        {
+            return Transforms[boneIndex].Item1 == null || Transforms[boneIndex].Item2 == null
+                                                       || Transforms[boneIndex].Item1.parent !=
+                                                       RecordedParents[boneIndex].Item1
+                                                       || Transforms[boneIndex].Item2.parent !=
+                                                       RecordedParents[boneIndex].Item2;
+        }
+
+        public void UpdateNow()
+        {
+            _update?.Invoke();
+        }
+    }
+}

Runtime/ArmatureAwase/ArmatureLockJob.cs.meta

@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: 38a41bc7e55d4f7c8efeafd6107487da
+timeCreated: 1709207112

Runtime/ArmatureAwase/ArmatureLockJobAccessor.cs

@@ -0,0 +1,96 @@
+#region
+
+using Unity.Collections;
+using Unity.Collections.LowLevel.Unsafe;
+using UnityEngine.Serialization;
+
+#endregion
+
+namespace nadena.dev.modular_avatar.core.armature_lock
+{
+    /// <summary>
+    ///  Abstractly, an armature lock job works by taking the local transforms of the base armature and target armature,
+    ///  deciding whether to abort updates, and if not, what the transforms should be set to, and writing out the
+    ///  results.
+    ///
+    ///  This struct handles these common inputs and outputs for different armature lock types.
+    /// </summary>
+    internal struct ArmatureLockJobAccessor
+    {
+        internal void Allocate(int nBones, int nWords)
+        {
+            _in_baseBone = new NativeArray<TransformState>(nBones, Allocator.Persistent);
+            _in_targetBone = new NativeArray<TransformState>(nBones, Allocator.Persistent);
+            _out_baseBone = new NativeArray<TransformState>(nBones, Allocator.Persistent);
+            _out_targetBone = new NativeArray<TransformState>(nBones, Allocator.Persistent);
+            _out_dirty_baseBone = new NativeArray<int>(nBones, Allocator.Persistent);
+            _out_dirty_targetBone = new NativeArray<int>(nBones, Allocator.Persistent);
+            _boneToJobIndex = new NativeArray<int>(nBones, Allocator.Persistent);
+            _abortFlag = new NativeArray<int>(nWords, Allocator.Persistent);
+            _didAnyWriteFlag = new NativeArray<int>(nWords, Allocator.Persistent);
+        }
+
+        internal void Destroy()
+        {
+            if (_in_baseBone.IsCreated) _in_baseBone.Dispose();
+            _in_baseBone = default;
+            if (_in_targetBone.IsCreated) _in_targetBone.Dispose();
+            _in_targetBone = default;
+            if (_out_baseBone.IsCreated) _out_baseBone.Dispose();
+            _out_baseBone = default;
+            if (_out_targetBone.IsCreated) _out_targetBone.Dispose();
+            _out_targetBone = default;
+            if (_out_dirty_baseBone.IsCreated) _out_dirty_baseBone.Dispose();
+            _out_dirty_baseBone = default;
+            if (_out_dirty_targetBone.IsCreated) _out_dirty_targetBone.Dispose();
+            _out_dirty_targetBone = default;
+            if (_boneToJobIndex.IsCreated) _boneToJobIndex.Dispose();
+            _boneToJobIndex = default;
+            if (_abortFlag.IsCreated) _abortFlag.Dispose();
+            _abortFlag = default;
+            if (_didAnyWriteFlag.IsCreated) _didAnyWriteFlag.Dispose();
+            _didAnyWriteFlag = default;
+        }
+
+        /// <summary>
+        /// Initial transform states
+        /// </summary>
+        public NativeArray<TransformState> _in_baseBone, _in_targetBone;
+
+        /// <summary>
+        /// Transform states to write out (if _out_dirty is set)
+        /// </summary>
+        public NativeArray<TransformState> _out_baseBone, _out_targetBone;
+        
+        /// <summary>
+        /// Flags indicating whether the given bone should be written back to its transform
+        /// </summary>
+        public NativeArray<int> _out_dirty_baseBone, _out_dirty_targetBone;
+
+        /// <summary>
+        /// Indexed by the job index (via _boneToJobIndex). If set to a nonzero value, none of the bones in this
+        /// particular job (e.g. a single MergeArmature component) will be committed.
+        ///
+        /// Note: This array is written simultaneously from multiple threads. Jobs may set this to 1, but otherwise
+        /// shouldn't read this value.
+        /// </summary>
+        [NativeDisableContainerSafetyRestriction] [NativeDisableParallelForRestriction]
+        public NativeArray<int> _abortFlag;
+
+        /// <summary>
+        /// Indexed by the job index (via _boneToJobIndex). Should be set to a nonzero value when any bone in the job
+        /// has changes that need to be written out.
+        ///
+        /// Note: This array is written simultaneously from multiple threads. Jobs may set this to 1, but otherwise
+        /// shouldn't read this value.
+        /// </summary>
+        [NativeDisableContainerSafetyRestriction] [NativeDisableParallelForRestriction]
+        public NativeArray<int> _didAnyWriteFlag;
+
+        /// <summary>
+        /// Maps from bone index to job index.
+        /// </summary>
+        [FormerlySerializedAs("_statusWordIndex")]
+        public NativeArray<int> _boneToJobIndex;
+    }
+}

Runtime/ArmatureAwase/ArmatureLockJobAccessor.cs.meta

@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: 73bef51200bd478c9e22761598e22d16
+timeCreated: 1709116462

Runtime/ArmatureAwase/ArmatureLockOperator.cs

@@ -0,0 +1,445 @@
+#region
+
+using System;
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Linq;
+using Unity.Burst;
+using Unity.Collections;
+using Unity.Collections.LowLevel.Unsafe;
+using Unity.Jobs;
+using UnityEditor;
+using UnityEngine;
+using UnityEngine.Jobs;
+using UnityEngine.Profiling;
+
+#endregion
+
+namespace nadena.dev.modular_avatar.core.armature_lock
+{
+    internal abstract class ArmatureLockOperator<T> : IDisposable where T : ArmatureLockOperator<T>, new()
+    {
+        internal static readonly T Instance = new T();
+
+        private static long LastHierarchyChange = 0;
+        private ArmatureLockJobAccessor _accessor;
+
+        private TransformAccessArray _baseBones, _targetBones;
+
+        private int _commitFilter;
+
+        private bool _isDisposed = false;
+        private bool _isInit = false, _isValid = false;
+
+        private ImmutableList<ArmatureLockJob> _jobs = ImmutableList<ArmatureLockJob>.Empty;
+        private JobHandle _lastJob;
+        private List<ArmatureLockJob> _requestedJobs = new List<ArmatureLockJob>();
+        private long LastCheckedHierarchy = -1;
+
+        static ArmatureLockOperator()
+        {
+            Instance = new T();
+#if UNITY_EDITOR
+            EditorApplication.delayCall += StaticInit;
+#endif
+        }
+
+        protected ArmatureLockOperator()
+        {
+#if UNITY_EDITOR
+            AssemblyReloadEvents.beforeAssemblyReload += () => DeferDestroy.DestroyImmediate(this);
+#endif
+        }
+
+        protected abstract bool WritesBaseBones { get; }
+
+        public void Dispose()
+        {
+            if (_isDisposed) return;
+            _isDisposed = true;
+
+            if (!_isInit) return;
+
+            _lastJob.Complete();
+            DeferDestroy.DeferDestroyObj(_baseBones);
+            DeferDestroy.DeferDestroyObj(_targetBones);
+            DerivedDispose();
+            _accessor.Destroy();
+        }
+
+#if UNITY_EDITOR
+        protected static void StaticInit()
+        {
+            EditorApplication.hierarchyChanged += () => { LastHierarchyChange += 1; };
+            UpdateLoopController.UpdateCallbacks += Instance.Update;
+            ArmatureLockConfig.instance.OnGlobalEnableChange += Instance.Invalidate;
+
+            EditorApplication.playModeStateChanged += (change) =>
+            {
+                // If we allow ourselves to simply enter play mode without a final update, any movement applied by
+                // automatically leaving animation preview mode won't be applied, leaving any outfits in the wrong pose. 
+                if (change == PlayModeStateChange.ExitingEditMode)
+                {
+                    Instance.Update();
+                }
+            };
+        }
+#endif
+
+        /// <summary>
+        /// Initialize the lock operator with a particular list of transforms.
+        /// </summary>
+        /// <param name="transforms"></param>
+        protected abstract void Reinit(List<(Transform, Transform)> transforms, List<int> problems);
+
+        /// <summary>
+        /// Computes the new positions and status words for a given range of bones.
+        /// </summary>
+        /// <param name="accessor"></param>
+        /// <param name="startBone"></param>
+        /// <param name="endBone"></param>
+        /// <returns></returns>
+        protected abstract JobHandle Compute(ArmatureLockJobAccessor accessor, int? jobIndex, JobHandle dependency);
+
+        public ArmatureLockJob RegisterLock(IEnumerable<(Transform, Transform)> transforms)
+        {
+            ArmatureLockJob job = null;
+            job = new ArmatureLockJob(
+                transforms.ToImmutableList(),
+                () => RemoveJob(job),
+                () => UpdateSingle(job)
+            );
+
+            _requestedJobs.Add(job);
+            Invalidate();
+
+            return job;
+        }
+
+        private void Invalidate()
+        {
+            _isValid = false;
+        }
+
+        private void MaybeRevalidate()
+        {
+            if (!_isValid)
+            {
+                // Do an update to make sure all the old jobs are in sync first, before we reset our state.
+                if (_isInit) SingleUpdate(null);
+                Reset();
+            }
+        }
+
+        private void Reset()
+        {
+            if (_isDisposed) return;
+
+            _lastJob.Complete();
+
+            if (_isInit)
+            {
+                _accessor.Destroy();
+                _baseBones.Dispose();
+                _targetBones.Dispose();
+            }
+
+            _isInit = true;
+
+            // TODO: toposort?
+            int[] boneToJobIndex = null;
+
+            List<int> problems = new List<int>();
+            do
+            {
+                var failed = problems.Select(p => _jobs[boneToJobIndex[p]]).Distinct().ToList();
+                foreach (var job in failed)
+                {
+                    job.IsValid = false;
+                    _requestedJobs.Remove(job);
+                }
+
+                problems.Clear();
+
+                _jobs = _requestedJobs.ToImmutableList();
+
+                _accessor.Destroy();
+                if (_baseBones.isCreated) _baseBones.Dispose();
+                if (_targetBones.isCreated) _targetBones.Dispose();
+
+                _baseBones = _targetBones = default;
+
+                var bones = _jobs.SelectMany(j => j.Transforms).ToList();
+                boneToJobIndex = _jobs.SelectMany((i, j) => Enumerable.Repeat(j, i.Transforms.Count)).ToArray();
+
+                var baseBones = bones.Select(t => t.Item1).ToArray();
+                var targetBones = bones.Select(t => t.Item2).ToArray();
+
+                _accessor.Allocate(
+                    bones.Count,
+                    _jobs.Count
+                );
+
+                _baseBones = new TransformAccessArray(baseBones);
+                _targetBones = new TransformAccessArray(targetBones);
+
+                Reinit(_jobs.SelectMany(j => j.Transforms).ToList(), problems);
+            } while (problems.Count > 0);
+
+            _isValid = true;
+        }
+
+        public void Update()
+        {
+            InternalUpdate();
+        }
+
+        private void UpdateSingle(ArmatureLockJob job)
+        {
+            var index = _jobs.IndexOf(job);
+            if (index < 0) return;
+
+            InternalUpdate(index);
+        }
+
+        private void InternalUpdate(int? jobIndex = null)
+        {
+            if (_isDisposed) return;
+
+            MaybeRevalidate();
+
+            SingleUpdate(jobIndex);
+        }
+
+        private long CycleStartHierarchyIndex = -1;
+        private int _nextCheckIndex = 0;
+
+        private void SingleUpdate(int? jobIndex)
+        {
+            if (!_isInit || _jobs.Count == 0) return;
+
+            Profiler.BeginSample("InternalUpdate");
+            _lastJob.Complete();
+
+            for (int i = 0; i < _jobs.Count; i++)
+            {
+                _accessor._abortFlag[i] = 0;
+                _accessor._didAnyWriteFlag[i] = 0;
+            }
+
+            _lastJob = ReadTransforms(jobIndex);
+            _lastJob = Compute(_accessor, jobIndex, _lastJob);
+
+            if (LastCheckedHierarchy != LastHierarchyChange)
+            {
+                Profiler.BeginSample("Recheck");
+
+                int startCheckIndex = _nextCheckIndex;
+                do
+                {
+                    if (_nextCheckIndex == 0)
+                    {
+                        CycleStartHierarchyIndex = LastHierarchyChange;
+                    }
+
+                    var job = _jobs[_nextCheckIndex % _jobs.Count];
+                    _nextCheckIndex = (1 + _nextCheckIndex) % _jobs.Count;
+
+                    if (job.HierarchyChanged)
+                    {
+                        job.IsValid = false;
+                        Invalidate();
+                    }
+                } while (_nextCheckIndex != startCheckIndex && !_lastJob.IsCompleted);
+
+                if (_nextCheckIndex == 0)
+                {
+                    LastCheckedHierarchy = CycleStartHierarchyIndex;
+                }
+
+                Profiler.EndSample();
+            }
+
+            // Before committing, do a spot check of any bones that moved, to see if their parents changed.
+            // This is needed because the hierarchyChanged event fires after Update ...
+
+            _lastJob.Complete();
+            Profiler.BeginSample("Revalidate dirty bones");
+            int boneBase = 0;
+            bool anyDirty = false;
+            for (int job = 0; job < _jobs.Count; job++)
+            {
+                int curBoneBase = boneBase;
+                boneBase += _jobs[job].Transforms.Count;
+                if (_accessor._didAnyWriteFlag[job] == 0) continue;
+
+                for (int b = curBoneBase; b < boneBase; b++)
+                {
+                    if (_accessor._out_dirty_targetBone[b] != 0 || _accessor._out_dirty_baseBone[b] != 0)
+                    {
+                        anyDirty = true;
+
+                        if (_jobs[job].BoneChanged(b - curBoneBase))
+                        {
+                            _accessor._abortFlag[job] = 1;
+                            _jobs[job].IsValid = false;
+                            break;
+                        }
+                    }
+                }
+            }
+
+            Profiler.EndSample();
+
+            if (anyDirty)
+            {
+                _lastJob = CommitTransforms(jobIndex, _lastJob);
+                _lastJob.Complete();
+            }
+
+            for (int i = 0; i < _jobs.Count; i++)
+            {
+                if (_accessor._abortFlag[i] != 0)
+                {
+                    Invalidate();
+                }
+                else
+                {
+                    _jobs[i].MarkLoop();
+                }
+
+                _jobs[i].WroteAny = _accessor._didAnyWriteFlag[i] != 0;
+            }
+
+            if (!_isValid)
+            {
+                Reset();
+            }
+
+            Profiler.EndSample();
+        }
+
+        private void RemoveJob(ArmatureLockJob job)
+        {
+            if (_requestedJobs.Remove(job)) Invalidate();
+        }
+
+        protected abstract void DerivedDispose();
+
+        #region Job logic
+
+        [BurstCompile]
+        struct ReadTransformsJob : IJobParallelForTransform
+        {
+            public NativeArray<TransformState> _bone;
+            public NativeArray<TransformState> _bone2;
+
+            [ReadOnly] public NativeArray<int> _boneToJobIndex;
+
+            [NativeDisableContainerSafetyRestriction] [NativeDisableParallelForRestriction]
+            public NativeArray<int> _abortFlag;
+
+            [BurstCompile]
+            public void Execute(int index, TransformAccess transform)
+            {
+#if UNITY_2021_1_OR_NEWER
+                if (!transform.isValid)
+                {
+                    _abortFlag[_boneToJobIndex[index]] = 1;
+                    return;
+                }
+#endif
+
+                _bone[index] = _bone2[index] = new TransformState
+                {
+                    localPosition = transform.localPosition,
+                    localRotation = transform.localRotation,
+                    localScale = transform.localScale
+                };
+            }
+        }
+
+        JobHandle ReadTransforms(int? jobIndex)
+        {
+            var baseRead = new ReadTransformsJob()
+            {
+                _bone = _accessor._in_baseBone,
+                _bone2 = _accessor._out_baseBone,
+                _boneToJobIndex = _accessor._boneToJobIndex,
+                _abortFlag = _accessor._abortFlag
+            }.ScheduleReadOnly(_baseBones, 32);
+
+            var targetRead = new ReadTransformsJob()
+            {
+                _bone = _accessor._in_targetBone,
+                _bone2 = _accessor._out_targetBone,
+                _boneToJobIndex = _accessor._boneToJobIndex,
+                _abortFlag = _accessor._abortFlag
+            }.ScheduleReadOnly(_targetBones, 32, baseRead);
+
+            return JobHandle.CombineDependencies(baseRead, targetRead);
+        }
+
+        [BurstCompile]
+        struct CommitTransformsJob : IJobParallelForTransform
+        {
+            [ReadOnly] public NativeArray<TransformState> _boneState;
+            [ReadOnly] public NativeArray<int> _dirtyBoneFlag;
+            [ReadOnly] public NativeArray<int> _boneToJobIndex;
+
+            [NativeDisableContainerSafetyRestriction] [NativeDisableParallelForRestriction] [ReadOnly]
+            public NativeArray<int> _abortFlag;
+
+            public int jobIndexFilter;
+
+            [BurstCompile]
+            public void Execute(int index, TransformAccess transform)
+            {
+#if UNITY_2021_1_OR_NEWER
+                if (!transform.isValid) return;
+#endif
+
+                var jobIndex = _boneToJobIndex[index];
+                if (jobIndexFilter >= 0 && jobIndex != jobIndexFilter) return;
+                if (_abortFlag[jobIndex] != 0) return;
+                if (_dirtyBoneFlag[index] == 0) return;
+
+                transform.localPosition = _boneState[index].localPosition;
+                transform.localRotation = _boneState[index].localRotation;
+                transform.localScale = _boneState[index].localScale;
+            }
+        }
+
+        JobHandle CommitTransforms(int? jobIndex, JobHandle prior)
+        {
+            JobHandle job = new CommitTransformsJob()
+            {
+                _boneState = _accessor._out_targetBone,
+                _dirtyBoneFlag = _accessor._out_dirty_targetBone,
+                _boneToJobIndex = _accessor._boneToJobIndex,
+                _abortFlag = _accessor._abortFlag,
+                jobIndexFilter = jobIndex ?? -1
+            }.Schedule(_targetBones, prior);
+
+            if (WritesBaseBones)
+            {
+                var job2 = new CommitTransformsJob()
+                {
+                    _boneState = _accessor._out_baseBone,
+                    _dirtyBoneFlag = _accessor._out_dirty_baseBone,
+                    _boneToJobIndex = _accessor._boneToJobIndex,
+                    _abortFlag = _accessor._abortFlag,
+                    jobIndexFilter = jobIndex ?? -1
+                }.Schedule(_baseBones, prior);
+
+                return JobHandle.CombineDependencies(job, job2);
+            }
+            else
+            {
+                return job;
+            }
+        }
+
+        #endregion
+    }
+}

Runtime/ArmatureAwase/ArmatureLockOperator.cs.meta

@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: ad2138add6244aa19ae24ccc42389efb
+timeCreated: 1709207125

Runtime/ArmatureAwase/BidirectionalArmatureLock.cs

@@ -1,233 +1,107 @@
 #region
 
-using System;
 using System.Collections.Generic;
-using nadena.dev.modular_avatar.JacksonDunstan.NativeCollections;
 using Unity.Burst;
 using Unity.Collections;
+using Unity.Collections.LowLevel.Unsafe;
 using Unity.Jobs;
 using UnityEngine;
-using UnityEngine.Jobs;
 
 #endregion
 
 namespace nadena.dev.modular_avatar.core.armature_lock
 {
-    internal class BidirectionalArmatureLock : ArmatureLock, IDisposable
+    internal class BidirectionalArmatureLockOperator : ArmatureLockOperator<BidirectionalArmatureLockOperator>
     {
-        private bool _disposed;
+        private NativeArray<TransformState> SavedState;
-        private TransformAccessArray _baseBoneAccess, _mergeBoneAccess;
+        protected override bool WritesBaseBones => true;
-        private readonly Transform[] _baseBones, _mergeBones, _baseParentBones, _mergeParentBones;
 
-        private NativeArray<TransformState> BaseBones, MergeBones, SavedMerge;
+        protected override void Reinit(List<(Transform, Transform)> transforms, List<int> problems)
-        private NativeArray<bool> ShouldWriteBase, ShouldWriteMerge;
-        private NativeIntPtr WroteAny;
-
-        private JobHandle LastOp;
-        private JobHandle LastPrepare;
-
-        public BidirectionalArmatureLock(IReadOnlyList<(Transform, Transform)> bones)
         {
-            _baseBones = new Transform[bones.Count];
+            if (SavedState.IsCreated) SavedState.Dispose();
-            _mergeBones = new Transform[bones.Count];
-            _baseParentBones = new Transform[bones.Count];
-            _mergeParentBones = new Transform[bones.Count];
 
-            BaseBones = new NativeArray<TransformState>(_baseBones.Length, Allocator.Persistent);
+            SavedState = new NativeArray<TransformState>(transforms.Count, Allocator.Persistent);
-            MergeBones = new NativeArray<TransformState>(_baseBones.Length, Allocator.Persistent);
-            SavedMerge = new NativeArray<TransformState>(_baseBones.Length, Allocator.Persistent);
 
-            for (int i = 0; i < _baseBones.Length; i++)
+            for (int i = 0; i < transforms.Count; i++)
             {
-                var (mergeBone, baseBone) = bones[i];
+                var (baseBone, mergeBone) = transforms[i];
-                _baseBones[i] = baseBone;
+                SavedState[i] = TransformState.FromTransform(mergeBone);
-                _mergeBones[i] = mergeBone;
-                _baseParentBones[i] = baseBone.parent;
-                _mergeParentBones[i] = mergeBone.parent;
 
-                var mergeState = TransformState.FromTransform(mergeBone);
+                if (TransformState.Differs(TransformState.FromTransform(baseBone), SavedState[i]))
-                SavedMerge[i] = mergeState;
+                {
-                MergeBones[i] = mergeState;
+                    problems.Add(i);
-                BaseBones[i] = TransformState.FromTransform(baseBone);
+                }
             }
+        }
 
-            _baseBoneAccess = new TransformAccessArray(_baseBones);
+        protected override JobHandle Compute(ArmatureLockJobAccessor accessor, int? jobIndex, JobHandle dependency)
-            _mergeBoneAccess = new TransformAccessArray(_mergeBones);
+        {
+            return new ComputeOperator()
+            {
+                base_in = accessor._in_baseBone,
+                merge_in = accessor._in_targetBone,
+                base_out = accessor._out_baseBone,
+                merge_out = accessor._out_targetBone,
 
-            ShouldWriteBase = new NativeArray<bool>(_baseBones.Length, Allocator.Persistent);
+                SavedState = SavedState,
-            ShouldWriteMerge = new NativeArray<bool>(_baseBones.Length, Allocator.Persistent);
+                baseDirty = accessor._out_dirty_baseBone,
-            WroteAny = new NativeIntPtr(Allocator.Persistent);
+                mergeDirty = accessor._out_dirty_targetBone,
+                boneToJobIndex = accessor._boneToJobIndex,
+                wroteAny = accessor._didAnyWriteFlag,
+
+                singleJobIndex = jobIndex ?? -1
+            }.Schedule(accessor._in_baseBone.Length, 16, dependency);
+        }
+
+        protected override void DerivedDispose()
+        {
+            SavedState.Dispose();
         }
 
         [BurstCompile]
-        struct Compute : IJobParallelForTransform
+        private struct ComputeOperator : IJobParallelFor
         {
-            public NativeArray<TransformState> BaseBones, SavedMerge;
+            public int singleJobIndex;
 
-            [WriteOnly] public NativeArray<TransformState> MergeBones;
+            public NativeArray<TransformState> base_in, merge_in, base_out, merge_out;
 
-            [WriteOnly] public NativeArray<bool> ShouldWriteBase, ShouldWriteMerge;
+            public NativeArray<TransformState> SavedState;
 
-            [WriteOnly] public NativeIntPtr.Parallel WroteAny;
+            [WriteOnly] public NativeArray<int> baseDirty, mergeDirty;
+            [ReadOnly] public NativeArray<int> boneToJobIndex;
 
-            public void Execute(int index, TransformAccess mergeTransform)
+            [NativeDisableContainerSafetyRestriction] [NativeDisableParallelForRestriction] [WriteOnly]
+            public NativeArray<int> wroteAny;
+
+            [BurstCompile]
+            public void Execute(int index)
             {
-                var baseBone = BaseBones[index];
+                var jobIndex = boneToJobIndex[index];
-                var mergeBone = new TransformState()
-                {
-                    localPosition = mergeTransform.localPosition,
-                    localRotation = mergeTransform.localRotation,
-                    localScale = mergeTransform.localScale,
-                };
-                MergeBones[index] = mergeBone;
 
-                var saved = SavedMerge[index];
+                if (singleJobIndex != -1 && jobIndex != singleJobIndex) return;
+
+                var baseBone = base_in[index];
+                var mergeBone = merge_in[index];
+                var saved = SavedState[index];
 
                 if (TransformState.Differs(saved, mergeBone))
                 {
-                    ShouldWriteBase[index] = true;
+                    baseDirty[index] = 1;
-                    ShouldWriteMerge[index] = false;
+                    mergeDirty[index] = 0;
 
-                    var mergeToBase = mergeBone;
+                    SavedState[index] = base_out[index] = merge_in[index];
-                    BaseBones[index] = mergeToBase;
+
-                    SavedMerge[index] = mergeBone;
+                    wroteAny[jobIndex] = 1;
-                    WroteAny.SetOne();
                 }
                 else if (TransformState.Differs(saved, baseBone))
                 {
-                    ShouldWriteMerge[index] = true;
+                    mergeDirty[index] = 1;
-                    ShouldWriteBase[index] = false;
+                    baseDirty[index] = 0;
 
-                    MergeBones[index] = baseBone;
+                    SavedState[index] = merge_out[index] = base_in[index];
-                    SavedMerge[index] = baseBone;
+
-                    WroteAny.SetOne();
+                    wroteAny[jobIndex] = 1;
                 }
-                else
-                {
-                    ShouldWriteBase[index] = false;
-                    ShouldWriteMerge[index] = false;
-                }
-            }
-        }
-
-        [BurstCompile]
-        struct Commit : IJobParallelForTransform
-        {
-            [ReadOnly] public NativeArray<TransformState> BoneState;
-            [ReadOnly] public NativeArray<bool> ShouldWrite;
-
-            public void Execute(int index, TransformAccess transform)
-            {
-                if (ShouldWrite[index])
-                {
-                    var boneState = BoneState[index];
-
-                    transform.localPosition = boneState.localPosition;
-                    transform.localRotation = boneState.localRotation;
-                    transform.localScale = boneState.localScale;
-                }
-            }
-        }
-
-        public override void Dispose()
-        {
-            if (_disposed) return;
-
-            LastOp.Complete();
-
-            // work around crashes caused by destroying TransformAccessArray from within Undo processing
-            DeferDestroy.DeferDestroyObj(_baseBoneAccess);
-            DeferDestroy.DeferDestroyObj(_mergeBoneAccess);
-            BaseBones.Dispose();
-            MergeBones.Dispose();
-            SavedMerge.Dispose();
-            ShouldWriteBase.Dispose();
-            ShouldWriteMerge.Dispose();
-            WroteAny.Dispose();
-
-            _disposed = true;
-        }
-
-        public override void Prepare()
-        {
-            if (_disposed) return;
-
-            LastOp.Complete();
-
-            WroteAny.Value = 0;
-
-            var readBase = new ReadBone()
-            {
-                _state = BaseBones,
-            }.Schedule(_baseBoneAccess);
-
-            LastOp = LastPrepare = new Compute()
-            {
-                BaseBones = BaseBones,
-                MergeBones = MergeBones,
-                SavedMerge = SavedMerge,
-                ShouldWriteBase = ShouldWriteBase,
-                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++)
-            {
-                if (_baseBones[i] == null || _mergeBones[i] == null || _baseParentBones[i] == null ||
-                    _mergeParentBones[i] == null)
-                {
-                    return false;
-                }
-
-                if (_baseBones[i].parent != _baseParentBones[i] || _mergeBones[i].parent != _mergeParentBones[i])
-                {
-                    return false;
-                }
-            }
-
-            return true;
-        }
-
-        public override bool IsStable()
-        {
-            Prepare();
-            if (!CheckConsistency()) return false;
-            LastPrepare.Complete();
-
-            return WroteAny.Value == 0;
-        }
-
-        public override LockResult Execute()
-        {
-            if (!CheckConsistency()) return LockResult.Failed;
-
-            var commitBase = new Commit()
-            {
-                BoneState = BaseBones,
-                ShouldWrite = ShouldWriteBase,
-            }.Schedule(_baseBoneAccess, LastPrepare);
-            var commitMerge = new Commit()
-            {
-                BoneState = MergeBones,
-                ShouldWrite = ShouldWriteMerge,
-            }.Schedule(_mergeBoneAccess, LastPrepare);
-
-            commitBase.Complete();
-            commitMerge.Complete();
-
-            if (WroteAny.Value == 0)
-            {
-                return LockResult.NoOp;
-            }
-            else
-            {
-                return LockResult.Success;
             }
         }
     }

Runtime/ArmatureAwase/DeferDestroy.cs

@@ -1,7 +1,7 @@
 #region
 
 using System;
-using UnityEngine;
+using UnityEditor;
 
 #endregion
 
@@ -33,7 +33,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                 return;
             }
 #if UNITY_EDITOR
-            UnityEditor.EditorApplication.delayCall += () => obj.Dispose();
+            EditorApplication.delayCall += () => obj.Dispose();
 #endif
         }
     }

Runtime/ArmatureAwase/LockResult.cs

@@ -1,9 +0,0 @@
-namespace nadena.dev.modular_avatar.core.armature_lock
-{
-    internal enum LockResult
-    {
-        Failed,
-        Success,
-        NoOp,
-    }
-}

Runtime/ArmatureAwase/LockResult.cs.meta

@@ -1,3 +0,0 @@
-fileFormatVersion: 2
-guid: 361faa0a05e34f7b8fbd1b2ae73d27bf
-timeCreated: 1693713933

Runtime/ArmatureAwase/OnewayArmatureLock.cs

@@ -1,58 +1,132 @@
 #region
 
-using System;
 using System.Collections.Generic;
-using nadena.dev.modular_avatar.JacksonDunstan.NativeCollections;
 using Unity.Burst;
 using Unity.Collections;
+using Unity.Collections.LowLevel.Unsafe;
 using Unity.Jobs;
 using UnityEngine;
-using UnityEngine.Jobs;
 #if UNITY_EDITOR
-using UnityEditor;
 #endif
 
 #endregion
 
 namespace nadena.dev.modular_avatar.core.armature_lock
 {
-    internal class OnewayArmatureLock : ArmatureLock, IDisposable
+    internal class OnewayArmatureLockOperator : ArmatureLockOperator<OnewayArmatureLockOperator>
     {
+        private Transform[] _baseBones, _mergeBones, _baseParentBones, _mergeParentBones;
+        private NativeArray<BoneStaticData> _boneStaticData;
+        public NativeArray<TransformState> _mergeSavedState;
+
+        private List<(Transform, Transform)> _transforms;
+        protected override bool WritesBaseBones => false;
+
+        protected override void Reinit(List<(Transform, Transform)> transforms, List<int> problems)
+        {
+            if (_boneStaticData.IsCreated) _boneStaticData.Dispose();
+            if (_mergeSavedState.IsCreated) _mergeSavedState.Dispose();
+            
+            _transforms = transforms;
+
+            _boneStaticData = new NativeArray<BoneStaticData>(transforms.Count, Allocator.Persistent);
+
+            _baseBones = new Transform[_transforms.Count];
+            _mergeBones = new Transform[_transforms.Count];
+            _baseParentBones = new Transform[_transforms.Count];
+            _mergeParentBones = new Transform[_transforms.Count];
+            _mergeSavedState = new NativeArray<TransformState>(_transforms.Count, Allocator.Persistent);
+
+            for (int i = 0; i < transforms.Count; i++)
+            {
+                var (baseBone, mergeBone) = transforms[i];
+                var mergeParent = mergeBone.parent;
+                var baseParent = baseBone.parent;
+
+                if (mergeParent == null || baseParent == null)
+                {
+                    problems.Add(i);
+                    continue;
+                }
+
+                if (SmallScale(mergeParent.localScale) || SmallScale(mergeBone.localScale) ||
+                    SmallScale(baseBone.localScale))
+                {
+                    problems.Add(i);
+                    continue;
+                }
+
+                _baseBones[i] = baseBone;
+                _mergeBones[i] = mergeBone;
+                _baseParentBones[i] = baseParent;
+                _mergeParentBones[i] = mergeParent;
+
+                _mergeSavedState[i] = TransformState.FromTransform(mergeBone);
+
+                // We want to emulate the hierarchy:
+                // baseParent
+                //  - baseBone 
+                //    - v_mergeBone 
+                //
+                // However our hierarchy actually is:
+                // mergeParent
+                //   - mergeBone
+                //
+                // Our question is: What is the local affine transform of mergeBone -> mergeParent space, given a new
+                // baseBone -> baseParent affine transform?
+
+                // First, relative to baseBone, what is the local affine transform of mergeBone?
+                var mat_l = baseBone.worldToLocalMatrix * mergeBone.localToWorldMatrix;
+                // We also find parent -> mergeParent
+                var mat_r = mergeParent.worldToLocalMatrix * baseParent.localToWorldMatrix;
+                // Now we can multiply:
+                // (baseParent -> mergeParent) * (baseBone -> baseParent) * (mergeBone -> baseBone)
+                //  = (baseParent -> mergeParent) * (mergeBone -> baseParent)
+                //  = (mergeBone -> mergeParent)
+
+                _boneStaticData[i] = new BoneStaticData()
+                {
+                    _mat_l = mat_r,
+                    _mat_r = mat_l
+                };
+            }
+        }
+
+        private bool SmallScale(Vector3 scale)
+        {
+            var epsilon = 0.000001f;
+
+            return (scale.x < epsilon || scale.y < epsilon || scale.z < epsilon);
+        }
+
+        protected override JobHandle Compute(ArmatureLockJobAccessor accessor, int? jobIndex, JobHandle dependency)
+        {
+            return new ComputePosition()
+            {
+                _baseState = accessor._in_baseBone,
+                _mergeState = accessor._in_targetBone,
+                _mergeSavedState = _mergeSavedState,
+                _boneStatic = _boneStaticData,
+                _fault = accessor._abortFlag,
+                _wroteAny = accessor._didAnyWriteFlag,
+                _wroteBone = accessor._out_dirty_targetBone,
+                jobIndexLimit = jobIndex ?? -1,
+                _boneToJobIndex = accessor._boneToJobIndex,
+                _outputState = accessor._out_targetBone,
+            }.Schedule(accessor._in_baseBone.Length, 32, dependency);
+        }
+
+        protected override void DerivedDispose()
+        {
+            if (_boneStaticData.IsCreated) _boneStaticData.Dispose();
+            if (_mergeSavedState.IsCreated) _mergeSavedState.Dispose();
+        }
+
         struct BoneStaticData
         {
             public Matrix4x4 _mat_l, _mat_r;
         }
 
-        private NativeArray<BoneStaticData> _boneStaticData;
-        private NativeArray<TransformState> _mergeSavedState;
-        private NativeArray<TransformState> _baseState, _mergeState;
-
-        private NativeIntPtr _fault, _wroteAny;
-
-        private readonly Transform[] _baseBones, _mergeBones, _baseParentBones, _mergeParentBones;
-        private TransformAccessArray _baseBonesAccessor, _mergeBonesAccessor;
-
-        private bool _disposed;
-        private JobHandle LastOp, LastPrepare;
-
-        [BurstCompile]
-        struct WriteBone : IJobParallelForTransform
-        {
-            [ReadOnly] public NativeIntPtr _fault, _shouldWrite;
-
-            [ReadOnly] public NativeArray<TransformState> _values;
-
-            public void Execute(int index, TransformAccess transform)
-            {
-                if (_fault.Value == 0 && _shouldWrite.Value != 0)
-                {
-                    var val = _values[index];
-                    transform.localPosition = val.localPosition;
-                    transform.localRotation = val.localRotation;
-                    transform.localScale = val.localScale;
-                }
-            }
-        }
 
         [BurstCompile]
         struct ComputePosition : IJobParallelFor
@@ -63,11 +137,23 @@ namespace nadena.dev.modular_avatar.core.armature_lock
             [ReadOnly] public NativeArray<TransformState> _baseState;
 
             public NativeArray<TransformState> _mergeSavedState;
+            public NativeArray<TransformState> _outputState;
+            public NativeArray<int> _wroteBone;
 
-            public NativeIntPtr.Parallel _fault, _wroteAny;
+            public int jobIndexLimit;
+
+            [ReadOnly] public NativeArray<int> _boneToJobIndex;
+
+            // job indexed
+            [NativeDisableContainerSafetyRestriction] [NativeDisableParallelForRestriction]
+            public NativeArray<int> _fault, _wroteAny;
 
             public void Execute(int index)
             {
+                var jobIndex = _boneToJobIndex[index];
+
+                if (jobIndexLimit >= 0 && jobIndex >= jobIndexLimit) return;
+
                 var boneStatic = _boneStatic[index];
                 var mergeState = _mergeState[index];
                 var baseState = _baseState[index];
@@ -79,8 +165,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
                 if (TransformState.Differs(mergeSaved, mergeState))
                 {
-                    TransformState.Differs(mergeSaved, mergeState);
+                    _fault[jobIndex] = 1;
-                    _fault.Increment();
                 }
 
                 var relTransform = boneStatic._mat_l * Matrix4x4.TRS(basePos, baseRot, baseScale) * boneStatic._mat_r;
@@ -98,219 +183,12 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
                 if (TransformState.Differs(mergeSaved, newState))
                 {
-                    _wroteAny.SetOne();
+                    _wroteAny[jobIndex] = 1;
+                    _wroteBone[index] = 1;
                     _mergeSavedState[index] = newState;
+                    _outputState[index] = newState;
                 }
             }
         }
-
-        public OnewayArmatureLock(IReadOnlyList<(Transform, Transform)> mergeToBase)
-        {
-            _boneStaticData = new NativeArray<BoneStaticData>(mergeToBase.Count, Allocator.Persistent);
-            _mergeSavedState = new NativeArray<TransformState>(mergeToBase.Count, Allocator.Persistent);
-            _baseState = new NativeArray<TransformState>(mergeToBase.Count, Allocator.Persistent);
-            _mergeState = new NativeArray<TransformState>(mergeToBase.Count, Allocator.Persistent);
-
-            _fault = new NativeIntPtr(Allocator.Persistent);
-            _wroteAny = new NativeIntPtr(Allocator.Persistent);
-
-            _baseBones = new Transform[mergeToBase.Count];
-            _mergeBones = new Transform[mergeToBase.Count];
-            _baseParentBones = new Transform[mergeToBase.Count];
-            _mergeParentBones = new Transform[mergeToBase.Count];
-
-            try
-            {
-                for (int i = 0; i < mergeToBase.Count; i++)
-                {
-                    var (mergeBone, baseBone) = mergeToBase[i];
-                    var mergeParent = mergeBone.parent;
-                    var baseParent = baseBone.parent;
-
-                    if (mergeParent == null || baseParent == null)
-                    {
-                        throw new Exception("Can't handle root objects");
-                    }
-
-                    if (SmallScale(mergeParent.localScale) || SmallScale(mergeBone.localScale) ||
-                        SmallScale(baseBone.localScale))
-                    {
-                        throw new Exception("Can't handle near-zero scale bones");
-                    }
-
-                    _baseBones[i] = baseBone;
-                    _mergeBones[i] = mergeBone;
-                    _baseParentBones[i] = baseParent;
-                    _mergeParentBones[i] = mergeParent;
-
-                    _baseState[i] = TransformState.FromTransform(baseBone);
-                    _mergeSavedState[i] = _mergeState[i] = TransformState.FromTransform(mergeBone);
-
-                    // We want to emulate the hierarchy:
-                    // baseParent
-                    //  - baseBone 
-                    //    - v_mergeBone 
-                    //
-                    // However our hierarchy actually is:
-                    // mergeParent
-                    //   - mergeBone
-                    //
-                    // Our question is: What is the local affine transform of mergeBone -> mergeParent space, given a new
-                    // baseBone -> baseParent affine transform?
-
-                    // First, relative to baseBone, what is the local affine transform of mergeBone?
-                    var mat_l = baseBone.worldToLocalMatrix * mergeBone.localToWorldMatrix;
-                    // We also find parent -> mergeParent
-                    var mat_r = mergeParent.worldToLocalMatrix * baseParent.localToWorldMatrix;
-                    // Now we can multiply:
-                    // (baseParent -> mergeParent) * (baseBone -> baseParent) * (mergeBone -> baseBone)
-                    //  = (baseParent -> mergeParent) * (mergeBone -> baseParent)
-                    //  = (mergeBone -> mergeParent)
-
-                    _boneStaticData[i] = new BoneStaticData()
-                    {
-                        _mat_l = mat_r,
-                        _mat_r = mat_l
-                    };
-                }
-            }
-            catch (Exception e)
-            {
-                _boneStaticData.Dispose();
-                _mergeSavedState.Dispose();
-                _baseState.Dispose();
-                _mergeState.Dispose();
-                _fault.Dispose();
-                _wroteAny.Dispose();
-
-                throw e;
-            }
-
-            _baseBonesAccessor = new TransformAccessArray(_baseBones);
-            _mergeBonesAccessor = new TransformAccessArray(_mergeBones);
-
-            EnableAssemblyReloadCallback = true;
-        }
-
-        private bool SmallScale(Vector3 scale)
-        {
-            var epsilon = 0.000001f;
-
-            return (scale.x < epsilon || scale.y < epsilon || scale.z < epsilon);
-        }
-
-        public override void Prepare()
-        {
-            if (_disposed) return;
-
-            LastOp.Complete();
-
-            _baseBonesAccessor.SetTransforms(_baseBones);
-            _mergeBonesAccessor.SetTransforms(_mergeBones);
-
-            _fault.Value = 0;
-            _wroteAny.Value = 0;
-
-            var jobReadBase = new ReadBone
-            {
-                _state = _baseState
-            }.Schedule(_baseBonesAccessor);
-            var jobReadMerged = new ReadBone
-            {
-                _state = _mergeState
-            }.Schedule(_mergeBonesAccessor);
-            var readAll = JobHandle.CombineDependencies(jobReadBase, jobReadMerged);
-            LastOp = LastPrepare = new ComputePosition
-            {
-                _boneStatic = _boneStaticData,
-                _mergeState = _mergeState,
-                _baseState = _baseState,
-                _mergeSavedState = _mergeSavedState,
-                _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++)
-            {
-                if (_baseBones[i] == null || _mergeBones[i] == null || _baseParentBones[i] == null ||
-                    _mergeParentBones[i] == null)
-                {
-                    return false;
-                }
-
-                if (_baseBones[i].parent != _baseParentBones[i] || _mergeBones[i].parent != _mergeParentBones[i])
-                {
-                    return false;
-                }
-            }
-
-            return true;
-        }
-
-        public override bool IsStable()
-        {
-            Prepare();
-            if (!CheckConsistency()) return false;
-
-            LastPrepare.Complete();
-
-            return _fault.Value == 0 && _wroteAny.Value == 0;
-        }
-
-        /// <summary>
-        /// Executes the armature lock job.
-        /// </summary>
-        /// <returns>True if successful, false if cached data was invalidated and needs recreating</returns>
-        public override LockResult Execute()
-        {
-            if (!CheckConsistency()) return LockResult.Failed;
-
-            var commit = new WriteBone()
-            {
-                _fault = _fault,
-                _values = _mergeSavedState,
-                _shouldWrite = _wroteAny
-            }.Schedule(_mergeBonesAccessor, LastPrepare);
-
-            commit.Complete();
-
-            if (_fault.Value != 0)
-            {
-                return LockResult.Failed;
-            }
-            else if (_wroteAny.Value == 0)
-            {
-                return LockResult.NoOp;
-            }
-            else
-            {
-                return LockResult.Success;
-            }
-        }
-
-        public override void Dispose()
-        {
-            if (_disposed) return;
-
-            LastOp.Complete();
-            _boneStaticData.Dispose();
-            _mergeSavedState.Dispose();
-            _baseState.Dispose();
-            _mergeState.Dispose();
-            _fault.Dispose();
-            _wroteAny.Dispose();
-            // work around crashes caused by destroying TransformAccessArray from within Undo processing
-            DeferDestroy.DeferDestroyObj(_baseBonesAccessor);
-            DeferDestroy.DeferDestroyObj(_mergeBonesAccessor);
-            _disposed = true;
-
-            EnableAssemblyReloadCallback = false;
-        }
     }
 }

Runtime/ArmatureAwase/ReadBone.cs

@@ -1,22 +0,0 @@
-using Unity.Burst;
-using Unity.Collections;
-using UnityEngine.Jobs;
-
-namespace nadena.dev.modular_avatar.core.armature_lock
-{
-    [BurstCompile]
-    internal struct ReadBone : IJobParallelForTransform
-    {
-        public NativeArray<TransformState> _state;
-
-        public void Execute(int index, TransformAccess transform)
-        {
-            _state[index] = new TransformState
-            {
-                localPosition = transform.localPosition,
-                localRotation = transform.localRotation,
-                localScale = transform.localScale
-            };
-        }
-    }
-}

Runtime/ArmatureAwase/ReadBone.cs.meta

@@ -1,3 +0,0 @@
-fileFormatVersion: 2
-guid: df590c12d16249608a9d8a8204b154bf
-timeCreated: 1693712551

Runtime/ArmatureAwase/TransformState.cs

@@ -1,7 +1,12 @@
-using System.Runtime.CompilerServices;
+#region
+
+using System.Runtime.CompilerServices;
 using Unity.Burst;
+using UnityEditor;
 using UnityEngine;
 
+#endregion
+
 namespace nadena.dev.modular_avatar.core.armature_lock
 {
     internal struct TransformState
@@ -14,7 +19,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         public Quaternion localRotation;
         public Vector3 localScale;
 
-        public static TransformState FromTransform(Transform mergeBone)
+        internal static TransformState FromTransform(Transform mergeBone)
         {
             return new TransformState
             {
@@ -24,10 +29,10 @@ namespace nadena.dev.modular_avatar.core.armature_lock
             };
         }
 
-        public void ToTransform(Transform bone)
+        internal void ToTransform(Transform bone)
         {
 #if UNITY_EDITOR
-            UnityEditor.Undo.RecordObject(bone, UnityEditor.Undo.GetCurrentGroupName());
+            Undo.RecordObject(bone, Undo.GetCurrentGroupName());
 #endif
             bone.localPosition = localPosition;
             bone.localRotation = localRotation;
@@ -36,7 +41,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
         [BurstCompile]
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public static bool Differs(TransformState self, TransformState other)
+        internal static bool Differs(TransformState self, TransformState other)
         {
             var deltaMergePos = (self.localPosition - other.localPosition).sqrMagnitude;
             var deltaMergeRot = self.localRotation * Quaternion.Inverse(other.localRotation);

Runtime/ArmatureAwase/Unity2019Compat.cs

@@ -0,0 +1,20 @@
+#if !UNITY_2021_1_OR_NEWER
+
+using Unity.Jobs;
+using UnityEngine.Jobs;
+
+namespace nadena.dev.modular_avatar.core.armature_lock
+{
+    internal static class Unity2019Compat
+    {
+        internal static JobHandle ScheduleReadOnly<T>(this T task, TransformAccessArray transforms, int batchCount,
+            JobHandle dependsOn = default)
+            where T : struct, IJobParallelForTransform
+        {
+            return task.Schedule(transforms, dependsOn);
+        }
+    }
+}
+
+
+#endif

Runtime/ArmatureAwase/Unity2019Compat.cs.meta

@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: dda11d07446e441d8313a99f53903d99
+timeCreated: 1709287006

Runtime/ArmatureAwase/UpdateLoopController.cs

@@ -1,33 +1,37 @@
-using System;
+#region
+
+using System;
+using System.Collections.Generic;
+using Unity.Jobs;
+using UnityEditor;
+
+#endregion
 
 namespace nadena.dev.modular_avatar.core.armature_lock
 {
     internal static class UpdateLoopController
     {
-        internal static event Action OnArmatureLockPrepare;
+        internal static event Action UpdateCallbacks;
-        internal static event Action OnArmatureLockUpdate;
         internal static event Action OnMoveIndependentlyUpdate;
 
 #if UNITY_EDITOR
-        [UnityEditor.InitializeOnLoadMethod]
+        [InitializeOnLoadMethod]
         private static void Init()
         {
-            UnityEditor.EditorApplication.update += () =>
+            EditorApplication.update += Update;
-            {
+        }
-                if (ArmatureLockConfig.instance.GlobalEnable)
-                {
-                    OnArmatureLockPrepare?.Invoke();
-                    OnArmatureLockUpdate?.Invoke();
-                }
 
-                OnMoveIndependentlyUpdate?.Invoke();
+        private static List<JobHandle> jobs = new List<JobHandle>();
-            };
+
+        private static void Update()
+        {
+            if (ArmatureLockConfig.instance.GlobalEnable)
+            {
+                UpdateCallbacks?.Invoke();
+            }
+
+            OnMoveIndependentlyUpdate?.Invoke();
         }
 #endif
-
-        internal static void InvokeArmatureLockPrepare()
-        {
-            OnArmatureLockPrepare?.Invoke();
-        }
     }
 }

Runtime/ModularAvatarMergeArmature.cs

@@ -22,13 +22,16 @@
  * SOFTWARE.
  */
 
+#region
+
 using System;
 using System.Collections.Generic;
 using nadena.dev.modular_avatar.core.armature_lock;
 using UnityEngine;
-using UnityEngine.Analytics;
 using UnityEngine.Serialization;
 
+#endregion
+
 namespace nadena.dev.modular_avatar.core
 {
     [Serializable]
@@ -109,28 +112,26 @@ namespace nadena.dev.modular_avatar.core
             SetLockMode();
         }
 
-        private void SetLockMode()
+        internal void SetLockMode()
         {
             if (this == null) return;
 
             if (_lockController == null)
             {
                 _lockController = ArmatureLockController.ForMerge(this, GetBonesForLock);
+                _lockController.WhenUnstable += OnUnstableLock;
             }
 
-            if (_lockController.Mode != LockMode)
+            _lockController.Mode = LockMode;
-            {
-                _lockController.Mode = LockMode;
-
-                if (!_lockController.IsStable())
-                {
-                    _lockController.Mode = LockMode = ArmatureLockMode.NotLocked;
-                }
-            }
 
             _lockController.Enabled = enabled;
         }
 
+        private void OnUnstableLock()
+        {
+            _lockController.Mode = LockMode = ArmatureLockMode.NotLocked;
+        }
+
         private void MigrateLockConfig()
         {
             if (LockMode == ArmatureLockMode.Legacy)
@@ -190,7 +191,7 @@ namespace nadena.dev.modular_avatar.core
                     var baseChild = FindCorrespondingBone(t, baseBone);
                     if (baseChild != null)
                     {
-                        mergeBones.Add((t, baseChild));
+                        mergeBones.Add((baseChild, t));
                         ScanHierarchy(t, baseChild);
                     }
                 }

Runtime/nadena.dev.modular-avatar.core.asmdef

@@ -1,5 +1,6 @@
 {
     "name": "nadena.dev.modular-avatar.core",
+    "rootNamespace": "",
     "references": [
         "Unity.Burst",
         "nadena.dev.ndmf.runtime"
@@ -13,7 +14,8 @@
         "VRCSDK3A.dll",
         "VRC.Dynamics.dll",
         "VRC.SDK3.Dynamics.Contact.dll",
-        "VRC.SDK3.Dynamics.PhysBone.dll"
+        "VRC.SDK3.Dynamics.PhysBone.dll",
+        "System.Collections.Immutable.dll"
     ],
     "autoReferenced": true,
     "defineConstraints": [],