# compression.py - Mercurial utility functions for compression

#

# This software may be used and distributed according to the terms of the

# GNU General Public License version 2 or any later version.

from __future__ import absolute_import, print_function

import bz2

import collections

import zlib

from .. import (

    error,

    i18n,

    pycompat,

)

from . import (

    stringutil,

)

safehasattr = pycompat.safehasattr

_ = i18n._

# compression code

SERVERROLE = 'server'

CLIENTROLE = 'client'

compewireprotosupport = collections.namedtuple(r'compenginewireprotosupport',

                                               (r'name', r'serverpriority',

                                                r'clientpriority'))

class propertycache(object):

    def __init__(self, func):

        self.func = func

        self.name = func.__name__

    def __get__(self, obj, type=None):

        result = self.func(obj)

        self.cachevalue(obj, result)

        return result

    def cachevalue(self, obj, value):

        # __dict__ assignment required to bypass __setattr__ (eg: repoview)

        obj.__dict__[self.name] = value

class compressormanager(object):

    """Holds registrations of various compression engines.

    This class essentially abstracts the differences between compression

    engines to allow new compression formats to be added easily, possibly from

    extensions.

    Compressors are registered against the global instance by calling its

    ``register()`` method.

    """

    def __init__(self):

        self._engines = {}

        # Bundle spec human name to engine name.

        self._bundlenames = {}

        # Internal bundle identifier to engine name.

        self._bundletypes = {}

        # Revlog header to engine name.

        self._revlogheaders = {}

        # Wire proto identifier to engine name.

        self._wiretypes = {}

    def __getitem__(self, key):

        return self._engines[key]

    def __contains__(self, key):

        return key in self._engines

    def __iter__(self):

        return iter(self._engines.keys())

    def register(self, engine):

        """Register a compression engine with the manager.

        The argument must be a ``compressionengine`` instance.

        """

        if not isinstance(engine, compressionengine):

            raise ValueError(_('argument must be a compressionengine'))

        name = engine.name()

        if name in self._engines:

            raise error.Abort(_('compression engine %s already registered') %

                              name)

        bundleinfo = engine.bundletype()

        if bundleinfo:

            bundlename, bundletype = bundleinfo

            if bundlename in self._bundlenames:

                raise error.Abort(_('bundle name %s already registered') %

                                  bundlename)

            if bundletype in self._bundletypes:

                raise error.Abort(_('bundle type %s already registered by %s') %

                                  (bundletype, self._bundletypes[bundletype]))

            # No external facing name declared.

            if bundlename:

                self._bundlenames[bundlename] = name

            self._bundletypes[bundletype] = name

        wiresupport = engine.wireprotosupport()

        if wiresupport:

            wiretype = wiresupport.name

            if wiretype in self._wiretypes:

                raise error.Abort(_('wire protocol compression %s already '

                                    'registered by %s') %

                                  (wiretype, self._wiretypes[wiretype]))

            self._wiretypes[wiretype] = name

        revlogheader = engine.revlogheader()

        if revlogheader and revlogheader in self._revlogheaders:

            raise error.Abort(_('revlog header %s already registered by %s') %

                              (revlogheader, self._revlogheaders[revlogheader]))

        if revlogheader:

            self._revlogheaders[revlogheader] = name

        self._engines[name] = engine

    @property

    def supportedbundlenames(self):

        return set(self._bundlenames.keys())

    @property

    def supportedbundletypes(self):

        return set(self._bundletypes.keys())

    def forbundlename(self, bundlename):

        """Obtain a compression engine registered to a bundle name.

        Will raise KeyError if the bundle type isn't registered.

        Will abort if the engine is known but not available.

        """

        engine = self._engines[self._bundlenames[bundlename]]

        if not engine.available():

            raise error.Abort(_('compression engine %s could not be loaded') %

                              engine.name())

        return engine

    def forbundletype(self, bundletype):

        """Obtain a compression engine registered to a bundle type.

        Will raise KeyError if the bundle type isn't registered.

        Will abort if the engine is known but not available.

        """

        engine = self._engines[self._bundletypes[bundletype]]

        if not engine.available():

            raise error.Abort(_('compression engine %s could not be loaded') %

                              engine.name())

        return engine

    def supportedwireengines(self, role, onlyavailable=True):

        """Obtain compression engines that support the wire protocol.

        Returns a list of engines in prioritized order, most desired first.

        If ``onlyavailable`` is set, filter out engines that can't be

        loaded.

        """

        assert role in (SERVERROLE, CLIENTROLE)

        attr = 'serverpriority' if role == SERVERROLE else 'clientpriority'

        engines = [self._engines[e] for e in self._wiretypes.values()]

        if onlyavailable:

            engines = [e for e in engines if e.available()]

        def getkey(e):

            # Sort first by priority, highest first. In case of tie, sort

            # alphabetically. This is arbitrary, but ensures output is

            # stable.

            w = e.wireprotosupport()

            return -1 * getattr(w, attr), w.name

        return list(sorted(engines, key=getkey))

    def forwiretype(self, wiretype):

        engine = self._engines[self._wiretypes[wiretype]]

        if not engine.available():

            raise error.Abort(_('compression engine %s could not be loaded') %

                              engine.name())

        return engine

    def forrevlogheader(self, header):

        """Obtain a compression engine registered to a revlog header.

        Will raise KeyError if the revlog header value isn't registered.

        """

        return self._engines[self._revlogheaders[header]]

compengines = compressormanager()

