~garyvdm/bzr/loggraphviz

# Copyright (C) 2009, 2010 Canonical Ltd

#

# This program is free software; you can redistribute it and/or

# modify it under the terms of the GNU General Public License

# as published by the Free Software Foundation; either version 2

# of the License, or (at your option) any later version.

#

# This program is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

# GNU General Public License for more details.

#

# You should have received a copy of the GNU General Public License

# along with this program; if not, write to the Free Software

# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.

"""

Layout a revision graph for visual representation, allowing for filtering and

collapsible branches.

Usage example

=============

.. python::

  bi = loggraphviz.BranchInfo('branch-label', tree, branch)

  graph_viz = loggraphviz.GraphVizLoader([bi], bi, False)

  graph_viz.load()

  state = loggraphviz.GraphVizFilterState(graph_viz)

  computed = graph_viz.compute_viz(state)

"""

import gc

from itertools import izip

from bzrlib import errors

from bzrlib.bzrdir import BzrDir

from bzrlib.transport.local import LocalTransport

from bzrlib.revision import NULL_REVISION, CURRENT_REVISION

from bzrlib.graph import (

    Graph,

    StackedParentsProvider,

    KnownGraph,

    DictParentsProvider,

    )

class BranchInfo(object):

    """Wrapper for a branch, it's working tree, if available, and a label."""

    def __init__(self, label, tree, branch, index=None):

        self.label = label

        self.tree = tree

        self.branch = branch

        self.index = index

    def __hash__(self):

        return self.branch.base.__hash__()

    def __eq__(self, other):

        if isinstance(other, BranchInfo):

            return self.branch.base.__eq__(other.branch.base)

        return False

class RevisionData(object):

    """

    Container for data for a revision in the graph that gets calculated

    when the graph is loaded.

    """

    # Instance of this object are typically named "rev".

    __slots__ = ["index", "_merge_sort_node", "branch", "_revno_str",

                 "merges", "merged_by", 'branch_id', 'color']

    def __init__(self, index, merge_sort_node):

        """Create a new RevisionData instance."""

        self.index = index

        self._merge_sort_node = merge_sort_node

        self.branch = None

        self._revno_str = None

        self.merges = []

        """Revision indexes that this revision merges"""

        self.merged_by = None

        """Revision index that merges this revision."""

        self.branch_id = self._merge_sort_node.revno[0:-1]

        self.color = reduce(lambda x, y: x + y, self.branch_id, 0)

    revid = property(lambda self: self._merge_sort_node.key)

    merge_depth = property(lambda self: self._merge_sort_node.merge_depth)

    revno_sequence = property(lambda self: self._merge_sort_node.revno)

    end_of_merge = property(lambda self: self._merge_sort_node.end_of_merge)

    def get_revno_str(self):

        if self._revno_str is None:

            self._revno_str = ".".join(["%d" % (revno)

                                for revno in self.revno_sequence])

            if self.revid.startswith(CURRENT_REVISION):

                self._revno_str += " ?"

        return self._revno_str

    revno_str = property(get_revno_str)

    def __repr__(self):

        return "%s <%s %s>" % (self.__class__.__name__, self.revno_str,

                              self.revid)

class BranchLine(object):

    """Container for data for a branch line, aka merge line."""

    __slots__ = ["branch_id", "revs", "merges", "merged_by",

                 "color", "merge_depth"]

    def __init__(self, branch_id):

        self.branch_id = branch_id

        self.revs = []

        self.merges = []

        self.merged_by = []

        self.merge_depth = 0

    def __repr__(self):

        return "%s <%s>" % (self.__class__.__name__, self.branch_id)

class GhostRevisionError(errors.InternalBzrError):

    _fmt = "{%(revision_id)s} is a ghost."

    def __init__(self, revision_id):

        errors.BzrError.__init__(self)

        self.revision_id = revision_id

