~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/_btree_serializer_pyx.pyx

Major code cleanup.

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# Copyright (C) 2008, 2009, 2010 Canonical Ltd
 
2
#
 
3
# This program is free software; you can redistribute it and/or modify
 
4
# it under the terms of the GNU General Public License as published by
 
5
# the Free Software Foundation; either version 2 of the License, or
 
6
# (at your option) any later version.
 
7
#
 
8
# This program is distributed in the hope that it will be useful,
 
9
# but WITHOUT ANY WARRANTY; without even the implied warranty of
 
10
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
11
# GNU General Public License for more details.
 
12
#
 
13
# You should have received a copy of the GNU General Public License
 
14
# along with this program; if not, write to the Free Software
 
15
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 
16
#
 
17
 
 
18
"""Pyrex extensions to btree node parsing."""
 
19
 
 
20
#python2.4 support
 
21
cdef extern from "python-compat.h":
 
22
    pass
 
23
 
 
24
cdef extern from "stdlib.h":
 
25
    ctypedef unsigned size_t
 
26
 
 
27
cdef extern from "Python.h":
 
28
    ctypedef int Py_ssize_t # Required for older pyrex versions
 
29
    ctypedef struct PyObject:
 
30
        pass
 
31
    int PyList_Append(object lst, object item) except -1
 
32
 
 
33
    char *PyString_AsString(object p) except NULL
 
34
    object PyString_FromStringAndSize(char *, Py_ssize_t)
 
35
    PyObject *PyString_FromStringAndSize_ptr "PyString_FromStringAndSize" (char *, Py_ssize_t)
 
36
    object PyString_FromFormat(char *, ...)
 
37
    int PyString_CheckExact(object s)
 
38
    int PyString_CheckExact_ptr "PyString_CheckExact" (PyObject *)
 
39
    Py_ssize_t PyString_Size(object p)
 
40
    Py_ssize_t PyString_GET_SIZE_ptr "PyString_GET_SIZE" (PyObject *)
 
41
    char * PyString_AS_STRING_ptr "PyString_AS_STRING" (PyObject *)
 
42
    char * PyString_AS_STRING(object)
 
43
    Py_ssize_t PyString_GET_SIZE(object)
 
44
    int PyString_AsStringAndSize_ptr(PyObject *, char **buf, Py_ssize_t *len)
 
45
    void PyString_InternInPlace(PyObject **)
 
46
    int PyTuple_CheckExact(object t)
 
47
    object PyTuple_New(Py_ssize_t n_entries)
 
48
    void PyTuple_SET_ITEM(object, Py_ssize_t offset, object) # steals the ref
 
49
    Py_ssize_t PyTuple_GET_SIZE(object t)
 
50
    PyObject *PyTuple_GET_ITEM_ptr_object "PyTuple_GET_ITEM" (object tpl, int index)
 
51
    void Py_INCREF(object)
 
52
    void Py_DECREF_ptr "Py_DECREF" (PyObject *)
 
53
    void *PyMem_Malloc(size_t nbytes)
 
54
    void PyMem_Free(void *)
 
55
    void memset(void *, int, size_t)
 
56
 
 
57
cdef extern from "string.h":
 
58
    void *memcpy(void *dest, void *src, size_t n)
 
59
    void *memchr(void *s, int c, size_t n)
 
60
    int memcmp(void *s1, void *s2, size_t n)
 
61
    # GNU extension
 
62
    # void *memrchr(void *s, int c, size_t n)
 
63
    int strncmp(char *s1, char *s2, size_t n)
 
64
    unsigned long strtoul(char *s1, char **out, int base)
 
65
    long long strtoll(char *s1, char **out, int base)
 
66
 
 
67
 
 
68
# It seems we need to import the definitions so that the pyrex compiler has
 
69
# local names to access them.
 
70
from _static_tuple_c cimport StaticTuple, \
 
71
    import_static_tuple_c, StaticTuple_New, \
 
72
    StaticTuple_Intern, StaticTuple_SET_ITEM, StaticTuple_CheckExact, \
 
73
    StaticTuple_GET_SIZE, StaticTuple_GET_ITEM
 
74
# This tells the test infrastructure that StaticTuple is a class, so we don't
 
75
# have to worry about exception checking.
 
76
## extern cdef class StaticTuple
 
77
import sys
 
78
 
 
79
 
 
80
# TODO: Find some way to import this from _dirstate_helpers
 
81
cdef void* _my_memrchr(void *s, int c, size_t n): # cannot_raise
 
82
    # memrchr seems to be a GNU extension, so we have to implement it ourselves
 
83
    # It is not present in any win32 standard library
 
84
    cdef char *pos
 
85
    cdef char *start
 
86
 
 
87
    start = <char*>s
 
88
    pos = start + n - 1
 
89
    while pos >= start:
 
90
        if pos[0] == c:
 
91
            return <void*>pos
 
92
        pos = pos - 1
 
93
    return NULL
 
94
 
 
95
 
 
96
# TODO: Import this from _dirstate_helpers when it is merged
 
97
cdef object safe_string_from_size(char *s, Py_ssize_t size):
 
98
    if size < 0:
 
99
        raise AssertionError(
 
100
            'tried to create a string with an invalid size: %d @0x%x'
 
101
            % (size, <int>s))
 
102
    return PyString_FromStringAndSize(s, size)
 
103
 
 
104
 
 
105
cdef object safe_interned_string_from_size(char *s, Py_ssize_t size):
 
106
    cdef PyObject *py_str
 
107
    if size < 0:
 
108
        raise AssertionError(
 
109
            'tried to create a string with an invalid size: %d @0x%x'
 
110
            % (size, <int>s))
 
111
    py_str = PyString_FromStringAndSize_ptr(s, size)
 
112
    PyString_InternInPlace(&py_str)
 
113
    result = <object>py_str
 
114
    # Casting a PyObject* to an <object> triggers an INCREF from Pyrex, so we
 