class compressionengine(object):

    """Base class for compression engines.

    Compression engines must implement the interface defined by this class.

    """

    def name(self):

        """Returns the name of the compression engine.

        This is the key the engine is registered under.

        This method must be implemented.

        """

        raise NotImplementedError()

    def available(self):

        """Whether the compression engine is available.

        The intent of this method is to allow optional compression engines

        that may not be available in all installations (such as engines relying

        on C extensions that may not be present).

        """

        return True

    def bundletype(self):

        """Describes bundle identifiers for this engine.

        If this compression engine isn't supported for bundles, returns None.

        If this engine can be used for bundles, returns a 2-tuple of strings of

        the user-facing "bundle spec" compression name and an internal

        identifier used to denote the compression format within bundles. To

        exclude the name from external usage, set the first element to ``None``.

        If bundle compression is supported, the class must also implement

        ``compressstream`` and `decompressorreader``.

        The docstring of this method is used in the help system to tell users

        about this engine.

        """

        return None

    def wireprotosupport(self):

        """Declare support for this compression format on the wire protocol.

        If this compression engine isn't supported for compressing wire

        protocol payloads, returns None.

        Otherwise, returns ``compenginewireprotosupport`` with the following

        fields:

        * String format identifier

        * Integer priority for the server

        * Integer priority for the client

        The integer priorities are used to order the advertisement of format

        support by server and client. The highest integer is advertised

        first. Integers with non-positive values aren't advertised.

        The priority values are somewhat arbitrary and only used for default

        ordering. The relative order can be changed via config options.

        If wire protocol compression is supported, the class must also implement

        ``compressstream`` and ``decompressorreader``.

        """

        return None

    def revlogheader(self):

        """Header added to revlog chunks that identifies this engine.

        If this engine can be used to compress revlogs, this method should

        return the bytes used to identify chunks compressed with this engine.

        Else, the method should return ``None`` to indicate it does not

        participate in revlog compression.

        """

        return None

    def compressstream(self, it, opts=None):

        """Compress an iterator of chunks.

        The method receives an iterator (ideally a generator) of chunks of

        bytes to be compressed. It returns an iterator (ideally a generator)

        of bytes of chunks representing the compressed output.

        Optionally accepts an argument defining how to perform compression.

        Each engine treats this argument differently.

        """

        raise NotImplementedError()

    def decompressorreader(self, fh):

        """Perform decompression on a file object.

        Argument is an object with a ``read(size)`` method that returns

        compressed data. Return value is an object with a ``read(size)`` that

        returns uncompressed data.

        """

        raise NotImplementedError()

    def revlogcompressor(self, opts=None):

        """Obtain an object that can be used to compress revlog entries.

        The object has a ``compress(data)`` method that compresses binary

        data. This method returns compressed binary data or ``None`` if

        the data could not be compressed (too small, not compressible, etc).

        The returned data should have a header uniquely identifying this

        compression format so decompression can be routed to this engine.

        This header should be identified by the ``revlogheader()`` return

        value.

        The object has a ``decompress(data)`` method that decompresses

        data. The method will only be called if ``data`` begins with

        ``revlogheader()``. The method should return the raw, uncompressed

        data or raise a ``StorageError``.

        The object is reusable but is not thread safe.

        """

        raise NotImplementedError()

class _CompressedStreamReader(object):

    def __init__(self, fh):

        if safehasattr(fh, 'unbufferedread'):

            self._reader = fh.unbufferedread

        else:

            self._reader = fh.read

        self._pending = []

        self._pos = 0

        self._eof = False

    def _decompress(self, chunk):

        raise NotImplementedError()

    def read(self, l):

        buf = []

        while True:

            while self._pending:

                if len(self._pending[0]) > l + self._pos:

                    newbuf = self._pending[0]

                    buf.append(newbuf[self._pos:self._pos + l])

                    self._pos += l

                    return ''.join(buf)

                newbuf = self._pending.pop(0)

                if self._pos:

                    buf.append(newbuf[self._pos:])

                    l -= len(newbuf) - self._pos

                else:

                    buf.append(newbuf)

                    l -= len(newbuf)

                self._pos = 0

            if self._eof:

                return ''.join(buf)

            chunk = self._reader(65536)

            self._decompress(chunk)

            if not chunk and not self._pending and not self._eof:

                # No progress and no new data, bail out

                return ''.join(buf)

class _GzipCompressedStreamReader(_CompressedStreamReader):

    def __init__(self, fh):

        super(_GzipCompressedStreamReader, self).__init__(fh)

        self._decompobj = zlib.decompressobj()

    def _decompress(self, chunk):

        newbuf = self._decompobj.decompress(chunk)

        if newbuf:

            self._pending.append(newbuf)

        d = self._decompobj.copy()

        try:

            d.decompress('x')

            d.flush()

            if d.unused_data == 'x':

                self._eof = True

        except zlib.error:

            pass

class _BZ2CompressedStreamReader(_CompressedStreamReader):

    def __init__(self, fh):

        super(_BZ2CompressedStreamReader, self).__init__(fh)

        self._decompobj = bz2.BZ2Decompressor()

    def _decompress(self, chunk):

        newbuf = self._decompobj.decompress(chunk)

        if newbuf:

            self._pending.append(newbuf)

        try:

            while True:

                newbuf = self._decompobj.decompress('')

                if newbuf:

                    self._pending.append(newbuf)

                else:

                    break

        except EOFError:

            self._eof = True

class _TruncatedBZ2CompressedStreamReader(_BZ2CompressedStreamReader):

    def __init__(self, fh):

        super(_TruncatedBZ2CompressedStreamReader, self).__init__(fh)

        newbuf = self._decompobj.decompress('BZ')

        if newbuf:

            self._pending.append(newbuf)

class _ZstdCompressedStreamReader(_CompressedStreamReader):