Experiments/icepack/pyice.py

#!/usr/bin/python3
# facebook t38065718
'''
PyICE is a memory mapped implementation of cached byte code.  Multiple Python
modules are compiled into a single memory mapped file.  The actual code objects
are stored in the memory mapped file.

PyICE will produce a file which is called an ice pack.  An ice pack consists of
multiple compiled Python byte codes.  An ice pack will further reduce memory
usage by folding constants across all of the modules.
'''

import io
import math
import mmap
import os
import re
import struct
import sys
import time
try:
    import _pyice
except ImportError:
    _pyice = None

from collections import deque, OrderedDict
from importlib.util import spec_from_file_location, decode_source
from os import path
from types import CodeType

'''
IcePack file format:

ICEPACK[byte version]

4 bytes - TIMESTAMP

Section offsets, an array of int32s that provide the offset to each section:
    Modules: A sorted tree of the module names provided by this icepack.
    Code Objects: The metadata for each code object.
    Strings: A table of strings which are referred to in the icepack
    Bytes: A table of bytes objects which are referred to in the icepack
    Ints: A table of integer values which don't fit in 24-bit values
    Big Ints: A table of integer values which don't fit in 32-bit values
    Floats: A table of floating point numbers in the icepack
    Complexs: A table of complex numbers
    Tuples: A table of tuple objects, composed from this or other tables
    FrozenSet: A table of frozen set objects

Sections
========

Modules
-------

A sorted list of the modules/packages.  Top-level packages/modules are listed
first, and then an offset to any child packages is provided.  If there are no
children then the offset will be 0.

Each section is organized as:
[uint32 module count], ([uint32 name reference], [uint32 code id],
                        [uint32 is_package], [uint32 str filename ref]
                        [uint32 child offset])*

The module count is the number of name reference / child offsets.  The
name reference is a reference into the string table.  The child offset is
the absolute offset from the start of the icepack to where anyh child modules
exist.

Code Objects
------------

Code objects are variable sized based upon the number of variables, etc...
which they reference.  Therefore this is broken into two parts.  The first part
is an array of uint32's which is the absolute offset to the code object.

[uint32 absolute offset]*

Following are the code objects themselves, laid out as:

[uint32 co_bytes offset to bytes table]
[uint32 co_argcount]
[uint32 co_kwonlyargcount]
[uint32 co_nlocals]
[uint32 co_stacksize]
[uint32 co_flags]
[uint32 co_firstlineno]
[uint32 co_name reference to str table]
[uint32 co_filename reference to str table]
[uint32 co_lnotab reference to bytes table]

[uint32 length] [uint32 co_cellvars reference to str table]*
[uint32 length] [uint32 co_freevars reference to str table]*
[uint32 length] [uint32 co_names reference to str table]*
[uint32 length] [uint32 co_varnames reference to str table]*
[uint32 length] [uint32 co_consts object references]*

Strings
-------
Strings are variable sized based upon the length of the string.
Therefore this is broken into two parts.  The first parts is an array of int32's
which is the absolute offset to the string contents.  Then the arrays of utf8
encoded bytes represent the string.

[uint32 count] number of strings
[uint32 absolute offset]*

[uint32 length] [byte utf8 encoded]*

Bytes
-----
Bytes are variable sized based upon the length of the byte string.
Therefore this is broken into two parts.  The first parts is an array of int32's
which is the absolute offset to the byte object.  Then the arrays of bytes
follow.

[uint32 count] number of byte objects
[uint32 absolute offset]*

[uint32 length] [byte]*

Ints
----
Ints are fixed 32-bit values.  They are encoded as an array of the integers.
They are preceeded by a uint32 count.

Big Ints
--------
Tuples are variable sized based upon the number of bytes they contain.
Therefore this is broken into two parts.  The first parts is an array of u
uint32's which is the absolute offset to the big int representation.  Then
the big ints follow.

[uint32 count] number of strings
[uint32 absolute offset]*

Big int's are stored as:

[uint32 length] [byte array]*

The byte array is in the format returned by int.to_bytes(), serialized as
little endian, and signed.

Floats
------
Floats are fixed sized double values.  They are encoded as an array of the
floating point value which can simply be indexed.  They are preceeded by a
uint32 count.

Complexes
---------
Complex's are fixed sized pairs of double values.  They are encoded as an array
of the real followed by the imaginary value stored as a double.  They are
preceeded by a int32 count.

Tuples
------
Tuples are variable sized based upon the number of items they contain.
Therefore this is broken into two parts.  The first parts is an array of
int32's which is the absolute offset to the tuple object.  Then the tuple
objects follow.

[uint32 count] number of tuples
[uint32 absolute offset]*

Tuples them selves are encoded as:

[uint32 length] [uint32 object references]*

Frozen Sets
-----------
Frozen sets are variable sized based upon the number of items they contain.
Therefore this is broken into two parts.  The first parts is an array of
int32's which is the absolute offset to the set object.  Then the set
objects follow.

[uint32 count] number of sets
[uint32 absolute offset]*

Sets them selves are encoded as:

[uint32 length] [uint32 object references]*

Object References
=================
Occasionally we need to embed an object reference to an untyped object.
This occurs in the co_const array on objects as well as for the elements
of a tuple.

References are encoded using a 32-bit value indicating the table to to
get the value from and the index into that table.  The low byte is used
to indicate the type of value and the value is encoded in the high 24
bits.

0: None
1: Named constant
    Upper 24-bits:
        0 = False
        1 = True
        2 = Ellipsis
2: small int, upper 24 bits is the int value
3: int32, upper 24-bits is a table reference
4: large int, upper 24-bits is a table reference
5: bytes, upper 24-bits is a table reference
6: str, upper 24-bits is a table reference
7: float, upper 24-bits is a table reference
8: complex, upper 24-bits is a table reference
9: tuple, upper 24-bits is a table reference
10: code, upper 24-bits is a table reference
'''

