1#! /usr/bin/env python
2
3"""
4Integer number allocator.
5
6Basically, these keep track of a set of allocatable values in
7some range (you provide min and max) and let you allocate out of
8the range and return values into the range.
9
10You may pick a value using "next since last time", or "next
11available after provided value".  Note that next-after will
12wrap around as needed (modular arithmetic style).
13
14The free lists are thread-locked so that this code can be used
15with threads.
16
17    >>> a = NumAlloc(5, 10) # note closed interval: 5..10 inclusive
18    >>> a
19    NumAlloc(5, 10)
20    >>> a.avail
21    [[5, 10]]
22    >>> a.alloc()
23    5
24    >>> a.avail
25    [[6, 10]]
26    >>> a.alloc(8)
27    8
28    >>> a.avail
29    [[6, 7], [9, 10]]
30    >>> a.free(5)
31    >>> a.avail
32    [[5, 7], [9, 10]]
33    >>> a.free(8)
34    >>> a.avail
35    [[5, 10]]
36
37Attempting to free a value that is already free is an error:
38
39    >>> a.free(5)
40    Traceback (most recent call last):
41       ...
42    ValueError: free: 5 already available
43
44You can, however, free a value that is outside the min/max
45range.  You can also free multiple values at once:
46
47    >>> a.free_multi([0, 1, 2, 4])
48    >>> a.avail
49    [[0, 2], [4, 10]]
50    >>> a.free_multi([3, 12])
51    >>> a.avail
52    [[0, 10], [12, 12]]
53
54Note that this changes the min/max values:
55
56    >>> a
57    NumAlloc(0, 12)
58
59To prevent adding values outside the min/max range, create the
60NumArray with autoextend=False, or set .autoextend=False at any
61time:
62
63    >>> a.autoextend = False
64    >>> a
65    NumAlloc(0, 12, autoextend=False)
66    >>> a.free(13)
67    Traceback (most recent call last):
68       ...
69    ValueError: free: 13 is outside range limit
70
71You can create an empty range, which is really only useful once
72you free values into it:
73
74    >>> r = NumAlloc(0, -1)
75    >>> r
76    NumAlloc(0, -1)
77    >>> r.alloc() is None
78    True
79    >>> r.free_multi(range(50))
80    >>> r
81    NumAlloc(0, 49)
82
83Note that r.alloc() starts from where you last left off, even if
84you've freed a value:
85
86    >>> r.alloc()
87    0
88    >>> r.free(0)
89    >>> r.alloc()
90    1
91
92Of course, in multithreaded code you can't really depend on this
93since it will race other threads.  Still, it generally makes for
94efficient allocation.  To force allocation to start from the
95range's minimum, provide the minimum (e.g., r.min_val) as an
96argument to r.alloc():
97
98    >>> r.alloc()
99    2
100    >>> r.alloc(r.min_val)
101    0
102
103Providing a number to alloc() tries to allocate that number,
104but wraps around to the next one if needed:
105
106    >>> r.alloc(49)
107    49
108    >>> r.alloc(49)
109    3
110    >>> r.alloc(99999)
111    4
112    >>> r.avail
113    [[5, 48]]
114
115There is currently no way to find all allocated values, although
116the obvious method (going through r.avail) will work.  Any iterator
117would not be thread-safe.
118"""
119
120import threading
121
122class NumAlloc(object):
123    """
124    Number allocator object.
125    """
126    def __init__(self, min_val, max_val, autoextend=True):
127        self.min_val = min_val
128        self.max_val = max_val
129        if min_val <= max_val:
130            self.avail = [[min_val, max_val]]
131        else:
132            self.avail = []
133        self.autoextend = autoextend
134        self.last = None
135        self.lock = threading.Lock()
136
137    def __repr__(self):
138        myname = self.__class__.__name__
139        if self.autoextend:
140            ae = ''
141        else:
142            ae = ', autoextend=False'
143        return '{0}({1}, {2}{3})'.format(myname, self.min_val, self.max_val, ae)
144
145    def _find_block(self, val):
146        """
147        Find the block that contains val, or that should contain val.
148        Remember that self.avail is a list of avaliable ranges of
149        the form [[min1, max1], [min2, max2], ..., [minN, maxN]]
150        where max1 < min2, max2 < min3, ..., < minN.
151
152        The input value either falls into one of the available
153        blocks, or falls into a gap between two available blocks.
154        We want to know which block it goes in, or if it goes
155        between two, which block it comes before.
156
157        We can do a binary search to find this block.  When we
158        find it, return its index and its values.
159
160        If we find that val is not in a block, return the position
161        where the value should go, were it to be put into a new
162        block by itself.  E.g., suppose val is 17, and there is a
163        block [14,16] and a block [18,20]. We would make this
164        [14,16],[17,17],[18,20] by inserting [17,17] between them.
165        (Afterward, we will want to fuse all three blocks to make
166        [14,18].  However, if we insert as block 0, e.g., if the
167        list starts with [18,20] and we insert to get
168        [17,17][18,20], we really end up just modifying block 0 to
169        [17,20].  Or, if we insert as the new final block, we
170        might end up modifying the last block.)
171        """
172        low = 0
173        high = len(self.avail) - 1
174        while low <= high:
175            mid = low + ((high - low) // 2)
176            pair = self.avail[mid]
177            if val < pair[0]:
178                # must go before block mid
179                high = mid - 1
180            elif val > pair[1]:
181                # must go after block mid
182                low = mid + 1
183            else:
184                # val >= first and val <= last, so we found it
185                return mid, pair
186        # Low > high: no block actually contains val, or
187        # there are no blocks at all.  If there are no blocks,
188        # return block #0 and None.  Otherwise return the
189        return low, None
190
191    def alloc(self, val=None):
192        """
193        Get new available value.
194
195        If val is None, we start from the most recently
196        allocated value, plus 1.
197
198        If val is a numeric value, we start from that value.
199        Hence, since the range is min_val..max_val, you can
200        provide min_val to take the first available value.
201
202        This may return None, if no values are still available.
203        """
204        with self.lock:
205            if val is None:
206                val = self.last + 1 if self.last is not None else self.min_val
207            if val is None or val > self.max_val or val < self.min_val:
208                val = self.min_val
209            i, pair = self._find_block(val)
210            if pair is None:
211                # Value is is not available.  The next
212                # available value that is greater than val
213                # is in the block right after block i.
214                # If there is no block after i, the next
215                # available value is in block 0.  If there
216                # is no block 0, there are no available
217                # values.
218                nblocks = len(self.avail)
219                i += 1
220                if i >= nblocks:
221                    if nblocks == 0:
222                        return None
223                    i = 0
224                pair = self.avail[i]
225                val = pair[0]
226            # Value val is available - take it.
227            #
228            # There are four special cases to handle.
229            #
230            # 1. pair[0] < val < pair[1]: split the pair.
231            # 2. pair[0] == val < pair[1]: increase pair[0].
232            # 3. pair[0] == val == pair[1]: delete the pair
233            # 4. pair[0] < val == pair[1]: decrease pair[1].
234            assert pair[0] <= val <= pair[1]
235            if pair[0] == val:
236                # case 2 or 3: Take the left edge or delete the pair.
237                if val == pair[1]:
238                    del self.avail[i]
239                else:
240                    pair[0] = val + 1
241            else:
242                # case 1 or 4: split the pair or take the right edge.
243                if val == pair[1]:
244                    pair[1] = val - 1
245                else:
246                    newpair = [val + 1, pair[1]]
247                    pair[1] = val - 1
248                    self.avail.insert(i + 1, newpair)
249            self.last = val
250            return val
251
252    def free(self, val):
253        "Free one value"
254        self._free_multi('free', [val])
255
256    def free_multi(self, values):
257        "Free many values (provide any iterable)"
258        values = list(values)
259        values.sort()
260        self._free_multi('free_multi', values)
261
262    def _free_multi(self, how, values):
263        """
264        Free a (sorted) list of values.
265        """
266        if len(values) == 0:
267            return
268        with self.lock:
269            while values:
270                # Take highest value, and any contiguous lower values.
271                # Note that it can be significantly faster this way
272                # since coalesced ranges make for shorter copies.
273                highval = values.pop()
274                val = highval
275                while len(values) and values[-1] == val - 1:
276                    val = values.pop()
277                self._free_range(how, val, highval)
278
279    def _maybe_increase_max(self, how, val):
280        """
281        If needed, widen our range to include new high val -- i.e.,
282        possibly increase self.max_val.  Do nothing if this is not a
283        new all time high; fail if we have autoextend disabled.
284        """
285        if val <= self.max_val:
286            return
287        if self.autoextend:
288            self.max_val = val
289            return
290        raise ValueError('{0}: {1} is outside range limit'.format(how, val))
291
292    def _maybe_decrease_min(self, how, val):
293        """
294        If needed, widen our range to include new low val -- i.e.,
295        possibly decrease self.min_val.  Do nothing if this is not a
296        new all time low; fail if we have autoextend disabled.
297        """
298        if val >= self.min_val:
299            return
300        if self.autoextend:
301            self.min_val = val
302            return
303        raise ValueError('{0}: {1} is outside range limit'.format(how, val))
304
305    def _free_range(self, how, val, highval):
306        """
307        Free the range [val..highval].  Note, val==highval it's just
308        a one-element range.
309
310        The lock is already held.
311        """
312        # Find the place to store the lower value.
313        # We should never find an actual pair here.
314        i, pair = self._find_block(val)
315        if pair:
316            raise ValueError('{0}: {1} already available'.format(how, val))
317        # If we're freeing a range, check that the high val
318        # does not span into the *next* range, either.
319        if highval > val and i < len(self.avail):
320            if self.avail[i][0] <= highval:
321                raise ValueError('{0}: {2} (from {{1}..{2}) already '
322                                 'available'.format(how, val, highval))
323
324        # We'll need to insert a block and perhaps fuse it
325        # with blocks before and/or after.  First, check
326        # whether there *is* a before and/or after, and find
327        # their corresponding edges and whether we abut them.
328        if i > 0:
329            abuts_below = self.avail[i - 1][1] + 1 == val
330        else:
331            abuts_below = False
332        if i < len(self.avail):
333            abuts_above = self.avail[i][0] - 1 == highval
334        else:
335            abuts_above = False
336        # Now there are these four cases:
337        # 1. abuts below and above: fuse the two blocks.
338        # 2. abuts below only: adjust previous (i-1'th) block
339        # 3. abuts above only: adjust next (i'th) block
340        # 4. doesn't abut: insert new block
341        if abuts_below:
342            if abuts_above:
343                # case 1
344                self.avail[i - 1][1] = self.avail[i][1]
345                del self.avail[i]
346            else:
347                # case 2
348                self._maybe_increase_max(how, highval)
349                self.avail[i - 1][1] = highval
350        else:
351            if abuts_above:
352                # case 3
353                self._maybe_decrease_min(how, val)
354                self.avail[i][0] = val
355            else:
356                # case 4
357                self._maybe_decrease_min(how, val)
358                self._maybe_increase_max(how, highval)
359                newblock = [val, highval]
360                self.avail.insert(i, newblock)
361
362if __name__ == '__main__':
363    import doctest
364    import sys
365
366    doctest.testmod()
367    if sys.version_info[0] >= 3:
368        xrange = range
369    # run some worst case tests
370    # NB: coalesce is terribly slow when done bottom up
371    r = NumAlloc(0, 2**16 - 1)
372    for i in xrange(r.min_val, r.max_val, 2):
373        r.alloc(i)
374    print('worst case alloc: len(r.avail) = {0}'.format(len(r.avail)))
375    for i in xrange(r.max_val - 1, r.min_val, -2):
376        r.free(i)
377    print('free again; len(r.avail) should be 1; is {0}'.format(len(r.avail)))
378    if len(r.avail) != 1:
379        sys.exit('failure')
380