~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

#

"""Pyrex extensions to btree node parsing."""

#python2.4 support

cdef extern from "python-compat.h":

    pass

cdef extern from "stdlib.h":

    ctypedef unsigned size_t

cdef extern from "Python.h":

    ctypedef int Py_ssize_t # Required for older pyrex versions

    ctypedef struct PyObject:

        pass

    int PyList_Append(object lst, object item) except -1

    char *PyString_AsString(object p) except NULL

    object PyString_FromStringAndSize(char *, Py_ssize_t)

    PyObject *PyString_FromStringAndSize_ptr "PyString_FromStringAndSize" (char *, Py_ssize_t)

    object PyString_FromFormat(char *, ...)

    int PyString_CheckExact(object s)

    int PyString_CheckExact_ptr "PyString_CheckExact" (PyObject *)

    Py_ssize_t PyString_Size(object p)

    Py_ssize_t PyString_GET_SIZE_ptr "PyString_GET_SIZE" (PyObject *)

    char * PyString_AS_STRING_ptr "PyString_AS_STRING" (PyObject *)

    char * PyString_AS_STRING(object)

    Py_ssize_t PyString_GET_SIZE(object)

    int PyString_AsStringAndSize_ptr(PyObject *, char **buf, Py_ssize_t *len)

    void PyString_InternInPlace(PyObject **)

    int PyTuple_CheckExact(object t)

    object PyTuple_New(Py_ssize_t n_entries)

    void PyTuple_SET_ITEM(object, Py_ssize_t offset, object) # steals the ref

    Py_ssize_t PyTuple_GET_SIZE(object t)

    PyObject *PyTuple_GET_ITEM_ptr_object "PyTuple_GET_ITEM" (object tpl, int index)

    void Py_INCREF(object)

    void Py_DECREF_ptr "Py_DECREF" (PyObject *)

    void *PyMem_Malloc(size_t nbytes)

    void PyMem_Free(void *)

    void memset(void *, int, size_t)

cdef extern from "string.h":

    void *memcpy(void *dest, void *src, size_t n)

    void *memchr(void *s, int c, size_t n)

    int memcmp(void *s1, void *s2, size_t n)

    # GNU extension

    # void *memrchr(void *s, int c, size_t n)

    int strncmp(char *s1, char *s2, size_t n)

    unsigned long strtoul(char *s1, char **out, int base)

    long long strtoll(char *s1, char **out, int base)

# It seems we need to import the definitions so that the pyrex compiler has

# local names to access them.

from _static_tuple_c cimport StaticTuple, \

    import_static_tuple_c, StaticTuple_New, \

    StaticTuple_Intern, StaticTuple_SET_ITEM, StaticTuple_CheckExact, \

    StaticTuple_GET_SIZE, StaticTuple_GET_ITEM

# This tells the test infrastructure that StaticTuple is a class, so we don't

# have to worry about exception checking.

## extern cdef class StaticTuple

import sys

# TODO: Find some way to import this from _dirstate_helpers

cdef void* _my_memrchr(void *s, int c, size_t n): # cannot_raise

    # memrchr seems to be a GNU extension, so we have to implement it ourselves

    # It is not present in any win32 standard library

    cdef char *pos

    cdef char *start

    start = <char*>s

    pos = start + n - 1

    while pos >= start:

        if pos[0] == c:

            return <void*>pos

        pos = pos - 1

    return NULL

# TODO: Import this from _dirstate_helpers when it is merged

cdef object safe_string_from_size(char *s, Py_ssize_t size):

    if size < 0:

        raise AssertionError(

            'tried to create a string with an invalid size: %d @0x%x'

            % (size, <int>s))

    return PyString_FromStringAndSize(s, size)

cdef object safe_interned_string_from_size(char *s, Py_ssize_t size):

    cdef PyObject *py_str

    if size < 0:

        raise AssertionError(

            'tried to create a string with an invalid size: %d @0x%x'

            % (size, <int>s))

    py_str = PyString_FromStringAndSize_ptr(s, size)

    PyString_InternInPlace(&py_str)

    result = <object>py_str

    # Casting a PyObject* to an <object> triggers an INCREF from Pyrex, so we

    # DECREF it to avoid geting immortal strings

    Py_DECREF_ptr(py_str)

    return result