CODE_FORMAT = struct.Struct(
                'I' +  # code index in the bytes table
                'I' +  # co_argcount
                'I' +  # co_kwonlyargcount
                'I' +  # co_nlocals
                'I' +  # co_stacksize
                'I' +  # co_flags
                'I' +  # co_firstlineno
                'I' +  # co_name index in str table
                'I' +  # co_filename index in str table
                'I' +  # co_nlotab index in bytes table
                'I' +  # cellvars index in string array table
                'I' +  # freevars index in string array table
                'I' +  # names index in string array table
                'I' +  # varnames index in string array table
                'I' +  # consts index in tuple table
                ''
              )

INT = struct.Struct('i')
UINT = struct.Struct('I')
FLOAT = struct.Struct('d')
COMPLEX = struct.Struct('dd')
TIMESTAMP_OFFSET = 8
SECTION_OFFSET = 12
SECTION_FORMAT = struct.Struct('IIIIIIIIII')
MODULE_ENTRY = struct.Struct('IIIII')

EXTENSION = '.icepack'

class ModuleInfo:
    def __init__(self, code=-1, filename=None, is_package=False):
        self.children = OrderedDict()
        self.code = code
        self.filename = filename
        self.child_offset = 0
        self.is_package = is_package


OBJECT_TYPE_NONE = 0x00
OBJECT_TYPE_NAMED_CONSTANT = 0x01
OBJECT_TYPE_INT32 = 0x03
OBJECT_TYPE_BIGINT = 0x04
OBJECT_TYPE_BYTES = 0x05
OBJECT_TYPE_STR = 0x06
OBJECT_TYPE_FLOAT = 0x07
OBJECT_TYPE_COMPLEX = 0x08
OBJECT_TYPE_TUPLE = 0x09
OBJECT_TYPE_CODE = 0x0A
OBJECT_TYPE_FROZENSET = 0x0B

# Named constants
OBJECT_FALSE = 0x0001
OBJECT_TRUE = 0x0101
OBJECT_ELLIPSIS = 0x0201

def _float_equals(a, b):
    if math.isnan(a) and math.isnan(b):
        return True
    elif (a == 0 and b == 0 and
          math.copysign(1, a) != math.copysign(1, b)):
        return False
    else:
        return a == b

class PyObjectValue:
    '''handles equality with slightly different semantics than normal
Python ==.  Disallows equality between conflicting types (e.g. 0 != 0.0 !=
False).  Allows NaN == NaN and +0.0 != -0.0 for both floats and components
of complex numbers'''
    def __init__(self, value):
        assert type(value) is not ObjectValue
        if type(value) == tuple:
            self.value = tuple(ObjectValue(v) for v in value)
        elif type(value) == frozenset:
            self.value = frozenset(ObjectValue(v) for v in value)
        else:
            self.value = value
        self.hash = hash(self.value) ^ hash(type(self.value))

    def __repr__(self):
        return 'ObjectValue(' + repr(self.value) + ')'

    def __hash__(self):
        return self.hash

    def __eq__(self, other):
        if type(other) is not ObjectValue:
            return False
        elif type(self.value) is not type(other.value):
            return False
        elif type(self.value) == float:
            return _float_equals(self.value, other.value)
        elif type(self.value) == complex:
            return (_float_equals(self.value.real, other.value.real) and
                    _float_equals(self.value.imag, other.value.imag))

        return self.value == other.value


if _pyice is None:
    ObjectValue = PyObjectValue
    class IcePackError(Exception):
        pass
else:
    # Use the C accelerator version of ObjectValue
    ObjectValue = _pyice.CObjectValue
    IcePackError = _pyice.IcePackError




def _align_file(file, align=8):
    len = file.tell()
    padding = (((len + align - 1) & (~(align - 1))) - len)
    file.write(b'\x00' * padding)


