fix: a bunch of bugs related to invalid jobs in ArmatureAwase (#985)

Closes: #839, #856

Runtime/ArmatureAwase/AllocationMap.cs

@@ -2,9 +2,7 @@
 
 using System;
 using System.Collections.Generic;
-using System.Linq;
 using System.Runtime.CompilerServices;
-using UnityEngine;
 
 #endregion
 
@@ -52,7 +50,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                 _onDispose = null;
             }
         }
-
+        
         /// <summary>
         /// A list of allocated (and unallocated) segments.
         ///
@@ -62,6 +60,8 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         /// </summary>
         List<Segment> segments = new List<Segment>();
 
+        public event Action<ISegment> OnSegmentDispose;
+
         public ISegment Allocate(int requested)
         {
             int needed = Math.Max(1, requested);
@@ -114,6 +114,8 @@ namespace nadena.dev.modular_avatar.core.armature_lock
             var s = inputSegment as Segment;
             if (s == null) throw new ArgumentException("Passed a foreign segment???");
 
+            OnSegmentDispose?.Invoke(inputSegment);
+
             int index = segments.BinarySearch(s, Comparer<Segment>.Create((a, b) => a._offset.CompareTo(b._offset)));
             if (index < 0 || segments[index] != s)
             {

Runtime/ArmatureAwase/ArmatureLockController.cs

@@ -181,6 +181,12 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                 return false;
             }
 
+#if UNITY_EDITOR
+            if (xforms.Count == 0 || EditorUtility.IsPersistent(xforms[0].Item1))
+                // Bail out if we're trying to lock a prefab...
+                return true;
+#endif
+
             try
             {
                 switch (Mode)

Runtime/ArmatureAwase/ArmatureLockJob.cs

@@ -59,7 +59,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
             {
                 var transitioned = (_isValid && !value);
                 _isValid = value;
-
+                
                 if (transitioned)
                 {
 #if UNITY_EDITOR

Runtime/ArmatureAwase/ArmatureLockOperator.cs

@@ -371,17 +371,16 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
             _lastJob.Complete();
             Profiler.BeginSample("Revalidate dirty bones");
-            int boneBase = 0;
             bool anyDirty = false;
             for (int job = 0; job < _jobs.Count; job++)
             {
                 if (accessor._abortFlag[job]) continue;
-                
-                int curBoneBase = boneBase;
-                boneBase += _jobs[job].Transforms.Count;
                 if (!accessor._didAnyWriteFlag[job]) continue;
 
-                for (int b = curBoneBase; b < boneBase; b++)
+                var curBoneBase = _jobs[job].Segment.Offset;
+                var boneEnd = curBoneBase + _jobs[job].Segment.Length;
+
+                for (var b = curBoneBase; b < boneEnd; b++)
                 {
                     if (accessor._out_dirty_targetBone[b] || accessor._out_dirty_baseBone[b])
                     {
@@ -459,7 +458,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                 if (!_boneInUse[index]) return;
 
 #if UNITY_2021_1_OR_NEWER
-                if (!transform.isValid)
+                if (!transform.isValid && _boneInUse[index])
                 {
                     _abortFlag[_boneToJobIndex[index]] = true;
                     return;

Runtime/ArmatureAwase/NativeMemoryManager.cs

@@ -96,6 +96,8 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                 Array = new NativeArray<bool>(1, Allocator.Persistent)
             };
             arrays.Add(InUseMask);
+
+            _allocationMap.OnSegmentDispose += seg => { SetInUseMask(seg.Offset, seg.Length, false); };
         }
 
         public NativeArrayRef<T> CreateArray<T>() where T : unmanaged
@@ -183,6 +185,8 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
         private void Defragment()
         {
+            SetInUseMask(0, _allocatedLength, false);
+            
             _allocationMap.Defragment((src, dst, length) =>
             {
                 foreach (var array in arrays)
@@ -190,11 +194,12 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                     array.MemMove(src, dst, length);
                 }
 
+                SetInUseMask(dst, length, true);
+
                 OnSegmentMove?.Invoke(src, dst, length);
             });
         }
-
-
+        
         private void ResizeNativeArrays(int minimumLength)
         {
             int targetLength = Math.Max((int)(1.5 * _allocatedLength), minimumLength);

Runtime/ArmatureAwase/OnewayArmatureLock.cs

@@ -176,6 +176,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                             _mergeParentState[index]))
                     {
                         _boneStatic[index] = state;
+                        _mergeSavedState[index] = mergeState;
                     }
                     else
                     {