115
    # DECREF it to avoid geting immortal strings
 
116
    Py_DECREF_ptr(py_str)
 
117
    return result
 
118
 
 
119
# This sets up the StaticTuple C_API functionality
 
120
import_static_tuple_c()
 
121
 
 
122
 
 
123
cdef class BTreeLeafParser:
 
124
    """Parse the leaf nodes of a BTree index.
 
125
 
 
126
    :ivar bytes: The PyString object containing the uncompressed text for the
 
127
        node.
 
128
    :ivar key_length: An integer describing how many pieces the keys have for
 
129
        this index.
 
130
    :ivar ref_list_length: An integer describing how many references this index
 
131
        contains.
 
132
    :ivar keys: A PyList of keys found in this node.
 
133
 
 
134
    :ivar _cur_str: A pointer to the start of the next line to parse
 
135
    :ivar _end_str: A pointer to the end of bytes
 
136
    :ivar _start: Pointer to the location within the current line while
 
137
        parsing.
 
138
    :ivar _header_found: True when we have parsed the header for this node
 
139
    """
 
140
 
 
141
    cdef object bytes
 
142
    cdef int key_length
 
143
    cdef int ref_list_length
 
144
    cdef object keys
 
145
 
 
146
    cdef char * _cur_str
 
147
    cdef char * _end_str
 
148
    # The current start point for parsing
 
149
    cdef char * _start
 
150
 
 
151
    cdef int _header_found
 
152
 
 
153
    def __init__(self, bytes, key_length, ref_list_length):
 
154
        self.bytes = bytes
 
155
        self.key_length = key_length
 
156
        self.ref_list_length = ref_list_length
 
157
        self.keys = []
 
158
        self._cur_str = NULL
 
159
        self._end_str = NULL
 
160
        self._header_found = 0
 
161
        # keys are tuples
 
162
 
 
163
    cdef extract_key(self, char * last):
 
164
        """Extract a key.
 
165
 
 
166
        :param last: points at the byte after the last byte permitted for the
 
167
            key.
 
168
        """
 
169
        cdef char *temp_ptr
 
170
        cdef int loop_counter
 
171
        cdef StaticTuple key
 
172
 
 
173
        key = StaticTuple_New(self.key_length)
 
174
        for loop_counter from 0 <= loop_counter < self.key_length:
 
175
            # grab a key segment
 
176
            temp_ptr = <char*>memchr(self._start, c'\0', last - self._start)
 
177
            if temp_ptr == NULL:
 
178
                if loop_counter + 1 == self.key_length:
 
179
                    # capture to last
 
180
                    temp_ptr = last
 
181
                else:
 
182
                    # Invalid line
 
183
                    failure_string = ("invalid key, wanted segment from " +
 
184
                        repr(safe_string_from_size(self._start,
 
185
                                                   last - self._start)))
 
186
                    raise AssertionError(failure_string)
 
187
            # capture the key string
 
188
            if (self.key_length == 1
 
189
                and (temp_ptr - self._start) == 45
 
190
                and strncmp(self._start, 'sha1:', 5) == 0):
 
191
                key_element = safe_string_from_size(self._start,
 
192
                                                    temp_ptr - self._start)
 
193
            else:
 
194
                key_element = safe_interned_string_from_size(self._start,
 
195
                                                         temp_ptr - self._start)
 
196
            # advance our pointer
 
197
            self._start = temp_ptr + 1
 
198
            Py_INCREF(key_element)
 
199
            StaticTuple_SET_ITEM(key, loop_counter, key_element)
 
200
        key = StaticTuple_Intern(key)
 
201
        return key
 
202
 
 
203
    cdef int process_line(self) except -1:
 
204
        """Process a line in the bytes."""
 
205
        cdef char *last
 
206
        cdef char *temp_ptr
 
207
        cdef char *ref_ptr
 
208
        cdef char *next_start
 
209
        cdef int loop_counter
 
210
        cdef Py_ssize_t str_len
 
211
 
 
212
        self._start = self._cur_str
 
213
        # Find the next newline
 
214
        last = <char*>memchr(self._start, c'\n', self._end_str - self._start)
 
215
        if last == NULL:
 
216
            # Process until the end of the file
 
217
            last = self._end_str
 
218
            self._cur_str = self._end_str
 
219
        else:
 
220
            # And the next string is right after it
 
221
            self._cur_str = last + 1
 
222
            # The last character is right before the '\n'
 
223
 
 
224
        if last == self._start:
 
225
            # parsed it all.
 
226
            return 0
 
227
        if last < self._start:
 
228
            # Unexpected error condition - fail
 
229
            raise AssertionError("last < self._start")
 
230
        if 0 == self._header_found:
 
231
            # The first line in a leaf node is the header "type=leaf\n"
 
232
            if strncmp("type=leaf", self._start, last - self._start) == 0:
 
233
                self._header_found = 1
 
234
                return 0
 
235
            else:
 
236
                raise AssertionError('Node did not start with "type=leaf": %r'
 
237
                    % (safe_string_from_size(self._start, last - self._start)))
 
238
 
 
239
        key = self.extract_key(last)
 
240
        # find the value area
 
241
        temp_ptr = <char*>_my_memrchr(self._start, c'\0', last - self._start)
 
242
        if temp_ptr == NULL:
 
243
            # Invalid line
 
244
            raise AssertionError("Failed to find the value area")
 
245
        else:
 
246
            # Because of how conversions were done, we ended up with *lots* of
 
247
            # values that are identical. These are all of the 0-length nodes
 
248
            # that are referred to by the TREE_ROOT (and likely some other
 
249
            # directory nodes.) For example, bzr has 25k references to
 
250
            # something like '12607215 328306 0 0', which ends up consuming 1MB
 
251
            # of memory, just for those strings.
 
252
            str_len = last - temp_ptr - 1
 
253
            if (str_len > 4
 
254
                and strncmp(" 0 0", last - 4, 4) == 0):
 
