~bzr-pqm/bzr/bzr.dev

# Copyright (C) 2008, 2009, 2010 Canonical Ltd
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

"""Compiled extensions for doing compression."""

#python2.4 support
cdef extern from "python-compat.h":
    pass


cdef extern from "Python.h":
    ctypedef struct PyObject:
        pass
    ctypedef int Py_ssize_t # Required for older pyrex versions
    int PyString_CheckExact(object)
    char * PyString_AS_STRING(object)
    Py_ssize_t PyString_GET_SIZE(object)
    object PyString_FromStringAndSize(char *, Py_ssize_t)


cdef extern from *:
    ctypedef unsigned long size_t
    void * malloc(size_t) nogil
    void * realloc(void *, size_t) nogil
    void free(void *) nogil
    void memcpy(void *, void *, size_t) nogil


cdef extern from "delta.h":
    struct source_info:
        void *buf
        unsigned long size
        unsigned long agg_offset
    struct delta_index:
        pass
    ctypedef enum delta_result:
        DELTA_OK
        DELTA_OUT_OF_MEMORY
        DELTA_INDEX_NEEDED
        DELTA_SOURCE_EMPTY
        DELTA_SOURCE_BAD
        DELTA_BUFFER_EMPTY
        DELTA_SIZE_TOO_BIG
    delta_result create_delta_index(source_info *src,
                                    delta_index *old,
                                    delta_index **fresh) nogil
    delta_result create_delta_index_from_delta(source_info *delta,
                                               delta_index *old,
                                               delta_index **fresh) nogil
    void free_delta_index(delta_index *index) nogil
    delta_result create_delta(delta_index *indexes,
                              void *buf, unsigned long bufsize,
                              unsigned long *delta_size,
                              unsigned long max_delta_size,
                              void **delta_data) nogil
    unsigned long get_delta_hdr_size(unsigned char **datap,
                                     unsigned char *top) nogil
    unsigned long sizeof_delta_index(delta_index *index)
    Py_ssize_t DELTA_SIZE_MIN


cdef void *safe_malloc(size_t count) except NULL:
    cdef void *result
    result = malloc(count)
    if result == NULL:
        raise MemoryError('Failed to allocate %d bytes of memory' % (count,))
    return result


cdef void *safe_realloc(void * old, size_t count) except NULL:
    cdef void *result
    result = realloc(old, count)
    if result == NULL:
        raise MemoryError('Failed to reallocate to %d bytes of memory'
                          % (count,))
    return result


cdef int safe_free(void **val) except -1:
    assert val != NULL
    if val[0] != NULL:
        free(val[0])
        val[0] = NULL

def make_delta_index(source):
    return DeltaIndex(source)


cdef object _translate_delta_failure(delta_result result):
    if result == DELTA_OUT_OF_MEMORY:
        return MemoryError("Delta function failed to allocate memory")
    elif result == DELTA_INDEX_NEEDED:
        return ValueError("Delta function requires delta_index param")
    elif result == DELTA_SOURCE_EMPTY:
        return ValueError("Delta function given empty source_info param")
    elif result == DELTA_SOURCE_BAD:
        return RuntimeError("Delta function given invalid source_info param")
    elif result == DELTA_BUFFER_EMPTY:
        return ValueError("Delta function given empty buffer params")
    return AssertionError("Unrecognised delta result code: %d" % result)


