# changegroup.py - Mercurial changegroup manipulation functions

#

#  Copyright 2006 Matt Mackall <mpm@selenic.com>

#

# 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

import os

import struct

import tempfile

import weakref

from .i18n import _

from .node import (

    hex,

    nullrev,

    short,

)

from . import (

    dagutil,

    discovery,

    error,

    mdiff,

    phases,

    pycompat,

    util,

)

_CHANGEGROUPV1_DELTA_HEADER = "20s20s20s20s"

_CHANGEGROUPV2_DELTA_HEADER = "20s20s20s20s20s"

_CHANGEGROUPV3_DELTA_HEADER = ">20s20s20s20s20sH"

def readexactly(stream, n):

    '''read n bytes from stream.read and abort if less was available'''

    s = stream.read(n)

    if len(s) < n:

        raise error.Abort(_("stream ended unexpectedly"

                           " (got %d bytes, expected %d)")

                          % (len(s), n))

    return s

def getchunk(stream):

    """return the next chunk from stream as a string"""

    d = readexactly(stream, 4)

    l = struct.unpack(">l", d)[0]

    if l <= 4:

        if l:

            raise error.Abort(_("invalid chunk length %d") % l)

        return ""

    return readexactly(stream, l - 4)

def chunkheader(length):

    """return a changegroup chunk header (string)"""

    return struct.pack(">l", length + 4)

def closechunk():

    """return a changegroup chunk header (string) for a zero-length chunk"""

    return struct.pack(">l", 0)

def writechunks(ui, chunks, filename, vfs=None):

    """Write chunks to a file and return its filename.

    The stream is assumed to be a bundle file.

    Existing files will not be overwritten.

    If no filename is specified, a temporary file is created.

    """

    fh = None

    cleanup = None

    try:

        if filename:

            if vfs:

                fh = vfs.open(filename, "wb")

            else:

                # Increase default buffer size because default is usually

                # small (4k is common on Linux).

                fh = open(filename, "wb", 131072)

        else:

            fd, filename = tempfile.mkstemp(prefix="hg-bundle-", suffix=".hg")

            fh = os.fdopen(fd, pycompat.sysstr("wb"))

        cleanup = filename

        for c in chunks:

            fh.write(c)

        cleanup = None

        return filename

    finally:

        if fh is not None:

            fh.close()

        if cleanup is not None:

            if filename and vfs:

                vfs.unlink(cleanup)

            else:

                os.unlink(cleanup)