255
                # This drops peak mem for bzr.dev from 87.4MB => 86.2MB
 
256
                # For Launchpad 236MB => 232MB
 
257
                value = safe_interned_string_from_size(temp_ptr + 1, str_len)
 
258
            else:
 
259
                value = safe_string_from_size(temp_ptr + 1, str_len)
 
260
            # shrink the references end point
 
261
            last = temp_ptr
 
262
 
 
263
        if self.ref_list_length:
 
264
            ref_lists = StaticTuple_New(self.ref_list_length)
 
265
            loop_counter = 0
 
266
            while loop_counter < self.ref_list_length:
 
267
                ref_list = []
 
268
                # extract a reference list
 
269
                loop_counter = loop_counter + 1
 
270
                if last < self._start:
 
271
                    raise AssertionError("last < self._start")
 
272
                # find the next reference list end point:
 
273
                temp_ptr = <char*>memchr(self._start, c'\t', last - self._start)
 
274
                if temp_ptr == NULL:
 
275
                    # Only valid for the last list
 
276
                    if loop_counter != self.ref_list_length:
 
277
                        # Invalid line
 
278
                        raise AssertionError(
 
279
                            "invalid key, loop_counter != self.ref_list_length")
 
280
                    else:
 
281
                        # scan to the end of the ref list area
 
282
                        ref_ptr = last
 
283
                        next_start = last
 
284
                else:
 
285
                    # scan to the end of this ref list
 
286
                    ref_ptr = temp_ptr
 
287
                    next_start = temp_ptr + 1
 
288
                # Now, there may be multiple keys in the ref list.
 
289
                while self._start < ref_ptr:
 
290
                    # loop finding keys and extracting them
 
291
                    temp_ptr = <char*>memchr(self._start, c'\r',
 
292
                                             ref_ptr - self._start)
 
293
                    if temp_ptr == NULL:
 
294
                        # key runs to the end
 
295
                        temp_ptr = ref_ptr
 
296
 
 
297
                    PyList_Append(ref_list, self.extract_key(temp_ptr))
 
298
                ref_list = StaticTuple_Intern(StaticTuple(*ref_list))
 
299
                Py_INCREF(ref_list)
 
300
                StaticTuple_SET_ITEM(ref_lists, loop_counter - 1, ref_list)
 
301
                # prepare for the next reference list
 
302
                self._start = next_start
 
303
            node_value = StaticTuple(value, ref_lists)
 
304
        else:
 
305
            if last != self._start:
 
306
                # unexpected reference data present
 
307
                raise AssertionError("unexpected reference data present")
 
308
            node_value = StaticTuple(value, StaticTuple())
 
309
        PyList_Append(self.keys, StaticTuple(key, node_value))
 
310
        return 0
 
311
 
 
312
    def parse(self):
 
313
        cdef Py_ssize_t byte_count
 
314
        if not PyString_CheckExact(self.bytes):
 
315
            raise AssertionError('self.bytes is not a string.')
 
316
        byte_count = PyString_Size(self.bytes)
 
317
        self._cur_str = PyString_AsString(self.bytes)
 
318
        # This points to the last character in the string
 
319
        self._end_str = self._cur_str + byte_count
 
320
        while self._cur_str < self._end_str:
 
321
            self.process_line()
 
322
        return self.keys
 
323
 
 
324
 
 
325
def _parse_leaf_lines(bytes, key_length, ref_list_length):
 
326
    parser = BTreeLeafParser(bytes, key_length, ref_list_length)
 
327
    return parser.parse()
 
328
 
 
329
 
 
330
# TODO: We can go from 8 byte offset + 4 byte length to a simple lookup,
 
331
#       because the block_offset + length is likely to be repeated. However,
 
332
#       the big win there is to cache across pages, and not just one page
 
333
#       Though if we did cache in a page, we could certainly use a short int.
 
334
#       And this goes from 40 bytes to 30 bytes.
 
335
#       One slightly ugly option would be to cache block offsets in a global.
 
336
#       However, that leads to thread-safety issues, etc.
 
337
ctypedef struct gc_chk_sha1_record:
 
338
    long long block_offset
 
339
    unsigned int block_length
 
340
    unsigned int record_start
 
341
    unsigned int record_end
 
342
    char sha1[20]
 
343
 
 
344
 
 
345
cdef int _unhexbuf[256]
 
346
cdef char *_hexbuf
 
347
_hexbuf = '0123456789abcdef'
 
348
 
 
349
cdef _populate_unhexbuf():
 
350
    cdef int i
 
351
    for i from 0 <= i < 256:
 
352
        _unhexbuf[i] = -1
 
353
    for i from 0 <= i < 10: # 0123456789 => map to the raw number
 
354
        _unhexbuf[(i + c'0')] = i
 
355
    for i from 10 <= i < 16: # abcdef => 10, 11, 12, 13, 14, 15, 16
 
356
        _unhexbuf[(i - 10 + c'a')] = i
 
357
    for i from 10 <= i < 16: # ABCDEF => 10, 11, 12, 13, 14, 15, 16
 
358
        _unhexbuf[(i - 10 + c'A')] = i
 
359
_populate_unhexbuf()
 
360
 
 
361
 
 
362
cdef int _unhexlify_sha1(char *as_hex, char *as_bin): # cannot_raise
 
363
    """Take the hex sha1 in as_hex and make it binary in as_bin
 
364
    
 
365
    Same as binascii.unhexlify, but working on C strings, not Python objects.
 
366
    """
 
367
    cdef int top
 
368
    cdef int bot
 
369
    cdef int i, j
 
370
    cdef char *cur
 
371
    
 
372
    # binascii does this using isupper() and tolower() and ?: syntax. I'm
 
373
    # guessing a simple lookup array should be faster.
 
374
    j = 0
 
375
    for i from 0 <= i < 20:
 
376
        top = _unhexbuf[<unsigned char>(as_hex[j])]
 
377
        j = j + 1
 