class GraphVizLoader(object):

    """

    Loads graph for branches and provides computed layout for visual

    representation.

    """

    # Most list/dicts related to revisions are unfiltered. When we do a graph

    # layout, we filter these revisions. A revision may be filter out because:

    # * It's branch is hidden (or collapsed).

    # * We have a specified file_id(s), and the revision does not touch the

    #   file_id(s).

    # * We have a search, and the revision does not match the search.

    #

    # The main list of unfiltered revisions is self.revisions. A revisions

    # index in revisions are normally called index. The main list of filtered

    # revisions is filtered_revs. Revision indexes in this list are called

    # f_index.

    def __init__(self, branches, primary_bi, no_graph):

        self.branches = branches

        """List of BranchInfo for each branch."""

        self.primary_bi = primary_bi

        self.no_graph = no_graph

        self.repos = []

        self.local_repo_copies = []

        """A list of repositories that revisions will be attempted to be loaded        

        from first."""

        self.revid_head_info = {}

        """Dict with a keys of head revid and value of

            (list of (branch, label),

             list of revids that are unique to this head)

            We can't just store the BranchInfo, because the label may have

            have addition text in it, e.g. "Working Tree", "Pending Merges"

        """

        self.revid_branch_info = {}

        self.ghosts = set()

        self.revisions = []

        """List of RevisionInfo from merge_sort."""

        self.revid_rev = {}

        self.graph_children = {}

        self.tags = {}      # map revid -> tags set

    def load(self):

        # Get a unique list of repositories. If the url is the same,

        # we consider it the same repositories

        repo_urls = set()

        for bi in self.branches:

            repo = bi.branch.repository

            if repo.base not in repo_urls:

                repo_urls.add(repo.base)

                self.repos.append(repo)

        no_local_repos = True

        for repo in self.repos:

            if repo_is_local(repo):

                no_local_repos = False

        if no_local_repos:

            self.load_current_dir_repo()

        self.repos.sort(key=lambda repo: not repo_is_local(repo))

        self.lock_read_branches()

        try:

            head_revids, graph_parents = self.load_graph_parents()

            self.process_graph_parents(head_revids, graph_parents)

            self.compute_head_info()

            del self.graph

            if not self.no_graph:

                self.compute_branch_lines()

                self.compute_merge_info()

            self.load_tags()

        finally:

            self.unlock_branches()

    def load_current_dir_repo(self):

        # There are no local repositories. Try open the repository

        # of the current directory, and try load revisions data from

        # this before trying from remote repositories. This makes

        # the common use case of viewing a remote branch that is

        # related to the current branch much faster, because most

        # of the revision can be loaded from the local repository.

        try:

            bzrdir, relpath = BzrDir.open_containing(u".")

            repo = bzrdir.find_repository()

            self.repos.add(repo)

            self.local_repo_copies.append(repo)

        except Exception:

            pass

    def update_ui(self):

        pass

    def throbber_show(self):

        pass

    def throbber_hide(self):

        pass

    def lock_read_branches(self):

        for bi in self.branches:

            bi.branch.lock_read()

        for repo in self.repos:

            repo.lock_read()

    def unlock_branches(self):

        for bi in self.branches:

            bi.branch.unlock()

        for repo in self.repos:

            repo.unlock()

    #def lock_read_repos(self):

    #    for repo in self.repos.itervalues():

    #        repo.lock_read()

    #

    #def unlock_repos(self):

    #    for repo in self.repos.itervalues():

    #        repo.unlock()

    def load_tags(self):

        self.tags = {}

        for bi in self.branches:

            # revid to tags map

            branch_tags = bi.branch.tags.get_reverse_tag_dict()

            for revid, tags in branch_tags.iteritems():

                if revid in self.tags:

                    self.tags[revid].update(set(tags))

                else:

                    self.tags[revid] = set(tags)

    def append_head_info(self, revid, branch_info, tag):

        if not revid == NULL_REVISION:

            if not revid in self.revid_head_info:

                self.revid_head_info[revid] = ([], [])

            self.revid_head_info[revid][0].append((branch_info, tag))

            # So that early calls to get_revid_branch work

            self.revid_branch_info[revid] = branch_info

    def load_branch_heads(self, bi):

        branch_heads = []

        def append_head_info(revid, branch_info, label):

            self.append_head_info(revid, branch_info, label)

            branch_heads.append(revid)

        if len(self.branches) > 0:

            label = bi.label

        else:

            label = None

        branch_last_revision = bi.branch.last_revision()

        append_head_info(branch_last_revision, bi, bi.label)

        self.update_ui()

        if bi.tree:

            parent_ids = bi.tree.get_parent_ids()

            if parent_ids:

                # first parent is last revision of the tree

                revid = parent_ids[0]

                if revid != branch_last_revision:

                    # working tree is out of date

                    if label:

                        wt_label = "%s - Working Tree" % label

                    else:

                        wt_label = "Working Tree"

                    append_head_info(revid, bi, wt_label)

                # other parents are pending merges

                for revid in parent_ids[1:]:

                    if label:

                        pm_label = "%s - Pending Merge" % label

                    else:

                        pm_label = "Pending Merge"

                    append_head_info(revid, bi, pm_label)

            self.update_ui()

        return branch_heads, branch_heads, ()

    def load_graph_parents(self):

        """Load the heads of the graph, and the graph parents"""

        extra_parents = {}

        branches_heads = []

        def load_branch_heads(bi, insert_at_begin=False):

            load_heads, sort_heads, extra_parents_ = self.load_branch_heads(bi)

            for key, parents in extra_parents_:

                extra_parents[key] = parents

            if insert_at_begin:

                branches_heads.insert(0, (load_heads, sort_heads))

            else:

                branches_heads.append((load_heads, sort_heads))

        for bi in self.branches:

            # Don't do the primary branch, as that will be inserted later at

            # the first position.

            if bi != self.primary_bi:

                load_branch_heads(bi)

        if len(branches_heads) >= 2:

            head_revids = [revid for load_heads, sort_heads in branches_heads

                                 for revid in load_heads]

            head_revs = self.load_revisions(head_revids)

            get_max_timestamp = lambda branch_heads: max(

                [head_revs[revid].timestamp for revid in branch_heads[0]])

            branches_heads.sort(key=get_max_timestamp, reverse=True)

        if self.primary_bi:

            load_branch_heads(self.primary_bi, True)

        load_heads = [revid for load_heads_, sort_heads_ in branches_heads

                      for revid in load_heads_]

        sort_heads = [revid for load_heads_, sort_heads_ in branches_heads

                      for revid in sort_heads_]

        parents_providers = [repo._make_parents_provider() \

                             for repo in self.repos]

        parents_providers.append(DictParentsProvider(extra_parents))

        self.graph = Graph(StackedParentsProvider(parents_providers))

        return sort_heads, self.graph.iter_ancestry(sort_heads)

    def process_graph_parents(self, head_revids, graph_parents_iter):

        graph_parents = {}

        self.ghosts = set()

        for (revid, parent_revids) in graph_parents_iter:

            if revid == NULL_REVISION:

                continue

            if parent_revids is None:

                #Ghost

                graph_parents[revid] = ()

                self.ghosts.add(revid)

            elif parent_revids == (NULL_REVISION,):

                graph_parents[revid] = ()

            else:

                graph_parents[revid] = parent_revids

            if len(graph_parents) % 100 == 0:

                self.update_ui()

        graph_parents["top:"] = head_revids

        if len(graph_parents) > 0:

            enabled = gc.isenabled()

            if enabled:

                gc.disable()

            def make_kg():

                return KnownGraph(graph_parents)

            self.known_graph = make_kg()

            merge_sorted_revisions = self.known_graph.merge_sort('top:')

            # So far, we are a bit faster than the pure-python code. But the

            # last step hurts. Specifically, we take

            #   377ms KnownGraph(self.graph_parents)

            #   263ms kg.merge_sort() [640ms combined]

            #  1322ms self.revisions = [...]

            # vs 

            #  1152ms tsort.merge_sort(self.graph_parents)

            #   691ms self.revisions = [...]

            #

            # It is a gc thing... :(

            # Adding gc.disable() / gc.enable() around this whole loop changes

            # things to be:

            #   100ms   KnownGraph(self.graph_parents)

            #    77ms   kg.merge_sort() [177ms combined]

            #   174ms   self.revisions = [...]

            # vs

            #   639ms   tsort.merge_sort(self.graph_parents)

            #   150ms   self.revisions = [...]

            # Also known as "wow that's a lot faster". This is because KG()

            # creates a bunch of Python objects, then merge_sort() creates a

            # bunch more. And then self.revisions() creates another whole set.

            # And all of these are moderately long lived, so you have a *lot*

            # of allocations without removals (which triggers the gc checker

            # over and over again.) And they probably don't live in cycles

            # anyway, so you can skip it for now, and just run at the end.

            # self.revisions *is* a little bit slower. Probably because pyrex

            # MergeSortNodes use long integers rather than PyIntObject and thus

            # create them on-the-fly.

            # Get rid of the 'top:' revision

            merge_sorted_revisions.pop(0)

            self.revisions = [RevisionData(index, node)

                for index, node in enumerate(merge_sorted_revisions)]

            if enabled:

                gc.enable()

        else:

            self.revisions = ()

        self.revid_rev = {}

        self.revno_rev = {}

        self.max_mainline_revno = 0

        for rev in self.revisions:

            self.max_mainline_revno = max(self.max_mainline_revno, 

                                          rev.revno_sequence[0])

            self.revid_rev[rev.revid] = rev

            self.revno_rev[rev.revno_sequence] = rev

    def branch_id_sort_key(self, x):

        merge_depth = self.branch_lines[x].merge_depth

        # Note: This greatly affects the layout of the graph.

        #

        # Branch line that have a smaller merge depth should be to the left

        # of those with bigger merge depths.

        #

        # For branch lines that have the same parent in the mainline -

        # those with bigger branch numbers to be to the rights. E.g. for

        # the following dag, you want the graph to appear as on the left,

        # not as on the right:

        #

        # 3     F_       F

        #       | \      |\

        # 1.2.1 |  E     | E

        #       |  |     | \

        # 2     D  |     D_|

        #       |\ |     | +_

        # 1.1.2 | C|     | | C

        #       | |/     |  \|

        # 1.1.1 | B      |   B

        #       |/       | /

        # 1     A        A

        #

        # Otherwise, those with a greater mainline parent revno should

        # appear to the left.

        if len(x) == 0:

            return (merge_depth)

        else:

            return (merge_depth, -x[0], x[1])

    def compute_branch_lines(self):

        self.branch_lines = {}

        """A list of branch lines (aka merge lines).

        For a revisions, the revision number less the least significant

        digit is the branch_id, and used as the key for the dict. Hence

        revision with the same revision number less the least significant

        digit are considered to be in the same branch line. e.g.: for

        revisions 290.12.1 and 290.12.2, the branch_id would be 290.12,

        and these two revisions will be in the same branch line.

        """

        self.branch_ids = []

        """List of branch ids, sorted in the order that the branches will

        be shown, from left to right on the graph."""

        for rev in self.revisions:

            branch_line = None

            if rev.branch_id not in self.branch_lines:

                branch_line = BranchLine(rev.branch_id)

                self.branch_lines[rev.branch_id] = branch_line

            else:

                branch_line = self.branch_lines[rev.branch_id]

            branch_line.revs.append(rev)

            branch_line.merge_depth = max(rev.merge_depth,

                                          branch_line.merge_depth)

            rev.branch = branch_line

        self.branch_ids = self.branch_lines.keys()

        self.branch_ids.sort(key=self.branch_id_sort_key)

    def compute_merge_info(self):

        def set_merged_by(rev, merged_by, merged_by_rev, do_branches=False):

            if merged_by is None:

                return

            if merged_by_rev is None:

                merged_by_rev = self.revisions[merged_by]

            rev.merged_by = merged_by

            merged_by_rev.merges.append(rev.index)

            if do_branches:

                branch_id = rev.branch_id

                branch_merged_by = self.branch_lines[branch_id].merged_by

                merged_by_branch_id = self.revisions[merged_by].branch_id

                merged_by_branch_merges = \

                    self.branch_lines[merged_by_branch_id].merges

                if not branch_id in merged_by_branch_merges:

                    merged_by_branch_merges.append(branch_id)

                if not merged_by_branch_id in branch_merged_by:

                    branch_merged_by.append(merged_by_branch_id)

        for rev in self.revisions:

            parents = [self.revid_rev[parent] for parent in

                       self.known_graph.get_parent_keys(rev.revid)]

            if len(parents) > 0:

                if rev.branch_id == parents[0].branch_id:

                    set_merged_by(parents[0], rev.merged_by, None)

            for parent in parents[1:]:

                if rev.merge_depth <= parent.merge_depth:

                    set_merged_by(parent, rev.index, rev, do_branches=True)

    def compute_head_info(self):

        def get_revid_head(heads):

            map = {}

            for i in xrange(len(heads)):

                prev_revids = [revid for revid, head in heads[:i]]

                unique_ancestors = self.graph.find_unique_ancestors(

                    heads[i][0], prev_revids)

                for ancestor_revid in unique_ancestors:

                    map[ancestor_revid] = heads[i][1]

            return map

        if len(self.branches) > 1:

            head_revid_branch_info = sorted(

                [(revid, branch_info)

                 for revid, (head_info, ur) in self.revid_head_info.iteritems()

                 for (branch_info, tag) in head_info],

                key=lambda x: not repo_is_local(x[1].branch.repository))

            self.revid_branch_info = get_revid_head(head_revid_branch_info)

        else:

            self.revid_branch_info = {}

        head_count = 0

        for head_info, ur in self.revid_head_info.itervalues():

            head_count += len(head_info)

        if head_count > 1:

            # Populate unique revisions for heads

            for revid, (head_info, ur) in self.revid_head_info.iteritems():

                rev = None

                if revid in self.revid_rev:

                    rev = self.revid_rev[revid]

                if rev and rev.merged_by:

                    # This head has been merged.

                    # d

                    # |\

                    # b c

                    # |/

                    # a

                    # if revid == c,then we want other_revids = [b]

                    merged_by_revid = self.revisions[rev.merged_by].revid

                    other_revids = [self.known_graph

                                    .get_parent_keys(merged_by_revid)[0]]

                else:

                    other_revids = [other_revid for other_revid \

                        in self.revid_head_info.iterkeys() \

                        if not other_revid == revid]

                ur.append(revid)

                ur.extend([revid for revid \

                    in self.graph.find_unique_ancestors(revid, other_revids) \

                    if not revid == NULL_REVISION and revid in self.revid_rev])

                ur.sort(key=lambda x: self.revid_rev[x].index)

    def compute_viz(self, state):

        # Overview:

        # Work out which revision need to be displayed.

        # Create ComputedGraphViz and ComputedRevisionData objects

        # Assign columns for branches, and lines that go between branches.

        #   These are intermingled, because some of the lines need to come

        #   before it's branch, and others need to come after. Other lines

        #   (such a the line from the last rev in a branch) are treated a

        #   special cases.

        # Return ComputedGraphViz object

        gc_enabled = gc.isenabled()

        gc.disable()

        try:

            computed = ComputedGraphViz(self)

            computed.filtered_revs = [ComputedRevisionData(rev) for rev in

                                      state.get_filtered_revisions()]

            c_revisions = computed.revisions

            for f_index, c_rev in enumerate(computed.filtered_revs):

                c_revisions[c_rev.rev.index] = c_rev

                c_rev.f_index = f_index

            for (revid, (head_info,

                         unique_revids)) in self.revid_head_info.iteritems():

                for unique_revid in unique_revids:

                    rev = self.revid_rev[unique_revid]

                    c_rev = c_revisions[rev.index]

                    if c_rev is not None:

                        c_rev.branch_labels.extend(head_info)

                        break

        finally:

            if gc_enabled:

                gc.enable()

        if self.no_graph:

            for c_rev in computed.filtered_revs:

                c_rev.col_index = c_rev.rev.merge_depth * 0.5

            return computed

        # This will hold a tuple of (child, parent, col_index, direct) for each

        # line that needs to be drawn. If col_index is not none, then the line

        # is drawn along that column, else the the line can be drawn directly

        # between the child and parent because either the child and parent are

        # in the same branch line, or the child and parent are 1 row apart.

        lines = []

        lines_by_column = []

        def branch_line_col_search_order(start_col_index):

            for col_index in range(start_col_index, len(lines_by_column)):

                yield col_index

            #for col_index in range(parent_col_index-1, -1, -1):

            #    yield col_index

        def line_col_search_order(parent_col_index, child_col_index):

            if parent_col_index is not None and child_col_index is not None:

                max_index = max(parent_col_index, child_col_index)

                min_index = min(parent_col_index, child_col_index)

                # First yield the columns between the child and parent.

                for col_index in range(max_index, min_index - 1, -1):

                    yield col_index

            elif child_col_index is not None:

                max_index = child_col_index

                min_index = child_col_index

                yield child_col_index

            elif parent_col_index is not None:

                max_index = parent_col_index

                min_index = parent_col_index

                yield parent_col_index

            else:

                max_index = 0

                min_index = 0

                yield 0

            i = 1

            # then yield the columns on either side.

            while max_index + i < len(lines_by_column) or \

                  min_index - i > -1:

                if max_index + i < len(lines_by_column):

                    yield max_index + i

                #if min_index - i > -1:

                #    yield min_index - i

                i += 1

        def is_col_free_for_range(col_index, child_f_index, parent_f_index):

            return not any(

                range_overlaps(child_f_index, parent_f_index,

                               line_child_f_index, line_parent_f_index)

                for line_child_f_index, line_parent_f_index

                in lines_by_column[col_index])

        def find_free_column(col_search_order, child_f_index, parent_f_index):

            for col_index in col_search_order:

                if is_col_free_for_range(col_index,

                                             child_f_index, parent_f_index):

                    break

            else:

                # No free columns found. Add an empty one on the end.

                col_index = len(lines_by_column)

                lines_by_column.append([])

            return col_index

        def append_line(child, parent, direct, col_index=None):

            lines.append((child, parent, col_index, direct))

            if col_index is not None:

                lines_by_column[int(round(col_index))].append(

                    (child.f_index, parent.f_index))

        def find_visible_parent(c_rev, parent, twisty_hidden_parents):

            if c_revisions[parent.index] is not None:

                return (c_rev, c_revisions[parent.index], True)

            else:

                if parent.index in twisty_hidden_parents:

                    # no need to draw a line if there is a twisty,

                    # except if this is the last in the branch.

                    return None

                # The parent was not visible. Search for a ancestor

                # that is. Stop searching if we make a hop, i.e. we

                # go away from our branch, and we come back to it.

                has_seen_different_branch = False

                while c_revisions[parent.index] is None:

                    if not parent.branch_id == c_rev.rev.branch_id:

                        has_seen_different_branch = True

                    # find grand parent.

                    g_parent_ids = (self.known_graph

                                    .get_parent_keys(parent.revid))

                    if len(g_parent_ids) == 0:

                        return None

                    else:

                        parent = self.revid_rev[g_parent_ids[0]]

                    if (has_seen_different_branch and

                        parent.branch_id == branch_id):

                        # We have gone away and come back to our

                        # branch. Stop.

                        return None

                if parent:

                    # Not Direct

                    return (c_rev, c_revisions[parent.index], False)

        def append_branch_parent_lines(branch_rev_visible_parents):

            groups = group_overlapping(branch_rev_visible_parents)

            for parents, start, end, group_key in groups:

                # Since all parents go from the same branch line to the

                # same branch line, we can use the col indexes of the

                # parent.

                if end - start == 1:

                    col_index = None

                else:

                    col_search_order = line_col_search_order(

                        parents[0][1].col_index, parents[0][0].col_index)                        

                    col_index = find_free_column(col_search_order,

                                                 start, end)

                col_offset_increment = 1.0 / len(parents)

                for i, (c_rev, parent_c_rev, direct) in enumerate(parents):

                    if col_index is None:

                        col_index_offset = None

                    else:

                        col_index_offset = (col_index - 0.5 +

                                            (i * col_offset_increment) +

                                            (col_offset_increment / 2))

                    append_line(c_rev, parent_c_rev,

                                direct, col_index_offset)

        for branch_id in self.branch_ids:

            if not branch_id in state.branch_line_state:

                continue

            branch_line = self.branch_lines[branch_id]

            branch_revs = [c_revisions[rev.index]

                           for rev in branch_line.revs

                           if c_revisions[rev.index] is not None]

            if not branch_revs:

                continue

            branch_rev_visible_parents_post = []

            branch_rev_visible_parents_pre = []

            # Lists of ([(c_rev, parent_c_rev, is_direct)],

            #            start, end, group_key]

            last_c_rev = branch_revs[-1]

            last_rev_left_parents = (self.known_graph

                                     .get_parent_keys(last_c_rev.rev.revid))

            if last_rev_left_parents: