/brz/remove-bazaar

# Copyright (C) 2008, 2009 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

"""Core compression logic for compressing streams of related files."""

from itertools import izip

from cStringIO import StringIO

import struct

import time

import zlib

try:

    import pylzma

except ImportError:

    pylzma = None

from bzrlib import (

    annotate,

    debug,

    diff,

    errors,

    graph as _mod_graph,

    osutils,

    pack,

    patiencediff,

    trace,

    )

from bzrlib.graph import Graph

from bzrlib.knit import _DirectPackAccess

from bzrlib.osutils import (

    contains_whitespace,

    sha_string,

    split_lines,

    )

from bzrlib.btree_index import BTreeBuilder

from bzrlib.lru_cache import LRUSizeCache

from bzrlib.tsort import topo_sort

from bzrlib.versionedfile import (

    adapter_registry,

    AbsentContentFactory,

    ChunkedContentFactory,

    FulltextContentFactory,

    VersionedFiles,

    )

_USE_LZMA = False and (pylzma is not None)

_NO_LABELS = True

_FAST = False

# osutils.sha_string('')

_null_sha1 = 'da39a3ee5e6b4b0d3255bfef95601890afd80709'

def encode_base128_int(val):

    """Convert an integer into a 7-bit lsb encoding."""

    bytes = []

    count = 0

    while val >= 0x80:

        bytes.append(chr((val | 0x80) & 0xFF))

        val >>= 7

    bytes.append(chr(val))

    return ''.join(bytes)

def decode_base128_int(bytes):

    """Decode an integer from a 7-bit lsb encoding."""

    offset = 0

    val = 0

    shift = 0

    bval = ord(bytes[offset])

    while bval >= 0x80:

        val |= (bval & 0x7F) << shift

        shift += 7

        offset += 1

        bval = ord(bytes[offset])

    val |= bval << shift

    offset += 1

    return val, offset

def sort_gc_optimal(parent_map):

    """Sort and group the keys in parent_map into groupcompress order.

    groupcompress is defined (currently) as reverse-topological order, grouped by

    the key prefix.

    :return: A sorted-list of keys

    """

    # groupcompress ordering is approximately reverse topological,

    # properly grouped by file-id.

    per_prefix_map = {}

    for item in parent_map.iteritems():

        key = item[0]

        if isinstance(key, str) or len(key) == 1:

            prefix = ''

        else:

            prefix = key[0]

        try:

            per_prefix_map[prefix].append(item)

        except KeyError:

            per_prefix_map[prefix] = [item]

    present_keys = []

    for prefix in sorted(per_prefix_map):

        present_keys.extend(reversed(topo_sort(per_prefix_map[prefix])))

    return present_keys

class GroupCompressBlockEntry(object):

    """Track the information about a single object inside a GC group.

    This is generally just the dumb data structure.

    """

    def __init__(self, key, type, sha1, start, length):

        self.key = key

        self.type = type # delta, fulltext, external?

        self.sha1 = sha1 # Sha1 of content

        self.start = start # Byte offset to start of data

        self.length = length # Length of content

    def __repr__(self):

        return '%s(%s, %s, %s, %s, %s)' % (

            self.__class__.__name__,

            self.key, self.type, self.sha1, self.start, self.length

            )

    @property

    def end(self):

        return self.start + self.length

# The max zlib window size is 32kB, so if we set 'max_size' output of the

# decompressor to the requested bytes + 32kB, then we should guarantee

# num_bytes coming out.

_ZLIB_DECOMP_WINDOW = 32*1024