378
        bot = _unhexbuf[<unsigned char>(as_hex[j])]
 
379
        j = j + 1
 
380
        if top == -1 or bot == -1:
 
381
            return 0
 
382
        as_bin[i] = <unsigned char>((top << 4) + bot);
 
383
    return 1
 
384
 
 
385
 
 
386
def _py_unhexlify(as_hex):
 
387
    """For the test infrastructure, just thunks to _unhexlify_sha1"""
 
388
    if len(as_hex) != 40 or not PyString_CheckExact(as_hex):
 
389
        raise ValueError('not a 40-byte hex digest')
 
390
    as_bin = PyString_FromStringAndSize(NULL, 20)
 
391
    if _unhexlify_sha1(PyString_AS_STRING(as_hex), PyString_AS_STRING(as_bin)):
 
392
        return as_bin
 
393
    return None
 
394
 
 
395
 
 
396
cdef void _hexlify_sha1(char *as_bin, char *as_hex): # cannot_raise
 
397
    cdef int i, j
 
398
    cdef char c
 
399
 
 
400
    j = 0
 
401
    for i from 0 <= i < 20:
 
402
        c = as_bin[i]
 
403
        as_hex[j] = _hexbuf[(c>>4)&0xf]
 
404
        j = j + 1
 
405
        as_hex[j] = _hexbuf[(c)&0xf]
 
406
        j = j + 1
 
407
 
 
408
 
 
409
def _py_hexlify(as_bin):
 
410
    """For test infrastructure, thunk to _hexlify_sha1"""
 
411
    if len(as_bin) != 20 or not PyString_CheckExact(as_bin):
 
412
        raise ValueError('not a 20-byte binary digest')
 
413
    as_hex = PyString_FromStringAndSize(NULL, 40)
 
414
    _hexlify_sha1(PyString_AS_STRING(as_bin), PyString_AS_STRING(as_hex))
 
415
    return as_hex
 
416
 
 
417
 
 
418
cdef int _key_to_sha1(key, char *sha1): # cannot_raise
 
419
    """Map a key into its sha1 content.
 
420
 
 
421
    :param key: A tuple of style ('sha1:abcd...',)
 
422
    :param sha1: A char buffer of 20 bytes
 
423
    :return: 1 if this could be converted, 0 otherwise
 
424
    """
 
425
    cdef char *c_val
 
426
    cdef PyObject *p_val
 
427
 
 
428
    if StaticTuple_CheckExact(key) and StaticTuple_GET_SIZE(key) == 1:
 
429
        p_val = <PyObject *>StaticTuple_GET_ITEM(key, 0)
 
430
    elif (PyTuple_CheckExact(key) and PyTuple_GET_SIZE(key) == 1):
 
431
        p_val = PyTuple_GET_ITEM_ptr_object(key, 0)
 
432
    else:
 
433
        # Not a tuple or a StaticTuple
 
434
        return 0
 
435
    if (PyString_CheckExact_ptr(p_val) and PyString_GET_SIZE_ptr(p_val) == 45):
 
436
        c_val = PyString_AS_STRING_ptr(p_val)
 
437
    else:
 
438
        return 0
 
439
    if strncmp(c_val, 'sha1:', 5) != 0:
 
440
        return 0
 
441
    if not _unhexlify_sha1(c_val + 5, sha1):
 
442
        return 0
 
443
    return 1
 
444
 
 
445
 
 
446
def _py_key_to_sha1(key):
 
447
    """Map a key to a simple sha1 string.
 
448
 
 
449
    This is a testing thunk to the C function.
 
450
    """
 
451
    as_bin_sha = PyString_FromStringAndSize(NULL, 20)
 
452
    if _key_to_sha1(key, PyString_AS_STRING(as_bin_sha)):
 
453
        return as_bin_sha
 
454
    return None
 
455
 
 
456
 
 
457
cdef StaticTuple _sha1_to_key(char *sha1):
 
458
    """Compute a ('sha1:abcd',) key for a given sha1."""
 
459
    cdef StaticTuple key
 
460
    cdef object hexxed
 
461
    cdef char *c_buf
 
462
    hexxed = PyString_FromStringAndSize(NULL, 45)
 
463
    c_buf = PyString_AS_STRING(hexxed)
 
464
    memcpy(c_buf, 'sha1:', 5)
 
465
    _hexlify_sha1(sha1, c_buf+5)
 
466
    key = StaticTuple_New(1)
 
467
    Py_INCREF(hexxed)
 
468
    StaticTuple_SET_ITEM(key, 0, hexxed)
 
469
    # This is a bit expensive. To parse 120 keys takes 48us, to return them all
 
470
    # can be done in 66.6us (so 18.6us to build them all).
 
471
    # Adding simple hash() here brings it to 76.6us (so computing the hash
 
472
    # value of 120keys is 10us), Intern is 86.9us (another 10us to look and add
 
473
    # them to the intern structure.)
 
474
    # However, since we only intern keys that are in active use, it is probably
 
475
    # a win. Since they would have been read from elsewhere anyway.
 
476
    # We *could* hang the PyObject form off of the gc_chk_sha1_record for ones
 
477
    # that we have deserialized. Something to think about, at least.
 
478
    key = StaticTuple_Intern(key)
 
479
    return key
 
480
 
 
481
 
 
482
def _py_sha1_to_key(sha1_bin):
 
483
    """Test thunk to check the sha1 mapping."""
 
484
    if not PyString_CheckExact(sha1_bin) or PyString_GET_SIZE(sha1_bin) != 20:
 
485
        raise ValueError('sha1_bin must be a str of exactly 20 bytes')
 
486
    return _sha1_to_key(PyString_AS_STRING(sha1_bin))
 
487
 
 
488
 
 
489
cdef unsigned int _sha1_to_uint(char *sha1): # cannot_raise
 
490
    cdef unsigned int val
 
491
    # Must be in MSB, because that is how the content is sorted
 
