← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/_btree_serializer_pyx.pyx

  • Committer: Robert Collins
  • Date: 2009-09-07 03:08:30 UTC
  • mto: This revision was merged to the branch mainline in revision 4690.
  • Revision ID: robertc@robertcollins.net-20090907030830-rf59kt28d550eauj
Milestones language tightning, internal consistency.

Show diffs side-by-side

added added

removed removed

Lines of Context:
bzrlib/_btree_serializer_pyx.pyx
 
1
# Copyright (C) 2008, 2009 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
    int PyString_CheckExact(object s)
 
37
    int PyString_CheckExact_ptr "PyString_CheckExact" (PyObject *)
 
38
    Py_ssize_t PyString_Size(object p)
 
39
    Py_ssize_t PyString_GET_SIZE_ptr "PyString_GET_SIZE" (PyObject *)
 
40
    char * PyString_AS_STRING_ptr "PyString_AS_STRING" (PyObject *)
 
41
    int PyString_AsStringAndSize_ptr(PyObject *, char **buf, Py_ssize_t *len)
 
42
    void PyString_InternInPlace(PyObject **)
 
43
    int PyTuple_CheckExact(object t)
 
44
    object PyTuple_New(Py_ssize_t n_entries)
 
45
    void PyTuple_SET_ITEM(object, Py_ssize_t offset, object) # steals the ref
 
46
    Py_ssize_t PyTuple_GET_SIZE(object t)
 
47
    PyObject *PyTuple_GET_ITEM_ptr_object "PyTuple_GET_ITEM" (object tpl, int index)
 
48
    void Py_INCREF(object)
 
49
    void Py_DECREF_ptr "Py_DECREF" (PyObject *)
 
50
 
 
51
cdef extern from "string.h":
 
52
    void *memcpy(void *dest, void *src, size_t n)
 
53
    void *memchr(void *s, int c, size_t n)
 
54
    # GNU extension
 
55
    # void *memrchr(void *s, int c, size_t n)
 
56
    int strncmp(char *s1, char *s2, size_t n)
 
57
 
 
58
 
 
59
# TODO: Find some way to import this from _dirstate_helpers
 
60
cdef void* _my_memrchr(void *s, int c, size_t n):
 
61
    # memrchr seems to be a GNU extension, so we have to implement it ourselves
 
62
    # It is not present in any win32 standard library
 
63
    cdef char *pos
 
64
    cdef char *start
 
65
 
 
66
    start = <char*>s
 
67
    pos = start + n - 1
 
68
    while pos >= start:
 
69
        if pos[0] == c:
 
70
            return <void*>pos
 
71
        pos = pos - 1
 
72
    return NULL
 
73
 
 
74
# TODO: Import this from _dirstate_helpers when it is merged
 
75
cdef object safe_string_from_size(char *s, Py_ssize_t size):
 
76
    if size < 0:
 
77
        raise AssertionError(
 
78
            'tried to create a string with an invalid size: %d @0x%x'
 
79
            % (size, <int>s))
 
80
    return PyString_FromStringAndSize(s, size)
 
81
 
 
82
 
 
83
cdef object safe_interned_string_from_size(char *s, Py_ssize_t size):
 
84
    cdef PyObject *py_str
 
85
    if size < 0:
 
86
        raise AssertionError(
 
87
            'tried to create a string with an invalid size: %d @0x%x'
 
88
            % (size, <int>s))
 
89
    py_str = PyString_FromStringAndSize_ptr(s, size)
 
90
    PyString_InternInPlace(&py_str)
 
91
    result = <object>py_str
 
92
    # Casting a PyObject* to an <object> triggers an INCREF from Pyrex, so we
 
93
    # DECREF it to avoid geting immortal strings
 
94
    Py_DECREF_ptr(py_str)
 
95
    return result
 
96
 
 
97
 
 
98
cdef class BTreeLeafParser:
 
99
    """Parse the leaf nodes of a BTree index.
 
100
 
 
101
    :ivar bytes: The PyString object containing the uncompressed text for the
 
102
        node.
 
103
    :ivar key_length: An integer describing how many pieces the keys have for
 
104
        this index.
 
105
    :ivar ref_list_length: An integer describing how many references this index
 
106
        contains.
 
107
    :ivar keys: A PyList of keys found in this node.
 
108
 
 
109
    :ivar _cur_str: A pointer to the start of the next line to parse
 
110
    :ivar _end_str: A pointer to the end of bytes
 
111
    :ivar _start: Pointer to the location within the current line while
 
112
        parsing.
 
113
    :ivar _header_found: True when we have parsed the header for this node
 
114
    """
 