class _TypeTable:
    section_count = 1

    def add(self, value) -> bool:
        raise NotImplementedError(str(type(self)))

    def write(self, value, outfile):
        raise NotImplementedError(str(type(self)))

    def write_table(self, value, outfile):
        raise NotImplementedError(str(type(self)))

    @staticmethod
    def add_const_to_table(table, value) -> bool:
        if value not in table:
            table[value] = len(table)
            return True
        return False

    @staticmethod
    def write_simple_table(table, outfile, get_bytes, align=4):
        values = [(y, x) for x, y in table.items()]
        values.sort()
        outfile.write(UINT.pack(len(values)))
        _align_file(outfile, align)
        for _index, value in values:
            outfile.write(get_bytes(value))

    @staticmethod
    def write_variable_len_table(table, outfile, get_bytes_and_len):
        '''writes the string table, in the format: (offset)* (utf8 value)*'''
        values = [(y, x, get_bytes_and_len(x)) for x, y in table.items()]
        values.sort()
        outfile.write(UINT.pack(len(values)))
        offset = outfile.tell() + len(table) * 4
        for _index, _value, (encoded, _length) in values:
            outfile.write(UINT.pack(offset))
            offset += len(encoded) + 4

        for _index, _value, (encoded, length) in values:
            outfile.write(UINT.pack(length))
            outfile.write(encoded)


class _ConstantTable(_TypeTable):
    section_count = 0

    def __init__(self, id):
        self.id = id

    def add(self, value) -> bool:
        return False

    def write(self, value, outfile):
        outfile.write(UINT.pack(self.id))

    def write_table(self, value, outfile):
        return ()


class _BoolTable(_TypeTable):
    section_count = 0

    def add(self, value) -> bool:
        return False

    def write(self, value, outfile):
        outfile.write(UINT.pack(OBJECT_TRUE if value else OBJECT_FALSE))

    def write_table(self, value, outfile):
        return ()


class _SimpleTable(_TypeTable):
    object_type: int

    def __init__(self):
        self.table = {}

    def add(self, value) -> bool:
        return self.add_const_to_table(self.table, value)

    def write(self, value, outfile):
        str_id = self.table[value]
        outfile.write(UINT.pack(str_id << 8 | self.object_type))

    def __getitem__(self, value):
        return self.table[value]


class _IntTable(_TypeTable):
    section_count = 2

    def __init__(self):
        self.int_table = {}
        self.bigint_table = {}

    def add(self, value):
        if -1 << 31 <= value <= (1 << 31) - 1:
            return self.add_const_to_table(self.int_table, value)
        else:
            return self.add_const_to_table(self.bigint_table, value)

    def write(self, value, outfile):
        if -1 << 31 <= value <= (1 << 31) - 1:
            index = self.int_table[value]
            outfile.write(UINT.pack(OBJECT_TYPE_INT32 | index << 8))
        else:
            index = self.bigint_table[value]
            outfile.write(UINT.pack(OBJECT_TYPE_BIGINT | index << 8))

    @staticmethod
    def bigint_to_bytes_and_len(value):
        length = (value.bit_length()+7)//8
        try:
            res = value.to_bytes(length, 'little', signed=True)
        except OverflowError:
            # Positive numbers like 0x80000000 will fail because bit_length
            # doesn't account for the need for a sign bit
            res = value.to_bytes(length + 1, 'little', signed=True)
        return res, len(res)

    def write_table(self, maker, outfile):
        int_offset = outfile.tell()
        self.write_simple_table(self.int_table, outfile, INT.pack)
        bigint_offset = outfile.tell()
        self.write_variable_len_table(self.bigint_table, outfile,
                                      self.bigint_to_bytes_and_len)
        return int_offset, bigint_offset


class _BytesTable(_SimpleTable):
    object_type = OBJECT_TYPE_BYTES

    def write_table(self, maker, outfile):
        offset = outfile.tell()
        self.write_variable_len_table(self.table, outfile,
                                      lambda value: (value, len(value)))
        return offset,


class _StrTable(_SimpleTable):
    object_type = OBJECT_TYPE_STR

    def write_table(self, maker, outfile):
        offset = outfile.tell()
        self.write_variable_len_table(self.table, outfile,
                                      self.get_bytes_and_len)
        return offset,

    def get_bytes_and_len(self, value):
        # Null terminate strings for making life easier in dealing
        # w/ the module table lookups
        res = value.encode('utf8', 'surrogatepass')

        return res + b'\0', len(res)


class _ObjectValueTable(_SimpleTable):
    def add(self, value):
        return super().add(ObjectValue(value))

    def write(self, value, outfile):
        super().write(ObjectValue(value), outfile)

    def __getitem__(self, value):
        return super().__getitem__(ObjectValue(value))


class _FloatTable(_ObjectValueTable):
    object_type = OBJECT_TYPE_FLOAT

    def write_table(self, maker, outfile):
        offset = outfile.tell()
        self.write_simple_table(self.table, outfile,
                                lambda value: FLOAT.pack(value.value), 8)
        return offset,


class _ComplexTable(_ObjectValueTable):
    object_type = OBJECT_TYPE_COMPLEX

    def write_table(self, maker, outfile):
        offset = outfile.tell()
        def writer(value):
            return COMPLEX.pack(value.value.real, value.value.imag)
        self.write_simple_table(self.table, outfile, writer, 8)
        return offset,


