1 files changed, 13 insertions, 2 deletions

diff --git a/subprojects/store-query/src/main/java/tools/refinery/store/query/dnf/FunctionalQuery.java b/subprojects/store-query/src/main/java/tools/refinery/store/query/dnf/FunctionalQuery.java
index 5a32b1ba..225f6844 100644
--- a/subprojects/store-query/src/main/java/tools/refinery/store/query/dnf/FunctionalQuery.java
+++ b/subprojects/store-query/src/main/java/tools/refinery/store/query/dnf/FunctionalQuery.java
@@ -5,6 +5,7 @@
  */
 package tools.refinery.store.query.dnf;
 
+import tools.refinery.store.query.InvalidQueryException;
 import tools.refinery.store.query.literal.CallPolarity;
 import tools.refinery.store.query.term.Aggregator;
 import tools.refinery.store.query.term.AssignedValue;
@@ -26,14 +27,14 @@ public final class FunctionalQuery<T> extends Query<T> {
                         var parameter = parameters.get(i);
                         var parameterType = parameter.tryGetType();
                         if (parameterType.isPresent()) {
-                                throw new IllegalArgumentException("Expected parameter %s of %s to be a node variable, got %s instead"
+                                throw new InvalidQueryException("Expected parameter %s of %s to be a node variable, got %s instead"
                                                 .formatted(parameter, dnf, parameterType.get().getName()));
                         }
                 }
                 var outputParameter = parameters.get(outputIndex);
                 var outputParameterType = outputParameter.tryGetType();
                 if (outputParameterType.isEmpty() || !outputParameterType.get().equals(type)) {
-                        throw new IllegalArgumentException("Expected parameter %s of %s to be %s, but got %s instead".formatted(
+                        throw new InvalidQueryException("Expected parameter %s of %s to be %s, but got %s instead".formatted(
                                         outputParameter, dnf, type, outputParameterType.map(Class::getName).orElse("node")));
                 }
                 this.type = type;
@@ -54,6 +55,16 @@ public final class FunctionalQuery<T> extends Query<T> {
                 return null;
         }
 
+        @Override
+        protected FunctionalQuery<T> withDnfInternal(Dnf newDnf) {
+                return newDnf.asFunction(type);
+        }
+
+        @Override
+        public FunctionalQuery<T> withDnf(Dnf newDnf) {
+                return (FunctionalQuery<T>) super.withDnf(newDnf);
+        }
+
         public AssignedValue<T> call(List<NodeVariable> arguments) {
                 return targetVariable -> {
                         var argumentsWithTarget = new ArrayList<Variable>(arguments.size() + 1);