class GroupCompressBlock(object):

    """An object which maintains the internal structure of the compressed data.

    This tracks the meta info (start of text, length, type, etc.)

    """

    # Group Compress Block v1 Zlib

    GCB_HEADER = 'gcb1z\n'

    GCB_LZ_HEADER = 'gcb1l\n'

    def __init__(self):

        # map by key? or just order in file?

        self._entries = {}

        self._compressor_name = None

        self._z_header_length = None

        self._header_length = None

        self._z_header = None

        self._z_content = None

        self._z_content_decompressor = None

        self._z_content_length = None

        self._content_length = None

        self._content = None

    def __len__(self):

        return self._content_length + self._header_length

    def _parse_header(self):

        """Parse the header part of the block."""

        assert self._z_header is not None

        if self._z_header == '':

            # Nothing to process

            self._z_header = None

            return

        if self._compressor_name == 'lzma':

            header = pylzma.decompress(self._z_header)

        else:

            assert self._compressor_name == 'zlib'

            header = zlib.decompress(self._z_header)

        self._z_header = None # We have consumed the header

        lines = header.split('\n')

        del header

        info_dict = {}

        for line in lines:

            if not line: #End of record

                if not info_dict:

                    break

                self.add_entry(**info_dict)

                info_dict = {}

                continue

            key, value = line.split(':', 1)

            if key == 'key':

                value = tuple(map(intern, value.split('\x00')))

            elif key in ('start', 'length'):

                value = int(value)

            elif key == 'type':

                value = intern(value)

            info_dict[key] = value

    def _ensure_content(self, num_bytes=None):

        """Make sure that content has been expanded enough.

        :param num_bytes: Ensure that we have extracted at least num_bytes of

            content. If None, consume everything

        """

        # TODO: If we re-use the same content block at different times during

        #       get_record_stream(), it is possible that the first pass will

        #       get inserted, triggering an extract/_ensure_content() which

        #       will get rid of _z_content. And then the next use of the block

        #       will try to access _z_content (to send it over the wire), and

        #       fail because it is already extracted. Consider never releasing

        #       _z_content because of this.

        if num_bytes is None:

            num_bytes = self._content_length

        if self._content_length is not None:

            assert num_bytes <= self._content_length

        if self._content is None:

            assert self._z_content is not None

            if self._z_content == '':

                self._content = ''

            elif self._compressor_name == 'lzma':

                # We don't do partial lzma decomp yet

                self._content = pylzma.decompress(self._z_content)

            else:

                # Start a zlib decompressor

                assert self._compressor_name == 'zlib'

                if num_bytes is None:

                    self._content = zlib.decompress(self._z_content)

                else:

                    self._z_content_decompressor = zlib.decompressobj()

                    # Seed the decompressor with the uncompressed bytes, so

                    # that the rest of the code is simplified

                    self._content = self._z_content_decompressor.decompress(

                        self._z_content, num_bytes + _ZLIB_DECOMP_WINDOW)

                # Any bytes remaining to be decompressed will be in the

                # decompressors 'unconsumed_tail'

        # Do we have enough bytes already?

        if num_bytes is not None and len(self._content) >= num_bytes:

            return

        if num_bytes is None and self._z_content_decompressor is None:

            # We must have already decompressed everything

            return

        # If we got this far, and don't have a decompressor, something is wrong

        assert self._z_content_decompressor is not None

        remaining_decomp = self._z_content_decompressor.unconsumed_tail

        if num_bytes is None:

            if remaining_decomp:

                # We don't know how much is left, but we'll decompress it all

                self._content += self._z_content_decompressor.decompress(

                    remaining_decomp)

                # Note: There what I consider a bug in zlib.decompressobj

                #       If you pass back in the entire unconsumed_tail, only

                #       this time you don't pass a max-size, it doesn't

                #       change the unconsumed_tail back to None/''.

                #       However, we know we are done with the whole stream

                self._z_content_decompressor = None

            self._content_length = len(self._content)

        else:

            # If we have nothing left to decomp, we ran out of decomp bytes

            assert remaining_decomp

            needed_bytes = num_bytes - len(self._content)

            # We always set max_size to 32kB over the minimum needed, so that

            # zlib will give us as much as we really want.

            # TODO: If this isn't good enough, we could make a loop here,

            #       that keeps expanding the request until we get enough

            self._content += self._z_content_decompressor.decompress(

                remaining_decomp, needed_bytes + _ZLIB_DECOMP_WINDOW)

            assert len(self._content) >= num_bytes

            if not self._z_content_decompressor.unconsumed_tail:

                # The stream is finished

                self._z_content_decompressor = None

    def _parse_bytes(self, bytes):

        """Read the various lengths from the header.

        This also populates the various 'compressed' buffers.

        :return: The position in bytes just after the last newline

        """

        # At present, there are 4 lengths to be read, we have 2 integers for

        # the length of the compressed and uncompressed header, and 2 integers

        # for the compressed and uncompressed content

        # 14 bytes can represent > 1TB, so to avoid checking too far, cap the

        # search to 14 bytes.

        pos = bytes.index('\n', 6, 20)

        self._z_header_length = int(bytes[6:pos])

        pos += 1

        pos2 = bytes.index('\n', pos, pos + 14)

        self._header_length = int(bytes[pos:pos2])

        end_of_z_lengths = pos2

        pos2 += 1

        # Older versions don't have the content lengths, if we want to preserve

        # backwards compatibility, we could try/except over these, and allow

        # them to be skipped

        try:

            pos = bytes.index('\n', pos2, pos2 + 14)

            self._z_content_length = int(bytes[pos2:pos])

            pos += 1

            pos2 = bytes.index('\n', pos, pos + 14)

            self._content_length = int(bytes[pos:pos2])

            pos = pos2 + 1

            assert len(bytes) == (pos + self._z_header_length +

                                  self._z_content_length)

            pos2 = pos + self._z_header_length

            self._z_header = bytes[pos:pos2]

            self._z_content = bytes[pos2:]

            assert len(self._z_content) == self._z_content_length

        except ValueError:

            # This is the older form, which did not encode its content length

            pos = end_of_z_lengths + 1

            pos2 = pos + self._z_header_length

            self._z_header = bytes[pos:pos2]

            self._z_content = bytes[pos2:]

            self._z_content_length = len(self._z_content)

    @classmethod

    def from_bytes(cls, bytes):

        out = cls()

        if bytes[:6] not in (cls.GCB_HEADER, cls.GCB_LZ_HEADER):

            raise ValueError('bytes did not start with %r' % (cls.GCB_HEADER,))

        if bytes[4] == 'z':

            out._compressor_name = 'zlib'

        elif bytes[4] == 'l':

            out._compressor_name = 'lzma'

        else:

            raise ValueError('unknown compressor: %r' % (bytes,))

        out._parse_bytes(bytes)

        if not _NO_LABELS:

            out._parse_header()

        return out

    def extract(self, key, start, end, sha1=None):

        """Extract the text for a specific key.

        :param key: The label used for this content

        :param sha1: TODO (should we validate only when sha1 is supplied?)

        :return: The bytes for the content

        """

        if start == end == 0:

            return ''

        self._ensure_content(end)

        # The bytes are 'f' or 'd' for the type, then a variable-length

        # base128 integer for the content size, then the actual content

        # We know that the variable-length integer won't be longer than 5

        # bytes (it takes 5 bytes to encode 2^32)

        c = self._content[start]

        if c == 'f':

            type = 'fulltext'

        else:

            if c != 'd':

                raise ValueError('Unknown content control code: %s'

                                 % (c,))

            type = 'delta'

        content_len, len_len = decode_base128_int(

                            self._content[start + 1:start + 6])

        content_start = start + 1 + len_len

        if end != content_start + content_len:

            raise ValueError('end != len according to field header'

                ' %s != %s' % (end, content_start + content_len))

        content = self._content[content_start:end]

        if c == 'f':

            bytes = content

        elif c == 'd':

            bytes = _groupcompress_pyx.apply_delta(self._content, content)

        return bytes

    def add_entry(self, key, type, sha1, start, length):

        """Add new meta info about an entry.

        :param key: The key for the new content

        :param type: Whether this is a delta or fulltext entry (external?)

        :param sha1: sha1sum of the fulltext of this entry

        :param start: where the encoded bytes start

        :param length: total number of bytes in the encoded form

        :return: The entry?

        """

        entry = GroupCompressBlockEntry(key, type, sha1, start, length)

        if key in self._entries:

            raise ValueError('Duplicate key found: %s' % (key,))

        self._entries[key] = entry

        return entry

    def set_content(self, content):

        """Set the content of this block."""

        self._content_length = len(content)

        self._content = content

        self._z_content = None

        self._z_header_length = None

    def to_bytes(self):

        """Encode the information into a byte stream."""

        compress = zlib.compress

        if _USE_LZMA:

            compress = pylzma.compress

        chunks = []

        for key in sorted(self._entries):

            entry = self._entries[key]

            chunk = ('key:%s\n'

                     'sha1:%s\n'

                     'type:%s\n'

                     'start:%s\n'

                     'length:%s\n'

                     '\n'

                     ) % ('\x00'.join(entry.key),

                          entry.sha1,

                          entry.type,

                          entry.start,

                          entry.length,

                          )

            chunks.append(chunk)

        bytes = ''.join(chunks)

        info_len = len(bytes)

        z_header_bytes = compress(bytes)

        del bytes, chunks

        z_header_len = len(z_header_bytes)

        # TODO: we may want to have the header compressed in the same chain

        #       as the data, or we may not, evaulate it

        #       having them compressed together is probably a win for

        #       revisions and the 'inv' portion of chk inventories. As the

        #       label in the header is duplicated in the text.

        #       For chk pages and real bytes, I would guess this is not

        #       true.

        if _NO_LABELS:

            z_header_bytes = ''

            z_header_len = 0

            info_len = 0

        if self._z_content is not None:

            content_len = self._content_length

            z_content_len = self._z_content_length

            z_content_bytes = self._z_content

        else:

            assert self._content is not None

            content_len = self._content_length

            z_content_bytes = compress(self._content)

            self._z_content = z_content_bytes

            z_content_len = len(z_content_bytes)

            self._z_content_length = z_content_len

        if _USE_LZMA:

            header = self.GCB_LZ_HEADER

        else:

            header = self.GCB_HEADER

        chunks = [header,

                  '%d\n%d\n%d\n%d\n' % (z_header_len, info_len,

                                        z_content_len, content_len)

                 ]

        chunks.append(z_header_bytes)

        chunks.append(z_content_bytes)

        return ''.join(chunks)