class _SequenceTable(_ObjectValueTable):
    def __init__(self, object_type, maker):
        super().__init__()
        self.maker = maker
        self.object_type = object_type

    def add(self, value) -> bool:
        # TODO: Fixme
        if super().add(value):
            for x in value:
                self.maker.add_const(x)
            return True
        return False

    def write_table(self, maker, outfile):
        table_offset = outfile.tell()
        values = [(y, x) for x, y in self.table.items()]
        values.sort()
        outfile.write(UINT.pack(len(values)))
        offset = outfile.tell() + len(self.table) * 4
        for _index, value in values:
            outfile.write(UINT.pack(offset))
            offset += 4 + len(value.value) * 4

        for _index, value in values:
            maker.write_array(value.value)
        return table_offset,


class _CodeTable(_TypeTable):
    def __init__(self, maker):
        self.table = {}
        self.maker = maker

    @property
    def count(self):
        return len(self.table)

    def code_id(self, code):
        '''Python has very weird equality semantics around code objects, so that
it doesn't compare the filename or the co_lnotab fields to determine if they're
equal.  That results in oddities such as:

>>> def f():
...     x = 1
...     x = 2
...
>>> g = f
>>> def f():
...     x = 1
...
...     x = 2
...
>>> f.__code__ == g.__code__
True

So we always consider object identify to ensure we don't merge code objects
together and end up with wrong filename or line number information.
'''
        return id(code), code

    def __getitem__(self, code):
        return self.table[self.code_id(code)]

    def write(self, value, outfile):
        code_id = self.code_id(value)
        outfile.write(UINT.pack(self.table[code_id] << 8 | OBJECT_TYPE_CODE))
        pass

    def add(self, value):
        code_id = self.code_id(value)
        if code_id not in self.table:
            self.table[code_id] = len(self.table)
            self.maker.add_const(value.co_code)
            self.maker.enqueue_code(value)

    def write_table(self, maker, outfile):
        table_offset = outfile.tell()
        code_offsets = []
        # write space for offsets to code objects
        outfile.write(INT.pack(len(self.table)))
        code_start = outfile.tell()
        outfile.truncate(code_start + 4 * len(self.table))
        outfile.seek(0, 2)
        codes = [(y, x) for (_, x), y in self.table.items()]
        codes.sort()
        for _i, code in codes:
            code_offsets.append(outfile.tell())
            header = CODE_FORMAT.pack(
                                      maker.bytes[code.co_code],
                                      code.co_argcount,
                                      code.co_kwonlyargcount,
                                      code.co_nlocals,
                                      code.co_stacksize,
                                      code.co_flags,
                                      code.co_firstlineno,
                                      maker.strs[code.co_name],
                                      maker.strs[code.co_filename],
                                      maker.bytes[code.co_lnotab],
                                      maker.get_tuple_id(code.co_cellvars),
                                      maker.get_tuple_id(code.co_freevars),
                                      maker.get_tuple_id(code.co_names),
                                      maker.get_tuple_id(code.co_varnames),
                                      maker.get_tuple_id(code.co_consts),
            )
            outfile.write(header)

        outfile.seek(code_start)
        for offset in code_offsets:
            outfile.write(UINT.pack(offset))

        outfile.seek(0, io.SEEK_END)
        return table_offset,