115
 
 
116
    cdef object bytes
 
117
    cdef int key_length
 
118
    cdef int ref_list_length
 
119
    cdef object keys
 
120
 
 
121
    cdef char * _cur_str
 
122
    cdef char * _end_str
 
123
    # The current start point for parsing
 
124
    cdef char * _start
 
125
 
 
126
    cdef int _header_found
 
127
 
 
128
    def __init__(self, bytes, key_length, ref_list_length):
 
129
        self.bytes = bytes
 
130
        self.key_length = key_length
 
131
        self.ref_list_length = ref_list_length
 
132
        self.keys = []
 
133
        self._cur_str = NULL
 
134
        self._end_str = NULL
 
135
        self._header_found = 0
 
136
 
 
137
    cdef extract_key(self, char * last):
 
138
        """Extract a key.
 
139
 
 
140
        :param last: points at the byte after the last byte permitted for the
 
141
            key.
 
142
        """
 
143
        cdef char *temp_ptr
 
144
        cdef int loop_counter
 
145
        # keys are tuples
 
146
        key = PyTuple_New(self.key_length)
 
147
        for loop_counter from 0 <= loop_counter < self.key_length:
 
148
            # grab a key segment
 
149
            temp_ptr = <char*>memchr(self._start, c'\0', last - self._start)
 
150
            if temp_ptr == NULL:
 
151
                if loop_counter + 1 == self.key_length:
 
152
                    # capture to last
 
153
                    temp_ptr = last
 
154
                else:
 
155
                    # Invalid line
 
156
                    failure_string = ("invalid key, wanted segment from " +
 
157
                        repr(safe_string_from_size(self._start,
 
158
                                                   last - self._start)))
 
159
                    raise AssertionError(failure_string)
 
160
            # capture the key string
 
161
            # TODO: Consider using PyIntern_FromString, the only caveat is that
 
162
            # it assumes a NULL-terminated string, so we have to check if
 
163
            # temp_ptr[0] == c'\0' or some other char.
 
164
            key_element = safe_interned_string_from_size(self._start,
 
165
                                                         temp_ptr - self._start)
 
166
            # advance our pointer
 
167
            self._start = temp_ptr + 1
 
168
            Py_INCREF(key_element)
 
169
            PyTuple_SET_ITEM(key, loop_counter, key_element)
 
170
        return key
 
171
 
 
172
    cdef int process_line(self) except -1:
 
173
        """Process a line in the bytes."""
 
174
        cdef char *last
 
175
        cdef char *temp_ptr
 
176
        cdef char *ref_ptr
 
177
        cdef char *next_start
 
178
        cdef int loop_counter
 
179
 
 
180
        self._start = self._cur_str
 
181
        # Find the next newline
 
182
        last = <char*>memchr(self._start, c'\n', self._end_str - self._start)
 
183
        if last == NULL:
 
184
            # Process until the end of the file
 
185
            last = self._end_str
 
186
            self._cur_str = self._end_str
 
187
        else:
 
188
            # And the next string is right after it
 
189
            self._cur_str = last + 1
 
190
            # The last character is right before the '\n'
 
191
            last = last
 
192
 
 
193
        if last == self._start:
 
194
            # parsed it all.
 
195
            return 0
 
196
        if last < self._start:
 
197
            # Unexpected error condition - fail
 
198
            return -1
 
199
        if 0 == self._header_found:
 
200
            # The first line in a leaf node is the header "type=leaf\n"
 
201
            if strncmp("type=leaf", self._start, last - self._start) == 0:
 
202
                self._header_found = 1
 
203
                return 0
 
204
            else:
 
205
                raise AssertionError('Node did not start with "type=leaf": %r'
 
206
                    % (safe_string_from_size(self._start, last - self._start)))
 
207
                return -1
 
208
 
 
209
        key = self.extract_key(last)
 
210
        # find the value area
 
211
        temp_ptr = <char*>_my_memrchr(self._start, c'\0', last - self._start)
 
212
        if temp_ptr == NULL:
 
213
            # Invalid line
 
214
            return -1
 
215
        else:
 
216
            # capture the value string
 
217
            value = safe_string_from_size(temp_ptr + 1, last - temp_ptr - 1)
 
218
            # shrink the references end point
 
219
            last = temp_ptr
 
220
        if self.ref_list_length:
 
221
            ref_lists = []
 
222
            loop_counter = 0
 
223
            while loop_counter < self.ref_list_length:
 
224
                ref_list = []
 
225
                # extract a reference list
 