class _LazyGroupCompressFactory(object):

    """Yield content from a GroupCompressBlock on demand."""

    def __init__(self, key, parents, manager, start, end, first):

        """Create a _LazyGroupCompressFactory

        :param key: The key of just this record

        :param parents: The parents of this key (possibly None)

        :param gc_block: A GroupCompressBlock object

        :param start: Offset of the first byte for this record in the

            uncompressd content

        :param end: Offset of the byte just after the end of this record

            (ie, bytes = content[start:end])

        :param first: Is this the first Factory for the given block?

        """

        self.key = key

        self.parents = parents

        self.sha1 = None

        # Note: This attribute coupled with Manager._factories creates a

        #       reference cycle. Perhaps we would rather use a weakref(), or

        #       find an appropriate time to release the ref. After the first

        #       get_bytes_as call? After Manager.get_record_stream() returns

        #       the object?

        self._manager = manager

        self._bytes = None

        self.storage_kind = 'groupcompress-block'

        if not first:

            self.storage_kind = 'groupcompress-block-ref'

        self._first = first

        self._start = start

        self._end = end

    def __repr__(self):

        return '%s(%s, first=%s)' % (self.__class__.__name__,

            self.key, self._first)

    def get_bytes_as(self, storage_kind):

        if storage_kind == self.storage_kind:

            if self._first:

                # wire bytes, something...

                return self._manager._wire_bytes()

            else:

                return ''

        if storage_kind in ('fulltext', 'chunked'):

            if self._bytes is None:

                # Grab and cache the raw bytes for this entry

                # and break the ref-cycle with _manager since we don't need it

                # anymore

                self._manager._prepare_for_extract()

                block = self._manager._block

                self._bytes = block.extract(self.key, self._start, self._end)

                # XXX: It seems the smart fetch extracts inventories and chk

                #      pages as fulltexts to find the next chk pages, but then

                #      passes them down to be inserted as a

                #      groupcompress-block, so this is not safe to do. Perhaps

                #      we could just change the storage kind to "fulltext" at

                #      that point?

                # self._manager = None

            if storage_kind == 'fulltext':

                return self._bytes

            else:

                return [self._bytes]

        raise errors.UnavailableRepresentation(self.key, storage_kind,

                                               self.storage_kind)