class IceMaker:
    '''Generates an ice pack from a set of modules and saves it to a file like
object.'''

    def __init__(self, outfile):
        '''Creates a new IceMaker which will save the contents of the provided
modules to outfile which should be a seekable file-like object'''
        self.outfile = outfile
        self.consts = set()
        self.queue = deque()
        self.modules = {}       # all modules, e.g. a.b.c -> ModuleInfo
        self.timestamp = 0

        self.codes = _CodeTable(self)
        self.strs = _StrTable()
        self.tuples = _SequenceTable(OBJECT_TYPE_TUPLE, self)
        self.bytes = _BytesTable()
        self.ints = _IntTable()
        self.floats = _FloatTable()
        self.complexes = _ComplexTable()
        self.frozensets = _SequenceTable(OBJECT_TYPE_FROZENSET, self)
        self.type_handlers = {
            type(None): _ConstantTable(0),
            type(Ellipsis): _ConstantTable(OBJECT_ELLIPSIS),
            bool: _BoolTable(),
            int: self.ints,
            bytes: self.bytes,
            str: self.strs,
            float: self.floats,
            complex: self.complexes,
            CodeType: self.codes,
            ObjectValue: None,
            tuple: self.tuples,
            frozenset: self.frozensets,
        }

        # Ensure we have an empty code for namespace modules
        self.empty_code = compile('', '', 'exec', dont_inherit=True)
        self.enqueue_code(self.empty_code)

    def add_module(self, code, name, filename, is_package=False, timestamp=0):
        if timestamp > self.timestamp:
            self.timestamp = timestamp
        self.enqueue_code(code)
        self.modules[name] = ModuleInfo(self.codes[code],
                                        filename, is_package)
        self.add_const(filename)
        for name_part in name.split('.'):
            self.add_const(name_part)
        self.process()

    def enqueue_code(self, code):
        self.add_const(code)

        # TODO: Consider the order of serialization.  Currently we're doing
        # breadth first, which means children won't be near their parents,
        # probably resulting in more pages being read in when not all modules
        # are used.  Switching to an appendleft here would result in children
        # being closer to their parents, but in opposite of the order they
        # appear in co_consts.  And the order in co_consts appears to be the
        # order they're referred to in code, so we could end up with extra
        # seeking (which might not really matter)
        self.queue.append(code)

    def write_str(self, value):
        self.outfile.write(UINT.pack(self.strs[value]))

    def add_const(self, const):
        handler = self.type_handlers[type(const)]
        return handler.add(const)

    def process(self):
        while self.queue:
            code = self.queue.popleft()

            self.add_const(code.co_filename)
            self.add_const(code.co_name)
            self.add_const(code.co_lnotab)
            self.add_const(code.co_names)
            self.add_const(code.co_varnames)
            self.add_const(code.co_cellvars)
            self.add_const(code.co_freevars)
            self.add_const(code.co_consts)

    def write_object_value(self, value):
        if type(value.value) == tuple:
            tuple_id = self.tuples.table[value]
            self.outfile.write(UINT.pack(tuple_id << 8 | OBJECT_TYPE_TUPLE))
        elif type(value.value) == frozenset:
            set_id = self.frozensets.table[value]
            self.outfile.write(UINT.pack(set_id << 8 | OBJECT_TYPE_FROZENSET))
        else:
            self.write_reference(value.value)

    def write_reference(self, value):
        if type(value) is ObjectValue:
            self.write_object_value(value)
        else:
            handler = self.type_handlers[type(value)]
            handler.write(value, self.outfile)

    def write_array(self, arr):
        self.outfile.write(UINT.pack(len(arr)))
        for value in arr:
            self.write_reference(value)

    def write_str_array(self, arr):
        self.outfile.write(UINT.pack(len(arr)))
        for value in arr:
            self.write_str(value)

    def write(self):
        # Write the tables
        sections = [
            self.codes,
            self.strs,
            self.bytes,
            self.ints,
            self.floats,
            self.complexes,
            self.tuples,
            self.frozensets,
        ]
        sec_count = sum(section.section_count for section in sections) + 1
        self.outfile.write(b'ICEPACK\x00')    # write header and version

        self.outfile.write(UINT.pack(self.timestamp))

        # Get space for section offsets
        offset_start = self.outfile.tell()
        self.outfile.write(b'\0\0\0\0' * (sec_count))

        # Then write the sections
        offsets = []

        _align_file(self.outfile)
        offsets.append(self.outfile.tell())
        self.write_modules()

        for section in sections:
            _align_file(self.outfile)
            offsets.extend(section.write_table(self, self.outfile))

        # Then update the section offsets
        self.outfile.seek(offset_start)

        for offset in offsets:
            self.outfile.write(UINT.pack(offset))
        self.outfile.seek(0, io.SEEK_END)

    def make_module_tree(self):
        tree = ModuleInfo()
        all_modules = list(self.modules.items())
        all_modules.sort()
        # First build a tree of all modules...
        for mod_name, mod_info in all_modules:
            cur = tree
            names = mod_name.split('.')
            for name in names[:-1]:
                next = cur.children.get(name)
                if next is None:

                    ns_module = ModuleInfo(self.codes[self.empty_code], '',
                                           is_package=True)
                    next = cur.children[name] = ns_module
                cur = next
            cur.children[names[-1]] = mod_info

        self.calculate_module_relative_offsets(tree)
        return tree

    def calculate_module_relative_offsets(self, tree, offset=0):
        '''Recurses through the tree and calculates the relative offsets for
where their children will live'''

        # space for the # of children, and their names/code/child offset
        offset += 4 + len(tree.children) * MODULE_ENTRY.size
        for _name, item in tree.children.items():
            if item.children:
                item.child_offset = offset
                offset = self.calculate_module_relative_offsets(item,
                                                                offset)

        return offset

    def write_modules(self):
        '''We write out the module table as a sorted tree that we can binary
search.  The format at each level is:
count, (name_offset, children_offset)*
If a level has no children 0 is written'''
        self.write_module(self.make_module_tree(), self.outfile.tell())

    def write_module(self, tree, base_offset):
        # Write this entry
        self.outfile.write(UINT.pack(len(tree.children)))
        for name, item in tree.children.items():
            self.write_str(name)

            if item.code == -1:
                self.outfile.write(UINT.pack(0xffffffff))
            else:
                self.outfile.write(UINT.pack(item.code))
            self.outfile.write(UINT.pack(1 if item.is_package else 0))
            self.write_str(item.filename)
            if item.child_offset == 0:
                self.outfile.write(b'\0\0\0\0')
            else:
                self.outfile.write(UINT.pack(base_offset + item.child_offset))

        # Then write the children
        for item in tree.children.values():
            if item.children:
                self.write_module(item, base_offset)

    def get_tuple_id(self, value):
        return self.tuples[value]