226
                loop_counter = loop_counter + 1
 
227
                if last < self._start:
 
228
                    return -1
 
229
                # find the next reference list end point:
 
230
                temp_ptr = <char*>memchr(self._start, c'\t', last - self._start)
 
231
                if temp_ptr == NULL:
 
232
                    # Only valid for the last list
 
233
                    if loop_counter != self.ref_list_length:
 
234
                        # Invalid line
 
235
                        return -1
 
236
                        raise AssertionError("invalid key")
 
237
                    else:
 
238
                        # scan to the end of the ref list area
 
239
                        ref_ptr = last
 
240
                        next_start = last
 
241
                else:
 
242
                    # scan to the end of this ref list
 
243
                    ref_ptr = temp_ptr
 
244
                    next_start = temp_ptr + 1
 
245
                # Now, there may be multiple keys in the ref list.
 
246
                while self._start < ref_ptr:
 
247
                    # loop finding keys and extracting them
 
248
                    temp_ptr = <char*>memchr(self._start, c'\r',
 
249
                                             ref_ptr - self._start)
 
250
                    if temp_ptr == NULL:
 
251
                        # key runs to the end
 
252
                        temp_ptr = ref_ptr
 
253
                    PyList_Append(ref_list, self.extract_key(temp_ptr))
 
254
                PyList_Append(ref_lists, tuple(ref_list))
 
255
                # prepare for the next reference list
 
256
                self._start = next_start
 
257
            ref_lists = tuple(ref_lists)
 
258
            node_value = (value, ref_lists)
 
259
        else:
 
260
            if last != self._start:
 
261
                # unexpected reference data present
 
262
                return -1
 
263
            node_value = (value, ())
 
264
        PyList_Append(self.keys, (key, node_value))
 
265
        return 0
 
266
 
 
267
    def parse(self):
 
268
        cdef Py_ssize_t byte_count
 
269
        if not PyString_CheckExact(self.bytes):
 
270
            raise AssertionError('self.bytes is not a string.')
 
271
        byte_count = PyString_Size(self.bytes)
 
272
        self._cur_str = PyString_AsString(self.bytes)
 
273
        # This points to the last character in the string
 
274
        self._end_str = self._cur_str + byte_count
 
275
        while self._cur_str < self._end_str:
 
276
            self.process_line()
 
277
        return self.keys
 
278
 
 
279
 
 
280
def _parse_leaf_lines(bytes, key_length, ref_list_length):
 
281
    parser = BTreeLeafParser(bytes, key_length, ref_list_length)
 
282
    return parser.parse()
 
283
 
 
284
 
 
285
def _flatten_node(node, reference_lists):
 
286
    """Convert a node into the serialized form.
 
287
 
 
288
    :param node: A tuple representing a node:
 
289
        (index, key_tuple, value, references)
 
290
    :param reference_lists: Does this index have reference lists?
 
291
    :return: (string_key, flattened)
 
292
        string_key  The serialized key for referencing this node
 
293
        flattened   A string with the serialized form for the contents
 
294
    """
 
295
    cdef int have_reference_lists
 
296
    cdef Py_ssize_t flat_len
 
297
    cdef Py_ssize_t key_len
 
298
    cdef Py_ssize_t node_len
 
299
    cdef PyObject * val
 
300
    cdef char * value
 
301
    cdef Py_ssize_t value_len
 
302
    cdef char * out
 
303
    cdef Py_ssize_t refs_len
 
304
    cdef Py_ssize_t next_len
 
305
    cdef int first_ref_list
 
306
    cdef int first_reference
 
307
    cdef int i
 
308
    cdef PyObject *ref_bit
 
309
    cdef Py_ssize_t ref_bit_len
 
310
 
 
311
    if not PyTuple_CheckExact(node):
 
312
        raise TypeError('We expected a tuple() for node not: %s'
 
313
            % type(node))
 
314
    node_len = PyTuple_GET_SIZE(node)
 
315
    have_reference_lists = reference_lists
 
316
    if have_reference_lists:
 
317
        if node_len != 4:
 
318
            raise ValueError('With ref_lists, we expected 4 entries not: %s'
 
319
                % len(node))
 
320
    elif node_len < 3:
 
321
        raise ValueError('Without ref_lists, we need at least 3 entries not: %s'
 
322
            % len(node))
 
323
    # I don't expect that we can do faster than string.join()
 
324
    string_key = '\0'.join(<object>PyTuple_GET_ITEM_ptr_object(node, 1))
 
325
 
 
326
    # TODO: instead of using string joins, precompute the final string length,
 
