Question How can I write an enumerator for nested unmanaged enumerables?

I've written enumerators based on the existing NativeHashMap.Enumerator and UnsafeHashMap.Enumerator structs that work in burst, but I am running into trouble when I try to combine enumerators for nested unmanaged structs. Specifically I am running into a problem where the first level works, but the second level only returns the first element of the set because the enumerator it depends on is getting corrupted or reset.

In my example I have three structs: Parent, Child, and GrandChild. Parent contains a NativeHashMap of Child, and Child contains an UnsafeHashMap of GrandChild. I can iterate over a Parent's Chilren using TestEnumerator.NativeMap. I can also iterate over a Child's GrandChildren using TestEnumerator.UnsafeMap. My problem is that when I try to combine these two enumerators into TestEnumerator.AllGrandChildren, the UnsafeMap Enumerator gets corrupted each time that AllGrandChildren.MoveNext() is called, and only the first GrandChild from every Child is returned.

Code (CSharp):

1. using UnityEngine;
2. using UnityEngine;
3. using Unity.Collections;
4. using System.Collections.Generic;
5. using System;
6. using Unity.Collections.LowLevel.Unsafe;
7. using System.Collections;
8.  
9. public class NestedEnumeratorTest : MonoBehaviour
10. {
11.     void Start()
12.     {
13.         {
14.             TestChild child = new TestChild(Allocator.Temp);
15.  
16.             for (int i = 0; i < 100; i++)
17.             {
18.                 child.Add(new TestGrandChild(-1, i));
19.             }
20.  
21.             int count = 0;
22.             foreach (TestGrandChild grandchild in child.GrandChildren)
23.             {
24.                 count++;
25.             }
26.             Debug.Log("Number of grandchildren in child: " + count); // Logs 100
27.         }
28.  
29.         {
30.             TestParent parent = new TestParent(Allocator.Temp);
31.  
32.             for (int i = 0; i < 10; i++)
33.             {
34.                 TestChild child = new TestChild(Allocator.Temp);
35.  
36.                 for (int j = 0; j < 10; j++)
37.                 {
38.                     child.Add(new TestGrandChild(i, j));
39.                 }
40.  
41.                 parent.Add(child);
42.             }
43.  
44.             int count = 0;
45.             foreach (TestGrandChild grandchild in parent.AllGrandChildren)
46.             {
47.                 count++;
48.             }
49.             Debug.Log("Number of grandchildren in parent: " + count); // Logs 10
50.         }
51.     }
52. }
53.  
54. public struct TestEmumerable
55. {
56.     public struct AllGrandChildren : IEnumerable<TestGrandChild>
57.     {
58.         private NativeHashMap<int, TestChild>.Enumerator enumerator;
59.  
60.         public readonly TestEmumerator.AllGrandChildren Enumerator => new(enumerator);
61.  
62.         public AllGrandChildren(NativeHashMap<int, TestChild>.Enumerator enumerator)
63.         {
64.             this.enumerator = enumerator;
65.         }
66.  
67.         public IEnumerator<TestGrandChild> GetEnumerator()
68.         {
69.             return Enumerator;
70.         }
71.  
72.         IEnumerator IEnumerable.GetEnumerator()
73.         {
74.             throw new NotImplementedException();
75.         }
76.     }
77.  
78.     public struct NativeMap<Key, Value> : IEnumerable<Value>
79.     where Key : unmanaged, IEquatable<Key>
80.     where Value : unmanaged
81.     {
82.         private NativeHashMap<Key, Value>.Enumerator enumerator;
83.  
84.         public readonly TestEmumerator.NativeMap<Key, Value> Enumerator => new(enumerator);
85.  
86.         public NativeMap(NativeHashMap<Key, Value>.Enumerator enumerator)
87.         {
88.             this.enumerator = enumerator;
89.         }
90.  
91.         public IEnumerator<Value> GetEnumerator()
92.         {
93.             return Enumerator;
94.         }
95.  
96.         IEnumerator IEnumerable.GetEnumerator()
97.         {
98.             throw new NotImplementedException();
99.         }
100.     }
101.  
102.     public struct UnsafeMap<Key, Value> : IEnumerable<Value>
103.         where Key : unmanaged, IEquatable<Key>
104.         where Value : unmanaged
105.     {
106.         private UnsafeHashMap<Key, Value>.Enumerator enumerator;
107.  
108.         public readonly TestEmumerator.UnsafeMap<Key, Value> Enumerator => new(enumerator);
109.  
110.         public UnsafeMap(UnsafeHashMap<Key, Value>.Enumerator enumerator)
111.         {
112.             this.enumerator = enumerator;
113.         }
114.  
115.         public IEnumerator<Value> GetEnumerator()
116.         {
117.             return Enumerator;
118.         }
119.  
120.         IEnumerator IEnumerable.GetEnumerator()
121.         {
122.             throw new NotImplementedException();
123.         }
124.     }
125. }
126.  
127. public struct TestEmumerator
128. {
129.     public struct AllGrandChildren : IEnumerator<TestGrandChild>
130.     {
131.         private NativeMap<int, TestChild> children; // Why can't I use a custom type here?
132.         private UnsafeMap<int, TestGrandChild> grandchildren;
133.         private NativeReference<bool> started;
134.  
135.         public AllGrandChildren(NativeHashMap<int, TestChild>.Enumerator children)
136.         {
137.             this.children = new(children);
138.             grandchildren = new();
139.             started = new(false, Allocator.Temp);
140.         }
141.  
142.         public TestGrandChild Current => grandchildren.Current;
143.  
144.         object IEnumerator.Current => Current;
145.  
146.         public bool MoveNext()
147.         {
148.             bool moveNext;
149.  
150.             if (started.Value)
151.             {
152.                 moveNext = grandchildren.MoveNext();
153.             }
154.             else
155.             {
156.                 started.Value = true;
157.                 moveNext = false;
158.             }
159.  
160.             while (!moveNext && children.MoveNext())
161.             {
162.                 grandchildren = children.Current.GrandChildren.Enumerator;
163.                 moveNext = grandchildren.MoveNext();
164.             }
165.  
166.             return moveNext;
167.         }
168.  
169.         public void Reset()
170.         {
171.             if (started.Value)
172.             {
173.                 started.Value = false;
174.             }
175.  
176.             children.Reset();
177.         }
178.  
179.         public void Dispose()
180.         {
181.  
182.         }
183.     }
184.  
185.     public struct NativeMap<Key, Value> : IEnumerator<Value>, IDisposable
186.         where Key : unmanaged, IEquatable<Key>
187.         where Value : unmanaged
188.     {
189.         private NativeHashMap<Key, Value>.Enumerator enumerator;
190.  
191.         public NativeMap(NativeHashMap<Key, Value>.Enumerator enumerator)
192.         {
193.             this.enumerator = enumerator;
194.         }
195.  
196.         public Value Current => enumerator.Current.Value;
197.  
198.         object IEnumerator.Current => Current;
199.  
200.         public void Dispose()
201.         {
202.  
203.         }
204.  
205.         public bool MoveNext()
206.         {
207.             return enumerator.MoveNext();
208.         }
209.  
210.         public void Reset()
211.         {
212.             enumerator.Reset();
213.         }
214.     }
215.  
216.     public struct UnsafeMap<Key, Value> : IEnumerator<Value>, IDisposable
217.         where Key : unmanaged, IEquatable<Key>
218.         where Value : unmanaged
219.     {
220.         private UnsafeHashMap<Key, Value>.Enumerator enumerator;
221.  
222.         public UnsafeMap(UnsafeHashMap<Key, Value>.Enumerator enumerator)
223.         {
224.             this.enumerator = enumerator;
225.         }
226.  
227.         public unsafe Value Current => enumerator.Current.Value;
228.  
229.         object IEnumerator.Current => Current;
230.  
231.         public void Dispose()
232.         {
233.  
234.         }
235.  
236.         public bool MoveNext()
237.         {
238.             return enumerator.MoveNext();
239.         }
240.  
241.         public void Reset()
242.         {
243.             enumerator.Reset();
244.         }
245.     }
246. }
247.  
248. public struct TestParent
249. {
250.     NativeHashMap<int, TestChild> children;
251.     public TestEmumerable.NativeMap<int, TestChild> Children => new(children.GetEnumerator());
252.     public TestEmumerable.AllGrandChildren AllGrandChildren => new(children.GetEnumerator());
253.  
254.     public TestParent(Allocator allocator)
255.     {
256.         children = new(10, allocator);
257.     }
258.  
259.     public void Add(TestChild child)
260.     {
261.         children[children.Count] = child;
262.     }
263. }
264.  
265. public struct TestChild
266. {
267.     private UnsafeHashMap<int, TestGrandChild> grandchildren;
268.     public TestEmumerable.UnsafeMap<int, TestGrandChild> GrandChildren => new(grandchildren.GetEnumerator());
269.  
270.     public TestChild(Allocator allocator)
271.     {
272.         grandchildren = new(10, allocator);
273.     }
274.  
275.     public void Add(TestGrandChild grandchild)
276.     {
277.         grandchildren[grandchildren.Count] = grandchild;
278.     }
279. }
280.  
281. public struct TestGrandChild
282. {
283.     int parent;
284.     int child;
285.     public TestGrandChild(int parent, int child)
286.     {
287.         this.parent = parent;
288.         this.child = child;
289.     }
290. }
291.  

Edit: I solved my problem, but I don't know why my solution works. In TestEmumerator.AllGrandChildren, If I use NativeHashMap.Enumerator instead of my custom TestEnumerator.NativeMap type for the children property, everything starts working. Can someone help explain why this is?

 

it's because the way you have it set up, you're copying structs around by value.

For example:

Code (CSharp):

1. grandchildren = children.Current.GrandChildren.Enumerator;

This line here copies the enumerator and creates a new one. When you call

MoveNext

, it updates the state of the copy. To fix it, you either need to return references of these structs, or set the fields to the modified copies to retain the state changes.