492
    val = (((<unsigned int>(sha1[0]) & 0xff) << 24)
 
493
           | ((<unsigned int>(sha1[1]) & 0xff) << 16)
 
494
           | ((<unsigned int>(sha1[2]) & 0xff) << 8)
 
495
           | ((<unsigned int>(sha1[3]) & 0xff) << 0))
 
496
    return val
 
497
 
 
498
 
 
499
cdef _format_record(gc_chk_sha1_record *record):
 
500
    # This is inefficient to go from a logical state back to a
 
501
    # string, but it makes things work a bit better internally for now.
 
502
    if record.block_offset >= 0xFFFFFFFF:
 
503
        # %llu is what we really want, but unfortunately it was only added
 
504
        # in python 2.7... :(
 
505
        block_offset_str = str(record.block_offset)
 
506
        value = PyString_FromFormat('%s %u %u %u',
 
507
                                PyString_AS_STRING(block_offset_str),
 
508
                                record.block_length,
 
509
                                record.record_start, record.record_end)
 
510
    else:
 
511
        value = PyString_FromFormat('%lu %u %u %u',
 
512
                                    <unsigned long>record.block_offset,
 
513
                                    record.block_length,
 
514
                                    record.record_start, record.record_end)
 
515
    return value
 
516
 
 
517
 
 
518
cdef class GCCHKSHA1LeafNode:
 
519
    """Track all the entries for a given leaf node."""
 
520
 
 
521
    cdef gc_chk_sha1_record *records
 
522
    cdef public object last_key
 
523
    cdef gc_chk_sha1_record *last_record
 
524
    cdef public int num_records
 
525
    # This is the number of bits to shift to get to the interesting byte. A
 
526
    # value of 24 means that the very first byte changes across all keys.
 
527
    # Anything else means that there is a common prefix of bits that we can
 
528
    # ignore. 0 means that at least the first 3 bytes are identical, though
 
529
    # that is going to be very rare
 
530
    cdef public unsigned char common_shift
 
531
    # This maps an interesting byte to the first record that matches.
 
532
    # Equivalent to bisect.bisect_left(self.records, sha1), though only taking
 
533
    # into account that one byte.
 
534
    cdef unsigned char offsets[257]
 
535
 
 
536
    def __sizeof__(self):
 
537
        # :( Why doesn't Pyrex let me do a simple sizeof(GCCHKSHA1LeafNode)
 
538
        # like Cython? Explicitly enumerating everything here seems to leave my
 
539
        # size off by 2 (286 bytes vs 288 bytes actual). I'm guessing it is an
 
540
        # alignment/padding issue. Oh well- at least we scale properly with
 
541
        # num_records and are very close to correct, which is what I care
 
542
        # about.
 
543
        # If we ever decide to require cython:
 
544
        # return (sizeof(GCCHKSHA1LeafNode)
 
545
        #     + sizeof(gc_chk_sha1_record)*self.num_records)
 
546
        return (sizeof(PyObject) + sizeof(void*) + sizeof(int)
 
547
            + sizeof(gc_chk_sha1_record*) + sizeof(PyObject *)
 
548
            + sizeof(gc_chk_sha1_record*) + sizeof(char)
 
549
            + sizeof(unsigned char)*257
 
550
            + sizeof(gc_chk_sha1_record)*self.num_records)
 
551
 
 
552
    def __dealloc__(self):
 
553
        if self.records != NULL:
 
554
            PyMem_Free(self.records)
 
555
            self.records = NULL
 
556
 
 
557
    def __init__(self, bytes):
 
558
        self._parse_bytes(bytes)
 
559
        self.last_key = None
 
560
        self.last_record = NULL
 
561
 
 
562
    property min_key:
 
563
        def __get__(self):
 
564
            if self.num_records > 0:
 
565
                return _sha1_to_key(self.records[0].sha1)
 
566
            return None
 
567
 
 
568
    property max_key:
 
569
        def __get__(self):
 
570
            if self.num_records > 0:
 
571
                return _sha1_to_key(self.records[self.num_records-1].sha1)
 
572
            return None
 
573
 
 
574
    cdef StaticTuple _record_to_value_and_refs(self,
 
575
                                               gc_chk_sha1_record *record):
 
576
        """Extract the refs and value part of this record."""
 
577
        cdef StaticTuple value_and_refs
 
578
        cdef StaticTuple empty
 
579
        value_and_refs = StaticTuple_New(2)
 
580
        value = _format_record(record)
 
581
        Py_INCREF(value)
 
582
        StaticTuple_SET_ITEM(value_and_refs, 0, value)
 
583
        # Always empty refs
 
584
        empty = StaticTuple_New(0)
 
585
        Py_INCREF(empty)
 
586
        StaticTuple_SET_ITEM(value_and_refs, 1, empty)
 
587
        return value_and_refs
 
588
 
 
589
    cdef StaticTuple _record_to_item(self, gc_chk_sha1_record *record):
 
590
        """Turn a given record back into a fully fledged item.
 
591
        """
 
592
        cdef StaticTuple item
 
593
        cdef StaticTuple key
 
594
        cdef StaticTuple value_and_refs
 
595
        cdef object value
 
596
        key = _sha1_to_key(record.sha1)
 
597
        item = StaticTuple_New(2)
 
598
        Py_INCREF(key)
 
599
        StaticTuple_SET_ITEM(item, 0, key)
 
600
        value_and_refs = self._record_to_value_and_refs(record)
 
601
        Py_INCREF(value_and_refs)
 
602
        StaticTuple_SET_ITEM(item, 1, value_and_refs)
 
603
        return item
 
604
 
 
605
    cdef gc_chk_sha1_record* _lookup_record(self, char *sha1) except? NULL:
 
606
        """Find a gc_chk_sha1_record that matches the sha1 supplied."""
 
607
        cdef int lo, hi, mid, the_cmp
 
608
        cdef int offset
 
609
 
 
610
        # TODO: We can speed up misses by comparing this sha1 to the common
 