class cg1unpacker(object):

    """Unpacker for cg1 changegroup streams.

    A changegroup unpacker handles the framing of the revision data in

    the wire format. Most consumers will want to use the apply()

    method to add the changes from the changegroup to a repository.

    If you're forwarding a changegroup unmodified to another consumer,

    use getchunks(), which returns an iterator of changegroup

    chunks. This is mostly useful for cases where you need to know the

    data stream has ended by observing the end of the changegroup.

    deltachunk() is useful only if you're applying delta data. Most

    consumers should prefer apply() instead.

    A few other public methods exist. Those are used only for

    bundlerepo and some debug commands - their use is discouraged.

    """

    deltaheader = _CHANGEGROUPV1_DELTA_HEADER

    deltaheadersize = struct.calcsize(deltaheader)

    version = '01'

    _grouplistcount = 1 # One list of files after the manifests

    def __init__(self, fh, alg, extras=None):

        if alg is None:

            alg = 'UN'

        if alg not in util.compengines.supportedbundletypes:

            raise error.Abort(_('unknown stream compression type: %s')

                             % alg)

        if alg == 'BZ':

            alg = '_truncatedBZ'

        compengine = util.compengines.forbundletype(alg)

        self._stream = compengine.decompressorreader(fh)

        self._type = alg

        self.extras = extras or {}

        self.callback = None

    # These methods (compressed, read, seek, tell) all appear to only

    # be used by bundlerepo, but it's a little hard to tell.

    def compressed(self):

        return self._type is not None and self._type != 'UN'

    def read(self, l):

        return self._stream.read(l)

    def seek(self, pos):

        return self._stream.seek(pos)

    def tell(self):

        return self._stream.tell()

    def close(self):

        return self._stream.close()

    def _chunklength(self):

        d = readexactly(self._stream, 4)

        l = struct.unpack(">l", d)[0]

        if l <= 4:

            if l:

                raise error.Abort(_("invalid chunk length %d") % l)

            return 0

        if self.callback:

            self.callback()

        return l - 4

    def changelogheader(self):

        """v10 does not have a changelog header chunk"""

        return {}

    def manifestheader(self):

        """v10 does not have a manifest header chunk"""

        return {}

    def filelogheader(self):

        """return the header of the filelogs chunk, v10 only has the filename"""

        l = self._chunklength()

        if not l:

            return {}

        fname = readexactly(self._stream, l)

        return {'filename': fname}

    def _deltaheader(self, headertuple, prevnode):

        node, p1, p2, cs = headertuple

        if prevnode is None:

            deltabase = p1

        else:

            deltabase = prevnode

        flags = 0

        return node, p1, p2, deltabase, cs, flags

    def deltachunk(self, prevnode):

        l = self._chunklength()

        if not l:

            return {}

        headerdata = readexactly(self._stream, self.deltaheadersize)

        header = struct.unpack(self.deltaheader, headerdata)

        delta = readexactly(self._stream, l - self.deltaheadersize)

        node, p1, p2, deltabase, cs, flags = self._deltaheader(header, prevnode)

        return {'node': node, 'p1': p1, 'p2': p2, 'cs': cs,

                'deltabase': deltabase, 'delta': delta, 'flags': flags}

    def getchunks(self):

        """returns all the chunks contains in the bundle

        Used when you need to forward the binary stream to a file or another

        network API. To do so, it parse the changegroup data, otherwise it will

        block in case of sshrepo because it don't know the end of the stream.

        """

        # an empty chunkgroup is the end of the changegroup

        # a changegroup has at least 2 chunkgroups (changelog and manifest).

        # after that, changegroup versions 1 and 2 have a series of groups

        # with one group per file. changegroup 3 has a series of directory

        # manifests before the files.

        count = 0

        emptycount = 0

        while emptycount < self._grouplistcount:

            empty = True

            count += 1

            while True:

                chunk = getchunk(self)

                if not chunk:

                    if empty and count > 2:

                        emptycount += 1

                    break

                empty = False

                yield chunkheader(len(chunk))

                pos = 0

                while pos < len(chunk):

                    next = pos + 2**20

                    yield chunk[pos:next]

                    pos = next

            yield closechunk()

    def _unpackmanifests(self, repo, revmap, trp, prog, numchanges):

        # We know that we'll never have more manifests than we had

        # changesets.

        self.callback = prog(_('manifests'), numchanges)

        # no need to check for empty manifest group here:

        # if the result of the merge of 1 and 2 is the same in 3 and 4,

        # no new manifest will be created and the manifest group will

        # be empty during the pull

        self.manifestheader()

        repo.manifestlog._revlog.addgroup(self, revmap, trp)

        repo.ui.progress(_('manifests'), None)

        self.callback = None

    def apply(self, repo, tr, srctype, url, targetphase=phases.draft,

              expectedtotal=None):

        """Add the changegroup returned by source.read() to this repo.

        srctype is a string like 'push', 'pull', or 'unbundle'.  url is

        the URL of the repo where this changegroup is coming from.

        Return an integer summarizing the change to this repo:

        - nothing changed or no source: 0

        - more heads than before: 1+added heads (2..n)

        - fewer heads than before: -1-removed heads (-2..-n)

        - number of heads stays the same: 1

        """

        repo = repo.unfiltered()

        def csmap(x):

            repo.ui.debug("add changeset %s\n" % short(x))

            return len(cl)

        def revmap(x):

            return cl.rev(x)

        changesets = files = revisions = 0

        try:

            # The transaction may already carry source information. In this

            # case we use the top level data. We overwrite the argument

            # because we need to use the top level value (if they exist)

            # in this function.

            srctype = tr.hookargs.setdefault('source', srctype)

            url = tr.hookargs.setdefault('url', url)

            repo.hook('prechangegroup', throw=True, **tr.hookargs)

            # write changelog data to temp files so concurrent readers

            # will not see an inconsistent view

            cl = repo.changelog

            cl.delayupdate(tr)

            oldheads = set(cl.heads())

            trp = weakref.proxy(tr)

            # pull off the changeset group

            repo.ui.status(_("adding changesets\n"))

            clstart = len(cl)

            class prog(object):

                def __init__(self, step, total):

                    self._step = step

                    self._total = total

                    self._count = 1

                def __call__(self):

                    repo.ui.progress(self._step, self._count, unit=_('chunks'),

                                     total=self._total)

                    self._count += 1

            self.callback = prog(_('changesets'), expectedtotal)

            efiles = set()

            def onchangelog(cl, node):

                efiles.update(cl.readfiles(node))

            self.changelogheader()

            cgnodes = cl.addgroup(self, csmap, trp, addrevisioncb=onchangelog)

            efiles = len(efiles)

            if not cgnodes:

            clend = len(cl)

            changesets = clend - clstart

            repo.ui.progress(_('changesets'), None)

            self.callback = None

            # pull off the manifest group

            repo.ui.status(_("adding manifests\n"))

            self._unpackmanifests(repo, revmap, trp, prog, changesets)

            needfiles = {}

            if repo.ui.configbool('server', 'validate'):

                cl = repo.changelog

                ml = repo.manifestlog

                # validate incoming csets have their manifests

                for cset in xrange(clstart, clend):

                    mfnode = cl.changelogrevision(cset).manifest

                    mfest = ml[mfnode].readdelta()

                    # store file cgnodes we must see

                    for f, n in mfest.iteritems():

                        needfiles.setdefault(f, set()).add(n)

            # process the files

            repo.ui.status(_("adding file changes\n"))

            newrevs, newfiles = _addchangegroupfiles(

                repo, self, revmap, trp, efiles, needfiles)

            revisions += newrevs

            files += newfiles

            deltaheads = 0

            if oldheads:

                heads = cl.heads()

                deltaheads = len(heads) - len(oldheads)

                for h in heads:

                    if h not in oldheads and repo[h].closesbranch():

                        deltaheads -= 1

            htext = ""

            if deltaheads:

                htext = _(" (%+d heads)") % deltaheads

            repo.ui.status(_("added %d changesets"

                             " with %d changes to %d files%s\n")

                             % (changesets, revisions, files, htext))

            repo.invalidatevolatilesets()

            if changesets > 0:

                if 'node' not in tr.hookargs:

                    tr.hookargs['node'] = hex(cl.node(clstart))

                    tr.hookargs['node_last'] = hex(cl.node(clend - 1))

                    hookargs = dict(tr.hookargs)

                else:

                    hookargs = dict(tr.hookargs)

                    hookargs['node'] = hex(cl.node(clstart))

                    hookargs['node_last'] = hex(cl.node(clend - 1))

                repo.hook('pretxnchangegroup', throw=True, **hookargs)

            added = [cl.node(r) for r in xrange(clstart, clend)]

            if srctype in ('push', 'serve'):

                # Old servers can not push the boundary themselves.

                # New servers won't push the boundary if changeset already

                # exists locally as secret

                #

                # We should not use added here but the list of all change in

                # the bundle

                if repo.publishing():

                    phases.advanceboundary(repo, tr, phases.public, cgnodes)

                else:

                    # Those changesets have been pushed from the

                    # outside, their phases are going to be pushed

                    # alongside. Therefor `targetphase` is

                    # ignored.

                    phases.advanceboundary(repo, tr, phases.draft, cgnodes)

                    phases.retractboundary(repo, tr, phases.draft, added)

            elif srctype != 'strip':

                # publishing only alter behavior during push

                #

                # strip should not touch boundary at all

                phases.retractboundary(repo, tr, targetphase, added)

            if changesets > 0:

                def runhooks():

                    # These hooks run when the lock releases, not when the

                    # transaction closes. So it's possible for the changelog

                    # to have changed since we last saw it.

                    if clstart >= len(repo):

                        return

                    repo.hook("changegroup", **hookargs)

                    for n in added:

                        args = hookargs.copy()

                        args['node'] = hex(n)

                        del args['node_last']

                        repo.hook("incoming", **args)

                    newheads = [h for h in repo.heads()

                                if h not in oldheads]

                    repo.ui.log("incoming",

                                "%s incoming changes - new heads: %s\n",

                                len(added),

                                ', '.join([hex(c[:6]) for c in newheads]))

                tr.addpostclose('changegroup-runhooks-%020i' % clstart,

                                lambda tr: repo._afterlock(runhooks))

        finally:

            repo.ui.flush()

        # never return 0 here:

        if deltaheads < 0:

            ret = deltaheads - 1

        else:

            ret = deltaheads + 1

        return ret, added