class PyIceBreaker:
    def __init__(self, icepack, base_dir=''):
        self.file = icepack
        self.base_dir = base_dir
        self.map = mmap.mmap(self.file.fileno(),
                             length=0, access=mmap.ACCESS_READ)
        self.str_cache = {}
        self.bytes_cache = {}
        self.const_cache = {}
        self.code_cache = {}
        self.tuple_cache = {}
        self.int_cache = {}
        self.float_cache = {}
        self.complex_cache = {}
        self.bigint_cache = {}
        self.frozenset_cache = {}
        header = self.map[0:7]
        if header != b'ICEPACK':
            raise IcePackError('Invalid ice pack file: ' + repr(header))
        version = self.map[7]
        if version != 0:
            raise IcePackError('Unsupported IcePack version')

        ts_bytes = self.map[TIMESTAMP_OFFSET:SECTION_OFFSET]
        self.timestamp, = UINT.unpack(ts_bytes)

        sec_data = self.map[SECTION_OFFSET:
                            SECTION_OFFSET + SECTION_FORMAT.size]
        sections = SECTION_FORMAT.unpack(sec_data)

        (self.modules, self.code, self.strings, self.bytes, self.ints,
         self.bigints, self.floats, self.complexes, self.tuples,
         self.frozensets) = sections

        self.bytes_count, = UINT.unpack(self.map[self.bytes:self.bytes+4])
        self.str_count, = UINT.unpack(self.map[self.strings:self.strings+4])
        self.code_count, = UINT.unpack(self.map[self.code:self.code+4])
        self.int_count, = UINT.unpack(self.map[self.ints:self.ints+4])
        self.bigint_count, = UINT.unpack(self.map[self.bigints:self.bigints+4])
        self.float_count, = UINT.unpack(self.map[self.floats:self.floats+4])
        self.complex_count, = UINT.unpack(self.map[self.complexes:
                                                   self.complexes+4])
        self.tuple_count, = UINT.unpack(self.map[self.tuples:self.tuples+4])

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self.file.close()

    def read_none(self, const):
        if const != 0:
            raise ValueError('Invalid const value')
        return None

    def read_named_constant(self, const):
        if const <= 2:
            return (False, True, ...)[const]
        raise ValueError('Invalid const value')

    def read_bytes(self, index):
        if index > self.bytes_count:
            raise ValueError('Invalid bytes index')

        res = self.bytes_cache.get(index)
        if res is None:
            start = self.bytes + 4 + index * 4
            location, = UINT.unpack(self.map[start:start + 4])
            len, = UINT.unpack(self.map[location:location + 4])
            res = self.map[location + 4:location + 4 + len]
            self.bytes_cache[index] = res

        return res

    def read_str(self, index):
        if index > self.str_count:
            raise ValueError('Invalid str index ' + str(index))

        res = self.str_cache.get(index)
        if res is None:
            start = self.strings + 4 + index * 4
            location, = UINT.unpack(self.map[start:start + 4])
            len, = UINT.unpack(self.map[location:location + 4])
            utf8 = self.map[location + 4:location + 4 + len]
            self.str_cache[index] = res = utf8.decode('utf8',
                                                      'surrogatepass')
        return res

    def read_const(self, const):
        return _CONST_READERS[const & 0xff](self, const >> 8)

    def read_const_array(self, offset):
        count, = UINT.unpack(self.map[offset:offset + 4])
        o = offset + 4  # starting offset to the actual elements
        map = self.map
        return (self.read_const(UINT.unpack(map[o + i * 4:o + 4 + i * 4])[0])
                for i in range(count))

    def read_tuple(self, index):
        if index > self.tuple_count:
            raise ValueError('Invalid tuple index')

        res = self.tuple_cache.get(index)
        if res is None:
            start = self.tuples + 4 + index * 4
            location, = UINT.unpack(self.map[start:start + 4])
            res = tuple(self.read_const_array(location))
            self.tuple_cache[index] = res

        return res

    def read_frozenset(self, index):
        if index > self.tuple_count:
            raise ValueError('Invalid frozenset index')

        res = self.frozenset_cache.get(index)
        if res is None:
            start = self.frozensets + 4 + index * 4
            location, = UINT.unpack(self.map[start:start + 4])
            res = frozenset(self.read_const_array(location))
            self.frozenset_cache[index] = res

        return res

    def read_int(self, index):
        if index > self.int_count:
            raise ValueError('Invalid int index')

        res = self.int_cache.get(index)
        if res is None:
            start = self.ints + 4 + index * INT.size
            res, = INT.unpack(self.map[start:start + 4])
            self.int_cache[index] = res
        return res

    def read_float(self, index):
        if index > self.float_count:
            raise ValueError('Invalid float index')

        res = self.float_cache.get(index)
        if res is None:
            start = self.floats + 4 + 4 + index * FLOAT.size
            res, = FLOAT.unpack(self.map[start:start + FLOAT.size])
            self.float_cache[index] = res
        return res

    def read_complex(self, index):
        if index > self.complex_count:
            raise ValueError('Invalid complex index')

        res = self.complex_cache.get(index)
        if res is None:
            start = self.complexes + 4 + 4 + index * COMPLEX.size
            real, imag = COMPLEX.unpack(self.map[start:start + COMPLEX.size])
            res = self.complex_cache[index] = complex(real, imag)
        return res

    def read_bigint(self, index):
        if index > self.bigint_count:
            raise ValueError('Invalid bigint index')

        res = self.bigint_cache.get(index)
        if res is None:
            start = self.bigints + 4 + index * 4
            location, = UINT.unpack(self.map[start:start + 4])
            len, = UINT.unpack(self.map[location:location + 4])
            bigint_bytes = self.map[location + 4:location + 4 + len]
            res = int.from_bytes(bigint_bytes, 'little', signed=True)
            self.bigint_cache[index] = res
        return res

    def read_code(self, index):
        if index > self.code_count:
            raise ValueError('Invalid code index')
        start = self.code + 4 + index * 4
        location, = UINT.unpack(self.map[start:start + 4])
        header = self.map[location:location + CODE_FORMAT.size]
        (bytes, argcount, kwonlyargcount, nlocals, stacksize, flags,
         firstlineno, name, filename, lnotab, cellvars, freevars, names,
         varnames, consts) = CODE_FORMAT.unpack(header)
        code = self.get_code_buffer(bytes)
        cellvars = self.read_tuple(cellvars)
        freevars = self.read_tuple(freevars)
        names = self.read_tuple(names)
        varnames = self.read_tuple(varnames)
        consts = self.read_tuple(consts)
        fixed_fn = path.join(self.base_dir, self.read_str(filename))
        return CodeType(argcount, kwonlyargcount, nlocals, stacksize, flags,
                        code, consts, names, varnames,
                        fixed_fn, self.read_str(name),
                        firstlineno, self.read_bytes(lnotab), freevars,
                        cellvars)

    def get_code_buffer(self, index):
        start = self.bytes + 4 + index * 4
        location, = UINT.unpack(self.map[start:start + 4])
        len, = UINT.unpack(self.map[location:location + 4])
        return memoryview(self.map)[location + 4:location + 4 + len]

    def find_module(self, name):
        '''Finds a module in the module tree.  Returns a tuple of the code
and a bool indicating if the module is a package'''
        parts = name.split('.')
        cur = self.modules
        res = None
        for part in parts:
            if cur == 0:
                # Previous loop had no children
                return None
            count, = UINT.unpack(self.map[cur:cur+4])
            for i in range(count):
                start = cur + 4 + i * MODULE_ENTRY.size
                entry_bytes = self.map[start:start + MODULE_ENTRY.size]
                (iname, code, is_package, filename,
                 children) = MODULE_ENTRY.unpack(entry_bytes)
                if self.read_str(iname) == part:
                    cur = children
                    res = code
                    break
            else:
                res = None

        if res is not None:
            return (self.read_code(res), is_package, self.read_str(filename))

        return None

    def close(self):
        self.file.close()