611
        #       bits, and seeing if the common prefix matches, if not, we don't
 
612
        #       need to search for anything because it cannot match
 
613
        # Use the offset array to find the closest fit for this entry
 
614
        # follow that up with bisecting, since multiple keys can be in one
 
615
        # spot
 
616
        # Bisecting dropped us from 7000 comparisons to 582 (4.8/key), using
 
617
        # the offset array dropped us from 23us to 20us and 156 comparisions
 
618
        # (1.3/key)
 
619
        offset = self._offset_for_sha1(sha1)
 
620
        lo = self.offsets[offset]
 
621
        hi = self.offsets[offset+1]
 
622
        if hi == 255:
 
623
            # if hi == 255 that means we potentially ran off the end of the
 
624
            # list, so push it up to num_records
 
625
            # note that if 'lo' == 255, that is ok, because we can start
 
626
            # searching from that part of the list.
 
627
            hi = self.num_records
 
628
        local_n_cmp = 0
 
629
        while lo < hi:
 
630
            mid = (lo + hi) / 2
 
631
            the_cmp = memcmp(self.records[mid].sha1, sha1, 20)
 
632
            if the_cmp == 0:
 
633
                return &self.records[mid]
 
634
            elif the_cmp < 0:
 
635
                lo = mid + 1
 
636
            else:
 
637
                hi = mid
 
638
        return NULL
 
639
 
 
640
    def __contains__(self, key):
 
641
        cdef char sha1[20]
 
642
        cdef gc_chk_sha1_record *record
 
643
        if _key_to_sha1(key, sha1):
 
644
            # If it isn't a sha1 key, then it won't be in this leaf node
 
645
            record = self._lookup_record(sha1)
 
646
            if record != NULL:
 
647
                self.last_key = key
 
648
                self.last_record = record
 
649
                return True
 
650
        return False
 
651
 
 
652
    def __getitem__(self, key):
 
653
        cdef char sha1[20]
 
654
        cdef gc_chk_sha1_record *record
 
655
        record = NULL
 
656
        if self.last_record != NULL and key is self.last_key:
 
657
            record = self.last_record
 
658
        elif _key_to_sha1(key, sha1):
 
659
            record = self._lookup_record(sha1)
 
660
        if record == NULL:
 
661
            raise KeyError('key %r is not present' % (key,))
 
662
        return self._record_to_value_and_refs(record)
 
663
 
 
664
    def __len__(self):
 
665
        return self.num_records
 
666
 
 
667
    def all_keys(self):
 
668
        cdef int i
 
669
        result = []
 
670
        for i from 0 <= i < self.num_records:
 
671
            PyList_Append(result, _sha1_to_key(self.records[i].sha1))
 
672
        return result
 
673
 
 
674
    def all_items(self):
 
675
        cdef int i
 
676
        result = []
 
677
        for i from 0 <= i < self.num_records:
 
678
            item = self._record_to_item(&self.records[i])
 
679
            PyList_Append(result, item)
 
680
        return result
 
681
 
 
682
    cdef int _count_records(self, char *c_content, char *c_end): # cannot_raise
 
683
        """Count how many records are in this section."""
 
684
        cdef char *c_cur
 
685
        cdef int num_records
 
686
 
 
687
        c_cur = c_content
 
688
        num_records = 0
 
689
        while c_cur != NULL and c_cur < c_end:
 
690
            c_cur = <char *>memchr(c_cur, c'\n', c_end - c_cur);
 
691
            if c_cur == NULL:
 
692
                break
 
693
            c_cur = c_cur + 1
 
694
            num_records = num_records + 1
 
695
        return num_records
 
696
 
 
697
    cdef _parse_bytes(self, bytes):
 
698
        """Parse the string 'bytes' into content."""
 
699
        cdef char *c_bytes
 
700
        cdef char *c_cur
 
701
        cdef char *c_end
 
702
        cdef Py_ssize_t n_bytes
 
703
        cdef int num_records
 
704
        cdef int entry
 
705
        cdef gc_chk_sha1_record *cur_record
 
706
 
 
707
        if not PyString_CheckExact(bytes):
 
708
            raise TypeError('We only support parsing plain 8-bit strings.')
 
709
        # Pass 1, count how many records there will be
 
710
        n_bytes = PyString_GET_SIZE(bytes)
 
711
        c_bytes = PyString_AS_STRING(bytes)
 
712
        c_end = c_bytes + n_bytes
 
713
        if strncmp(c_bytes, 'type=leaf\n', 10):
 
714
            raise ValueError("bytes did not start with 'type=leaf\\n': %r"
 
715
                             % (bytes[:10],))
 
716
        c_cur = c_bytes + 10
 
717
        num_records = self._count_records(c_cur, c_end)
 
718
        # Now allocate the memory for these items, and go to town
 
719
        self.records = <gc_chk_sha1_record*>PyMem_Malloc(num_records *
 
720
            (sizeof(unsigned short) + sizeof(gc_chk_sha1_record)))
 
721
        self.num_records = num_records
 
722
        cur_record = self.records
 
723
        entry = 0
 
724
        while c_cur != NULL and c_cur < c_end and entry < num_records:
 
725
            c_cur = self._parse_one_entry(c_cur, c_end, cur_record)
 
726
            cur_record = cur_record + 1
 
727
            entry = entry + 1
 
728
        if (entry != self.num_records
 
729
            or c_cur != c_end
 
730
            or cur_record != self.records + self.num_records):
 
731
            raise ValueError('Something went wrong while parsing.')
 
732
        # Pass 3: build the offset map
 
733
        self._compute_common()
 
734
 
 
735
    cdef char *_parse_one_entry(self, char *c_cur, char *c_end,
 
736
                                gc_chk_sha1_record *cur_record) except NULL:
 
737
        """Read a single sha record from the bytes.
 
738
        :param c_cur: The pointer to the start of bytes
 
739
        :param cur_record: 
 
740
        """
 
741
        cdef char *c_next
 