class cg2unpacker(cg1unpacker):

    """Unpacker for cg2 streams.

    cg2 streams add support for generaldelta, so the delta header

    format is slightly different. All other features about the data

    remain the same.

    """

    deltaheader = _CHANGEGROUPV2_DELTA_HEADER

    deltaheadersize = struct.calcsize(deltaheader)

    version = '02'

    def _deltaheader(self, headertuple, prevnode):

        node, p1, p2, deltabase, cs = headertuple

        flags = 0

        return node, p1, p2, deltabase, cs, flags

class cg3unpacker(cg2unpacker):

    """Unpacker for cg3 streams.

    cg3 streams add support for exchanging treemanifests and revlog

    flags. It adds the revlog flags to the delta header and an empty chunk

    separating manifests and files.

    """

    deltaheader = _CHANGEGROUPV3_DELTA_HEADER

    deltaheadersize = struct.calcsize(deltaheader)

    version = '03'

    _grouplistcount = 2 # One list of manifests and one list of files

    def _deltaheader(self, headertuple, prevnode):

        node, p1, p2, deltabase, cs, flags = headertuple

        return node, p1, p2, deltabase, cs, flags

    def _unpackmanifests(self, repo, revmap, trp, prog, numchanges):

        super(cg3unpacker, self)._unpackmanifests(repo, revmap, trp, prog,

                                                  numchanges)

        for chunkdata in iter(self.filelogheader, {}):

            # If we get here, there are directory manifests in the changegroup

            d = chunkdata["filename"]

            repo.ui.debug("adding %s revisions\n" % d)

            dirlog = repo.manifestlog._revlog.dirlog(d)

            if not dirlog.addgroup(self, revmap, trp):

                raise error.Abort(_("received dir revlog group is empty"))

class headerlessfixup(object):

    def __init__(self, fh, h):

        self._h = h

        self._fh = fh

    def read(self, n):

        if self._h:

            d, self._h = self._h[:n], self._h[n:]

            if len(d) < n:

                d += readexactly(self._fh, n - len(d))

            return d

        return readexactly(self._fh, n)

class cg1packer(object):

    deltaheader = _CHANGEGROUPV1_DELTA_HEADER

    version = '01'

    def __init__(self, repo, bundlecaps=None):

        """Given a source repo, construct a bundler.

        bundlecaps is optional and can be used to specify the set of

        capabilities which can be used to build the bundle. While bundlecaps is

        unused in core Mercurial, extensions rely on this feature to communicate

        capabilities to customize the changegroup packer.

        """

        # Set of capabilities we can use to build the bundle.

        if bundlecaps is None:

            bundlecaps = set()