Updates binary VQL expression return values #28

viatra · Apr 17, 2019 · 438b5e8 · 438b5e8
1 parent 752d71c
commit 438b5e8
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/...cquerylabs.v4md/src/main/com/incquerylabs/v4md/expressions/BinaryQueryResultProvider.java b/...cquerylabs.v4md/src/main/com/incquerylabs/v4md/expressions/BinaryQueryResultProvider.java
@@ -17,8 +17,8 @@ public abstract class BinaryQueryResultProvider {
 	 * This method is responsible for calculating the results to be returned to MagicDraw.
 	 * @return A non-null, but possibly empty stream of results
 	 */
-	public abstract Stream<Element> getResults(ViatraQueryEngine engine, Object source);
-
+	public abstract Stream<Object> getResults(ViatraQueryEngine engine, Object source);
+	
 	/**
 	 * Returns the expected source object types for the given query
 	 */

diff --git a/com.incquerylabs.v4md/src/main/com/incquerylabs/v4md/expressions/BinaryVQLExpression.java b/com.incquerylabs.v4md/src/main/com/incquerylabs/v4md/expressions/BinaryVQLExpression.java
@@ -1,8 +1,8 @@
 package com.incquerylabs.v4md.expressions;
 
 import java.lang.reflect.Constructor;
-import java.util.HashSet;
-import java.util.Set;
+import java.util.ArrayList;
+import java.util.List;
 
 import org.eclipse.viatra.query.runtime.api.ViatraQueryEngine;
 
@@ -53,11 +53,12 @@ public BinaryVQLExpression(ParameterizedExpressionDescriptor descriptor, Express
 	public Object getValue(Element sourceParameter, ValueContext context) throws Exception {
 		ViatraQueryEngine engine = ViatraQueryAdapter.getAdapter(project).get().getEngine();
 
-		Set<Element> returnSet = new HashSet<Element>();
+		List<Object> returnSet = new ArrayList<>();
 		if (provider != null) {
 			provider.getResults(engine, sourceParameter).forEach(returnSet::add);
 		}
-		return returnSet;
+		// MagicDraw sometimes handles single-valued returns differently, so instead of 1-length lists we should return the element itself
+		return returnSet.size() == 1 ? returnSet.get(0) : returnSet;
 	}
 
 	@Override