742
        if strncmp(c_cur, 'sha1:', 5):
 
743
            raise ValueError('line did not start with sha1: %r'
 
744
                % (safe_string_from_size(c_cur, 10),))
 
745
        c_cur = c_cur + 5
 
746
        c_next = <char *>memchr(c_cur, c'\0', c_end - c_cur)
 
747
        if c_next == NULL or (c_next - c_cur != 40):
 
748
            raise ValueError('Line did not contain 40 hex bytes')
 
749
        if not _unhexlify_sha1(c_cur, cur_record.sha1):
 
750
            raise ValueError('We failed to unhexlify')
 
751
        c_cur = c_next + 1
 
752
        if c_cur[0] != c'\0':
 
753
            raise ValueError('only 1 null, not 2 as expected')
 
754
        c_cur = c_cur + 1
 
755
        cur_record.block_offset = strtoll(c_cur, &c_next, 10)
 
756
        if c_cur == c_next or c_next[0] != c' ':
 
757
            raise ValueError('Failed to parse block offset')
 
758
        c_cur = c_next + 1
 
759
        cur_record.block_length = strtoul(c_cur, &c_next, 10)
 
760
        if c_cur == c_next or c_next[0] != c' ':
 
761
            raise ValueError('Failed to parse block length')
 
762
        c_cur = c_next + 1
 
763
        cur_record.record_start = strtoul(c_cur, &c_next, 10)
 
764
        if c_cur == c_next or c_next[0] != c' ':
 
765
            raise ValueError('Failed to parse block length')
 
766
        c_cur = c_next + 1
 
767
        cur_record.record_end = strtoul(c_cur, &c_next, 10)
 
768
        if c_cur == c_next or c_next[0] != c'\n':
 
769
            raise ValueError('Failed to parse record end')
 
770
        c_cur = c_next + 1
 
771
        return c_cur
 
772
 
 
773
    cdef int _offset_for_sha1(self, char *sha1) except -1:
 
774
        """Find the first interesting 8-bits of this sha1."""
 
775
        cdef int this_offset
 
776
        cdef unsigned int as_uint
 
777
        as_uint = _sha1_to_uint(sha1)
 
778
        this_offset = (as_uint >> self.common_shift) & 0xFF
 
779
        return this_offset
 
780
 
 
781
    def _get_offset_for_sha1(self, sha1):
 
782
        return self._offset_for_sha1(PyString_AS_STRING(sha1))
 
783
 
 
784
    cdef _compute_common(self):
 
785
        cdef unsigned int first
 
786
        cdef unsigned int this
 
787
        cdef unsigned int common_mask
 
788
        cdef unsigned char common_shift
 
789
        cdef int i
 
790
        cdef int offset, this_offset
 
791
        cdef int max_offset
 
792
        # The idea with the offset map is that we should be able to quickly
 
793
        # jump to the key that matches a gives sha1. We know that the keys are
 
794
        # in sorted order, and we know that a lot of the prefix is going to be
 
795
        # the same across them.
 
796
        # By XORing the records together, we can determine what bits are set in
 
797
        # all of them
 
798
        if self.num_records < 2:
 
799
            # Everything is in common if you have 0 or 1 leaves
 
800
            # So we'll always just shift to the first byte
 
801
            self.common_shift = 24
 
802
        else:
 
803
            common_mask = 0xFFFFFFFF
 
804
            first = _sha1_to_uint(self.records[0].sha1)
 
805
            for i from 0 < i < self.num_records:
 
806
                this = _sha1_to_uint(self.records[i].sha1)
 
807
                common_mask = (~(first ^ this)) & common_mask
 
808
            common_shift = 24
 
809
            while common_mask & 0x80000000 and common_shift > 0:
 
810
                common_mask = common_mask << 1
 
811
                common_shift = common_shift - 1
 
812
            self.common_shift = common_shift
 
813
        offset = 0
 
814
        max_offset = self.num_records
 
815
        # We cap this loop at 254 records. All the other offsets just get
 
816
        # filled with 0xff as the singleton saying 'too many'.
 
817
        # It means that if we have >255 records we have to bisect the second
 
818
        # half of the list, but this is going to be very rare in practice.
 
819
        if max_offset > 255:
 
820
            max_offset = 255
 
821
        for i from 0 <= i < max_offset:
 
822
            this_offset = self._offset_for_sha1(self.records[i].sha1)
 
823
            while offset <= this_offset:
 
824
                self.offsets[offset] = i
 
825
                offset = offset + 1
 
826
        while offset < 257:
 
827
            self.offsets[offset] = max_offset
 
828
            offset = offset + 1
 
829
 
 
830
    def _get_offsets(self):
 
831
        cdef int i
 
832
        result = []
 
833
        for i from 0 <= i < 257:
 
834
            PyList_Append(result, self.offsets[i])
 
835
        return result
 
836
 
 
837
 
 
838
def _parse_into_chk(bytes, key_length, ref_list_length):
 
839
    """Parse into a format optimized for chk records."""
 
840
    assert key_length == 1
 
841
    assert ref_list_length == 0
 
842
    return GCCHKSHA1LeafNode(bytes)
 
843
 
 
844
 
 
845
def _flatten_node(node, reference_lists):
 
846
    """Convert a node into the serialized form.
 
847
 
 
848
    :param node: A tuple representing a node:
 
849
        (index, key_tuple, value, references)
 
850
    :param reference_lists: Does this index have reference lists?
 
851
    :return: (string_key, flattened)
 
852
        string_key  The serialized key for referencing this node
 
853
        flattened   A string with the serialized form for the contents
 
854
    """
 
855
    cdef int have_reference_lists
 
856
    cdef Py_ssize_t flat_len
 
857
    cdef Py_ssize_t key_len
 
858
    cdef Py_ssize_t node_len
 
859
    cdef char * value
 
860
    cdef Py_ssize_t value_len
 
861
    cdef char * out
 
