NativeContainer doesn't work in jobs - What's missing?

Hey!

I need some help on the NativeContainer I’ve written. It’s working fine from mainthread but values are not updated when used in a job. No idea what I’m missing here.
I’m using UnsafeList and have a [NativeContainer] tag.

[NativeContainer]
    [StructLayout(LayoutKind.Sequential)]
    public unsafe struct ArrayHashMap<TKey, TValue> : IDisposable
        where TKey : unmanaged
        where TValue : unmanaged
    {
        [NativeDisableUnsafePtrRestriction] internal byte* Keys;
        [NativeDisableUnsafePtrRestriction] internal byte* Values;

        [NativeDisableUnsafePtrRestriction] internal UnsafeList<int>* buckets;
        [NativeDisableUnsafePtrRestriction] internal UnsafeList<int>* next;
       
        internal int keyCapacity;
        internal int bucketCapacityMask;
        internal int allocatedIndexLength;

        internal Allocator m_AllocatorLabel;
       
#if ENABLE_UNITY_COLLECTIONS_CHECKS
        internal AtomicSafetyHandle m_Safety;
        [NativeSetClassTypeToNullOnSchedule] internal DisposeSentinel m_DisposeSentinel;
#endif
        public ArrayHashMap(int capacity, Allocator allocator = Allocator.Persistent)
        {
#if ENABLE_UNITY_COLLECTIONS_CHECKS
            DisposeSentinel.Create(out m_Safety, out m_DisposeSentinel, 0, allocator);
#endif

            m_AllocatorLabel = allocator;

            Keys = null;
            Values = null;
           
            next = UnsafeList<int>.Create(capacity, allocator);
            buckets = UnsafeList<int>.Create(capacity * 2, allocator);

            keyCapacity = 0;
            bucketCapacityMask = 0;
            allocatedIndexLength = 0;
        }

        public void SetArrays(NativeArray<TKey> keyArray, NativeArray<TValue> valueArray)
        {
            if (!keyArray.IsCreated || !valueArray.IsCreated)
                throw new Exception("Key or values are not created!");
            if (keyArray.Length != valueArray.Length)
                throw new Exception("Key and value length is not the same!");
            if (keyArray.Length == 0 || valueArray.Length == 0)
            {
                allocatedIndexLength = 0;
                return;
            }
           
            Keys = (byte*) keyArray.GetUnsafeReadOnlyPtr();
            Values = (byte*) valueArray.GetUnsafeReadOnlyPtr();

            int length = keyArray.Length;
            int bucketLength = length * 2;

            keyCapacity = length;
            bucketLength = math.ceilpow2(bucketLength);
            bucketCapacityMask = bucketLength - 1;

            Debug.Log($"Set next/buckets cap to {length}/{bucketLength}");
            next->Resize(length, NativeArrayOptions.UninitializedMemory);
            buckets->Resize(bucketLength, NativeArrayOptions.UninitializedMemory);

            for (int i = 0; i < length; i++)
                (*next)[i] = -1;
           
            for (int i = 0; i < bucketLength; i++)
                (*buckets)[i] = -1;
           
            allocatedIndexLength = length;
           
            Debug.Log($"SetArrays with allocatedIndexLength {allocatedIndexLength}");
           
            CalculateBuckets();
        }

        private void Clear()
        {
            // set all to -1
            UnsafeUtility.MemSet(buckets->Ptr, 0xff, (bucketCapacityMask + 1) * 4);
            UnsafeUtility.MemSet(next->Ptr, 0xff, (keyCapacity) * 4);
            next->Clear();
            buckets->Clear();
           
            allocatedIndexLength = 0;
        }


        private void CalculateBuckets()
        {
            //Debug.Log($"CalculateBuckets with length {allocatedIndexLength} nextCap: {next.Capacity} bucketsCap: {buckets.Capacity}");
            for (int i = 0; i < allocatedIndexLength; i++)
            {
                var bucketIndex = (*(TKey*) (Keys + i * sizeof(TKey))).GetHashCode() & bucketCapacityMask;
               
                (*next)[i] = (*buckets)[bucketIndex];
                (*buckets)[bucketIndex] = i;
            }
        }

        public void PrintValues()
        {
            Debug.Log($"PrintValues with length {allocatedIndexLength}");
            for (int i = 0; i < allocatedIndexLength; i++)
            {
                var key = (*(TKey*)(Keys + i * sizeof(TKey)));
                Debug.Log($"Key: {key}");
            }
            for (int i = 0; i < allocatedIndexLength; i++)
            {
                var value = (*(TValue*)(Values + i * sizeof(TValue)));
                Debug.Log($"value: {value}");
            }
            for (int i = 0; i < allocatedIndexLength; i++)
            {
                var nextValue = (*(int*)(next + i * sizeof(int)));
                Debug.Log($"nextValue: {nextValue}");
            }
            for (int i = 0; i < (bucketCapacityMask + 1); i++)
            {
                var bucketValue = (*(int*)(buckets + i * sizeof(int)));
                Debug.Log($"bucketValue: {bucketValue}");
            }
        }
       
        public bool TryGetFirstRefValue(TKey key, out byte* item, out ArrayHashMapIterator<TKey> it)           
        {
            it.key = key;           

            if (allocatedIndexLength <= 0)
            {
                it.EntryIndex = it.NextEntryIndex = -1;
                item = null;
                return false;
            }

            // First find the slot based on the hash           
            int bucket = key.GetHashCode() & bucketCapacityMask;
            it.EntryIndex = it.NextEntryIndex = (*buckets)[bucket];
            return TryGetNextRefValue(out item, ref it);
        }

        public bool TryGetNextRefValue(out byte* item, ref ArrayHashMapIterator<TKey> it)          
        {
            int entryIdx = it.NextEntryIndex;
            it.NextEntryIndex = -1;
            it.EntryIndex = -1;
            item = null;
          
            if (entryIdx < 0 || entryIdx >= keyCapacity)
            {
                return false;
            }
           
            //while (!Keys[entryIdx].Equals(it.key))
            while (!(*(TKey*) (Keys + entryIdx * sizeof(TKey))).Equals(it.key))
            {
                entryIdx = (*next)[entryIdx];
                if (entryIdx < 0 || entryIdx >= keyCapacity)
                {
                    return false;
                }
            }

            it.NextEntryIndex = (*next)[entryIdx];
            it.EntryIndex = entryIdx;

            // Read the value
            item = Values + entryIdx * sizeof(TValue);

            return true;
        }
       
        public bool ContainsKey(TKey key)
        {
            return TryGetFirstRefValue(key, out var temp0, out var temp1);
        }
       
        public ArrayHashMapEnumerator<TKey, TValue> GetValuesForKey(TKey key)
        {
            return new ArrayHashMapEnumerator<TKey, TValue> { Map = this, key = key, isFirst = true };
        }

        public void Dispose()
        {
#if ENABLE_UNITY_COLLECTIONS_CHECKS
            if (!UnsafeUtility.IsValidAllocator(m_AllocatorLabel))
                throw new InvalidOperationException("The NativeArray can not be Disposed because it was not allocated with a valid allocator.");

            DisposeSentinel.Dispose(ref m_Safety, ref m_DisposeSentinel);
#endif
           
            next->Dispose();
            buckets->Dispose();
        }
    }
   
    public unsafe struct ArrayHashMapEnumerator<TKey, TValue> where TKey : unmanaged where TValue : unmanaged
    {
        public ArrayHashMap<TKey, TValue> Map;

        public TKey key;
        public bool isFirst;
        private byte* value;

        private ArrayHashMapIterator<TKey> iterator;
        public ref TValue Current => ref UnsafeUtility.AsRef<TValue>(value);

        public bool MoveNext()
        {
            //Avoids going beyond the end of the collection.
            if (!isFirst)
                return Map.TryGetNextRefValue(out value, ref iterator);
           
            isFirst = false;
            return Map.TryGetFirstRefValue(key, out value, out iterator);

        }
    }
       
    public struct ArrayHashMapIterator<TKey>
        where TKey : struct
    {
        internal TKey key;
        internal int NextEntryIndex;
        internal int EntryIndex;

        /// <summary>
        /// Returns the entry index.
        /// </summary>
        /// <returns>The entry index.</returns>
        public int GetEntryIndex() => EntryIndex;
    }

