source: fedd/deter/ip_allocator.py @ 187a8f9

#!/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