cdef class DeltaIndex:

    # We need Pyrex 0.9.8+ to understand a 'list' definition, and this object
    # isn't performance critical
    # cdef readonly list _sources
    cdef readonly object _sources
    cdef source_info *_source_infos
    cdef delta_index *_index
    cdef public unsigned long _source_offset
    cdef readonly unsigned int _max_num_sources

    def __init__(self, source=None):
        self._sources = []
        self._index = NULL
        self._max_num_sources = 65000
        self._source_infos = <source_info *>safe_malloc(sizeof(source_info)
                                                        * self._max_num_sources)
        self._source_offset = 0

        if source is not None:
            self.add_source(source, 0)

    def __sizeof__(self):
        # We want to track the _source_infos allocations, but the referenced
        # void* are actually tracked in _sources itself.
        # XXX: Cython is capable of doing sizeof(class) and returning the size
        #      of the underlying struct. Pyrex (<= 0.9.9) refuses, so we need
        #      to do it manually. *sigh* Note that we might get it wrong
        #      because of alignment issues.
        cdef Py_ssize_t size
        # PyObject start, vtable *, 3 object pointers, 2 C ints
        size = ((sizeof(PyObject) + sizeof(void*) + 3*sizeof(PyObject*)
                 + sizeof(unsigned long)
                 + sizeof(unsigned int))
                + (sizeof(source_info) * self._max_num_sources)
                + sizeof_delta_index(self._index))
        return size

    def __repr__(self):
        return '%s(%d, %d)' % (self.__class__.__name__,
            len(self._sources), self._source_offset)

    def __dealloc__(self):
        if self._index != NULL:
            free_delta_index(self._index)
            self._index = NULL
        safe_free(<void **>&self._source_infos)

    def _has_index(self):
        return (self._index != NULL)

    def add_delta_source(self, delta, unadded_bytes):
        """Add a new delta to the source texts.

        :param delta: The text of the delta, this must be a byte string.
        :param unadded_bytes: Number of bytes that were added to the source
            that were not indexed.
        """
        cdef char *c_delta
        cdef Py_ssize_t c_delta_size
        cdef delta_index *index
        cdef delta_result res
        cdef unsigned int source_location
        cdef source_info *src
        cdef unsigned int num_indexes

        if not PyString_CheckExact(delta):
            raise TypeError('delta is not a str')

        source_location = len(self._sources)
        if source_location >= self._max_num_sources:
            self._expand_sources()
        self._sources.append(delta)
        c_delta = PyString_AS_STRING(delta)
        c_delta_size = PyString_GET_SIZE(delta)
        src = self._source_infos + source_location
        src.buf = c_delta
        src.size = c_delta_size
        src.agg_offset = self._source_offset + unadded_bytes
        with nogil:
            res = create_delta_index_from_delta(src, self._index, &index)
        if res != DELTA_OK:
            raise _translate_delta_failure(res)
        self._source_offset = src.agg_offset + src.size
        if index != self._index:
            free_delta_index(self._index)
            self._index = index

    def add_source(self, source, unadded_bytes):
        """Add a new bit of source text to the delta indexes.

        :param source: The text in question, this must be a byte string
        :param unadded_bytes: Assume there are this many bytes that didn't get
            added between this source and the end of the previous source.
        """
        cdef char *c_source
        cdef Py_ssize_t c_source_size
        cdef delta_index *index
        cdef delta_result res
        cdef unsigned int source_location
        cdef source_info *src
        cdef unsigned int num_indexes

        if not PyString_CheckExact(source):
            raise TypeError('source is not a str')

        source_location = len(self._sources)
        if source_location >= self._max_num_sources:
            self._expand_sources()
        if source_location != 0 and self._index == NULL:
            # We were lazy about populating the index, create it now
            self._populate_first_index()
        self._sources.append(source)
        c_source = PyString_AS_STRING(source)
        c_source_size = PyString_GET_SIZE(source)
        src = self._source_infos + source_location
        src.buf = c_source
        src.size = c_source_size

        src.agg_offset = self._source_offset + unadded_bytes
        self._source_offset = src.agg_offset + src.size
        # We delay creating the index on the first insert
        if source_location != 0:
            with nogil:
                res = create_delta_index(src, self._index, &index)
            if res != DELTA_OK:
                raise _translate_delta_failure(res)
            if index != self._index:
                free_delta_index(self._index)
                self._index = index

    cdef _populate_first_index(self):
        cdef delta_index *index
        cdef delta_result res
        if len(self._sources) != 1 or self._index != NULL:
            raise AssertionError('_populate_first_index should only be'
                ' called when we have a single source and no index yet')

        # We know that self._index is already NULL, so create_delta_index
        # will always create a new index unless there's a malloc failure
        with nogil:
            res = create_delta_index(&self._source_infos[0], NULL, &index)
        if res != DELTA_OK:
            raise _translate_delta_failure(res)
        self._index = index

    cdef _expand_sources(self):
        raise RuntimeError('if we move self._source_infos, then we need to'
                           ' change all of the index pointers as well.')
        self._max_num_sources = self._max_num_sources * 2
        self._source_infos = <source_info *>safe_realloc(self._source_infos,
                                                sizeof(source_info)
                                                * self._max_num_sources)

    def make_delta(self, target_bytes, max_delta_size=0):
        """Create a delta from the current source to the target bytes."""
        cdef char *target
        cdef Py_ssize_t target_size
        cdef void * delta
        cdef unsigned long delta_size
        cdef unsigned long c_max_delta_size
        cdef delta_result res

        if self._index == NULL:
            if len(self._sources) == 0:
                return None
            # We were just lazy about generating the index
            self._populate_first_index()

        if not PyString_CheckExact(target_bytes):
            raise TypeError('target is not a str')

        target = PyString_AS_STRING(target_bytes)
        target_size = PyString_GET_SIZE(target_bytes)

        # TODO: inline some of create_delta so we at least don't have to double
        #       malloc, and can instead use PyString_FromStringAndSize, to
        #       allocate the bytes into the final string
        c_max_delta_size = max_delta_size
        with nogil:
            res = create_delta(self._index, target, target_size,
                               &delta_size, c_max_delta_size, &delta)
        result = None
        if res == DELTA_OK:
            result = PyString_FromStringAndSize(<char *>delta, delta_size)
            free(delta)
        elif res != DELTA_SIZE_TOO_BIG:
            raise _translate_delta_failure(res)
        return result