TestSystem:

[AlwaysUpdateSystem]
    public partial class ArrayHashMapTest : SystemBase
    {
        public NativeArray<int> Keys;
        public NativeArray<int> Values;
        public ArrayHashMap<int, int> arrayHashMap;

        protected override void OnCreate()
        {
            Keys = new NativeArray<int>(100, Allocator.Persistent);
            Values = new NativeArray<int>(100, Allocator.Persistent);

            arrayHashMap = new ArrayHashMap<int, int>(1);
        }

        protected override void OnDestroy()
        {
            Keys.Dispose();
            Values.Dispose();
            arrayHashMap.Dispose();
        }

        protected override void OnUpdate()
        {
            Debug.Log("ArrayHashMapTest Update");
            for (int i = 0; i < 20; i++)
            {
                Keys[i] = i;
                Values[i] = 20;
            }
           
            Keys[20] = 1;
            Values[20] = 2;

            Keys[21] = 1;
            Values[21] = 3;

            Keys[22] = 1;
            Values[22] = 4;

            Keys[23] = 1;
            Values[23] = 5;

            for (int i = 24; i < 40; i++)
            {
                Keys[i] = i;
                Values[i] = 30;
            }

            Keys[41] = 41;
            Values[41] = 5;
           
            for (int i = 42; i < 100; i++)
            {
                Keys[i] = i;
                Values[i] = 40;
            }
           
            //arrayHashMap.SetArrays(Keys, Values);
           
            new CalculateBucketsJob<int, int>()
            {
                keys = Keys,
                values = Values,
                hashmap = arrayHashMap
            }.Schedule(Dependency).Complete();

            var enumerator = arrayHashMap.GetValuesForKey(1);
           
            while (enumerator.MoveNext())
            {
                Debug.Log("hashmap value; " + enumerator.Current);
            }
           
            arrayHashMap.PrintValues();
        }
    }
1 Like

I think your internal int values need to be pointers. Because right now they are being copied by value whenever you pass your container into jobs.

3 Likes

Ah, thanks a lot. That must be it.

edit: That was indeed the problem! Now I understand why the additional UnsafeX structs are in place and it’s not just in one NativeContainer struct.
Unity devs, if you are reading this, some explicit mention in a doc would be great.

Docs like: Unity - Scripting API: NativeContainerAttribute don’t mention it in any way and use blittable types. Honestly, not a good example to learn NativeContainers. Start by teaching the good habits around writing NativeContainers like writing an unsafe part of the container and just a pointer in the NativeContainer. If you understand this part, everything else clicks too.

Custom job types | Jobs | 0.50.0-preview.9 has an int* but that went honestly over my head while reading, mostly because the later NativeCounter implementation uses several ints for each thread. Much better example, just needs some preamble why an int* was used in the first place.

1 Like