IceBreaker = PyIceBreaker

if _pyice is not None:
    # Use the accelerator version if it's available
    class IceBreaker(_pyice.CIceBreaker):
        def __new__(cls, icepack, base_dir=''):
            map = mmap.mmap(icepack.fileno(),
                            length=0, access=mmap.ACCESS_READ)
            if (isinstance(base_dir, str) and base_dir and
                not base_dir.endswith(path.sep)):
                base_dir += path.sep
            self = super().__new__(cls, map, base_dir)
            self.map = map
            self.file = icepack
            return self

        def __enter__(self):
            return self

        def __exit__(self, type, value, traceback):
            # Can't close this until all of the code objects are freed
            super().__exit__(type, value, traceback)
            self.file.close()

_CONST_READERS = {
    OBJECT_TYPE_NONE: PyIceBreaker.read_none,
    OBJECT_TYPE_NAMED_CONSTANT: PyIceBreaker.read_named_constant,
    OBJECT_TYPE_INT32: PyIceBreaker.read_int,
    OBJECT_TYPE_BIGINT: PyIceBreaker.read_bigint,
    OBJECT_TYPE_BYTES: PyIceBreaker.read_bytes,
    OBJECT_TYPE_STR: PyIceBreaker.read_str,
    OBJECT_TYPE_FLOAT: PyIceBreaker.read_float,
    OBJECT_TYPE_COMPLEX: PyIceBreaker.read_complex,
    OBJECT_TYPE_TUPLE: PyIceBreaker.read_tuple,
    OBJECT_TYPE_CODE: PyIceBreaker.read_code,
    OBJECT_TYPE_FROZENSET: PyIceBreaker.read_frozenset,
}

