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