# This sets up the StaticTuple C_API functionality

import_static_tuple_c()

cdef class BTreeLeafParser:

    """Parse the leaf nodes of a BTree index.

    :ivar bytes: The PyString object containing the uncompressed text for the

        node.

    :ivar key_length: An integer describing how many pieces the keys have for

        this index.

    :ivar ref_list_length: An integer describing how many references this index

        contains.

    :ivar keys: A PyList of keys found in this node.

    :ivar _cur_str: A pointer to the start of the next line to parse

    :ivar _end_str: A pointer to the end of bytes

    :ivar _start: Pointer to the location within the current line while

        parsing.

    :ivar _header_found: True when we have parsed the header for this node

    """

    cdef object bytes

    cdef int key_length

    cdef int ref_list_length

    cdef object keys

    cdef char * _cur_str

    cdef char * _end_str

    # The current start point for parsing

    cdef char * _start

    cdef int _header_found

    def __init__(self, bytes, key_length, ref_list_length):

        self.bytes = bytes

        self.key_length = key_length

        self.ref_list_length = ref_list_length

        self.keys = []

        self._cur_str = NULL

        self._end_str = NULL

        self._header_found = 0

        # keys are tuples

    cdef extract_key(self, char * last):

        """Extract a key.

        :param last: points at the byte after the last byte permitted for the

            key.

        """

        cdef char *temp_ptr

        cdef int loop_counter

        cdef StaticTuple key

        key = StaticTuple_New(self.key_length)

        for loop_counter from 0 <= loop_counter < self.key_length:

            # grab a key segment

            temp_ptr = <char*>memchr(self._start, c'\0', last - self._start)

            if temp_ptr == NULL:

                if loop_counter + 1 == self.key_length:

                    # capture to last

                    temp_ptr = last

                else:

                    # Invalid line

                    failure_string = ("invalid key, wanted segment from " +

                        repr(safe_string_from_size(self._start,

                                                   last - self._start)))

                    raise AssertionError(failure_string)

            # capture the key string

            if (self.key_length == 1

                and (temp_ptr - self._start) == 45

                and strncmp(self._start, 'sha1:', 5) == 0):

                key_element = safe_string_from_size(self._start,

                                                    temp_ptr - self._start)

            else:

                key_element = safe_interned_string_from_size(self._start,

                                                         temp_ptr - self._start)

            # advance our pointer

            self._start = temp_ptr + 1

            Py_INCREF(key_element)

            StaticTuple_SET_ITEM(key, loop_counter, key_element)

        key = StaticTuple_Intern(key)

        return key

    cdef int process_line(self) except -1:

        """Process a line in the bytes."""

        cdef char *last

        cdef char *temp_ptr

        cdef char *ref_ptr

        cdef char *next_start

        cdef int loop_counter

        cdef Py_ssize_t str_len

        self._start = self._cur_str

        # Find the next newline

        last = <char*>memchr(self._start, c'\n', self._end_str - self._start)

        if last == NULL:

            # Process until the end of the file

            last = self._end_str

            self._cur_str = self._end_str

        else:

            # And the next string is right after it

            self._cur_str = last + 1

            # The last character is right before the '\n'

        if last == self._start:

            # parsed it all.

            return 0

        if last < self._start:

            # Unexpected error condition - fail

            raise AssertionError("last < self._start")

        if 0 == self._header_found:

            # The first line in a leaf node is the header "type=leaf\n"

            if strncmp("type=leaf", self._start, last - self._start) == 0:

                self._header_found = 1

                return 0

            else:

                raise AssertionError('Node did not start with "type=leaf": %r'

                    % (safe_string_from_size(self._start, last - self._start)))

        key = self.extract_key(last)

        # find the value area

        temp_ptr = <char*>_my_memrchr(self._start, c'\0', last - self._start)

        if temp_ptr == NULL:

            # Invalid line

            raise AssertionError("Failed to find the value area")

        else:

            # Because of how conversions were done, we ended up with *lots* of

            # values that are identical. These are all of the 0-length nodes

            # that are referred to by the TREE_ROOT (and likely some other

            # directory nodes.) For example, bzr has 25k references to

            # something like '12607215 328306 0 0', which ends up consuming 1MB

            # of memory, just for those strings.

            str_len = last - temp_ptr - 1

            if (str_len > 4

                and strncmp(" 0 0", last - 4, 4) == 0):

                # This drops peak mem for bzr.dev from 87.4MB => 86.2MB

                # For Launchpad 236MB => 232MB

                value = safe_interned_string_from_size(temp_ptr + 1, str_len)

            else:

                value = safe_string_from_size(temp_ptr + 1, str_len)

            # shrink the references end point

            last = temp_ptr

        if self.ref_list_length:

            ref_lists = StaticTuple_New(self.ref_list_length)

            loop_counter = 0

            while loop_counter < self.ref_list_length:

                ref_list = []

                # extract a reference list

                loop_counter = loop_counter + 1

                if last < self._start:

                    raise AssertionError("last < self._start")

                # find the next reference list end point:

                temp_ptr = <char*>memchr(self._start, c'\t', last - self._start)

                if temp_ptr == NULL:

                    # Only valid for the last list

                    if loop_counter != self.ref_list_length:

                        # Invalid line

                        raise AssertionError(

                            "invalid key, loop_counter != self.ref_list_length")

                    else:

                        # scan to the end of the ref list area

                        ref_ptr = last

                        next_start = last

                else:

                    # scan to the end of this ref list

                    ref_ptr = temp_ptr

                    next_start = temp_ptr + 1

                # Now, there may be multiple keys in the ref list.

                while self._start < ref_ptr:

                    # loop finding keys and extracting them

                    temp_ptr = <char*>memchr(self._start, c'\r',

                                             ref_ptr - self._start)

                    if temp_ptr == NULL:

                        # key runs to the end

                        temp_ptr = ref_ptr

                    PyList_Append(ref_list, self.extract_key(temp_ptr))

                ref_list = StaticTuple_Intern(StaticTuple(*ref_list))

                Py_INCREF(ref_list)

                StaticTuple_SET_ITEM(ref_lists, loop_counter - 1, ref_list)

                # prepare for the next reference list

                self._start = next_start

            node_value = StaticTuple(value, ref_lists)

        else:

            if last != self._start:

                # unexpected reference data present

                raise AssertionError("unexpected reference data present")

            node_value = StaticTuple(value, StaticTuple())

        PyList_Append(self.keys, StaticTuple(key, node_value))

        return 0

    def parse(self):

        cdef Py_ssize_t byte_count

        if not PyString_CheckExact(self.bytes):

            raise AssertionError('self.bytes is not a string.')

        byte_count = PyString_Size(self.bytes)

        self._cur_str = PyString_AsString(self.bytes)

        # This points to the last character in the string

        self._end_str = self._cur_str + byte_count

        while self._cur_str < self._end_str:

            self.process_line()

        return self.keys

