Script to test different git-remote-hg revision walking algorithms

grhg-walk-test

#!/usr/bin/env python2

from mercurial import hg, ui, dagop, smartset
import sys, time

class Marks:

    def __init__(self):
        self.marks = None

    def set(self, v):
        self.marks = v

    def is_marked(self, rev):
        return rev in self.marks

def title(text):
    print "\n%s\n%s" % (text, '-' * len(text))

def test_gitrange(a, b):
    def gitrange(repo, a, b):
        if a == b:
            return []

        pfunc = repo.changelog.parentrevs
        it = iter(xrange(b.rev(), -1, -1))

        positive = []
        pending = set([b.rev()])
        negative = set([a.rev()])

        def get_parents(rev):
            for p in pfunc(rev):
                if p == -1:
                    continue
                yield p

        while pending:
            cur = it.next()

            if cur in negative:
                negative.remove(cur)
                for p in get_parents(cur):
                    negative.add(p)
                    pending.discard(p)
            elif cur in pending:
                positive.append(cur)
                pending.remove(cur)
                for p in get_parents(cur):
                    if p not in negative:
                        pending.add(p)

        positive.reverse()
        return positive

    for rev in gitrange(repo, a, b):
        c = repo[rev]

        if marks.is_marked(c.hex()):
            continue

        yield c

def test_gitrange_orig(a, b):
    def gitrange_orig(repo, a, b):
        positive = []
        pending = set([b.rev()])
        negative = set([a.rev()])
        for cur in xrange(b.rev(), -1, -1):
            if not pending:
                break

            parents = [p for p in repo.changelog.parentrevs(cur) if p >= 0]

            if cur in pending:
                positive.append(cur)
                pending.remove(cur)
                for p in parents:
                    if p not in negative:
                        pending.add(p)
            elif cur in negative:
                negative.remove(cur)
                for p in parents:
                    if p not in pending:
                        negative.add(p)
                    else:
                        pending.discard(p)

        positive.reverse()
        return positive

    for rev in gitrange_orig(repo, a, b):
        c = repo[rev]

        if marks.is_marked(c.hex()):
            continue

        yield c

def test_gitrange_new(a, b):

    def gitrange_new(repo, a, b):
        pending = set([b.rev()])
        bases = set([a.rev()])
        negative = pending.intersection(bases)
        pending.difference_update(negative)

        if not pending:
            return []

        pfunc = repo.changelog.parentrevs
        it = iter(xrange(b.rev(), -1, -1))

        def get_parents(rev):
            for p in pfunc(rev):
                if p == -1:
                    continue
                yield p

        positive = []
        while pending:
            cur = it.next()

            def discard_parent(p):
                pending.discard(p)
                negative.add(p)
                bases.add(p)

            if cur in negative:
                negative.remove(cur)
                for p in get_parents(cur):
                    discard_parent(p)
            elif cur in pending:
                positive.append(cur)
                pending.remove(cur)
                for p in get_parents(cur):
                    if p in bases or p in negative:
                        discard_parent(p)
                    else:
                        pending.add(p)
            elif cur in bases:
                for p in get_parents(cur):
                    if p in pending or p in negative:
                        discard_parent(p)
                    else:
                        bases.add(p)

        positive.reverse()
        return positive

    for rev in gitrange_new(repo, a, b):
        c = repo[rev]

        if marks.is_marked(c.hex()):
            continue

        yield c

def test_revwalk(a, b):
    def revwalk(repo, a, b):
        positive = []
        pending = set()
        if not marks.is_marked(b.hex()):
            pending.add(b.rev())
        for cur in xrange(b.rev(), -1, -1):
            if not pending:
                break

            if cur in pending:
                positive.append(cur)
                pending.remove(cur)
                parents = [p for p in repo.changelog.parentrevs(cur) if p >= 0]
                for p in parents:
                    if not marks.is_marked(repo[p].hex()):
                        pending.add(p)

        positive.reverse()
        return positive

    for rev in revwalk(repo, a, b):
        yield repo[rev]

def test_internal(a, b):
    if a:
        revs = smartset.baseset(repo.changelog.findmissingrevs(common=[a.rev()], heads=[b.rev()]))
    else:
        b = smartset.baseset([b.rev()])
        revs = dagop.revancestors(repo, b)
    for rev in revs:
        c = repo[rev]

        if marks.is_marked(c.hex()):
            continue

        yield c

def test_reference(a, b):
    revset = 'ancestors(%d)' % b.rev()
    if a: revset += ' -ancestors(%d)' % a.rev()
    for rev in repo.revs(revset):
        yield repo[rev]

def run_test(func, a, b, m):
    a = repo[a]
    b = repo[b]
    marks.set(m)
    return list(func(a, b))

def setup():
    n = repo.revs('%s' % branch).last()
    cur = repo.revs('%s~100' % branch).last()

    cur_marks = last_marks = {}

    revs = repo.revs('ancestors(%d)' % cur)
    cur_marks = { repo[rev].hex(): rev for rev in revs }

    revs = repo.revs('ancestors(%d)' % n)
    last_marks = { repo[rev].hex(): rev for rev in revs }

    tests['all'] = { 'a': -1, 'b': n, 'marks': {} }
    tests['cur'] = { 'a': cur, 'b': n, 'marks': cur_marks }
    tests['last'] = { 'a': n, 'b': n, 'marks': last_marks }

    for name, test in tests.iteritems():
        print '%s: %d %% %d' % (name, test['b'], test['a'])

def check():
    for test in tests.values():
        test['nodes'] = run_test(test_reference, test['a'], test['b'], test['marks'])

    for algo_name, func in algos.iteritems():
        for test_name, test in tests.iteritems():
            algo_nodes = run_test(func, test['a'], test['b'], test['marks'])
            print '%s: %s: %s' % (algo_name, test_name, 'ok' if algo_nodes == test['nodes'] else 'bad')

def benchmark(reps):
    for algo_name, func in algos.iteritems():
        start = time.time()
        for i in xrange(reps):
            for test_name, test in tests.iteritems():
                algo_nodes = run_test(func, test['a'], test['b'], test['marks'])
        end = time.time()
        elapsed = (end - start) / reps
        print '%s: %0.2f' % (algo_name, elapsed)

if len(sys.argv) < 2:
    print 'usage: walk <url> [<branch>] [<times>]'
    exit(1)

url = sys.argv[1]
branch = sys.argv[2] if len(sys.argv) >= 3 else "default"
times = int(sys.argv[3]) if len(sys.argv) >= 4 else 10

myui = ui.ui()
repo = hg.repository(myui, url)

marks = Marks()
tests = {}

algos = {}

algos['gitrange'] = test_gitrange
algos['gitrange_orig'] = test_gitrange_orig
algos['gitrange_new'] = test_gitrange_new
algos['revwalk'] = test_revwalk
algos['internal'] = test_internal

title('Setup')
setup()

title('Check')
check()

title('Benchmark')
benchmark(times)