opti: fix perf regressions in new armature lock system (#729)

* opti: fix perf regressions in new armature lock system

... by avoiding reinitializing everything whenever any target bone moves.

* chore: fixing unity 2019 issues

Runtime/ArmatureAwase/AllocationMap.cs

@@ -0,0 +1,149 @@
+#region
+
+using System;
+using System.Collections.Generic;
+
+#endregion
+
+namespace nadena.dev.modular_avatar.core.armature_lock
+{
+    internal interface ISegment
+    {
+        AllocationMap.DefragmentCallback Defragment { get; set; }
+        int Offset { get; }
+        int Length { get; }
+    }
+
+    internal class AllocationMap
+    {
+        public delegate void DefragmentCallback(int oldOffset, int newOffset, int length);
+
+        // Visible for unit tests
+        internal class Segment : ISegment
+        {
+            public int _offset;
+            public int _length;
+            public bool _inUse;
+
+            public AllocationMap.DefragmentCallback Defragment { get; set; }
+            public int Offset => _offset;
+            public int Length => _length;
+
+            internal Segment(int offset, int length, bool inUse)
+            {
+                _offset = offset;
+                _length = length;
+                _inUse = inUse;
+            }
+        }
+
+        /// <summary>
+        /// A list of allocated (and unallocated) segments.
+        ///
+        /// Invariant: The last element (if any) is always inUse.
+        /// Invariant: No two consecutive elements are free (inUse = false).
+        /// 
+        /// </summary>
+        List<Segment> segments = new List<Segment>();
+
+        public ISegment Allocate(int requestedLength)
+        {
+            for (int i = 0; i < segments.Count; i++)
+            {
+                var segment = segments[i];
+                if (segment._inUse) continue;
+
+                if (segment._length == requestedLength)
+                {
+                    segment._inUse = true;
+                    return segment;
+                }
+
+                if (segment._length > requestedLength)
+                {
+                    var remaining = new Segment(
+                        segment._offset + requestedLength,
+                        segment._length - requestedLength,
+                        false
+                    );
+
+                    segment._length = requestedLength;
+                    segment._inUse = true;
+                    segments.Insert(i + 1, remaining);
+
+                    return segment;
+                }
+            }
+
+            // Add a new in-use segment at the end
+            var newSegment = new Segment(
+                segments.Count == 0 ? 0 : segments[segments.Count - 1]._offset + segments[segments.Count - 1]._length,
+                requestedLength,
+                true
+            );
+            segments.Add(newSegment);
+
+            return newSegment;
+        }
+
+        public void FreeSegment(ISegment inputSegment)
+        {
+            var s = inputSegment as Segment;
+            if (s == null) throw new ArgumentException("Passed a foreign segment???");
+
+            int index = segments.BinarySearch(s, Comparer<Segment>.Create((a, b) => a._offset.CompareTo(b._offset)));
+            if (index < 0 || segments[index] != s) throw new Exception("Segment not found in FreeSegment");
+
+            if (index == segments.Count - 1)
+            {
+                segments.RemoveAt(index);
+                return;
+            }
+
+            if (index + 1 < segments.Count)
+            {
+                var next = segments[index + 1];
+                if (!next._inUse)
+                {
+                    next._offset = s._offset;
+                    next._length += s._length;
+                    segments.RemoveAt(index);
+                    return;
+                }
+            }
+
+            // Replace with a fresh segment object to avoid any issues with leaking old references to the segment
+            segments[index] = new Segment(s._offset, s._length, false);
+        }
+
+        /// <summary>
+        /// Defragments all free space. When a segment is moved, the passed callback is called with the old and new offsets,
+        /// and then the callback associated with the segment (if any) is also invoked.
+        /// </summary>
+        /// <param name="callback"></param>
+        public void Defragment(AllocationMap.DefragmentCallback callback)
+        {
+            int offset = 0;
+
+            for (int i = 0; i < segments.Count; i++)
+            {
+                var seg = segments[i];
+                if (!seg._inUse)
+                {
+                    segments.RemoveAt(i);
+                    i--;
+                    continue;
+                }
+
+                if (seg._offset != offset)
+                {
+                    callback(seg._offset, offset, seg._length);
+                    seg.Defragment?.Invoke(seg._offset, offset, seg._length);
+                    seg._offset = offset;
+                }
+
+                offset += seg.Length;
+            }
+        } 
+    }
+}

Runtime/ArmatureAwase/AllocationMap.cs.meta

@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: 1341ee6059a7410abe5a1170cdbf6355
+timeCreated: 1709527531

Runtime/ArmatureAwase/ArmatureLockController.cs

@@ -196,8 +196,9 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                         break;
                 }
             }
-            catch (Exception)
+            catch (Exception e)
             {
+                Debug.LogException(e);
                 _job = null;
                 return false;
             }

Runtime/ArmatureAwase/ArmatureLockJob.cs

@@ -23,8 +23,12 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         internal ImmutableList<(Transform, Transform)> RecordedParents;
         internal ImmutableList<(Transform, Transform)> Transforms;
 