def _parse_leaf_lines(bytes, key_length, ref_list_length):

    parser = BTreeLeafParser(bytes, key_length, ref_list_length)

    return parser.parse()

# TODO: We can go from 8 byte offset + 4 byte length to a simple lookup,

#       because the block_offset + length is likely to be repeated. However,

#       the big win there is to cache across pages, and not just one page

#       Though if we did cache in a page, we could certainly use a short int.

#       And this goes from 40 bytes to 30 bytes.

#       One slightly ugly option would be to cache block offsets in a global.

#       However, that leads to thread-safety issues, etc.

ctypedef struct gc_chk_sha1_record:

    long long block_offset

    unsigned int block_length

    unsigned int record_start

    unsigned int record_end

    char sha1[20]

cdef int _unhexbuf[256]

cdef char *_hexbuf

_hexbuf = '0123456789abcdef'

cdef _populate_unhexbuf():

    cdef int i

    for i from 0 <= i < 256:

        _unhexbuf[i] = -1

    for i from 0 <= i < 10: # 0123456789 => map to the raw number

        _unhexbuf[(i + c'0')] = i

    for i from 10 <= i < 16: # abcdef => 10, 11, 12, 13, 14, 15, 16

        _unhexbuf[(i - 10 + c'a')] = i

    for i from 10 <= i < 16: # ABCDEF => 10, 11, 12, 13, 14, 15, 16

        _unhexbuf[(i - 10 + c'A')] = i

