1  using System;
  2  using System.Collections;
  3  using System.Collections.Generic;
  4  using System.Numerics;
  5  using System.Runtime.CompilerServices;
  6  
  7  namespace ARMeilleure.Common
  8  {
  9      unsafe class BitMap : IEnumerable<int>, IDisposable
 10      {
 11          private const int IntSize = 64;
 12          private const int IntMask = IntSize - 1;
 13  
 14          private int _count;
 15          private long* _masks;
 16          private readonly Allocator _allocator;
 17  
 18          public BitMap(Allocator allocator)
 19          {
 20              _allocator = allocator;
 21          }
 22  
 23          public BitMap(Allocator allocator, int capacity) : this(allocator)
 24          {
 25              EnsureCapacity(capacity);
 26          }
 27  
 28          public bool Set(int bit)
 29          {
 30              EnsureCapacity(bit + 1);
 31  
 32              int wordIndex = bit / IntSize;
 33              int wordBit = bit & IntMask;
 34  
 35              long wordMask = 1L << wordBit;
 36  
 37              if ((_masks[wordIndex] & wordMask) != 0)
 38              {
 39                  return false;
 40              }
 41  
 42              _masks[wordIndex] |= wordMask;
 43  
 44              return true;
 45          }
 46  
 47          public void Clear(int bit)
 48          {
 49              EnsureCapacity(bit + 1);
 50  
 51              int wordIndex = bit / IntSize;
 52              int wordBit = bit & IntMask;
 53  
 54              long wordMask = 1L << wordBit;
 55  
 56              _masks[wordIndex] &= ~wordMask;
 57          }
 58  
 59          public bool IsSet(int bit)
 60          {
 61              EnsureCapacity(bit + 1);
 62  
 63              int wordIndex = bit / IntSize;
 64              int wordBit = bit & IntMask;
 65  
 66              return (_masks[wordIndex] & (1L << wordBit)) != 0;
 67          }
 68  
 69          public int FindFirstUnset()
 70          {
 71              for (int index = 0; index < _count; index++)
 72              {
 73                  long mask = _masks[index];
 74  
 75                  if (mask != -1L)
 76                  {
 77                      return BitOperations.TrailingZeroCount(~mask) + index * IntSize;
 78                  }
 79              }
 80  
 81              return _count * IntSize;
 82          }
 83  
 84          public bool Set(BitMap map)
 85          {
 86              EnsureCapacity(map._count * IntSize);
 87  
 88              bool modified = false;
 89  
 90              for (int index = 0; index < _count; index++)
 91              {
 92                  long newValue = _masks[index] | map._masks[index];
 93  
 94                  if (_masks[index] != newValue)
 95                  {
 96                      _masks[index] = newValue;
 97  
 98                      modified = true;
 99                  }
100              }
101  
102              return modified;
103          }
104  
105          public bool Clear(BitMap map)
106          {
107              EnsureCapacity(map._count * IntSize);
108  
109              bool modified = false;
110  
111              for (int index = 0; index < _count; index++)
112              {
113                  long newValue = _masks[index] & ~map._masks[index];
114  
115                  if (_masks[index] != newValue)
116                  {
117                      _masks[index] = newValue;
118  
119                      modified = true;
120                  }
121              }
122  
123              return modified;
124          }
125  
126          private void EnsureCapacity(int size)
127          {
128              int count = (size + IntMask) / IntSize;
129  
130              if (count > _count)
131              {
132                  var oldMask = _masks;
133                  var oldSpan = new Span<long>(_masks, _count);
134  
135                  _masks = _allocator.Allocate<long>((uint)count);
136                  _count = count;
137  
138                  var newSpan = new Span<long>(_masks, _count);
139  
140                  oldSpan.CopyTo(newSpan);
141                  newSpan[oldSpan.Length..].Clear();
142  
143                  _allocator.Free(oldMask);
144              }
145          }
146  
147          public void Dispose()
148          {
149              if (_masks != null)
150              {
151                  _allocator.Free(_masks);
152  
153                  _masks = null;
154              }
155          }
156  
157          IEnumerator IEnumerable.GetEnumerator()
158          {
159              return GetEnumerator();
160          }
161  
162          IEnumerator<int> IEnumerable<int>.GetEnumerator()
163          {
164              return GetEnumerator();
165          }
166  
167          public Enumerator GetEnumerator()
168          {
169              return new Enumerator(this);
170          }
171  
172          public struct Enumerator : IEnumerator<int>
173          {
174              private long _index;
175              private long _mask;
176              private int _bit;
177              private readonly BitMap _map;
178  
179              public readonly int Current => (int)_index * IntSize + _bit;
180              readonly object IEnumerator.Current => Current;
181  
182              public Enumerator(BitMap map)
183              {
184                  _index = -1;
185                  _mask = 0;
186                  _bit = 0;
187                  _map = map;
188              }
189  
190              [MethodImpl(MethodImplOptions.AggressiveInlining)]
191              public bool MoveNext()
192              {
193                  if (_mask != 0)
194                  {
195                      _mask &= ~(1L << _bit);
196                  }
197  
198                  // Manually hoist these loads, because RyuJIT does not.
199                  long count = (uint)_map._count;
200                  long* masks = _map._masks;
201  
202                  while (_mask == 0)
203                  {
204                      if (++_index >= count)
205                      {
206                          return false;
207                      }
208  
209                      _mask = masks[_index];
210                  }
211  
212                  _bit = BitOperations.TrailingZeroCount(_mask);
213  
214                  return true;
215              }
216  
217              public readonly void Reset() { }
218  
219              public readonly void Dispose() { }
220          }
221      }
222  }