class Freezer:
    def __init__(self, output, modules, optimize, exclude, verbose):
        self.modules = modules
        self.optimize = optimize
        self.exclude = exclude
        self.verbose = verbose
        self.outfile = open(output, 'wb')
        self.maker = IceMaker(self.outfile)

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.outfile.close()

    def build_file(self, basedir, fullpath):
        dir, file = path.split(fullpath)
        if file == "__init__.py":
            relname = path.relpath(dir,  basedir)
            module_name = relname.replace('/', '.').replace('\\', '.')
            is_package = True
        else:
            relname = path.relpath(path.splitext(fullpath)[0], basedir)
            module_name = relname.replace('/', '.').replace('\\', '.')
            is_package = False
        relfn = path.relpath(fullpath, basedir)

        for exclusion in self.exclude:
            if re.match(exclusion, relfn):
                if self.verbose:
                    print('Skipping', relfn)
                break
        else:
            if self.verbose:
                print('Including', module_name, 'from', relfn)
            with open(fullpath, 'rb') as inp:
                try:
                    bytes = inp.read()
                    compiled = compile(bytes, relfn,
                                       'exec', dont_inherit=True,
                                       optimize=self.optimize)
                    timestamp = int(os.stat(fullpath).st_mtime)
                    self.maker.add_module(compiled, module_name, relfn,
                                     is_package, timestamp)
                except SyntaxError as se:
                    if self.verbose:
                        print('Ignoring module with error: ', relfn, se)

    def build_dir(self, dir):
        for dirpath, _dirnames, filenames in os.walk(dir):
            for filename in filenames:
                if not filename.endswith('.py'):
                    continue

                fullpath = path.join(dirpath, filename)
                self.build_file(dir, fullpath)

    def freeze(self):
        start = time.time()
        for module in self.modules:
            if path.isdir(module):
                if self.verbose:
                    print('Including directory', module)
                self.build_dir(module)
            else:
                self.build_file(path.dirname(module), module)

            self.maker.write()
        end = time.time()
        print('IcePack built in', end - start, 'seconds')


def main():
    args = parser.parse_args()
    if not args.modules:
        print('No modules specified!')
        sys.exit(1)

    with Freezer(args.output, args.modules, args.optimize,
                 args.exclude or (), args.verbose) as freezer:
        freezer.freeze()


class PyIceImporter:
    def __init__(self, import_path):
        self.path = import_path
        try:
            if (EXTENSION + '/') in import_path:
                # sys.path entry should be
                # 'path/to/compiled.icepack//relative/loc'
                components = import_path.split(EXTENSION + '/')
                pack_name = components[0] + EXTENSION
                if path.isfile(pack_name):
                    self.disk_loc = components[1]
                    self.breaker = IceBreaker(open(pack_name, 'rb'),
                                              self.disk_loc)
                    return
        except IcePackError as e:
            print('failed to load ice pack (invalid)', e)
        except OSError as e:
            print('failed to load ice pack: ' + str(e), e)

        raise ImportError()

    def find_spec(self, fullname, target=None):
        if '\x00' in fullname:
            # Invalid module name, return None, and let the import machinery
            # report the module as not found.
            return None

        mod_info = self.breaker.find_module(fullname)
        if mod_info is None:
            return None

        mod, is_package, filename = mod_info

        disk_loc = path.join(self.disk_loc, fullname.replace('.', '/'))
        if filename:
            file_path = path.join(self.disk_loc, filename)
            try:
                mtime = os.stat(file_path).st_mtime
                if int(mtime) > self.breaker.timestamp:
                    # the file on disk has been updated since the icepack was
                    # generated, prefer the on-disk version.
                    return None
            except OSError:
                # no file on disk, use the icepack
                pass
        else:
            # namespace package
            file_path = None

        if is_package:
            search = [self.path, disk_loc]
        else:
            search = None
        loader = PyIceLoader(mod, self, file_path, is_package)
        spec = spec_from_file_location(fullname, file_path,
                                       loader=loader,
                                       submodule_search_locations=search)
        if not file_path:
            spec.has_location = False
        return spec


class PyIceLoader:
    def __init__(self, code, importer, filename, is_package):
        self.code = code
        self.importer = importer
        self.path = filename
        self._is_package = is_package

    def create_module(self, spec):
        return None

    def exec_module(self, mod):
        # self.path is the empty string for namespace packages, which don't
        # get a __file__ attribute
        if self.path:
            mod.__file__ = self.path
        exec(self.code, mod.__dict__)

    def is_package(self, fullname):
        return self._is_package

    def get_code(self, x):
        return self.code

    def get_source(self, fullname):
        if not self.path:
            # matching behavior of _NamespaceLoader in _bootstrap_external
            return ''

        with open(self.path, 'rb') as f:
            return decode_source(f.read())

    def get_data(self, path):
        """Return the data from path as raw bytes."""
        with open(path, 'rb') as file:
            return file.read()

    def get_filename(self, fullname):
        return self.path


def install():
    sys.path_hooks.append(PyIceImporter)


def uninstall():
    sys.path_hooks.remove(PyIceImporter)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(
        description='''
    PyICE - Produces an icepack, a memory mappable pre-compiled modules.
    ''')
    parser.add_argument('--exclude', type=str, nargs='+',
                        dest='exclude',
                        help='Adds a module to be excluded in the icepack.')
    parser.add_argument('--optimize',
                        default=-1,
                        type=int,
                        action='store',
                        dest='optimize',
                        help='The optimization level (-1, 0, 1 or 2).')
    parser.add_argument('--verbose',
                        default=False,
                        action='store_true',
                        dest='verbose',
                        help='Enable verbose output.')
    parser.add_argument('modules',
                        nargs='*',
                        help='Directories or files to be included in the IcePack.')
    parser.add_argument('--output', type=str,
                        default='out' + EXTENSION,
                        help='The destination filename')

    main()