# 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 weakref

from .i18n import _

from .node import (

    hex,

    nullid,

    nullrev,

    short,

)

from .thirdparty import (

    attr,

)

from . import (

    dagutil,

    error,

    manifest,

    match as matchmod,

    mdiff,

    phases,

    pycompat,

    repository,

    revlog,

    util,

)

from .utils import (

    stringutil,

)

_CHANGEGROUPV1_DELTA_HEADER = struct.Struct("20s20s20s20s")

_CHANGEGROUPV2_DELTA_HEADER = struct.Struct("20s20s20s20s20s")

_CHANGEGROUPV3_DELTA_HEADER = struct.Struct(">20s20s20s20s20sH")

LFS_REQUIREMENT = 'lfs'

readexactly = util.readexactly

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 = pycompat.mkstemp(prefix="hg-bundle-", suffix=".hg")

            fh = os.fdopen(fd, r"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 = deltaheader.size

    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 = self.deltaheader.unpack(headerdata)

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

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

        return (node, p1, p2, cs, deltabase, delta, 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.

        """

        # For changegroup 1 and 2, we expect 3 parts: changelog, manifestlog,

        # and a list of filelogs. For changegroup 3, we expect 4 parts:

        # changelog, manifestlog, a list of tree manifestlogs, and a list of

        # filelogs.

        #

        # Changelog and manifestlog parts are terminated with empty chunks. The

        # tree and file parts are a list of entry sections. Each entry section

        # is a series of chunks terminating in an empty chunk. The list of these

        # entry sections is terminated in yet another empty chunk, so we know

        # we've reached the end of the tree/file list when we reach an empty

        # chunk that was proceeded by no non-empty chunks.

        parts = 0

        while parts < 2 + self._grouplistcount:

            noentries = True

            while True:

                chunk = getchunk(self)

                if not chunk:

                    # The first two empty chunks represent the end of the

                    # changelog and the manifestlog portions. The remaining

                    # empty chunks represent either A) the end of individual

                    # tree or file entries in the file list, or B) the end of

                    # the entire list. It's the end of the entire list if there

                    # were no entries (i.e. noentries is True).

                    if parts < 2:

                        parts += 1

                    elif noentries:

                        parts += 1

                    break

                noentries = 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):

        self.callback = prog.increment

        # 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()

        deltas = self.deltaiter()

        repo.manifestlog.addgroup(deltas, revmap, trp)

        prog.complete()

        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, **pycompat.strkwargs(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)

            progress = repo.ui.makeprogress(_('changesets'), unit=_('chunks'),

                                            total=expectedtotal)

            self.callback = progress.increment

            efiles = set()

            def onchangelog(cl, node):

                efiles.update(cl.readfiles(node))

            self.changelogheader()

            deltas = self.deltaiter()

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

            efiles = len(efiles)

            if not cgnodes:

                repo.ui.develwarn('applied empty changegroup',

                                  config='warn-empty-changegroup')

            clend = len(cl)

            changesets = clend - clstart

            progress.complete()

            self.callback = None

            # pull off the manifest group

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

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

            # changesets.

            progress = repo.ui.makeprogress(_('manifests'), unit=_('chunks'),

                                            total=changesets)

            self._unpackmanifests(repo, revmap, trp, progress)

            needfiles = {}

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

                cl = repo.changelog

                ml = repo.manifestlog

                # validate incoming csets have their manifests

                for cset in pycompat.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, **pycompat.strkwargs(hookargs))

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

            phaseall = None

            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():

                    targetphase = phaseall = phases.public

                else:

                    # closer target phase computation

                    # Those changesets have been pushed from the

                    # outside, their phases are going to be pushed

                    # alongside. Therefor `targetphase` is

                    # ignored.

                    targetphase = phaseall = phases.draft

            if added:

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

            if phaseall is not None:

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

            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", **pycompat.strkwargs(hookargs))

                    for n in added:

                        args = hookargs.copy()

                        args['node'] = hex(n)

                        del args['node_last']

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

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

                                if h not in oldheads]

                    repo.ui.log("incoming",

                                "%d 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

    def deltaiter(self):

        """

        returns an iterator of the deltas in this changegroup

        Useful for passing to the underlying storage system to be stored.

        """

        chain = None

        for chunkdata in iter(lambda: self.deltachunk(chain), {}):

            # Chunkdata: (node, p1, p2, cs, deltabase, delta, flags)

            yield chunkdata

            chain = chunkdata[0]

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 = deltaheader.size

    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 = deltaheader.size

    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):

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

        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)