_populate_unhexbuf()

cdef int _unhexlify_sha1(char *as_hex, char *as_bin): # cannot_raise

    """Take the hex sha1 in as_hex and make it binary in as_bin

    Same as binascii.unhexlify, but working on C strings, not Python objects.

    """

    cdef int top

    cdef int bot

    cdef int i, j

    cdef char *cur

    # binascii does this using isupper() and tolower() and ?: syntax. I'm

    # guessing a simple lookup array should be faster.

    j = 0

    for i from 0 <= i < 20:

        top = _unhexbuf[<unsigned char>(as_hex[j])]

        j = j + 1

        bot = _unhexbuf[<unsigned char>(as_hex[j])]

        j = j + 1

        if top == -1 or bot == -1:

            return 0

        as_bin[i] = <unsigned char>((top << 4) + bot);

    return 1

def _py_unhexlify(as_hex):

    """For the test infrastructure, just thunks to _unhexlify_sha1"""

    if len(as_hex) != 40 or not PyString_CheckExact(as_hex):

        raise ValueError('not a 40-byte hex digest')

    as_bin = PyString_FromStringAndSize(NULL, 20)

    if _unhexlify_sha1(PyString_AS_STRING(as_hex), PyString_AS_STRING(as_bin)):

        return as_bin

    return None

cdef void _hexlify_sha1(char *as_bin, char *as_hex): # cannot_raise

    cdef int i, j

    cdef char c

    j = 0

    for i from 0 <= i < 20:

        c = as_bin[i]

        as_hex[j] = _hexbuf[(c>>4)&0xf]

        j = j + 1

        as_hex[j] = _hexbuf[(c)&0xf]

        j = j + 1

def _py_hexlify(as_bin):

    """For test infrastructure, thunk to _hexlify_sha1"""

    if len(as_bin) != 20 or not PyString_CheckExact(as_bin):

        raise ValueError('not a 20-byte binary digest')

    as_hex = PyString_FromStringAndSize(NULL, 40)

    _hexlify_sha1(PyString_AS_STRING(as_bin), PyString_AS_STRING(as_hex))

    return as_hex

cdef int _key_to_sha1(key, char *sha1): # cannot_raise

    """Map a key into its sha1 content.

    :param key: A tuple of style ('sha1:abcd...',)

    :param sha1: A char buffer of 20 bytes

    :return: 1 if this could be converted, 0 otherwise

    """

    cdef char *c_val

    cdef PyObject *p_val

    if StaticTuple_CheckExact(key) and StaticTuple_GET_SIZE(key) == 1:

        p_val = <PyObject *>StaticTuple_GET_ITEM(key, 0)

    elif (PyTuple_CheckExact(key) and PyTuple_GET_SIZE(key) == 1):

        p_val = PyTuple_GET_ITEM_ptr_object(key, 0)

    else:

        # Not a tuple or a StaticTuple

        return 0

    if (PyString_CheckExact_ptr(p_val) and PyString_GET_SIZE_ptr(p_val) == 45):

        c_val = PyString_AS_STRING_ptr(p_val)

    else:

        return 0

    if strncmp(c_val, 'sha1:', 5) != 0:

        return 0

    if not _unhexlify_sha1(c_val + 5, sha1):

        return 0

    return 1

def _py_key_to_sha1(key):

    """Map a key to a simple sha1 string.

    This is a testing thunk to the C function.

    """

    as_bin_sha = PyString_FromStringAndSize(NULL, 20)

    if _key_to_sha1(key, PyString_AS_STRING(as_bin_sha)):

        return as_bin_sha

    return None