def make_delta(source_bytes, target_bytes):
    """Create a delta, this is a wrapper around DeltaIndex.make_delta."""
    di = DeltaIndex(source_bytes)
    return di.make_delta(target_bytes)


def apply_delta(source_bytes, delta_bytes):
    """Apply a delta generated by make_delta to source_bytes."""
    cdef char *source
    cdef Py_ssize_t source_size
    cdef char *delta
    cdef Py_ssize_t delta_size

    if not PyString_CheckExact(source_bytes):
        raise TypeError('source is not a str')
    if not PyString_CheckExact(delta_bytes):
        raise TypeError('delta is not a str')
    source = PyString_AS_STRING(source_bytes)
    source_size = PyString_GET_SIZE(source_bytes)
    delta = PyString_AS_STRING(delta_bytes)
    delta_size = PyString_GET_SIZE(delta_bytes)
    # Code taken from patch-delta.c, only brought here to give better error
    # handling, and to avoid double allocating memory
    if (delta_size < DELTA_SIZE_MIN):
        # XXX: Invalid delta block
        raise RuntimeError('delta_size %d smaller than min delta size %d'
                           % (delta_size, DELTA_SIZE_MIN))

    return _apply_delta(source, source_size, delta, delta_size)


cdef unsigned char *_decode_copy_instruction(unsigned char *bytes,
    unsigned char cmd, unsigned int *offset,
    unsigned int *length) nogil: # cannot_raise
    """Decode a copy instruction from the next few bytes.

    A copy instruction is a variable number of bytes, so we will parse the
    bytes we care about, and return the new position, as well as the offset and
    length referred to in the bytes.

    :param bytes: Pointer to the start of bytes after cmd
    :param cmd: The command code
    :return: Pointer to the bytes just after the last decode byte
    """
    cdef unsigned int off, size, count
    off = 0
    size = 0
    count = 0
    if (cmd & 0x01):
        off = bytes[count]
        count = count + 1
    if (cmd & 0x02):
        off = off | (bytes[count] << 8)
        count = count + 1
    if (cmd & 0x04):
        off = off | (bytes[count] << 16)
        count = count + 1
    if (cmd & 0x08):
        off = off | (bytes[count] << 24)
        count = count + 1
    if (cmd & 0x10):
        size = bytes[count]
        count = count + 1
    if (cmd & 0x20):
        size = size | (bytes[count] << 8)
        count = count + 1
    if (cmd & 0x40):
        size = size | (bytes[count] << 16)
        count = count + 1
    if (size == 0):
        size = 0x10000
    offset[0] = off
    length[0] = size
    return bytes + count