327
    #       and then malloc a single string and copy everything in.
 
328
 
 
329
    # TODO: We probably want to use PySequenceFast, because we have lists and
 
330
    #       tuples, but we aren't sure which we will get.
 
331
 
 
332
    # line := string_key NULL flat_refs NULL value LF
 
333
    # string_key := BYTES (NULL BYTES)*
 
334
    # flat_refs := ref_list (TAB ref_list)*
 
335
    # ref_list := ref (CR ref)*
 
336
    # ref := BYTES (NULL BYTES)*
 
337
    # value := BYTES
 
338
    refs_len = 0
 
339
    if have_reference_lists:
 
340
        # Figure out how many bytes it will take to store the references
 
341
        ref_lists = <object>PyTuple_GET_ITEM_ptr_object(node, 3)
 
342
        next_len = len(ref_lists) # TODO: use a Py function
 
343
        if next_len > 0:
 
344
            # If there are no nodes, we don't need to do any work
 
345
            # Otherwise we will need (len - 1) '\t' characters to separate
 
346
            # the reference lists
 
347
            refs_len = refs_len + (next_len - 1)
 
348
            for ref_list in ref_lists:
 
349
                next_len = len(ref_list)
 
350
                if next_len > 0:
 
351
                    # We will need (len - 1) '\r' characters to separate the
 
352
                    # references
 
353
                    refs_len = refs_len + (next_len - 1)
 
354
                    for reference in ref_list:
 
355
                        if not PyTuple_CheckExact(reference):
 
356
                            raise TypeError(
 
357
                                'We expect references to be tuples not: %s'
 
358
                                % type(reference))
 
359
                        next_len = PyTuple_GET_SIZE(reference)
 
360
                        if next_len > 0:
 
361
                            # We will need (len - 1) '\x00' characters to
 
362
                            # separate the reference key
 
363
                            refs_len = refs_len + (next_len - 1)
 
364
                            for i from 0 <= i < next_len:
 
365
                                ref_bit = PyTuple_GET_ITEM_ptr_object(reference, i)
 
366
                                if not PyString_CheckExact_ptr(ref_bit):
 
367
                                    raise TypeError('We expect reference bits'
 
368
                                        ' to be strings not: %s'
 
369
                                        % type(<object>ref_bit))
 
370
                                refs_len = refs_len + PyString_GET_SIZE_ptr(ref_bit)
 
371
 
 
372
    # So we have the (key NULL refs NULL value LF)
 
373
    key_len = PyString_Size(string_key)
 
374
    val = PyTuple_GET_ITEM_ptr_object(node, 2)
 
375
    if not PyString_CheckExact_ptr(val):
 
376
        raise TypeError('Expected a plain str for value not: %s'
 
377
                        % type(<object>val))
 
378
    value = PyString_AS_STRING_ptr(val)
 
379
    value_len = PyString_GET_SIZE_ptr(val)
 
380
    flat_len = (key_len + 1 + refs_len + 1 + value_len + 1)
 
381
    line = PyString_FromStringAndSize(NULL, flat_len)
 
382
    # Get a pointer to the new buffer
 
383
    out = PyString_AsString(line)
 
384
    memcpy(out, PyString_AsString(string_key), key_len)
 
385
    out = out + key_len
 
386
    out[0] = c'\0'
 
387
    out = out + 1
 
388
    if refs_len > 0:
 
389
        first_ref_list = 1
 
390
        for ref_list in ref_lists:
 
391
            if first_ref_list == 0:
 
392
                out[0] = c'\t'
 
393
                out = out + 1
 
394
            first_ref_list = 0
 
395
            first_reference = 1
 
396
            for reference in ref_list:
 
397
                if first_reference == 0:
 
398
                    out[0] = c'\r'
 
399
                    out = out + 1
 
400
                first_reference = 0
 
401
                next_len = PyTuple_GET_SIZE(reference)
 
402
                for i from 0 <= i < next_len:
 
403
                    if i != 0:
 
404
                        out[0] = c'\x00'
 
405
                        out = out + 1
 
406
                    ref_bit = PyTuple_GET_ITEM_ptr_object(reference, i)
 
407
                    ref_bit_len = PyString_GET_SIZE_ptr(ref_bit)
 
408
                    memcpy(out, PyString_AS_STRING_ptr(ref_bit), ref_bit_len)
 
409
                    out = out + ref_bit_len
 
410
    out[0] = c'\0'
 
411
    out = out  + 1
 
412
    memcpy(out, value, value_len)
 
413
    out = out + value_len
 
414
    out[0] = c'\n'
 
415
    return string_key, line