cdef StaticTuple _sha1_to_key(char *sha1):

    """Compute a ('sha1:abcd',) key for a given sha1."""

    cdef StaticTuple key

    cdef object hexxed

    cdef char *c_buf

    hexxed = PyString_FromStringAndSize(NULL, 45)

    c_buf = PyString_AS_STRING(hexxed)

    memcpy(c_buf, 'sha1:', 5)

    _hexlify_sha1(sha1, c_buf+5)

    key = StaticTuple_New(1)

    Py_INCREF(hexxed)

    StaticTuple_SET_ITEM(key, 0, hexxed)

    # This is a bit expensive. To parse 120 keys takes 48us, to return them all

    # can be done in 66.6us (so 18.6us to build them all).

    # Adding simple hash() here brings it to 76.6us (so computing the hash

    # value of 120keys is 10us), Intern is 86.9us (another 10us to look and add

    # them to the intern structure.)

    # However, since we only intern keys that are in active use, it is probably

    # a win. Since they would have been read from elsewhere anyway.

    # We *could* hang the PyObject form off of the gc_chk_sha1_record for ones

    # that we have deserialized. Something to think about, at least.

    key = StaticTuple_Intern(key)

    return key

def _py_sha1_to_key(sha1_bin):

    """Test thunk to check the sha1 mapping."""

    if not PyString_CheckExact(sha1_bin) or PyString_GET_SIZE(sha1_bin) != 20:

        raise ValueError('sha1_bin must be a str of exactly 20 bytes')

    return _sha1_to_key(PyString_AS_STRING(sha1_bin))

cdef unsigned int _sha1_to_uint(char *sha1): # cannot_raise

    cdef unsigned int val

    # Must be in MSB, because that is how the content is sorted

    val = (((<unsigned int>(sha1[0]) & 0xff) << 24)

           | ((<unsigned int>(sha1[1]) & 0xff) << 16)

           | ((<unsigned int>(sha1[2]) & 0xff) << 8)

           | ((<unsigned int>(sha1[3]) & 0xff) << 0))

    return val

cdef _format_record(gc_chk_sha1_record *record):

    # This is inefficient to go from a logical state back to a

    # string, but it makes things work a bit better internally for now.

    if record.block_offset >= 0xFFFFFFFF:

        # %llu is what we really want, but unfortunately it was only added

        # in python 2.7... :(

        block_offset_str = str(record.block_offset)

        value = PyString_FromFormat('%s %u %u %u',

                                PyString_AS_STRING(block_offset_str),

                                record.block_length,

                                record.record_start, record.record_end)

    else:

        value = PyString_FromFormat('%lu %u %u %u',

                                    <unsigned long>record.block_offset,

                                    record.block_length,

                                    record.record_start, record.record_end)

    return value