cdef object _apply_delta(char *source, Py_ssize_t source_size,
                         char *delta, Py_ssize_t delta_size):
    """common functionality between apply_delta and apply_delta_to_source."""
    cdef unsigned char *data, *top
    cdef unsigned char *dst_buf, *out, cmd
    cdef Py_ssize_t size
    cdef unsigned int cp_off, cp_size
    cdef int failed

    data = <unsigned char *>delta
    top = data + delta_size

    # now the result size
    size = get_delta_hdr_size(&data, top)
    result = PyString_FromStringAndSize(NULL, size)
    dst_buf = <unsigned char*>PyString_AS_STRING(result)

    failed = 0
    with nogil:
        out = dst_buf
        while (data < top):
            cmd = data[0]
            data = data + 1
            if (cmd & 0x80):
                # Copy instruction
                data = _decode_copy_instruction(data, cmd, &cp_off, &cp_size)
                if (cp_off + cp_size < cp_size or
                    cp_off + cp_size > <unsigned int>source_size or
                    cp_size > <unsigned int>size):
                    failed = 1
                    break
                memcpy(out, source + cp_off, cp_size)
                out = out + cp_size
                size = size - cp_size
            else:
                # Insert instruction
                if cmd == 0:
                    # cmd == 0 is reserved for future encoding
                    # extensions. In the mean time we must fail when
                    # encountering them (might be data corruption).
                    failed = 2
                    break
                if cmd > size:
                    failed = 3
                    break
                memcpy(out, data, cmd)
                out = out + cmd
                data = data + cmd
                size = size - cmd
    if failed:
        if failed == 1:
            raise ValueError('Something wrong with:'
                ' cp_off = %s, cp_size = %s'
                ' source_size = %s, size = %s'
                % (cp_off, cp_size, source_size, size))
        elif failed == 2:
            raise ValueError('Got delta opcode: 0, not supported')
        elif failed == 3:
            raise ValueError('Insert instruction longer than remaining'
                ' bytes: %d > %d' % (cmd, size))

    # sanity check
    if (data != top or size != 0):
        raise RuntimeError('Did not extract the number of bytes we expected'
            ' we were left with %d bytes in "size", and top - data = %d'
            % (size, <int>(top - data)))
        return None

    # *dst_size = out - dst_buf;
    if (out - dst_buf) != PyString_GET_SIZE(result):
        raise RuntimeError('Number of bytes extracted did not match the'
            ' size encoded in the delta header.')
    return result


def apply_delta_to_source(source, delta_start, delta_end):
    """Extract a delta from source bytes, and apply it."""
    cdef char *c_source
    cdef Py_ssize_t c_source_size
    cdef char *c_delta
    cdef Py_ssize_t c_delta_size
    cdef Py_ssize_t c_delta_start, c_delta_end

    if not PyString_CheckExact(source):
        raise TypeError('source is not a str')
    c_source_size = PyString_GET_SIZE(source)
    c_delta_start = delta_start
    c_delta_end = delta_end
    if c_delta_start >= c_source_size:
        raise ValueError('delta starts after source')
    if c_delta_end > c_source_size:
        raise ValueError('delta ends after source')
    if c_delta_start >= c_delta_end:
        raise ValueError('delta starts after it ends')

    c_delta_size = c_delta_end - c_delta_start
    c_source = PyString_AS_STRING(source)
    c_delta = c_source + c_delta_start
    # We don't use source_size, because we know the delta should not refer to
    # any bytes after it starts
    return _apply_delta(c_source, c_delta_start, c_delta, c_delta_size)


def encode_base128_int(val):
    """Convert an integer into a 7-bit lsb encoding."""
    cdef unsigned int c_val
    cdef Py_ssize_t count
    cdef unsigned int num_bytes
    cdef unsigned char c_bytes[8] # max size for 32-bit int is 5 bytes

    c_val = val
    count = 0
    while c_val >= 0x80 and count < 8:
        c_bytes[count] = <unsigned char>((c_val | 0x80) & 0xFF)
        c_val = c_val >> 7
        count = count + 1
    if count >= 8 or c_val >= 0x80:
        raise ValueError('encode_base128_int overflowed the buffer')
    c_bytes[count] = <unsigned char>(c_val & 0xFF)
    count = count + 1
    return PyString_FromStringAndSize(<char *>c_bytes, count)


def decode_base128_int(bytes):
    """Decode an integer from a 7-bit lsb encoding."""
    cdef int offset
    cdef int val
    cdef unsigned int uval
    cdef int shift
    cdef Py_ssize_t num_low_bytes
    cdef unsigned char *c_bytes

    offset = 0
    val = 0
    shift = 0
    if not PyString_CheckExact(bytes):
        raise TypeError('bytes is not a string')
    c_bytes = <unsigned char*>PyString_AS_STRING(bytes)
    # We take off 1, because we have to be able to decode the non-expanded byte
    num_low_bytes = PyString_GET_SIZE(bytes) - 1
    while (c_bytes[offset] & 0x80) and offset < num_low_bytes:
        val = val | ((c_bytes[offset] & 0x7F) << shift)
        shift = shift + 7
        offset = offset + 1
    if c_bytes[offset] & 0x80:
        raise ValueError('Data not properly formatted, we ran out of'
                         ' bytes before 0x80 stopped being set.')
    val = val | (c_bytes[offset] << shift)
    offset = offset + 1
    if val < 0:
        uval = <unsigned int> val
        return uval, offset
    return val, offset