#!/usr/local/bin/python class ip_allocator: def __init__(self, base=0, size=4294967296, min_alloc=256): self.base = base self.size = size self.min_alloc = min_alloc i = self.min_alloc self.free = { } self.allocated = { } while i < self.size: self.free[i] = set() self.allocated[i] = set() i *= 2 self.free[self.size] = set([self.base]) self.allocated[self.size] = set() def split(self, base, size): return ( base, base + size) def make_chunk(self, size): if not self.free[size * 2]: if size < self.size: if not self.make_chunk(size * 2): return False else: return False base = min(self.free[size*2]) self.free[size*2].discard(base) for b in self.split(base, size): self.free[size].add(b) return True def allocate(self, size): # Make sure size is a power of 2 at least as big as minalloc ns = self.min_alloc while ns < self.size: if ns == size: break if ns > size: size = ns break ns *= 2 else: size = self.size if self.free[size] or self.make_chunk(size): rv = min(self.free[size]) self.free[size].discard(rv) self.allocated[size].add(rv) return (rv, size) else: return None def deallocate(self, base, size): if base not in self.allocated[size]: raise KeyError("Not allocated") if base % (size * 2) == 0: if base + size in self.free[size]: self.allocated[size].discard(base) self.allocated[size * 2].add(base) self.free[size].discard(base+size) self.deallocate(base, size *2) else: self.allocated[size].discard(base) self.free[size].add(base) else: if base - size in self.free[size]: self.allocated[size].discard(base) self.allocated[size * 2].add(base - size) self.free[size].discard(base - size) self.deallocate(base-size, size * 2) else: self.allocated[size].discard(base) self.free[size].add(base) return True def __str__(self): rv = "" i=self.size while i >= self.min_alloc: rv += "%d: %s\n" % (i, self.free[i]) i /= 2 return rv