cdef class GCCHKSHA1LeafNode:

    """Track all the entries for a given leaf node."""

    cdef gc_chk_sha1_record *records

    cdef public object last_key

    cdef gc_chk_sha1_record *last_record

    cdef public int num_records

    # This is the number of bits to shift to get to the interesting byte. A

    # value of 24 means that the very first byte changes across all keys.

    # Anything else means that there is a common prefix of bits that we can

    # ignore. 0 means that at least the first 3 bytes are identical, though

    # that is going to be very rare

    cdef public unsigned char common_shift

    # This maps an interesting byte to the first record that matches.

    # Equivalent to bisect.bisect_left(self.records, sha1), though only taking

    # into account that one byte.

    cdef unsigned char offsets[257]

    def __sizeof__(self):

        # :( Why doesn't Pyrex let me do a simple sizeof(GCCHKSHA1LeafNode)

        # like Cython? Explicitly enumerating everything here seems to leave my

        # size off by 2 (286 bytes vs 288 bytes actual). I'm guessing it is an

        # alignment/padding issue. Oh well- at least we scale properly with

        # num_records and are very close to correct, which is what I care

        # about.

        # If we ever decide to require cython:

        # return (sizeof(GCCHKSHA1LeafNode)

        #     + sizeof(gc_chk_sha1_record)*self.num_records)

        return (sizeof(PyObject) + sizeof(void*) + sizeof(int)

            + sizeof(gc_chk_sha1_record*) + sizeof(PyObject *)

            + sizeof(gc_chk_sha1_record*) + sizeof(char)

            + sizeof(unsigned char)*257

            + sizeof(gc_chk_sha1_record)*self.num_records)

    def __dealloc__(self):

        if self.records != NULL:

            PyMem_Free(self.records)

            self.records = NULL

    def __init__(self, bytes):

        self._parse_bytes(bytes)

        self.last_key = None

        self.last_record = NULL

    property min_key:

        def __get__(self):

            if self.num_records > 0:

                return _sha1_to_key(self.records[0].sha1)

            return None

    property max_key:

        def __get__(self):

            if self.num_records > 0:

                return _sha1_to_key(self.records[self.num_records-1].sha1)

            return None

    cdef StaticTuple _record_to_value_and_refs(self,

                                               gc_chk_sha1_record *record):

        """Extract the refs and value part of this record."""

        cdef StaticTuple value_and_refs

        cdef StaticTuple empty

        value_and_refs = StaticTuple_New(2)

        value = _format_record(record)

        Py_INCREF(value)

        StaticTuple_SET_ITEM(value_and_refs, 0, value)

        # Always empty refs

        empty = StaticTuple_New(0)

        Py_INCREF(empty)

        StaticTuple_SET_ITEM(value_and_refs, 1, empty)

        return value_and_refs

    cdef StaticTuple _record_to_item(self, gc_chk_sha1_record *record):

        """Turn a given record back into a fully fledged item.

        """

        cdef StaticTuple item

        cdef StaticTuple key

        cdef StaticTuple value_and_refs

        cdef object value

        key = _sha1_to_key(record.sha1)

        item = StaticTuple_New(2)

        Py_INCREF(key)

        StaticTuple_SET_ITEM(item, 0, key)

        value_and_refs = self._record_to_value_and_refs(record)

        Py_INCREF(value_and_refs)

        StaticTuple_SET_ITEM(item, 1, value_and_refs)

        return item

    cdef gc_chk_sha1_record* _lookup_record(self, char *sha1) except? NULL:

        """Find a gc_chk_sha1_record that matches the sha1 supplied."""

        cdef int lo, hi, mid, the_cmp

        cdef int offset

        # TODO: We can speed up misses by comparing this sha1 to the common

        #       bits, and seeing if the common prefix matches, if not, we don't

        #       need to search for anything because it cannot match

        # Use the offset array to find the closest fit for this entry

        # follow that up with bisecting, since multiple keys can be in one

        # spot

        # Bisecting dropped us from 7000 comparisons to 582 (4.8/key), using

        # the offset array dropped us from 23us to 20us and 156 comparisions

        # (1.3/key)

        offset = self._offset_for_sha1(sha1)

        lo = self.offsets[offset]

        hi = self.offsets[offset+1]

        if hi == 255:

            # if hi == 255 that means we potentially ran off the end of the

            # list, so push it up to num_records

            # note that if 'lo' == 255, that is ok, because we can start

            # searching from that part of the list.

            hi = self.num_records

        local_n_cmp = 0

        while lo < hi:

            mid = (lo + hi) / 2

            the_cmp = memcmp(self.records[mid].sha1, sha1, 20)

            if the_cmp == 0:

                return &self.records[mid]

            elif the_cmp < 0:

                lo = mid + 1

            else:

                hi = mid

        return NULL

    def __contains__(self, key):

        cdef char sha1[20]

        cdef gc_chk_sha1_record *record

        if _key_to_sha1(key, sha1):

            # If it isn't a sha1 key, then it won't be in this leaf node

            record = self._lookup_record(sha1)

            if record != NULL:

                self.last_key = key

                self.last_record = record

                return True

        return False

    def __getitem__(self, key):

        cdef char sha1[20]

        cdef gc_chk_sha1_record *record

        record = NULL