862
    cdef Py_ssize_t refs_len
 
863
    cdef Py_ssize_t next_len
 
864
    cdef int first_ref_list
 
865
    cdef int first_reference
 
866
    cdef int i
 
867
    cdef Py_ssize_t ref_bit_len
 
868
 
 
869
    if not PyTuple_CheckExact(node) and not StaticTuple_CheckExact(node):
 
870
        raise TypeError('We expected a tuple() or StaticTuple() for node not: %s'
 
871
            % type(node))
 
872
    node_len = len(node)
 
873
    have_reference_lists = reference_lists
 
874
    if have_reference_lists:
 
875
        if node_len != 4:
 
876
            raise ValueError('With ref_lists, we expected 4 entries not: %s'
 
877
                % len(node))
 
878
    elif node_len < 3:
 
879
        raise ValueError('Without ref_lists, we need at least 3 entries not: %s'
 
880
            % len(node))
 
881
    # TODO: We can probably do better than string.join(), namely
 
882
    #       when key has only 1 item, we can just grab that string
 
883
    #       And when there are 2 items, we could do a single malloc + len() + 1
 
884
    #       also, doing .join() requires a PyObject_GetAttrString call, which
 
885
    #       we could also avoid.
 
886
    # TODO: Note that pyrex 0.9.6 generates fairly crummy code here, using the
 
887
    #       python object interface, versus 0.9.8+ which uses a helper that
 
888
    #       checks if this supports the sequence interface.
 
889
    #       We *could* do more work on our own, and grab the actual items
 
890
    #       lists. For now, just ask people to use a better compiler. :)
 
891
    string_key = '\0'.join(node[1])
 
892
 
 
893
    # TODO: instead of using string joins, precompute the final string length,
 
894
    #       and then malloc a single string and copy everything in.
 
895
 
 
896
    # TODO: We probably want to use PySequenceFast, because we have lists and
 
897
    #       tuples, but we aren't sure which we will get.
 
898
 
 
899
    # line := string_key NULL flat_refs NULL value LF
 
900
    # string_key := BYTES (NULL BYTES)*
 
901
    # flat_refs := ref_list (TAB ref_list)*
 
902
    # ref_list := ref (CR ref)*
 
903
    # ref := BYTES (NULL BYTES)*
 
904
    # value := BYTES
 
905
    refs_len = 0
 
906
    if have_reference_lists:
 
907
        # Figure out how many bytes it will take to store the references
 
908
        ref_lists = node[3]
 
909
        next_len = len(ref_lists) # TODO: use a Py function
 
910
        if next_len > 0:
 
911
            # If there are no nodes, we don't need to do any work
 
912
            # Otherwise we will need (len - 1) '\t' characters to separate
 
913
            # the reference lists
 
914
            refs_len = refs_len + (next_len - 1)
 
915
            for ref_list in ref_lists:
 
916
                next_len = len(ref_list)
 
917
                if next_len > 0:
 
918
                    # We will need (len - 1) '\r' characters to separate the
 
919
                    # references
 
920
                    refs_len = refs_len + (next_len - 1)
 
921
                    for reference in ref_list:
 
922
                        if (not PyTuple_CheckExact(reference)
 
923
                            and not StaticTuple_CheckExact(reference)):
 
924
                            raise TypeError(
 
925
                                'We expect references to be tuples not: %s'
 
926
                                % type(reference))
 
927
                        next_len = len(reference)
 
928
                        if next_len > 0:
 
929
                            # We will need (len - 1) '\x00' characters to
 
930
                            # separate the reference key
 
931
                            refs_len = refs_len + (next_len - 1)
 
932
                            for ref_bit in reference:
 
933
                                if not PyString_CheckExact(ref_bit):
 
934
                                    raise TypeError('We expect reference bits'
 
935
                                        ' to be strings not: %s'
 
936
                                        % type(<object>ref_bit))
 
937
                                refs_len = refs_len + PyString_GET_SIZE(ref_bit)
 
938
 
 
939
    # So we have the (key NULL refs NULL value LF)
 
940
    key_len = PyString_Size(string_key)
 
941
    val = node[2]
 
942
    if not PyString_CheckExact(val):
 
943
        raise TypeError('Expected a plain str for value not: %s'
 
944
                        % type(val))
 
945
    value = PyString_AS_STRING(val)
 
946
    value_len = PyString_GET_SIZE(val)
 
947
    flat_len = (key_len + 1 + refs_len + 1 + value_len + 1)
 
948
    line = PyString_FromStringAndSize(NULL, flat_len)
 
949
    # Get a pointer to the new buffer
 
950
    out = PyString_AsString(line)
 
951
    memcpy(out, PyString_AsString(string_key), key_len)
 
952
    out = out + key_len
 
953
    out[0] = c'\0'
 
954
    out = out + 1
 
955
    if refs_len > 0:
 
956
        first_ref_list = 1
 
957
        for ref_list in ref_lists:
 
958
            if first_ref_list == 0:
 
959
                out[0] = c'\t'
 
960
                out = out + 1
 
961
            first_ref_list = 0
 
962
            first_reference = 1
 
963
            for reference in ref_list:
 
964
                if first_reference == 0:
 
965
                    out[0] = c'\r'
 
966
                    out = out + 1
 
967
                first_reference = 0
 
968
                next_len = len(reference)
 
969
                for i from 0 <= i < next_len:
 
970
                    if i != 0:
 
971
                        out[0] = c'\x00'
 
972
                        out = out + 1
 
973
                    ref_bit = reference[i]
 
974
                    ref_bit_len = PyString_GET_SIZE(ref_bit)
 
975
                    memcpy(out, PyString_AS_STRING(ref_bit), ref_bit_len)
 
976
                    out = out + ref_bit_len
 
977
    out[0] = c'\0'
 
978
    out = out  + 1
 
979
    memcpy(out, value, value_len)
 
980
    out = out + value_len
 
981
    out[0] = c'\n'
 
982
    return string_key, line