diff --git a/plugins/compare/bundles/org.eclipse.papyrus.compare.diagram.ide.ui/src/org/eclipse/papyrus/compare/diagram/ide/ui/contentmergeviewer/PapyrusTreeContentMergeViewer.java b/plugins/compare/bundles/org.eclipse.papyrus.compare.diagram.ide.ui/src/org/eclipse/papyrus/compare/diagram/ide/ui/contentmergeviewer/PapyrusTreeContentMergeViewer.java
index 597e7a1..70cd426 100644
--- a/plugins/compare/bundles/org.eclipse.papyrus.compare.diagram.ide.ui/src/org/eclipse/papyrus/compare/diagram/ide/ui/contentmergeviewer/PapyrusTreeContentMergeViewer.java
+++ b/plugins/compare/bundles/org.eclipse.papyrus.compare.diagram.ide.ui/src/org/eclipse/papyrus/compare/diagram/ide/ui/contentmergeviewer/PapyrusTreeContentMergeViewer.java
@@ -7,18 +7,18 @@
  * 
  * Contributors:
  *     Stefan Dirix - initial API and implementation
- *     Philip Langer - introduce caching
+ *     Philip Langer - introduce caching and improve selection
  *******************************************************************************/
 package org.eclipse.papyrus.compare.diagram.ide.ui.contentmergeviewer;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
+import com.google.common.collect.Sets;
 
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
@@ -195,6 +195,9 @@
 			leftContentComputations = null;
 			rightContentComputations = null;
 			cachedMapForSelection = null;
+		} else {
+			getLeftMergeViewer().refresh();
+			getRightMergeViewer().refresh();
 		}
 
 		IMergeViewerItem leftInitialItem = null;
@@ -262,13 +265,13 @@
 		Object itemValue = item.getSideValue(side);
 		IMergeViewerItem result;
 		if (cachedMapForSelection == null) {
-			cachedMapForSelection = new HashMap<>();
+			cachedMapForSelection = Maps.newHashMap();
 			result = null;
 		} else {
 			result = cachedMapForSelection.get(itemValue);
 		}
 