class _LazyGroupContentManager(object):

    """This manages a group of _LazyGroupCompressFactory objects."""

    def __init__(self, block):

        self._block = block

        # We need to preserve the ordering

        self._factories = []

        self._last_byte = 0

    def add_factory(self, key, parents, start, end):

        if not self._factories:

            first = True

        else:

            first = False

        # Note that this creates a reference cycle....

        factory = _LazyGroupCompressFactory(key, parents, self,

            start, end, first=first)

        self._last_byte = max(end, self._last_byte)

        self._factories.append(factory)

    def get_record_stream(self):

        """Get a record for all keys added so far."""

        for factory in self._factories:

            yield factory

            # Break the ref-cycle

            factory._bytes = None

            # XXX: this is not safe, the smart fetch code requests the content

            #      as both a 'fulltext', and then later on as a

            #      groupcompress-block. The iter_interesting_nodes code also is

            #      still buffering multiple records and returning them later.

            #      So that code would need to be updated to either re-fetch the

            #      original object, or buffer it somehow.

            # factory._manager = None

        # TODO: Consider setting self._factories = None after the above loop,

        #       as it will break the reference cycle

    def _trim_block(self, last_byte):

        """Create a new GroupCompressBlock, with just some of the content."""

        # None of the factories need to be adjusted, because the content is

        # located in an identical place. Just that some of the unreferenced

        # trailing bytes are stripped

        trace.mutter('stripping trailing bytes from groupcompress block'

                     ' %d => %d', self._block._content_length, last_byte)

        new_block = GroupCompressBlock()

        self._block._ensure_content(last_byte)

        new_block.set_content(self._block._content[:last_byte])

        self._block = new_block

    def _rebuild_block(self):

        """Create a new GroupCompressBlock with only the referenced texts."""

        compressor = GroupCompressor()

        tstart = time.time()

        old_length = self._block._content_length

        end_point = 0

        for factory in self._factories:

            bytes = factory.get_bytes_as('fulltext')

            (found_sha1, start_point, end_point, type,

             length) = compressor.compress(factory.key, bytes, factory.sha1)

            # Now update this factory with the new offsets, etc

            factory.sha1 = found_sha1

            factory._start = start_point

            factory._end = end_point

        self._last_byte = end_point

        new_block = compressor.flush()

        # TODO: Should we check that new_block really *is* smaller than the old

        #       block? It seems hard to come up with a method that it would

        #       expand, since we do full compression again. Perhaps based on a

        #       request that ends up poorly ordered?

        delta = time.time() - tstart

        self._block = new_block

        trace.mutter('creating new compressed block on-the-fly in %.3fs'

                     ' %d bytes => %d bytes', delta, old_length,

                     self._block._content_length)

    def _prepare_for_extract(self):

        """A _LazyGroupCompressFactory is about to extract to fulltext."""

        # We expect that if one child is going to fulltext, all will be. This

        # helps prevent all of them from extracting a small amount at a time.

        # Which in itself isn't terribly expensive, but resizing 2MB 32kB at a

        # time (self._block._content) is a little expensive.

        self._block._ensure_content(self._last_byte)

    def _check_rebuild_block(self):

        """Check to see if our block should be repacked."""

        total_bytes_used = 0

        last_byte_used = 0

        for factory in self._factories:

            total_bytes_used += factory._end - factory._start

            last_byte_used = max(last_byte_used, factory._end)

        # If we are using most of the bytes from the block, we have nothing

        # else to check (currently more that 1/2)

        if total_bytes_used * 2 >= self._block._content_length:

            return

        # Can we just strip off the trailing bytes? If we are going to be

        # transmitting more than 50% of the front of the content, go ahead

        if total_bytes_used * 2 > last_byte_used:

            self._trim_block(last_byte_used)

            return

        # We are using a small amount of the data, and it isn't just packed

        # nicely at the front, so rebuild the content.

        # Note: This would be *nicer* as a strip-data-from-group, rather than

        #       building it up again from scratch

        #       It might be reasonable to consider the fulltext sizes for

        #       different bits when deciding this, too. As you may have a small

        #       fulltext, and a trivial delta, and you are just trading around

        #       for another fulltext. If we do a simple 'prune' you may end up

        #       expanding many deltas into fulltexts, as well.

        #       If we build a cheap enough 'strip', then we could try a strip,

        #       if that expands the content, we then rebuild.

        self._rebuild_block()

    def _wire_bytes(self):

        """Return a byte stream suitable for transmitting over the wire."""

        self._check_rebuild_block()

        # The outer block starts with:

        #   'groupcompress-block\n'

        #   <length of compressed key info>\n

        #   <length of uncompressed info>\n

        #   <length of gc block>\n

        #   <header bytes>

        #   <gc-block>

        lines = ['groupcompress-block\n']

        # The minimal info we need is the key, the start offset, and the

        # parents. The length and type are encoded in the record itself.

        # However, passing in the other bits makes it easier.  The list of

        # keys, and the start offset, the length

        # 1 line key

        # 1 line with parents, '' for ()

        # 1 line for start offset

        # 1 line for end byte

        header_lines = []

        for factory in self._factories:

            key_bytes = '\x00'.join(factory.key)

            parents = factory.parents

            if parents is None:

                parent_bytes = 'None:'

            else:

                parent_bytes = '\t'.join('\x00'.join(key) for key in parents)

            record_header = '%s\n%s\n%d\n%d\n' % (

                key_bytes, parent_bytes, factory._start, factory._end)

            header_lines.append(record_header)

        header_bytes = ''.join(header_lines)

        del header_lines

        header_bytes_len = len(header_bytes)

        z_header_bytes = zlib.compress(header_bytes)

        del header_bytes

        z_header_bytes_len = len(z_header_bytes)

        block_bytes = self._block.to_bytes()

        lines.append('%d\n%d\n%d\n' % (z_header_bytes_len, header_bytes_len,

                                       len(block_bytes)))

        lines.append(z_header_bytes)

        lines.append(block_bytes)

        del z_header_bytes, block_bytes

        return ''.join(lines)

    @classmethod

    def from_bytes(cls, bytes):