hgext/narrow/narrowchangegroup.py

# narrowchangegroup.py - narrow clone changegroup creation and consumption
#
# Copyright 2017 Google, Inc.
#
# 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

from mercurial.i18n import _
from mercurial import (
    changegroup,
    extensions,
    node,
    util,
)

def setup():
    def generate(orig, self, commonrevs, clnodes, fastpathlinkrev, source):
        '''yield a sequence of changegroup chunks (strings)'''
        # Note: other than delegating to orig, the only deviation in
        # logic from normal hg's generate is marked with BEGIN/END
        # NARROW HACK.
        if not util.safehasattr(self, 'full_nodes'):
            # not sending a narrow bundle
            for x in orig(self, commonrevs, clnodes, fastpathlinkrev, source):
                yield x
            return

        repo = self._repo
        cl = repo.changelog
        mfl = repo.manifestlog
        mfrevlog = mfl._revlog

        clrevorder = {}
        mfs = {} # needed manifests
        fnodes = {} # needed file nodes
        changedfiles = set()

        # Callback for the changelog, used to collect changed files and manifest
        # nodes.
        # Returns the linkrev node (identity in the changelog case).
        def lookupcl(x):
            c = cl.read(x)
            clrevorder[x] = len(clrevorder)
            # BEGIN NARROW HACK
            #
            # Only update mfs if x is going to be sent. Otherwise we
            # end up with bogus linkrevs specified for manifests and
            # we skip some manifest nodes that we should otherwise
            # have sent.
            if x in self.full_nodes or cl.rev(x) in self.precomputed_ellipsis:
                n = c[0]
                # record the first changeset introducing this manifest version
                mfs.setdefault(n, x)
                # Set this narrow-specific dict so we have the lowest manifest
                # revnum to look up for this cl revnum. (Part of mapping
                # changelog ellipsis parents to manifest ellipsis parents)
                self.next_clrev_to_localrev.setdefault(cl.rev(x),
                                                       mfrevlog.rev(n))
            # We can't trust the changed files list in the changeset if the
            # client requested a shallow clone.
            if self.is_shallow:
                changedfiles.update(mfl[c[0]].read().keys())
            else:
                changedfiles.update(c[3])
            # END NARROW HACK
            # Record a complete list of potentially-changed files in
            # this manifest.
            return x

        self._verbosenote(_('uncompressed size of bundle content:\n'))
        size = 0
        for chunk in self.group(clnodes, cl, lookupcl, units=_('changesets')):
            size += len(chunk)
            yield chunk
        self._verbosenote(_('%8.i (changelog)\n') % size)

        # We need to make sure that the linkrev in the changegroup refers to
        # the first changeset that introduced the manifest or file revision.
        # The fastpath is usually safer than the slowpath, because the filelogs
        # are walked in revlog order.
        #
        # When taking the slowpath with reorder=None and the manifest revlog
        # uses generaldelta, the manifest may be walked in the "wrong" order.
        # Without 'clrevorder', we would get an incorrect linkrev (see fix in
        # cc0ff93d0c0c).
        #
        # When taking the fastpath, we are only vulnerable to reordering
        # of the changelog itself. The changelog never uses generaldelta, so
        # it is only reordered when reorder=True. To handle this case, we
        # simply take the slowpath, which already has the 'clrevorder' logic.
        # This was also fixed in cc0ff93d0c0c.
        fastpathlinkrev = fastpathlinkrev and not self._reorder
        # Treemanifests don't work correctly with fastpathlinkrev
        # either, because we don't discover which directory nodes to
        # send along with files. This could probably be fixed.
        fastpathlinkrev = fastpathlinkrev and (
            'treemanifest' not in repo.requirements)
        # Shallow clones also don't work correctly with fastpathlinkrev
        # because file nodes may need to be sent for a manifest even if they
        # weren't introduced by that manifest.
        fastpathlinkrev = fastpathlinkrev and not self.is_shallow

        for chunk in self.generatemanifests(commonrevs, clrevorder,
                fastpathlinkrev, mfs, fnodes, source):
            yield chunk
        # BEGIN NARROW HACK
        mfdicts = None
        if self.is_shallow:
            mfdicts = [(self._repo.manifestlog[n].read(), lr)
                       for (n, lr) in mfs.iteritems()]
        # END NARROW HACK
        mfs.clear()
        clrevs = set(cl.rev(x) for x in clnodes)

        if not fastpathlinkrev:
            def linknodes(unused, fname):
                return fnodes.get(fname, {})
        else:
            cln = cl.node
            def linknodes(filerevlog, fname):
                llr = filerevlog.linkrev
                fln = filerevlog.node
                revs = ((r, llr(r)) for r in filerevlog)
                return dict((fln(r), cln(lr)) for r, lr in revs if lr in clrevs)

        # BEGIN NARROW HACK
        #
        # We need to pass the mfdicts variable down into
        # generatefiles(), but more than one command might have
        # wrapped generatefiles so we can't modify the function
        # signature. Instead, we pass the data to ourselves using an
        # instance attribute. I'm sorry.
        self._mfdicts = mfdicts
        # END NARROW HACK
        for chunk in self.generatefiles(changedfiles, linknodes, commonrevs,
                                        source):
            yield chunk

        yield self.close()

        if clnodes:
            repo.hook('outgoing', node=node.hex(clnodes[0]), source=source)
    extensions.wrapfunction(changegroup.cg1packer, 'generate', generate)