-		if (result == null) {
+		if (result == null && itemValue != null) {
 			LinkedList<AbstractContentFunction> contentComputations;
 
 			switch (side) {
@@ -302,9 +305,55 @@
 				contentComputations.addAll(0, containerComputations);
 			}
 
-			// compute item
-			while (!contentComputations.isEmpty()) {
-				result = contentComputations.removeFirst().apply(itemValue, containers);
+			result = computeItemToBeSelected(itemValue, containers, contentComputations);
+		}
+
+		return result;
+	}
+
+	/**
+	 * Computes the item to be selected for the given item value, taking into account the containers, using
+	 * the content computations.
+	 * 
+	 * @param itemValue
+	 *            the item for which to base the selection.
+	 * @param containers
+	 *            the containers for that item.
+	 * @param contentComputations
+	 *            the computations for computing content.
+	 * @return the item to be selected.
+	 */
+	private IMergeViewerItem computeItemToBeSelected(Object itemValue, Collection<Object> containers,
+			LinkedList<AbstractContentFunction> contentComputations) {
+		// We try to limit the amount of time we spend looking for selection.
+		// Sometimes the object will not be present in the tree at all,
+		// causing the whole tree to be visited.
+		// So we keep track of which of the contains have not been visited,
+		// removing any containers for which there is already a selection computed.
+		IMergeViewerItem result = null;
+		Collection<Object> unvistedContainers = Sets.newHashSet(containers);
+		unvistedContainers.removeAll(cachedMapForSelection.keySet());
+		long start = System.currentTimeMillis();
+		while (!contentComputations.isEmpty()) {
+			AbstractContentFunction contentFunction = contentComputations.removeFirst();
+
+			// If we've visited all the containers and have been computing for more the 2 seconds,
+			// stop looking and select the container instead.
+			unvistedContainers.remove(contentFunction.getValue());
+			if (unvistedContainers.isEmpty() && (System.currentTimeMillis() - start) > 2 * 1000) {
+				break;
+			}
+
+			result = contentFunction.apply(itemValue, containers);
+			if (result != null) {
+				break;
+			}
+		}
+
+		// If we haven't found the selection, try to find a selection for the nearest container.
+		if (result == null) {
+			for (Object container : containers) {
+				result = cachedMapForSelection.get(container);
 				if (result != null) {
 					break;
 				}
@@ -356,7 +405,7 @@
 	private Collection<Object> getContainers(Object eObject) {
 		Collection<Object> containers;
 		if (eObject instanceof EObject) {
-			containers = new HashSet<>();
+			containers = Sets.newLinkedHashSet();
 			for (EObject container = ((EObject)eObject).eContainer(); container != null; container = container
 					.eContainer()) {
 				containers.add(container);
@@ -396,7 +445,7 @@
 		/**
 		 * The selections being cached.
 		 * 
-		 * @see PapyrusTreeContentMergeViewer#selections
+		 * @see PapyrusTreeContentMergeViewer#cachedMapForSelection
 		 */
 		protected final Map<Object, IMergeViewerItem> selections;
 
diff --git a/plugins/compare/bundles/org.eclipse.papyrus.compare.diagram.ide.ui/src/org/eclipse/papyrus/compare/diagram/ide/ui/contentmergeviewer/provider/PapyrusContextTester.java b/plugins/compare/bundles/org.eclipse.papyrus.compare.diagram.ide.ui/src/org/eclipse/papyrus/compare/diagram/ide/ui/contentmergeviewer/provider/PapyrusContextTester.java
index a10b688..4268fc2 100644
--- a/plugins/compare/bundles/org.eclipse.papyrus.compare.diagram.ide.ui/src/org/eclipse/papyrus/compare/diagram/ide/ui/contentmergeviewer/provider/PapyrusContextTester.java
+++ b/plugins/compare/bundles/org.eclipse.papyrus.compare.diagram.ide.ui/src/org/eclipse/papyrus/compare/diagram/ide/ui/contentmergeviewer/provider/PapyrusContextTester.java
@@ -1,5 +1,5 @@
 /*******************************************************************************
- * Copyright (c) 2016 EclipseSource Muenchen GmbH and others.
+ * Copyright (c) 2016, 2017 EclipseSource Muenchen GmbH and others.
  * All rights reserved. This program and the accompanying materials
  * are made available under the terms of the Eclipse Public License v1.0
  * which accompanies this distribution, and is available at
@@ -10,11 +10,13 @@
  *******************************************************************************/
 package org.eclipse.papyrus.compare.diagram.ide.ui.contentmergeviewer.provider;
 
+import static org.eclipse.papyrus.compare.diagram.ide.ui.internal.context.PapyrusContextUtil.isPapyrusContext;
+
 import java.util.Map;
+import java.util.WeakHashMap;
 
 import org.eclipse.emf.compare.Comparison;
 import org.eclipse.emf.compare.adapterfactory.context.AbstractContextTester;
-import org.eclipse.papyrus.compare.diagram.ide.ui.internal.context.PapyrusContextUtil;
 
 /**
  * Indicates whether we are in a Papyrus context.
@@ -24,12 +26,22 @@
 public class PapyrusContextTester extends AbstractContextTester {
 
 	/**
+	 * A weak cache of comparisons that have been already been tested.
+	 */
+	private final Map<Comparison, Boolean> cache = new WeakHashMap<>();
+
+	/**
 	 * {@inheritDoc}
 	 */
 	public boolean apply(Map<Object, Object> context) {
 		Comparison comparison = getComparison(context);
-		if (context != null) {
-			return PapyrusContextUtil.isPapyrusContext(comparison);
+		if (comparison != null) {
+			Boolean result = cache.get(comparison);
+			if (result == null) {
+				result = Boolean.valueOf(isPapyrusContext(comparison));
+				cache.put(comparison, result);
+			}
+			return result.booleanValue();
 		}
 		return false;
 	}
diff --git a/plugins/compare/bundles/org.eclipse.papyrus.compare.diagram.ide.ui/src/org/eclipse/papyrus/compare/diagram/ide/ui/internal/context/PapyrusContextUtil.java b/plugins/compare/bundles/org.eclipse.papyrus.compare.diagram.ide.ui/src/org/eclipse/papyrus/compare/diagram/ide/ui/internal/context/PapyrusContextUtil.java
index ffee9ad..75ba760 100644
--- a/plugins/compare/bundles/org.eclipse.papyrus.compare.diagram.ide.ui/src/org/eclipse/papyrus/compare/diagram/ide/ui/internal/context/PapyrusContextUtil.java
+++ b/plugins/compare/bundles/org.eclipse.papyrus.compare.diagram.ide.ui/src/org/eclipse/papyrus/compare/diagram/ide/ui/internal/context/PapyrusContextUtil.java
@@ -1,5 +1,5 @@
 /*******************************************************************************
- * Copyright (c) 2016 EclipseSource Muenchen GmbH and others.
+ * Copyright (c) 2016, 2017 EclipseSource Muenchen GmbH and others.
  * All rights reserved. This program and the accompanying materials
  * are made available under the terms of the Eclipse Public License v1.0
  * which accompanies this distribution, and is available at
@@ -36,11 +36,14 @@
  */
 public final class PapyrusContextUtil {
 
+	/** The value ".di". */
+	private static final String DI_FILE_EXTENSION = "." + DiModel.DI_FILE_EXTENSION; //$NON-NLS-1$
+
 	/** Predicate specifying whether the given URI string ends with the Papyrus DI file extension. */
 	private static final Predicate<String> ENDS_WITH_PAPYRUS_EXTENSION = new Predicate<String>() {
 		public boolean apply(String input) {
 			if (input != null) {
-				return input.endsWith(DiModel.DI_FILE_EXTENSION);
+				return input.endsWith(DI_FILE_EXTENSION);
 			}
 			return false;
 		}