1 files changed, 16 insertions, 18 deletions

diff --git a/plugins/infra/gmfdiag/org.eclipse.papyrus.infra.gmfdiag.modelexplorer/src/org/eclipse/papyrus/infra/gmfdiag/modelexplorer/queries/GetContainedDiagrams.java b/plugins/infra/gmfdiag/org.eclipse.papyrus.infra.gmfdiag.modelexplorer/src/org/eclipse/papyrus/infra/gmfdiag/modelexplorer/queries/GetContainedDiagrams.java
index b6ea0b02f61..1aa786b4f3d 100644
--- a/plugins/infra/gmfdiag/org.eclipse.papyrus.infra.gmfdiag.modelexplorer/src/org/eclipse/papyrus/infra/gmfdiag/modelexplorer/queries/GetContainedDiagrams.java
+++ b/plugins/infra/gmfdiag/org.eclipse.papyrus.infra.gmfdiag.modelexplorer/src/org/eclipse/papyrus/infra/gmfdiag/modelexplorer/queries/GetContainedDiagrams.java
@@ -12,10 +12,13 @@
  */
 package org.eclipse.papyrus.infra.gmfdiag.modelexplorer.queries;
 
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Iterator;
+import java.util.List;
 
 import org.eclipse.emf.ecore.EObject;
-import org.eclipse.emf.ecore.EStructuralFeature;
+import org.eclipse.emf.ecore.util.EcoreUtil;
 import org.eclipse.emf.facet.infra.query.core.exception.ModelQueryExecutionException;
 import org.eclipse.emf.facet.infra.query.core.java.IJavaModelQuery;
 import org.eclipse.emf.facet.infra.query.core.java.ParameterValueList;
@@ -23,26 +26,21 @@ import org.eclipse.gmf.runtime.notation.Diagram;
 import org.eclipse.papyrus.views.modelexplorer.NavigatorUtils;
 import org.eclipse.papyrus.views.modelexplorer.queries.AbstractEditorContainerQuery;
 
-import com.google.common.base.Function;
-import com.google.common.base.Predicate;
-
 /** Get the collection of all contained diagrams */
 public class GetContainedDiagrams extends AbstractEditorContainerQuery implements IJavaModelQuery<EObject, Collection<org.eclipse.gmf.runtime.notation.Diagram>> {
 
-	public Collection<org.eclipse.gmf.runtime.notation.Diagram> evaluate(final EObject context, final ParameterValueList parameterValues) throws ModelQueryExecutionException {
-		Predicate<EStructuralFeature.Setting> p = new Predicate<EStructuralFeature.Setting>() {
-
-			public boolean apply(EStructuralFeature.Setting arg0) {
-				return arg0.getEObject() instanceof Diagram;
+	public Collection<Diagram> evaluate(final EObject context, final ParameterValueList parameterValues) throws ModelQueryExecutionException {
+		List<Diagram> result = new ArrayList<Diagram>(3);
+		Iterator<EObject> roots = NavigatorUtils.getNotationRoots(context);
+		if (roots == null)
+			return result;
+		while (roots.hasNext()) {
+			EObject root = roots.next();
+			if (root instanceof Diagram) {
+				if (EcoreUtil.equals(((Diagram)root).getElement(), context))
+					result.add((Diagram)root);
 			}
-		};
-		Function<EStructuralFeature.Setting, Diagram> f = new Function<EStructuralFeature.Setting, Diagram>() {
-
-			public Diagram apply(EStructuralFeature.Setting arg0) {
-				return (Diagram)arg0.getEObject();
-			}
-
-		};
-		return NavigatorUtils.findFilterAndApply(context, p, f);
+		}
+		return result;
 	}
 }