-        internal ArmatureLockJob(ImmutableList<(Transform, Transform)> transforms, Action dispose, Action update)
+        internal ISegment Segment { get; private set; }
+
+        internal ArmatureLockJob(ISegment Segment, ImmutableList<(Transform, Transform)> transforms, Action dispose,
+            Action update)
         {
+            this.Segment = Segment;
             Transforms = transforms;
             RecordedParents = transforms.Select(((tuple, _) => (tuple.Item1.parent, tuple.Item2.parent)))
                 .ToImmutableList();

Runtime/ArmatureAwase/ArmatureLockJobAccessor.cs

@@ -17,45 +17,10 @@ namespace nadena.dev.modular_avatar.core.armature_lock
     /// </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;
+        public NativeArray<TransformState> _in_baseBone, _in_targetBone, _in_baseParentBone, _in_targetParentBone;
 
         /// <summary>
         /// Transform states to write out (if _out_dirty is set)
@@ -65,8 +30,14 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         /// <summary>
         /// Flags indicating whether the given bone should be written back to its transform
         /// </summary>
-        public NativeArray<int> _out_dirty_baseBone, _out_dirty_targetBone;
+        public NativeArray<bool> _out_dirty_baseBone, _out_dirty_targetBone;
 
+        /// <summary>
+        /// Indicates whether this bone index is associated with any job at all.
+        /// </summary>
+        [FormerlySerializedAs("_in_boneIsValid")]
+        public NativeArray<bool> _in_boneInUse;
+        
         /// <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.
@@ -75,7 +46,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         /// shouldn't read this value.
         /// </summary>
         [NativeDisableContainerSafetyRestriction] [NativeDisableParallelForRestriction]
-        public NativeArray<int> _abortFlag;
+        public NativeArray<bool> _abortFlag;
 
         /// <summary>
         /// Indexed by the job index (via _boneToJobIndex). Should be set to a nonzero value when any bone in the job
@@ -85,7 +56,7 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         /// shouldn't read this value.
         /// </summary>
         [NativeDisableContainerSafetyRestriction] [NativeDisableParallelForRestriction]
-        public NativeArray<int> _didAnyWriteFlag;
+        public NativeArray<bool> _didAnyWriteFlag;
 
         /// <summary>
         /// Maps from bone index to job index.

Runtime/ArmatureAwase/ArmatureLockOperator.cs

@@ -22,20 +22,57 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         internal static readonly T Instance = new T();
 
         private static long LastHierarchyChange = 0;
-        private ArmatureLockJobAccessor _accessor;
 
-        private TransformAccessArray _baseBones, _targetBones;
+        private TransformAccessArray _baseBones, _baseParentBones, _targetBones, _targetParentBones;
 
         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 List<ArmatureLockJob> _jobs = new List<ArmatureLockJob>();
         private long LastCheckedHierarchy = -1;
 
+        protected readonly NativeMemoryManager _memoryManager = new NativeMemoryManager();
+
+        private bool _transformAccessDirty = true;
+        private Transform[] _baseTransforms = Array.Empty<Transform>(), _targetTransforms = Array.Empty<Transform>();
+
+        private Transform[] _baseParentTransforms = Array.Empty<Transform>(),
+            _targetParentTransforms = Array.Empty<Transform>();
+
+        protected Transform[] BaseTransforms => _baseTransforms;
+        protected Transform[] TargetTransforms => _targetTransforms;
+
+        // Managed by _memoryManager
+        private NativeArrayRef<TransformState> _in_baseBone, _in_targetBone, _out_baseBone, _out_targetBone;
+        private NativeArrayRef<TransformState> _in_baseParentBone, _in_targetParentBone;
+
+        private NativeArrayRef<bool> _out_dirty_baseBone, _out_dirty_targetBone;
+        private NativeArrayRef<int> _boneToJobIndex;
+
+        // Not managed by _memoryManager (since they're not indexed by bone)
+        private NativeArray<bool> _abortFlag, _didAnyWriteFlag;
+
+        private ArmatureLockJobAccessor GetAccessor()
+        {
+            return new ArmatureLockJobAccessor()
+            {
+                _in_baseBone = _in_baseBone,
+                _in_targetBone = _in_targetBone,
+                _in_baseParentBone = _in_baseParentBone,
+                _in_targetParentBone = _in_targetParentBone,
+                _out_baseBone = _out_baseBone,
+                _out_targetBone = _out_targetBone,
+                _out_dirty_baseBone = _out_dirty_baseBone,
+                _out_dirty_targetBone = _out_dirty_targetBone,
+                _abortFlag = _abortFlag,
+                _didAnyWriteFlag = _didAnyWriteFlag,
+                _boneToJobIndex = _boneToJobIndex,
+                _in_boneInUse = _memoryManager.InUseMask,
+            };
+        }
+        
         static ArmatureLockOperator()
         {
             Instance = new T();
@@ -49,6 +86,18 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 #if UNITY_EDITOR
             AssemblyReloadEvents.beforeAssemblyReload += () => DeferDestroy.DestroyImmediate(this);
 #endif
+            _memoryManager.OnSegmentMove += MoveTransforms;
+
+            _in_baseBone = _memoryManager.CreateArray<TransformState>();
+            _in_targetBone = _memoryManager.CreateArray<TransformState>();
+            _out_baseBone = _memoryManager.CreateArray<TransformState>();
+            _out_targetBone = _memoryManager.CreateArray<TransformState>();
+            _in_baseParentBone = _memoryManager.CreateArray<TransformState>();
+            _in_targetParentBone = _memoryManager.CreateArray<TransformState>();
+
+            _out_dirty_baseBone = _memoryManager.CreateArray<bool>();
+            _out_dirty_targetBone = _memoryManager.CreateArray<bool>();
+            _boneToJobIndex = _memoryManager.CreateArray<int>();
         }
 
         protected abstract bool WritesBaseBones { get; }
@@ -57,14 +106,17 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         {
             if (_isDisposed) return;
             _isDisposed = true;
-
+            
-            if (!_isInit) return;
-
             _lastJob.Complete();
-            DeferDestroy.DeferDestroyObj(_baseBones);
+            if (_baseBones.isCreated) DeferDestroy.DeferDestroyObj(_baseBones);
-            DeferDestroy.DeferDestroyObj(_targetBones);
+            if (_targetBones.isCreated) DeferDestroy.DeferDestroyObj(_targetBones);
+            if (_baseParentBones.isCreated) DeferDestroy.DeferDestroyObj(_baseParentBones);
+            if (_targetParentBones.isCreated) DeferDestroy.DeferDestroyObj(_targetParentBones);
             DerivedDispose();
-            _accessor.Destroy();
+
+            _memoryManager.Dispose();
+            if (_abortFlag.IsCreated) _abortFlag.Dispose();
+            if (_didAnyWriteFlag.IsCreated) _didAnyWriteFlag.Dispose();
         }
 
 #if UNITY_EDITOR
@@ -72,7 +124,8 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         {
             EditorApplication.hierarchyChanged += () => { LastHierarchyChange += 1; };
             UpdateLoopController.UpdateCallbacks += Instance.Update;
-            ArmatureLockConfig.instance.OnGlobalEnableChange += Instance.Invalidate;
+            // TODO: On global enable, reset all jobs to init state?
+            //ArmatureLockConfig.instance.OnGlobalEnableChange += Instance.Invalidate;
 
             EditorApplication.playModeStateChanged += (change) =>
             {
@@ -87,10 +140,10 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 #endif
 
         /// <summary>
-        /// Initialize the lock operator with a particular list of transforms.
+        /// (Re-)initialize state for a single job
         /// </summary>
         /// <param name="transforms"></param>
-        protected abstract void Reinit(List<(Transform, Transform)> transforms, List<int> problems);
+        protected abstract bool SetupJob(ISegment segment);
 
         /// <summary>
         /// Computes the new positions and status words for a given range of bones.
@@ -103,92 +156,113 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
         public ArmatureLockJob RegisterLock(IEnumerable<(Transform, Transform)> transforms)
         {
+            if (_isDisposed) throw new ObjectDisposedException("ArmatureLockOperator");
+
+            var immutableTransforms = transforms.ToImmutableList();
+
+            var segment = _memoryManager.Allocate(immutableTransforms.Count());
+            
             ArmatureLockJob job = null;
             job = new ArmatureLockJob(
-                transforms.ToImmutableList(),
+                segment,
+                immutableTransforms,
                 () => RemoveJob(job),
                 () => UpdateSingle(job)
             );
 
-            _requestedJobs.Add(job);
+            EnsureTransformCapacity(_memoryManager.AllocatedLength);
-            Invalidate();
+
+            for (int i = 0; i < job.Transforms.Count(); i++)
+            {
+                var (baseBone, mergeBone) = job.Transforms[i];
+                _baseTransforms[i + segment.Offset] = baseBone;
+                _baseParentTransforms[i + segment.Offset] = baseBone.parent;
+                _targetTransforms[i + segment.Offset] = mergeBone;
+                _targetParentTransforms[i + segment.Offset] = mergeBone.parent;
+            }
+
+            int jobIndex = _jobs.IndexOf(null);
+            if (jobIndex >= 0)
+            {
+                _jobs[jobIndex] = job;
+            }
+            else
+            {
+                jobIndex = _jobs.Count();
+                _jobs.Add(job);
+            } 
+
+            EnsureJobFlagCapacity();
+
+            for (int i = 0; i < segment.Length; i++)
+            {
+                _boneToJobIndex.Array[segment.Offset + i] = jobIndex;
+            }
+
+            _transformAccessDirty = true;
+
+            bool ok = false;
+            try
+            {
+                ok = SetupJob(segment);
+            }
+            finally
+            {
+                if (!ok)
+                {
+                    // Initial setup failed; roll things back
+                    job.IsValid = false;
+                    RemoveJob(job);
+                }
+            }
 
             return job;
         }
 
-        private void Invalidate()
+        private void RemoveJob(ArmatureLockJob job)
         {
-            _isValid = false;
+            int index = _jobs.IndexOf(job);
+
+            if (index < 0) return;
+
+            _jobs[index] = null;
+
+            _memoryManager.Free(job.Segment);
         }
 
-        private void MaybeRevalidate()
+        private void EnsureJobFlagCapacity()
         {
-            if (!_isValid)
+            if (_abortFlag.IsCreated && _abortFlag.Length >= _jobs.Count) return;
-            {
+
-                // Do an update to make sure all the old jobs are in sync first, before we reset our state.
+            var priorLength = _abortFlag.Length;
-                if (_isInit) SingleUpdate(null);
+            
-                Reset();
+            if (_abortFlag.IsCreated) _abortFlag.Dispose();
-            }
+            if (_didAnyWriteFlag.IsCreated) _didAnyWriteFlag.Dispose();
+
+            int targetSize = Math.Max(Math.Max(16, _jobs.Count), (int)(priorLength * 1.5f));
+            _abortFlag = new NativeArray<bool>(targetSize, Allocator.Persistent);
+            _didAnyWriteFlag = new NativeArray<bool>(targetSize, Allocator.Persistent);
         }
 
-        private void Reset()
+        private void EnsureTransformCapacity(int targetLength)
         {
-            if (_isDisposed) return;
+            if (targetLength == _baseTransforms.Length) return;
 
-            _lastJob.Complete();
+            Array.Resize(ref _baseTransforms, targetLength);
-
+            Array.Resize(ref _baseParentTransforms, targetLength);
-            if (_isInit)
+            Array.Resize(ref _targetTransforms, targetLength);
-            {
+            Array.Resize(ref _targetParentTransforms, targetLength);
-                _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;
         }
 
+        private void MoveTransforms(int oldoffset, int newoffset, int length)
+        {
+            Array.Copy(_baseTransforms, oldoffset, _baseTransforms, newoffset, length);
+            Array.Copy(_baseParentTransforms, oldoffset, _baseParentTransforms, newoffset, length);
+            Array.Copy(_targetTransforms, oldoffset, _targetTransforms, newoffset, length);
+            Array.Copy(_targetParentTransforms, oldoffset, _targetParentTransforms, newoffset, length);
+            _transformAccessDirty = true;
+        }
+        
         public void Update()
         {
             InternalUpdate();
@@ -206,8 +280,6 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         {
             if (_isDisposed) return;
 
-            MaybeRevalidate();
-
             SingleUpdate(jobIndex);
         }
 
@@ -216,19 +288,54 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
         private void SingleUpdate(int? jobIndex)
         {
-            if (!_isInit || _jobs.Count == 0) return;
+            if (_jobs.Count == 0) return;
+
+            if (_isDisposed) return;
 
             Profiler.BeginSample("InternalUpdate");
             _lastJob.Complete();
 
+            EnsureJobFlagCapacity();
+
+            if (_transformAccessDirty)
+            {
+                Profiler.BeginSample("RecreateTransformAccess");
+
+                if (_baseBones.isCreated && _baseBones.length == _baseTransforms.Length)
+                {
+                    _baseBones.SetTransforms(_baseTransforms);
+                    _baseParentBones.SetTransforms(_baseParentTransforms);
+                    _targetBones.SetTransforms(_targetTransforms);
+                    _targetParentBones.SetTransforms(_targetParentTransforms);
+                }
+                else
+                {
+                    if (_baseBones.isCreated) _baseBones.Dispose();
+                    if (_targetBones.isCreated) _targetBones.Dispose();
+                    if (_baseParentBones.isCreated) _baseParentBones.Dispose();
+                    if (_targetParentBones.isCreated) _targetParentBones.Dispose();
+
+                    _baseBones = new TransformAccessArray(_baseTransforms);
+                    _baseParentBones = new TransformAccessArray(_baseParentTransforms);
+                    _targetBones = new TransformAccessArray(_targetTransforms);
+                    _targetParentBones = new TransformAccessArray(_targetParentTransforms);
+                }
+
+                _transformAccessDirty = false;
+
+                Profiler.EndSample();
+            }
+
+            var accessor = GetAccessor();
+            
             for (int i = 0; i < _jobs.Count; i++)
             {
-                _accessor._abortFlag[i] = 0;
+                accessor._abortFlag[i] = (_jobs[i] == null) || !_jobs[i].IsValid;
-                _accessor._didAnyWriteFlag[i] = 0;
+                accessor._didAnyWriteFlag[i] = false;
             }
 
             _lastJob = ReadTransforms(jobIndex);
-            _lastJob = Compute(_accessor, jobIndex, _lastJob);
+            _lastJob = Compute(accessor, jobIndex, _lastJob);
 
             if (LastCheckedHierarchy != LastHierarchyChange)
             {
@@ -245,10 +352,9 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                     var job = _jobs[_nextCheckIndex % _jobs.Count];
                     _nextCheckIndex = (1 + _nextCheckIndex) % _jobs.Count;
 
-                    if (job.HierarchyChanged)
+                    if (job != null && job.HierarchyChanged)
                     {
                         job.IsValid = false;
-                        Invalidate();
                     }
                 } while (_nextCheckIndex != startCheckIndex && !_lastJob.IsCompleted);
 
@@ -269,19 +375,21 @@ namespace nadena.dev.modular_avatar.core.armature_lock
             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] == 0) continue;
+                if (!accessor._didAnyWriteFlag[job]) continue;
 
                 for (int b = curBoneBase; b < boneBase; b++)
                 {
-                    if (_accessor._out_dirty_targetBone[b] != 0 || _accessor._out_dirty_baseBone[b] != 0)
+                    if (accessor._out_dirty_targetBone[b] || accessor._out_dirty_baseBone[b])
                     {
                         anyDirty = true;
 
                         if (_jobs[job].BoneChanged(b - curBoneBase))
                         {
-                            _accessor._abortFlag[job] = 1;
+                            accessor._abortFlag[job] = true;
                             _jobs[job].IsValid = false;
                             break;
                         }
@@ -299,96 +407,139 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
             for (int i = 0; i < _jobs.Count; i++)
             {
-                if (_accessor._abortFlag[i] != 0)
+                if (_jobs[i] == null) continue;
+
+                if (accessor._abortFlag[i])
                 {
-                    Invalidate();
+                    _jobs[i].IsValid = false;
                 }
                 else
                 {
                     _jobs[i].MarkLoop();
+                    _jobs[i].WroteAny = accessor._didAnyWriteFlag[i];
                 }
-
-                _jobs[i].WroteAny = _accessor._didAnyWriteFlag[i] != 0;
-            }
-
-            if (!_isValid)
-            {
-                Reset();
             }
 
             Profiler.EndSample();
         }
 
-        private void RemoveJob(ArmatureLockJob job)
+        protected virtual void DerivedDispose()
         {
-            if (_requestedJobs.Remove(job)) Invalidate();
+            // default no-op
         }
 
-        protected abstract void DerivedDispose();
-
         #region Job logic
 
+        [BurstCompile]
+        struct CopyTransformState : IJobParallelFor
+        {
+            [ReadOnly] public NativeArray<TransformState> _in;
+            [WriteOnly] public NativeArray<TransformState> _out;
+
+            public void Execute(int index)
+            {
+                _out[index] = _in[index];
+            }
+        }
+        
         [BurstCompile]
         struct ReadTransformsJob : IJobParallelForTransform
         {
-            public NativeArray<TransformState> _bone;
+            [WriteOnly] public NativeArray<TransformState> _bone;
-            public NativeArray<TransformState> _bone2;
 
             [ReadOnly] public NativeArray<int> _boneToJobIndex;
+            [ReadOnly] public NativeArray<bool> _boneInUse;
 
-            [NativeDisableContainerSafetyRestriction] [NativeDisableParallelForRestriction]
+            [NativeDisableContainerSafetyRestriction] [NativeDisableParallelForRestriction] [WriteOnly]
-            public NativeArray<int> _abortFlag;
+            public NativeArray<bool> _abortFlag;
 
             [BurstCompile]
             public void Execute(int index, TransformAccess transform)
             {
+                if (!_boneInUse[index]) return;
+
 #if UNITY_2021_1_OR_NEWER
                 if (!transform.isValid)
                 {
-                    _abortFlag[_boneToJobIndex[index]] = 1;
+                    _abortFlag[_boneToJobIndex[index]] = true;
                     return;
                 }
 #endif
 
-                _bone[index] = _bone2[index] = new TransformState
+                _bone[index] = new TransformState
                 {
                     localPosition = transform.localPosition,
                     localRotation = transform.localRotation,
-                    localScale = transform.localScale
+                    localScale = transform.localScale,
+                    localToWorldMatrix = transform.localToWorldMatrix,
                 };
             }
         }
 
         JobHandle ReadTransforms(int? jobIndex)
         {
+            var accessor = GetAccessor();
+            
             var baseRead = new ReadTransformsJob()
             {
-                _bone = _accessor._in_baseBone,
+                _bone = accessor._in_baseBone,
-                _bone2 = _accessor._out_baseBone,
+                _boneToJobIndex = accessor._boneToJobIndex,
-                _boneToJobIndex = _accessor._boneToJobIndex,
+                _abortFlag = accessor._abortFlag,
-                _abortFlag = _accessor._abortFlag
+                _boneInUse = accessor._in_boneInUse,
             }.ScheduleReadOnly(_baseBones, 32);
 
+            baseRead = new CopyTransformState()
+            {
+                _in = accessor._in_baseBone,
+                _out = accessor._out_baseBone
+            }.Schedule(accessor._in_baseBone.Length, 32, baseRead);
+
             var targetRead = new ReadTransformsJob()
             {
-                _bone = _accessor._in_targetBone,
+                _bone = accessor._in_targetBone,
-                _bone2 = _accessor._out_targetBone,
+                _boneToJobIndex = accessor._boneToJobIndex,
-                _boneToJobIndex = _accessor._boneToJobIndex,
+                _abortFlag = accessor._abortFlag,
-                _abortFlag = _accessor._abortFlag
+                _boneInUse = accessor._in_boneInUse,
-            }.ScheduleReadOnly(_targetBones, 32, baseRead);
+            }.ScheduleReadOnly(_targetBones, 32);
 
-            return JobHandle.CombineDependencies(baseRead, targetRead);
+            targetRead = new CopyTransformState()
+            {
+                _in = accessor._in_targetBone,
+                _out = accessor._out_targetBone
+            }.Schedule(accessor._in_targetBone.Length, 32, targetRead);
+
+            var baseParentRead = new ReadTransformsJob()
+            {
+                _bone = accessor._in_baseParentBone,
+                _boneToJobIndex = accessor._boneToJobIndex,
+                _abortFlag = accessor._abortFlag,
+                _boneInUse = accessor._in_boneInUse,
+            }.ScheduleReadOnly(_baseParentBones, 32);
+
+            var targetParentRead = new ReadTransformsJob()
+            {
+                _bone = accessor._in_targetParentBone,
+                _boneToJobIndex = accessor._boneToJobIndex,
+                _abortFlag = accessor._abortFlag,
+                _boneInUse = accessor._in_boneInUse,
+            }.ScheduleReadOnly(_targetParentBones, 32);
+
+            return JobHandle.CombineDependencies(
+                JobHandle.CombineDependencies(baseRead, targetRead),
+                JobHandle.CombineDependencies(baseParentRead, targetParentRead)
+            );
         }
 
         [BurstCompile]
         struct CommitTransformsJob : IJobParallelForTransform
         {
             [ReadOnly] public NativeArray<TransformState> _boneState;
-            [ReadOnly] public NativeArray<int> _dirtyBoneFlag;
+            [ReadOnly] public NativeArray<bool> _dirtyBoneFlag;
             [ReadOnly] public NativeArray<int> _boneToJobIndex;
+            [ReadOnly] public NativeArray<bool> _boneInUse;
 
             [NativeDisableContainerSafetyRestriction] [NativeDisableParallelForRestriction] [ReadOnly]
-            public NativeArray<int> _abortFlag;
+            public NativeArray<bool> _abortFlag;
 
             public int jobIndexFilter;
 
@@ -398,11 +549,12 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 #if UNITY_2021_1_OR_NEWER
                 if (!transform.isValid) return;
 #endif
+                if (!_boneInUse[index]) return;
 
                 var jobIndex = _boneToJobIndex[index];
                 if (jobIndexFilter >= 0 && jobIndex != jobIndexFilter) return;
-                if (_abortFlag[jobIndex] != 0) return;
+                if (_abortFlag[jobIndex]) return;
-                if (_dirtyBoneFlag[index] == 0) return;
+                if (!_dirtyBoneFlag[index]) return;
 
                 transform.localPosition = _boneState[index].localPosition;
                 transform.localRotation = _boneState[index].localRotation;
@@ -412,24 +564,28 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
         JobHandle CommitTransforms(int? jobIndex, JobHandle prior)
         {
+            var accessor = GetAccessor();
+            
             JobHandle job = new CommitTransformsJob()
             {
-                _boneState = _accessor._out_targetBone,
+                _boneState = accessor._out_targetBone,
-                _dirtyBoneFlag = _accessor._out_dirty_targetBone,
+                _dirtyBoneFlag = accessor._out_dirty_targetBone,
-                _boneToJobIndex = _accessor._boneToJobIndex,
+                _boneToJobIndex = accessor._boneToJobIndex,
-                _abortFlag = _accessor._abortFlag,
+                _abortFlag = accessor._abortFlag,
-                jobIndexFilter = jobIndex ?? -1
+                jobIndexFilter = jobIndex ?? -1,
+                _boneInUse = accessor._in_boneInUse,
             }.Schedule(_targetBones, prior);
 
             if (WritesBaseBones)
             {
                 var job2 = new CommitTransformsJob()
                 {
-                    _boneState = _accessor._out_baseBone,
+                    _boneState = accessor._out_baseBone,
-                    _dirtyBoneFlag = _accessor._out_dirty_baseBone,
+                    _dirtyBoneFlag = accessor._out_dirty_baseBone,
-                    _boneToJobIndex = _accessor._boneToJobIndex,
+                    _boneToJobIndex = accessor._boneToJobIndex,
-                    _abortFlag = _accessor._abortFlag,
+                    _abortFlag = accessor._abortFlag,
-                    jobIndexFilter = jobIndex ?? -1
+                    jobIndexFilter = jobIndex ?? -1,
+                    _boneInUse = accessor._in_boneInUse,
                 }.Schedule(_baseBones, prior);
 
                 return JobHandle.CombineDependencies(job, job2);

Runtime/ArmatureAwase/BidirectionalArmatureLock.cs

@@ -1,11 +1,9 @@
 #region
 
-using System.Collections.Generic;
 using Unity.Burst;
 using Unity.Collections;
 using Unity.Collections.LowLevel.Unsafe;
 using Unity.Jobs;
-using UnityEngine;
 
 #endregion
 
@@ -13,25 +11,31 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 {
     internal class BidirectionalArmatureLockOperator : ArmatureLockOperator<BidirectionalArmatureLockOperator>
     {
-        private NativeArray<TransformState> SavedState;
+        private NativeArrayRef<TransformState> SavedState;
         protected override bool WritesBaseBones => true;
 
-        protected override void Reinit(List<(Transform, Transform)> transforms, List<int> problems)
+        public BidirectionalArmatureLockOperator()
         {
-            if (SavedState.IsCreated) SavedState.Dispose();
+            SavedState = _memoryManager.CreateArray<TransformState>();
+        }
 
-            SavedState = new NativeArray<TransformState>(transforms.Count, Allocator.Persistent);
+        protected override bool SetupJob(ISegment segment)
-
+        {
-            for (int i = 0; i < transforms.Count; i++)
+            for (int i = 0; i < segment.Length; i++)
             {
-                var (baseBone, mergeBone) = transforms[i];
+                int bone = i + segment.Offset;
-                SavedState[i] = TransformState.FromTransform(mergeBone);
 
-                if (TransformState.Differs(TransformState.FromTransform(baseBone), SavedState[i]))
+                var baseBone = BaseTransforms[bone];
+                var targetBone = TargetTransforms[bone];
+
+                SavedState.Array[i] = TransformState.FromTransform(targetBone);
+                if (TransformState.Differs(TransformState.FromTransform(baseBone), SavedState.Array[i]))
                 {
-                    problems.Add(i);
+                    return false;
                 }
             }
+
+            return true;
         }
 
         protected override JobHandle Compute(ArmatureLockJobAccessor accessor, int? jobIndex, JobHandle dependency)
@@ -49,15 +53,12 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                 boneToJobIndex = accessor._boneToJobIndex,
                 wroteAny = accessor._didAnyWriteFlag,
 
+                boneInUse = accessor._in_boneInUse,
+
                 singleJobIndex = jobIndex ?? -1
             }.Schedule(accessor._in_baseBone.Length, 16, dependency);
         }
 
-        protected override void DerivedDispose()
-        {
-            SavedState.Dispose();
-        }
-
         [BurstCompile]
         private struct ComputeOperator : IJobParallelFor
         {
@@ -67,15 +68,19 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
             public NativeArray<TransformState> SavedState;
 
-            [WriteOnly] public NativeArray<int> baseDirty, mergeDirty;
+            [WriteOnly] public NativeArray<bool> baseDirty, mergeDirty;
             [ReadOnly] public NativeArray<int> boneToJobIndex;
 
             [NativeDisableContainerSafetyRestriction] [NativeDisableParallelForRestriction] [WriteOnly]
-            public NativeArray<int> wroteAny;
+            public NativeArray<bool> wroteAny;
+
+            [ReadOnly] public NativeArray<bool> boneInUse;
 
             [BurstCompile]
             public void Execute(int index)
             {
+                if (!boneInUse[index]) return;
+                
                 var jobIndex = boneToJobIndex[index];
 
                 if (singleJobIndex != -1 && jobIndex != singleJobIndex) return;
@@ -86,21 +91,21 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
                 if (TransformState.Differs(saved, mergeBone))
                 {
-                    baseDirty[index] = 1;
+                    baseDirty[index] = true;
-                    mergeDirty[index] = 0;
+                    mergeDirty[index] = false;
 
                     SavedState[index] = base_out[index] = merge_in[index];
 
-                    wroteAny[jobIndex] = 1;
+                    wroteAny[jobIndex] = true;
                 }
                 else if (TransformState.Differs(saved, baseBone))
                 {
-                    mergeDirty[index] = 1;
+                    mergeDirty[index] = true;
-                    baseDirty[index] = 0;
+                    baseDirty[index] = false;
 
                     SavedState[index] = merge_out[index] = base_in[index];
 
-                    wroteAny[jobIndex] = 1;
+                    wroteAny[jobIndex] = true;
                 }
             }
         }

Runtime/ArmatureAwase/NativeMemoryManager.cs

@@ -0,0 +1,219 @@
+#region
+
+using System;
+using System.Collections.Generic;
+using Unity.Collections;
+using Unity.Collections.LowLevel.Unsafe;
+
+#endregion
+
+namespace nadena.dev.modular_avatar.core.armature_lock
+{
+    internal class NativeArrayRef<T> : INativeArrayRef where T : unmanaged
+    {
+        internal NativeArray<T> Array;
+
+        public static implicit operator NativeArray<T>(NativeArrayRef<T> arrayRef) => arrayRef.Array;
+
+        public void Dispose()
+        {
+            Array.Dispose();
+        }
+
+        public void Resize(int n)
+        {
+            if (Array.Length == n) return;
+
+            var newArray = new NativeArray<T>(n, Allocator.Persistent);
+            
+            unsafe
+            {
+                UnsafeUtility.MemCpy(newArray.GetUnsafePtr(), Array.GetUnsafePtr(),
+                    Math.Min(n, Array.Length) * UnsafeUtility.SizeOf<T>());
+            }
+
+            /*
+            for (int i = 0; i < Math.Min(n, Array.Length); i++)
+            {
+                newArray[i] = Array[i];
+            }*/
+
+            Array.Dispose();
+
+            Array = newArray;
+        }
+
+        public void MemMove(int srcOffset, int dstOffset, int count)
+        {
+            if (srcOffset < 0 || dstOffset < 0
+                              || count < 0
+                              || srcOffset + count > Array.Length
+                              || dstOffset + count > Array.Length
+               )
+            {
+                throw new ArgumentOutOfRangeException();
+            }
+
+            
+            unsafe
+            {
+                UnsafeUtility.MemMove(((T*)Array.GetUnsafePtr()) + dstOffset, ((T*)Array.GetUnsafePtr()) + srcOffset,
+                    count * UnsafeUtility.SizeOf<T>());
+            }
+
+            /*
+            // We assume dstOffset < srcOffset
+            for (int i = 0; i < count; i++)
+            {
+                Array[dstOffset + i] = Array[srcOffset + i];
+            }*/
+        }
+    }
+
+    internal interface INativeArrayRef : IDisposable
+    {
+        void Resize(int n);
+        void MemMove(int srcOffset, int dstOffset, int count);
+    }
+
+    internal class NativeMemoryManager : IDisposable
+    {
+        private List<INativeArrayRef> arrays = new List<INativeArrayRef>();
+        public NativeArrayRef<bool> InUseMask { get; private set; }
+
+        public event AllocationMap.DefragmentCallback OnSegmentMove;
+
+        private int _allocatedLength = 1;
+        public int AllocatedLength => _allocatedLength;
+        private AllocationMap _allocationMap = new AllocationMap();
+        private bool _isDisposed;
+
+        public NativeMemoryManager()
+        {
+            // Bootstrap
+            InUseMask = new NativeArrayRef<bool>()
+            {
+                Array = new NativeArray<bool>(1, Allocator.Persistent)
+            };
+            arrays.Add(InUseMask);
+        }
+
+        public NativeArrayRef<T> CreateArray<T>() where T : unmanaged
+        {
+            if (_isDisposed)
+            {
+                throw new ObjectDisposedException(nameof(NativeMemoryManager));
+            }
+            
+            var arrayRef = new NativeArrayRef<T>()
+            {
+                Array = new NativeArray<T>(_allocatedLength, Allocator.Persistent)
+            };
+
+            arrays.Add(arrayRef);
+
+            return arrayRef;
+        }
+
+        public void Dispose()
+        {
+            if (_isDisposed) return;
+
+            _isDisposed = true;
+            
+            foreach (var array in arrays)
+            {
+                array.Dispose();
+            }
+        }
+
+        void SetInUseMask(int offset, int length, bool value)
+        {
+            unsafe
+            {
+                UnsafeUtility.MemSet((byte*)InUseMask.Array.GetUnsafePtr() + offset, value ? (byte)1 : (byte)0, length);
+            }
+
+            /*
+            for (int i = 0; i < length; i++)
+            {
+                try
+                {
+                    InUseMask.Array[offset + i] = value;
+                }
+                catch (Exception e)
+                {
+                    throw;
+                }
+            }*/
+        }
+
+        public ISegment Allocate(int requested)
+        {
+            if (_isDisposed)
+            {
+                throw new ObjectDisposedException(nameof(NativeMemoryManager));
+            }
+            
+            var segment = _allocationMap.Allocate(requested);
+
+            if (segment.Offset + segment.Length > _allocatedLength)
+            {
+                // Try defragmenting first.
+
+                // First, deallocate that segment we just created, since it'll be beyond the end of the array and break
+                // the memmove operations we'll be doing.
+                _allocationMap.FreeSegment(segment);
+                
+                Defragment();
+
+                segment = _allocationMap.Allocate(requested);
+            }
+
+            if (segment.Offset + segment.Length > _allocatedLength)
+            {
+                // We're still using more space than we have allocated, so allocate some more memory now
+                ResizeNativeArrays(segment.Offset + segment.Length);
+            }
+
+            SetInUseMask(segment.Offset, segment.Length, true);
+
+            return segment;
+        }
+
+        private void Defragment()
+        {
+            _allocationMap.Defragment((src, dst, length) =>
+            {
+                foreach (var array in arrays)
+                {
+                    array.MemMove(src, dst, length);
+                }
+
+                OnSegmentMove?.Invoke(src, dst, length);
+            });
+        }
+
+
+        private void ResizeNativeArrays(int minimumLength)
+        {
+            int targetLength = Math.Max((int)(1.5 * _allocatedLength), minimumLength);
+
+            foreach (var array in arrays)
+            {
+                array.Resize(targetLength);
+            }
+
+            SetInUseMask(_allocatedLength, targetLength - _allocatedLength, false);
+            _allocatedLength = targetLength;
+        }
+
+        public void Free(ISegment segment)
+        {
+            if (_isDisposed) return;
+            
+            _allocationMap.FreeSegment(segment);
+            SetInUseMask(segment.Offset, segment.Length, false);
+        }
+    }
+}

Runtime/ArmatureAwase/NativeMemoryManager.cs.meta

@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: b6602f940b944dc2b01f5f977fbc16a9
+timeCreated: 1709526526

Runtime/ArmatureAwase/OnewayArmatureLock.cs

@@ -15,54 +15,102 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 {
     internal class OnewayArmatureLockOperator : ArmatureLockOperator<OnewayArmatureLockOperator>
     {
-        private Transform[] _baseBones, _mergeBones, _baseParentBones, _mergeParentBones;
+        private NativeArrayRef<BoneStaticData> _boneStaticData;
-        private NativeArray<BoneStaticData> _boneStaticData;
+        private NativeArrayRef<TransformState> _mergeSavedState;
-        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)
+        public OnewayArmatureLockOperator()
         {
-            if (_boneStaticData.IsCreated) _boneStaticData.Dispose();
+            _boneStaticData = _memoryManager.CreateArray<BoneStaticData>();
-            if (_mergeSavedState.IsCreated) _mergeSavedState.Dispose();
+            _mergeSavedState = _memoryManager.CreateArray<TransformState>();
-            
+        }
-            _transforms = transforms;
 
-            _boneStaticData = new NativeArray<BoneStaticData>(transforms.Count, Allocator.Persistent);
+        protected override bool SetupJob(ISegment segment)
-
+        {
-            _baseBones = new Transform[_transforms.Count];
+            for (int i = 0; i < segment.Length; i++)
-            _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];
+                int bone = segment.Offset + i;
-                var mergeParent = mergeBone.parent;
-                var baseParent = baseBone.parent;
 
-                if (mergeParent == null || baseParent == null)
+                var baseState = TransformState.FromTransform(BaseTransforms[bone]);
+                var mergeState = TransformState.FromTransform(TargetTransforms[bone]);
+                var baseParentState = TransformState.FromTransform(BaseTransforms[bone].parent);
+                var mergeParentState = TransformState.FromTransform(TargetTransforms[bone].parent);
+
+                if (!new ComputePosition().SyncState(out var staticData, baseState, mergeState, baseParentState,
+                        mergeParentState))
                 {
-                    problems.Add(i);
+                    return false;
-                    continue;
                 }
 
-                if (SmallScale(mergeParent.localScale) || SmallScale(mergeBone.localScale) ||
+                _boneStaticData.Array[bone] = staticData;
-                    SmallScale(baseBone.localScale))
+                _mergeSavedState.Array[bone] = mergeState;
+            }
+
+            return true;
+        }
+
+        protected override JobHandle Compute(ArmatureLockJobAccessor accessor, int? jobIndex, JobHandle dependency)
+        {
+            return new ComputePosition()
+            {
+                _baseState = accessor._in_baseBone,
+                _mergeState = accessor._in_targetBone,
+                _baseParentState = accessor._in_baseParentBone,
+                _mergeParentState = accessor._in_baseParentBone,
+                _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,
+                _boneInUse = accessor._in_boneInUse,
+            }.Schedule(accessor._in_baseBone.Length, 32, dependency);
+        }
+        
+        struct BoneStaticData
+        {
+            public Matrix4x4 _mat_l, _mat_r;
+        }
+
+
+        [BurstCompile]
+        struct ComputePosition : IJobParallelFor
+        {
+            public NativeArray<BoneStaticData> _boneStatic;
+
+            [ReadOnly] public NativeArray<TransformState> _mergeState;
+            [ReadOnly] public NativeArray<TransformState> _baseState;
+
+            [ReadOnly] public NativeArray<TransformState> _mergeParentState;
+            [ReadOnly] public NativeArray<TransformState> _baseParentState;
+
+            public NativeArray<TransformState> _mergeSavedState;
+            public NativeArray<TransformState> _outputState;
+            public NativeArray<bool> _wroteBone;
+
+            public int jobIndexLimit;
+
+            [ReadOnly] public NativeArray<int> _boneToJobIndex;
+            [ReadOnly] public NativeArray<bool> _boneInUse;
+
+            // job indexed
+            [NativeDisableContainerSafetyRestriction] [NativeDisableParallelForRestriction]
+            public NativeArray<bool> _fault, _wroteAny;
+
+            public bool SyncState(out BoneStaticData result, TransformState baseState, TransformState mergeState,
+                TransformState baseParentState, TransformState mergeParentState)
+            {
+                if (SmallScale(mergeParentState.localScale) || SmallScale(mergeState.localScale) ||
+                    SmallScale(baseState.localScale))
                 {
-                    problems.Add(i);
+                    result = default;
-                    continue;
+                    return false;
                 }
 
-                _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 
@@ -76,80 +124,35 @@ namespace nadena.dev.modular_avatar.core.armature_lock
                 // baseBone -> baseParent affine transform?
 
                 // First, relative to baseBone, what is the local affine transform of mergeBone?
-                var mat_l = baseBone.worldToLocalMatrix * mergeBone.localToWorldMatrix;
+                var mat_l = baseState.worldToLocalMatrix * mergeState.localToWorldMatrix;
                 // We also find parent -> mergeParent
-                var mat_r = mergeParent.worldToLocalMatrix * baseParent.localToWorldMatrix;
+                var mat_r = mergeParentState.worldToLocalMatrix * baseParentState.localToWorldMatrix;
                 // Now we can multiply:
                 // (baseParent -> mergeParent) * (baseBone -> baseParent) * (mergeBone -> baseBone)
                 //  = (baseParent -> mergeParent) * (mergeBone -> baseParent)
                 //  = (mergeBone -> mergeParent)
 
-                _boneStaticData[i] = new BoneStaticData()
+                result = new BoneStaticData()
                 {
                     _mat_l = mat_r,
                     _mat_r = mat_l
                 };
+
+                return true;
             }
-        }
 
-        private bool SmallScale(Vector3 scale)
+            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,
+                var epsilon = 0.000001f;
-                _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;
-        }
-
-
-        [BurstCompile]
-        struct ComputePosition : IJobParallelFor
-        {
-            [ReadOnly] public NativeArray<BoneStaticData> _boneStatic;
-
-            [ReadOnly] public NativeArray<TransformState> _mergeState;
-            [ReadOnly] public NativeArray<TransformState> _baseState;
-
-            public NativeArray<TransformState> _mergeSavedState;
-            public NativeArray<TransformState> _outputState;
-            public NativeArray<int> _wroteBone;
-
-            public int jobIndexLimit;
-
-            [ReadOnly] public NativeArray<int> _boneToJobIndex;
-
-            // job indexed
-            [NativeDisableContainerSafetyRestriction] [NativeDisableParallelForRestriction]
-            public NativeArray<int> _fault, _wroteAny;
 
+                return (scale.x < epsilon || scale.y < epsilon || scale.z < epsilon);
+            }
+            
             public void Execute(int index)
             {
+                if (!_boneInUse[index]) return;
+                _wroteBone[index] = false;
+                
                 var jobIndex = _boneToJobIndex[index];
 
                 if (jobIndexLimit >= 0 && jobIndex >= jobIndexLimit) return;
@@ -165,7 +168,21 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
                 if (TransformState.Differs(mergeSaved, mergeState))
                 {
-                    _fault[jobIndex] = 1;
+                    // Reinitialize our transform matrices here, so we can continue to track on the next frame
+                    if (SyncState(out var state,
+                            _baseState[index],
+                            _mergeState[index],
+                            _baseParentState[index],
+                            _mergeParentState[index]))
+                    {
+                        _boneStatic[index] = state;
+                    }
+                    else
+                    {
+                        _fault[jobIndex] = true;
+                    }
+
+                    return;
                 }
 
                 var relTransform = boneStatic._mat_l * Matrix4x4.TRS(basePos, baseRot, baseScale) * boneStatic._mat_r;
@@ -183,8 +200,8 @@ namespace nadena.dev.modular_avatar.core.armature_lock
 
                 if (TransformState.Differs(mergeSaved, newState))
                 {
-                    _wroteAny[jobIndex] = 1;
+                    _wroteAny[jobIndex] = true;
-                    _wroteBone[index] = 1;
+                    _wroteBone[index] = true;
                     _mergeSavedState[index] = newState;
                     _outputState[index] = newState;
                 }

Runtime/ArmatureAwase/TransformState.cs

@@ -19,13 +19,19 @@ namespace nadena.dev.modular_avatar.core.armature_lock
         public Quaternion localRotation;
         public Vector3 localScale;
 
+        // Read on FromTransform, not written back in ToTransform
+        public Matrix4x4 localToWorldMatrix;
+
+        public Matrix4x4 worldToLocalMatrix => localToWorldMatrix.inverse;
+        
         internal static TransformState FromTransform(Transform mergeBone)
         {
             return new TransformState
             {
                 localPosition = mergeBone.localPosition,
                 localRotation = mergeBone.localRotation,
-                localScale = mergeBone.localScale
+                localScale = mergeBone.localScale,
+                localToWorldMatrix = mergeBone.localToWorldMatrix,
             };
         }
 

UnitTests~/ArmatureAwase.meta

@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: 1efcc3ec2e0142c880fb6d44f651e239
+timeCreated: 1709536475

UnitTests~/ArmatureAwase/AllocationMapTest.cs

@@ -0,0 +1,78 @@
+using System.Collections.Generic;
+using nadena.dev.modular_avatar.core.armature_lock;
+using NUnit.Framework;
+
+namespace UnitTests.ArmatureAwase
+{
+    public class AllocationMapTest
+    {
+        [Test]
+        public void Test()
+        {
+            AllocationMap map = new AllocationMap();
+            
+            ISegment s1 = map.Allocate(10);
+            AssertSegment(s1, 0, 10, true);
+            
+            ISegment s2 = map.Allocate(5);
+            AssertSegment(s2, 10, 5, true);
+            
+            map.FreeSegment(s1);
+            s1 = map.Allocate(5);
+            AssertSegment(s1, 0, 5, true);
+            
+            var s1a = map.Allocate(3);
+            AssertSegment(s1a, 5, 3, true);
+            
+            var s3 = map.Allocate(3);
+            AssertSegment(s3, 15, 3, true);
+
+            List<(ISegment, int, int, int)> segmentDefrags = new List<(ISegment, int, int, int)>();
+            List<(int, int, int)> globalDefrags = new List<(int, int, int)>();
+
+            s1.Defragment = (src, dst, length) => segmentDefrags.Add((s1, src, dst, length));
+            s1a.Defragment = (src, dst, length) => segmentDefrags.Add((s1a, src, dst, length));
+            s2.Defragment = (src, dst, length) => segmentDefrags.Add((s2, src, dst, length));
+            s3.Defragment = (src, dst, length) => segmentDefrags.Add((s3, src, dst, length));
+            
+            map.Defragment((src, dst, length) => globalDefrags.Add((src, dst, length)));
+            
+            Assert.AreEqual(segmentDefrags, new List<(ISegment, int, int, int)>()
+            {
+                (s2, 10, 8, 5),
+                (s3, 15, 13, 3),
+            });
+            
+            Assert.AreEqual(globalDefrags, new List<(int, int, int)>()
+            {
+                (10, 8, 5),
+                (15, 13, 3),
+            });
+        }
+
+        [Test]
+        public void SegmentCoalescing()
+        {
+            var map = new AllocationMap();
+            var s1 = map.Allocate(10);
+            var s2 = map.Allocate(10);
+            var s3 = map.Allocate(10);
+            
+            map.FreeSegment(s2);
+            map.FreeSegment(s1);
+            
+            var s4 = map.Allocate(20);
+            
+            AssertSegment(s4, 0, 20, true);
+        }
+
+        private void AssertSegment(ISegment segment, int offset, int length, bool inUse)
+        {
+            var s = segment as AllocationMap.Segment;
+            
+            Assert.AreEqual(offset, segment.Offset);
+            Assert.AreEqual(length, segment.Length);
+            Assert.AreEqual(inUse, s._inUse);
+        }
+    }
+}

UnitTests~/ArmatureAwase/AllocationMapTest.cs.meta

@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: 5d0c9fcd8a234b66b96dda10be362791
+timeCreated: 1709536483

UnitTests~/ArmatureAwase/NativeMemoryManagerTest.cs

@@ -0,0 +1,69 @@
+using System.Collections.Generic;
+using nadena.dev.modular_avatar.core.armature_lock;
+using NUnit.Framework;
+using Unity.Collections;
+
+namespace UnitTests.ArmatureAwase
+{
+    public class NativeMemoryManagerTest
+    {
+        [Test]
+        public void Test()
+        { 
+            var mm = new NativeMemoryManager();
+            var arr = mm.CreateArray<int>();
+
+            var s1 = mm.Allocate(8);
+            SetRange(arr, s1, 101);
+
+            var s2 = mm.Allocate(8);
+            SetRange(arr, s2, 102);
+            
+            mm.Free(s1);
+            AssertRange(mm.InUseMask, 0, 8, false);
+            AssertRange(mm.InUseMask, 8, 16, true);
+            AssertRange(mm.InUseMask, 16, -1, false);
+
+            List<(int, int, int)> defragOps = new List<(int, int, int)>();
+            mm.OnSegmentMove += (src, dst, length) => defragOps.Add((src, dst, length));
+            var s3 = mm.Allocate(16); // Forces reallocation/defragment
+            Assert.AreEqual(s2.Offset, 0);
+            Assert.AreEqual(defragOps, new List<(int, int, int)>()
+            {
+                (8, 0, 8),
+            });
+            SetRange(arr, s3, 103);
+            
+            AssertRange(arr, s2, 102);
+            
+            AssertRange(mm.InUseMask, s2, true);
+            AssertRange(mm.InUseMask, s3, true);
+            AssertRange(mm.InUseMask, s3.Offset, -1, false);
+            
+            mm.Dispose();
+            
+            Assert.IsFalse(arr.Array.IsCreated);
+        }
+
+        private void SetRange<T>(NativeArray<T> arr, ISegment segment, T value) where T : unmanaged
+        {
+            for (int i = 0; i < segment.Length; i++)
+            {
+                arr[i + segment.Offset] = value;
+            }
+        }
+        
+        private void AssertRange<T>(NativeArray<T> arr, ISegment segment, T value) where T : unmanaged
+        {
+            AssertRange<T>(arr, segment.Offset, segment.Offset + segment.Length, value);
+        }
+        
+        private void AssertRange<T>(NativeArray<T> arr, int start, int end, T value) where T : unmanaged
+        {
+            for (int i = start; i < end; i++)
+            {
+                Assert.AreEqual(value, arr[i]);
+            }
+        }
+    }
+}

UnitTests~/ArmatureAwase/NativeMemoryManagerTest.cs.meta

@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: d371b34b4f1e45f6b945509d26f48cee
+timeCreated: 1709536883