refactor: DNF query builder

5 files changed, 233 insertions, 238 deletions

diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/ViatraQueryableModelStore.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/ViatraQueryableModelStore.java
index a63abc3a..702eb659 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/ViatraQueryableModelStore.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/ViatraQueryableModelStore.java
@@ -5,9 +5,11 @@ import tools.refinery.store.model.ModelDiffCursor;
 import tools.refinery.store.model.ModelStore;
 import tools.refinery.store.model.ModelStoreImpl;
 import tools.refinery.store.model.representation.DataRepresentation;
-import tools.refinery.store.query.QueryableModel;
-import tools.refinery.store.query.QueryableModelStore;
-import tools.refinery.store.query.building.*;
+import tools.refinery.store.query.*;
+import tools.refinery.store.query.atom.DNFAtom;
+import tools.refinery.store.query.atom.DNFCallAtom;
+import tools.refinery.store.query.atom.EquivalenceAtom;
+import tools.refinery.store.query.atom.RelationViewAtom;
 import tools.refinery.store.query.viatra.internal.RawPatternMatcher;
 import tools.refinery.store.query.viatra.internal.ViatraQueryableModel;
 import tools.refinery.store.query.viatra.internal.pquery.DNF2PQuery;
@@ -23,10 +25,10 @@ public class ViatraQueryableModelStore implements QueryableModelStore {
 
         protected final Set<RelationView<?>> relationViews;
 
-        protected final Map<DNFPredicate, GenericQuerySpecification<RawPatternMatcher>> predicates;
+        protected final Map<DNF, GenericQuerySpecification<RawPatternMatcher>> predicates;
 
         public ViatraQueryableModelStore(ModelStore store, Set<RelationView<?>> relationViews,
-                                                                         Set<DNFPredicate> predicates) {
+                                                                         Set<DNF> predicates) {
                 this.store = store;
                 validateViews(store.getDataRepresentations(), relationViews);
                 this.relationViews = Collections.unmodifiableSet(relationViews);
@@ -35,7 +37,7 @@ public class ViatraQueryableModelStore implements QueryableModelStore {
         }
 
         public ViatraQueryableModelStore(Set<DataRepresentation<?, ?>> dataRepresentations,
-                                                                         Set<RelationView<?>> relationViews, Set<DNFPredicate> predicates) {
+                                                                         Set<RelationView<?>> relationViews, Set<DNF> predicates) {
                 this(new ModelStoreImpl(dataRepresentations), relationViews, predicates);
         }
 
@@ -49,44 +51,46 @@ public class ViatraQueryableModelStore implements QueryableModelStore {
                 }
         }
 
-        private void validatePredicates(Set<RelationView<?>> relationViews, Set<DNFPredicate> predicates) {
-                for (DNFPredicate dnfPredicate : predicates) {
+        private void validatePredicates(Set<RelationView<?>> relationViews, Set<DNF> predicates) {
+                for (DNF dnfPredicate : predicates) {
                         for (DNFAnd clause : dnfPredicate.getClauses()) {
-                                for (DNFAtom atom : clause.getConstraints()) {
-                                        if (atom instanceof RelationAtom relationAtom) {
-                                                validateRelationAtom(relationViews, dnfPredicate, relationAtom);
-                                        } else if (atom instanceof PredicateAtom predicateAtom) {
-                                                validatePredicateAtom(predicates, dnfPredicate, predicateAtom);
+                                for (DNFAtom atom : clause.constraints()) {
+                                        if (atom instanceof RelationViewAtom relationViewAtom) {
+                                                validateRelationAtom(relationViews, dnfPredicate, relationViewAtom);
+                                        } else if (atom instanceof DNFCallAtom queryCallAtom) {
+                                                validatePredicateAtom(predicates, dnfPredicate, queryCallAtom);
+                                        } else if (!(atom instanceof EquivalenceAtom)) {
+                                                throw new IllegalArgumentException("Unknown constraint: " + atom.toString());
                                         }
                                 }
                         }
                 }
         }
 
-        private void validateRelationAtom(Set<RelationView<?>> relationViews, DNFPredicate dnfPredicate,
-                                                                          RelationAtom relationAtom) {
-                if (!relationViews.contains(relationAtom.view())) {
+        private void validateRelationAtom(Set<RelationView<?>> relationViews, DNF dnfPredicate,
+                                                                          RelationViewAtom relationViewAtom) {
+                if (!relationViews.contains(relationViewAtom.getTarget())) {
                         throw new IllegalArgumentException(
                                         "%s %s contains reference to a view %s that is not in the model.".formatted(
-                                                        DNFPredicate.class.getSimpleName(), dnfPredicate.getUniqueName(),
-                                                        relationAtom.view().getName()));
+                                                        DNF.class.getSimpleName(), dnfPredicate.getUniqueName(),
+                                                        relationViewAtom.getTarget().getName()));
                 }
         }
 
-        private void validatePredicateAtom(Set<DNFPredicate> predicates, DNFPredicate dnfPredicate,
-                                                                           PredicateAtom predicateAtom) {
-                if (!predicates.contains(predicateAtom.getReferred())) {
+        private void validatePredicateAtom(Set<DNF> predicates, DNF dnfPredicate,
+                                                                           DNFCallAtom queryCallAtom) {
+                if (!predicates.contains(queryCallAtom.getTarget())) {
                         throw new IllegalArgumentException(
                                         "%s %s contains reference to a predicate %s that is not in the model.".formatted(
-                                                        DNFPredicate.class.getSimpleName(), dnfPredicate.getUniqueName(),
-                                                        predicateAtom.getReferred().getName()));
+                                                        DNF.class.getSimpleName(), dnfPredicate.getUniqueName(),
+                                                        queryCallAtom.getTarget().getName()));
                 }
         }
 
-        private Map<DNFPredicate, GenericQuerySpecification<RawPatternMatcher>> initPredicates(Set<DNFPredicate> predicates) {
-                Map<DNFPredicate, GenericQuerySpecification<RawPatternMatcher>> result = new HashMap<>();
+        private Map<DNF, GenericQuerySpecification<RawPatternMatcher>> initPredicates(Set<DNF> predicates) {
+                Map<DNF, GenericQuerySpecification<RawPatternMatcher>> result = new HashMap<>();
                 var dnf2PQuery = new DNF2PQuery();
-                for (DNFPredicate dnfPredicate : predicates) {
+                for (DNF dnfPredicate : predicates) {
                         GenericQuerySpecification<RawPatternMatcher> query = dnf2PQuery.translate(dnfPredicate).build();
                         result.put(dnfPredicate, query);
                 }
@@ -105,7 +109,7 @@ public class ViatraQueryableModelStore implements QueryableModelStore {
         }
 
         @Override
-        public Set<DNFPredicate> getPredicates() {
+        public Set<DNF> getPredicates() {
                 return predicates.keySet();
         }
 
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraQueryableModel.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraQueryableModel.java
index 1c15e8f9..cd9fb180 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraQueryableModel.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraQueryableModel.java
@@ -12,7 +12,7 @@ import tools.refinery.store.model.representation.DataRepresentation;
 import tools.refinery.store.model.representation.Relation;
 import tools.refinery.store.query.QueryableModel;
 import tools.refinery.store.query.QueryableModelStore;
-import tools.refinery.store.query.building.DNFPredicate;
+import tools.refinery.store.query.DNF;
 import tools.refinery.store.tuple.Tuple;
 import tools.refinery.store.tuple.TupleLike;
 
@@ -27,16 +27,16 @@ public class ViatraQueryableModel implements QueryableModel {
 
         protected final Model model;
 
-        protected final Map<DNFPredicate, GenericQuerySpecification<RawPatternMatcher>> predicates2PQuery;
+        protected final Map<DNF, GenericQuerySpecification<RawPatternMatcher>> predicates2PQuery;
 
         protected RelationalScope scope;
 
         protected AdvancedViatraQueryEngine engine;
 
-        protected Map<DNFPredicate, RawPatternMatcher> predicate2Matcher;
+        protected Map<DNF, RawPatternMatcher> predicate2Matcher;
 
         public ViatraQueryableModel(QueryableModelStore store, Model model,
-                                                                Map<DNFPredicate, GenericQuerySpecification<RawPatternMatcher>> predicates2PQuery) {
+                                                                Map<DNF, GenericQuerySpecification<RawPatternMatcher>> predicates2PQuery) {
                 this.store = store;
                 this.model = model;
                 this.predicates2PQuery = predicates2PQuery;
@@ -49,15 +49,15 @@ public class ViatraQueryableModel implements QueryableModel {
                 this.predicate2Matcher = initMatchers(this.engine, this.predicates2PQuery);
         }
 
-        private Map<DNFPredicate, RawPatternMatcher> initMatchers(
+        private Map<DNF, RawPatternMatcher> initMatchers(
                         AdvancedViatraQueryEngine engine,
-                        Map<DNFPredicate, GenericQuerySpecification<RawPatternMatcher>> predicates2pQuery) {
+                        Map<DNF, GenericQuerySpecification<RawPatternMatcher>> predicates2pQuery) {
                 // 1. prepare group
                 IQueryGroup queryGroup = GenericQueryGroup.of(Set.copyOf(predicates2pQuery.values()));
                 engine.prepareGroup(queryGroup, null);
 
                 // 2. then get all matchers
-                Map<DNFPredicate, RawPatternMatcher> result = new HashMap<>();
+                Map<DNF, RawPatternMatcher> result = new HashMap<>();
                 for (var entry : predicates2pQuery.entrySet()) {
                         var matcher = engine.getMatcher(entry.getValue());
                         result.put(entry.getKey(), matcher);
@@ -71,7 +71,7 @@ public class ViatraQueryableModel implements QueryableModel {
         }
 
         @Override
-        public Set<DNFPredicate> getPredicates() {
+        public Set<DNF> getPredicates() {
                 return store.getPredicates();
         }
 
@@ -117,7 +117,7 @@ public class ViatraQueryableModel implements QueryableModel {
                 return model.getSize(representation);
         }
 
-        protected RawPatternMatcher getMatcher(DNFPredicate predicate) {
+        protected RawPatternMatcher getMatcher(DNF predicate) {
                 var result = this.predicate2Matcher.get(predicate);
                 if (result == null) {
                         throw new IllegalArgumentException("Model does not contain predicate %s".formatted(predicate.getName()));
@@ -125,8 +125,8 @@ public class ViatraQueryableModel implements QueryableModel {
                         return result;
         }
 
-        protected void validateParameters(DNFPredicate predicate, Tuple parameters) {
-                int predicateArity = predicate.getVariables().size();
+        protected void validateParameters(DNF predicate, Tuple parameters) {
+                int predicateArity = predicate.getParameters().size();
                 int parameterArity = parameters.getSize();
                 if (parameterArity != predicateArity) {
                         throw new IllegalArgumentException(
@@ -136,45 +136,45 @@ public class ViatraQueryableModel implements QueryableModel {
         }
 
         @Override
-        public boolean hasResult(DNFPredicate predicate) {
+        public boolean hasResult(DNF predicate) {
                 return getMatcher(predicate).hasResult();
         }
 
         @Override
-        public boolean hasResult(DNFPredicate predicate, Tuple parameters) {
+        public boolean hasResult(DNF predicate, Tuple parameters) {
                 validateParameters(predicate, parameters);
                 return getMatcher(predicate).hasResult(parameters);
         }
 
         @Override
-        public Optional<TupleLike> oneResult(DNFPredicate predicate) {
+        public Optional<TupleLike> oneResult(DNF predicate) {
                 return getMatcher(predicate).oneResult();
         }
 
         @Override
-        public Optional<TupleLike> oneResult(DNFPredicate predicate, Tuple parameters) {
+        public Optional<TupleLike> oneResult(DNF predicate, Tuple parameters) {
                 validateParameters(predicate, parameters);
                 return getMatcher(predicate).oneResult(parameters);
         }
 
         @Override
-        public Stream<TupleLike> allResults(DNFPredicate predicate) {
+        public Stream<TupleLike> allResults(DNF predicate) {
                 return getMatcher(predicate).allResults();
         }
 
         @Override
-        public Stream<TupleLike> allResults(DNFPredicate predicate, Tuple parameters) {
+        public Stream<TupleLike> allResults(DNF predicate, Tuple parameters) {
                 validateParameters(predicate, parameters);
                 return getMatcher(predicate).allResults(parameters);
         }
 
         @Override
-        public int countResults(DNFPredicate predicate) {
+        public int countResults(DNF predicate) {
                 return getMatcher(predicate).countResults();
         }
 
         @Override
-        public int countResults(DNFPredicate predicate, Tuple parameters) {
+        public int countResults(DNF predicate, Tuple parameters) {
                 validateParameters(predicate, parameters);
                 return getMatcher(predicate).countResults(parameters);
 
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/pquery/DNF2PQuery.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/pquery/DNF2PQuery.java
index b6468e76..e3c586a0 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/pquery/DNF2PQuery.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/pquery/DNF2PQuery.java
@@ -11,130 +11,124 @@ import org.eclipse.viatra.query.runtime.matchers.psystem.basicenumerables.Positi
 import org.eclipse.viatra.query.runtime.matchers.psystem.basicenumerables.TypeConstraint;
 import org.eclipse.viatra.query.runtime.matchers.psystem.queries.PParameter;
 import org.eclipse.viatra.query.runtime.matchers.tuple.Tuples;
-import tools.refinery.store.query.building.*;
+import tools.refinery.store.query.*;
+import tools.refinery.store.query.atom.DNFAtom;
+import tools.refinery.store.query.atom.EquivalenceAtom;
+import tools.refinery.store.query.atom.DNFCallAtom;
+import tools.refinery.store.query.atom.RelationViewAtom;
 import tools.refinery.store.query.view.RelationView;
 
 import java.util.*;
 import java.util.stream.Collectors;
 
 public class DNF2PQuery {
-        private final Set<DNFPredicate> translating = new LinkedHashSet<>();
+        private final Set<DNF> translating = new LinkedHashSet<>();
 
-        private final Map<DNFPredicate, SimplePQuery> dnf2PQueryMap = new HashMap<>();
+        private final Map<DNF, SimplePQuery> dnf2PQueryMap = new HashMap<>();
 
         private final Map<RelationView<?>, RelationViewWrapper> view2WrapperMap = new HashMap<>();
 
-        public SimplePQuery translate(DNFPredicate predicate) {
-                if (translating.contains(predicate)) {
-                        var path = translating.stream().map(DNFPredicate::getName).collect(Collectors.joining(" -> "));
+        public SimplePQuery translate(DNF dnfQuery) {
+                if (translating.contains(dnfQuery)) {
+                        var path = translating.stream().map(DNF::getName).collect(Collectors.joining(" -> "));
                         throw new IllegalStateException("Circular reference %s -> %s detected".formatted(path,
-                                        predicate.getName()));
+                                        dnfQuery.getName()));
                 }
                 // We can't use computeIfAbsent here, because translating referenced queries calls this method in a reentrant
                 // way, which would cause a ConcurrentModificationException with computeIfAbsent.
-                var pQuery = dnf2PQueryMap.get(predicate);
+                var pQuery = dnf2PQueryMap.get(dnfQuery);
                 if (pQuery == null) {
-                        translating.add(predicate);
+                        translating.add(dnfQuery);
                         try {
-                                pQuery = doTranslate(predicate);
-                                dnf2PQueryMap.put(predicate, pQuery);
+                                pQuery = doTranslate(dnfQuery);
+                                dnf2PQueryMap.put(dnfQuery, pQuery);
                         } finally {
-                                translating.remove(predicate);
+                                translating.remove(dnfQuery);
                         }
                 }
                 return pQuery;
         }
 
-        private SimplePQuery doTranslate(DNFPredicate predicate) {
-                var query = new SimplePQuery(predicate.getName());
+        private SimplePQuery doTranslate(DNF dnfQuery) {
+                var pQuery = new SimplePQuery(dnfQuery.getUniqueName());
 
                 Map<Variable, PParameter> parameters = new HashMap<>();
-                for (Variable variable : predicate.getVariables()) {
-                        parameters.put(variable, new PParameter(variable.getName()));
+                for (Variable variable : dnfQuery.getParameters()) {
+                        parameters.put(variable, new PParameter(variable.getUniqueName()));
                 }
 
                 List<PParameter> parameterList = new ArrayList<>();
-                for (var param : predicate.getVariables()) {
+                for (var param : dnfQuery.getParameters()) {
                         parameterList.add(parameters.get(param));
                 }
-                query.setParameters(parameterList);
+                pQuery.setParameters(parameterList);
 
-                for (DNFAnd clause : predicate.getClauses()) {
-                        PBody body = new PBody(query);
+                for (DNFAnd clause : dnfQuery.getClauses()) {
+                        PBody body = new PBody(pQuery);
                         List<ExportedParameter> symbolicParameters = new ArrayList<>();
-                        for (var param : predicate.getVariables()) {
-                                PVariable pVar = body.getOrCreateVariableByName(param.getName());
+                        for (var param : dnfQuery.getParameters()) {
+                                PVariable pVar = body.getOrCreateVariableByName(param.getUniqueName());
                                 symbolicParameters.add(new ExportedParameter(body, pVar, parameters.get(param)));
                         }
                         body.setSymbolicParameters(symbolicParameters);
-                        query.addBody(body);
-                        for (DNFAtom constraint : clause.getConstraints()) {
+                        pQuery.addBody(body);
+                        for (DNFAtom constraint : clause.constraints()) {
                                 translateDNFAtom(constraint, body);
                         }
                 }
 
-                return query;
+                return pQuery;
         }
 
         private void translateDNFAtom(DNFAtom constraint, PBody body) {
-                if (constraint instanceof EquivalenceAtom equivalence) {
-                        translateEquivalenceAtom(equivalence, body);
-                }
-                if (constraint instanceof RelationAtom relation) {
-                        translateRelationAtom(relation, body);
-                }
-                if (constraint instanceof PredicateAtom predicate) {
-                        translatePredicateAtom(predicate, body);
+                if (constraint instanceof EquivalenceAtom equivalenceAtom) {
+                        translateEquivalenceAtom(equivalenceAtom, body);
+                } else if (constraint instanceof RelationViewAtom relationViewAtom) {
+                        translateRelationViewAtom(relationViewAtom, body);
+                } else if (constraint instanceof DNFCallAtom dnfCallAtom) {
+                        translateDNFCallAtom(dnfCallAtom, body);
+                } else {
+                        throw new IllegalArgumentException("Unknown constraint: " + constraint.toString());
                 }
         }
 
         private void translateEquivalenceAtom(EquivalenceAtom equivalence, PBody body) {
-                PVariable varSource = body.getOrCreateVariableByName(equivalence.getLeft().getName());
-                PVariable varTarget = body.getOrCreateVariableByName(equivalence.getRight().getName());
-                if (equivalence.isPositive())
+                PVariable varSource = body.getOrCreateVariableByName(equivalence.left().getUniqueName());
+                PVariable varTarget = body.getOrCreateVariableByName(equivalence.right().getUniqueName());
+                if (equivalence.positive()) {
                         new Equality(body, varSource, varTarget);
-                else
+                } else {
                         new Inequality(body, varSource, varTarget);
+                }
         }
 
-        private void translateRelationAtom(RelationAtom relation, PBody body) {
-                if (relation.substitution().size() != relation.view().getArity()) {
-                        throw new IllegalArgumentException("Arity (%d) does not match parameter numbers (%d)".formatted(
-                                        relation.view().getArity(), relation.substitution().size()));
-                }
-                Object[] variables = new Object[relation.substitution().size()];
-                for (int i = 0; i < relation.substitution().size(); i++) {
-                        variables[i] = body.getOrCreateVariableByName(relation.substitution().get(i).getName());
+        private void translateRelationViewAtom(RelationViewAtom relationViewAtom, PBody body) {
+                int arity = relationViewAtom.getSubstitution().size();
+                Object[] variables = new Object[arity];
+                for (int i = 0; i < arity; i++) {
+                        var variable = relationViewAtom.getSubstitution().get(i);
+                        variables[i] = body.getOrCreateVariableByName(variable.getUniqueName());
                 }
-                new TypeConstraint(body, Tuples.flatTupleOf(variables), wrapView(relation.view()));
+                new TypeConstraint(body, Tuples.flatTupleOf(variables), wrapView(relationViewAtom.getTarget()));
         }
 
         private RelationViewWrapper wrapView(RelationView<?> relationView) {
                 return view2WrapperMap.computeIfAbsent(relationView, RelationViewWrapper::new);
         }
 
-        private void translatePredicateAtom(PredicateAtom predicate, PBody body) {
-                Object[] variables = new Object[predicate.getSubstitution().size()];
-                for (int i = 0; i < predicate.getSubstitution().size(); i++) {
-                        variables[i] = body.getOrCreateVariableByName(predicate.getSubstitution().get(i).getName());
+        private void translateDNFCallAtom(DNFCallAtom queryCallAtom, PBody body) {
+                int arity = queryCallAtom.getSubstitution().size();
+                Object[] variables = new Object[arity];
+                for (int i = 0; i < arity; i++) {
+                        var variable = queryCallAtom.getSubstitution().get(i);
+                        variables[i] = body.getOrCreateVariableByName(variable.getUniqueName());
                 }
                 var variablesTuple = Tuples.flatTupleOf(variables);
-                var translatedReferred = translate(predicate.getReferred());
-                if (predicate.isPositive()) {
-                        if (predicate.isTransitive()) {
-                                if (predicate.getSubstitution().size() != 2) {
-                                        throw new IllegalArgumentException("Transitive Predicate Atoms must be binary.");
-                                }
-                                new BinaryTransitiveClosure(body, variablesTuple, translatedReferred);
-                        } else {
-                                new PositivePatternCall(body, variablesTuple, translatedReferred);
-                        }
-                } else {
-                        if (predicate.isTransitive()) {
-                                throw new IllegalArgumentException("Transitive Predicate Atoms cannot be negative.");
-                        } else {
-                                new NegativePatternCall(body, variablesTuple, translatedReferred);
-                        }
+                var translatedReferred = translate(queryCallAtom.getTarget());
+                switch (queryCallAtom.getKind()) {
+                case POSITIVE -> new PositivePatternCall(body, variablesTuple, translatedReferred);
+                case TRANSITIVE -> new BinaryTransitiveClosure(body, variablesTuple, translatedReferred);
+                case NEGATIVE -> new NegativePatternCall(body, variablesTuple, translatedReferred);
                 }
         }
 }
diff --git a/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/tests/QueryTest.java b/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/tests/QueryTest.java
index d6213b02..8984cb2c 100644
--- a/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/tests/QueryTest.java
+++ b/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/tests/QueryTest.java
@@ -3,9 +3,11 @@ package tools.refinery.store.query.viatra.tests;
 import org.junit.jupiter.api.Test;
 import tools.refinery.store.model.representation.Relation;
 import tools.refinery.store.model.representation.TruthValue;
-import tools.refinery.store.query.QueryableModel;
-import tools.refinery.store.query.QueryableModelStore;
-import tools.refinery.store.query.building.*;
+import tools.refinery.store.query.*;
+import tools.refinery.store.query.atom.CallKind;
+import tools.refinery.store.query.atom.EquivalenceAtom;
+import tools.refinery.store.query.atom.DNFCallAtom;
+import tools.refinery.store.query.atom.RelationViewAtom;
 import tools.refinery.store.query.viatra.ViatraQueryableModelStore;
 import tools.refinery.store.query.view.FilteredRelationView;
 import tools.refinery.store.query.view.KeyOnlyRelationView;
@@ -25,10 +27,11 @@ class QueryTest {
                 Relation<Boolean> asset = new Relation<>("Asset", 1, false);
                 RelationView<Boolean> personView = new KeyOnlyRelationView(person);
 
-                List<Variable> parameters = List.of(new Variable("p1"));
-                RelationAtom personRelationAtom = new RelationAtom(personView, parameters);
-                DNFAnd clause = new DNFAnd(Collections.emptySet(), List.of(personRelationAtom));
-                DNFPredicate predicate = new DNFPredicate("TypeConstraint", parameters, List.of(clause));
+                var p1 = new Variable("p1");
+                DNF predicate = DNF.builder("TypeConstraint")
+                                .parameters(p1)
+                                .clause(new RelationViewAtom(personView, p1))
+                                .build();
 
                 QueryableModelStore store = new ViatraQueryableModelStore(Set.of(person, asset), Set.of(personView),
                                 Set.of(predicate));
@@ -54,14 +57,14 @@ class QueryTest {
 
                 Variable p1 = new Variable("p1");
                 Variable p2 = new Variable("p2");
-                List<Variable> parameters = Arrays.asList(p1, p2);
-
-                RelationAtom personRelationAtom1 = new RelationAtom(personView, List.of(p1));
-                RelationAtom personRelationAtom2 = new RelationAtom(personView, List.of(p2));
-                RelationAtom friendRelationAtom = new RelationAtom(friendMustView, Arrays.asList(p1, p2));
-                DNFAnd clause = new DNFAnd(Collections.emptySet(),
-                                Arrays.asList(personRelationAtom1, personRelationAtom2, friendRelationAtom));
-                DNFPredicate predicate = new DNFPredicate("RelationConstraint", parameters, List.of(clause));
+                DNF predicate = DNF.builder("RelationConstraint")
+                                .parameters(p1, p2)
+                                .clause(
+                                                new RelationViewAtom(personView, p1),
+                                                new RelationViewAtom(personView, p1),
+                                                new RelationViewAtom(friendMustView, p1, p2)
+                                )
+                                .build();
 
                 QueryableModelStore store = new ViatraQueryableModelStore(Set.of(person, friend),
                                 Set.of(personView, friendMustView), Set.of(predicate));
@@ -93,15 +96,15 @@ class QueryTest {
 
                 Variable p1 = new Variable("p1");
                 Variable p2 = new Variable("p2");
-                List<Variable> parameters = Arrays.asList(p1, p2);
-
-                RelationAtom personRelationAtom1 = new RelationAtom(personView, List.of(p1));
-                RelationAtom personRelationAtom2 = new RelationAtom(personView, List.of(p2));
-                RelationAtom friendRelationAtom1 = new RelationAtom(friendMustView, Arrays.asList(p1, p2));
-                RelationAtom friendRelationAtom2 = new RelationAtom(friendMustView, Arrays.asList(p2, p1));
-                DNFAnd clause = new DNFAnd(Collections.emptySet(),
-                                Arrays.asList(personRelationAtom1, personRelationAtom2, friendRelationAtom1, friendRelationAtom2));
-                DNFPredicate predicate = new DNFPredicate("RelationConstraint", parameters, List.of(clause));
+                DNF predicate = DNF.builder("RelationConstraint")
+                                .parameters(p1, p2)
+                                .clause(
+                                                new RelationViewAtom(personView, p1),
+                                                new RelationViewAtom(personView, p2),
+                                                new RelationViewAtom(friendMustView, p1, p2),
+                                                new RelationViewAtom(friendMustView, p2, p1)
+                                )
+                                .build();
 
                 QueryableModelStore store = new ViatraQueryableModelStore(Set.of(person, friend),
                                 Set.of(personView, friendMustView), Set.of(predicate));
@@ -141,14 +144,14 @@ class QueryTest {
 
                 Variable p1 = new Variable("p1");
                 Variable p2 = new Variable("p2");
-                List<Variable> parameters = List.of(p1);
-
-                RelationAtom personRelationAtom1 = new RelationAtom(personView, List.of(p1));
-                RelationAtom personRelationAtom2 = new RelationAtom(personView, List.of(p2));
-                RelationAtom friendRelationAtom = new RelationAtom(friendMustView, Arrays.asList(p1, p2));
-                DNFAnd clause = new DNFAnd(Set.of(p2),
-                                Arrays.asList(personRelationAtom1, personRelationAtom2, friendRelationAtom));
-                DNFPredicate predicate = new DNFPredicate("RelationConstraint", parameters, List.of(clause));
+                DNF predicate = DNF.builder("RelationConstraint")
+                                .parameters(p1)
+                                .clause(
+                                                new RelationViewAtom(personView, p1),
+                                                new RelationViewAtom(personView, p2),
+                                                new RelationViewAtom(friendMustView, p1, p2)
+                                )
+                                .build();
 
                 QueryableModelStore store = new ViatraQueryableModelStore(Set.of(person, friend),
                                 Set.of(personView, friendMustView), Set.of(predicate));
@@ -182,24 +185,19 @@ class QueryTest {
 
                 Variable p1 = new Variable("p1");
                 Variable p2 = new Variable("p2");
-                List<Variable> parameters = Arrays.asList(p1, p2);
-
-                // Person-Person friendship
-                RelationAtom personRelationAtom1 = new RelationAtom(personView, List.of(p1));
-                RelationAtom personRelationAtom2 = new RelationAtom(personView, List.of(p2));
-                RelationAtom friendRelationAtom1 = new RelationAtom(friendMustView, Arrays.asList(p1, p2));
-                DNFAnd clause1 = new DNFAnd(Collections.emptySet(),
-                                Arrays.asList(personRelationAtom1, personRelationAtom2, friendRelationAtom1));
-
-                // Animal-Animal friendship
-                RelationAtom animalRelationAtom1 = new RelationAtom(animalView, List.of(p1));
-                RelationAtom animalRelationAtom2 = new RelationAtom(animalView, List.of(p2));
-                RelationAtom friendRelationAtom2 = new RelationAtom(friendMustView, Arrays.asList(p1, p2));
-                DNFAnd clause2 = new DNFAnd(Collections.emptySet(),
-                                Arrays.asList(animalRelationAtom1, animalRelationAtom2, friendRelationAtom2));
-
-                // No friendship between species
-                DNFPredicate predicate = new DNFPredicate("Or", parameters, Arrays.asList(clause1, clause2));
+                DNF predicate = DNF.builder("Or")
+                                .parameters(p1, p2)
+                                .clause(
+                                                new RelationViewAtom(personView, p1),
+                                                new RelationViewAtom(personView, p2),
+                                                new RelationViewAtom(friendMustView, p1, p2)
+                                )
+                                .clause(
+                                                new RelationViewAtom(animalView, p1),
+                                                new RelationViewAtom(animalView, p2),
+                                                new RelationViewAtom(friendMustView, p1, p2)
+                                )
+                                .build();
 
                 QueryableModelStore store = new ViatraQueryableModelStore(Set.of(person, animal, friend),
                                 Set.of(personView, animalView, friendMustView), Set.of(predicate));
@@ -226,14 +224,14 @@ class QueryTest {
 
                 Variable p1 = new Variable("p1");
                 Variable p2 = new Variable("p2");
-                List<Variable> parameters = Arrays.asList(p1, p2);
-
-                RelationAtom personRelationAtom1 = new RelationAtom(personView, List.of(p1));
-                RelationAtom personRelationAtom2 = new RelationAtom(personView, List.of(p2));
-                EquivalenceAtom equivalenceAtom = new EquivalenceAtom(true, p1, p2);
-                DNFAnd clause = new DNFAnd(Collections.emptySet(),
-                                Arrays.asList(personRelationAtom1, personRelationAtom2, equivalenceAtom));
-                DNFPredicate predicate = new DNFPredicate("Equality", parameters, List.of(clause));
+                DNF predicate = DNF.builder("Equality")
+                                .parameters(p1, p2)
+                                .clause(
+                                                new RelationViewAtom(personView, p1),
+                                                new RelationViewAtom(personView, p2),
+                                                new EquivalenceAtom(p1, p2)
+                                )
+                                .build();
 
                 QueryableModelStore store = new ViatraQueryableModelStore(Set.of(person), Set.of(personView), Set.of(predicate));
                 QueryableModel model = store.createModel();
@@ -258,16 +256,16 @@ class QueryTest {
                 Variable p1 = new Variable("p1");
                 Variable p2 = new Variable("p2");
                 Variable p3 = new Variable("p3");
-                List<Variable> parameters = Arrays.asList(p1, p2, p3);
-
-                RelationAtom personRelationAtom1 = new RelationAtom(personView, List.of(p1));
-                RelationAtom personRelationAtom2 = new RelationAtom(personView, List.of(p2));
-                RelationAtom friendRelationAtom1 = new RelationAtom(friendMustView, Arrays.asList(p1, p3));
-                RelationAtom friendRelationAtom2 = new RelationAtom(friendMustView, Arrays.asList(p2, p3));
-                EquivalenceAtom inequivalenceAtom = new EquivalenceAtom(false, p1, p2);
-                DNFAnd clause = new DNFAnd(Collections.emptySet(), Arrays.asList(personRelationAtom1, personRelationAtom2,
-                                friendRelationAtom1, friendRelationAtom2, inequivalenceAtom));
-                DNFPredicate predicate = new DNFPredicate("Inequality", parameters, List.of(clause));
+                DNF predicate = DNF.builder("Inequality")
+                                .parameters(p1, p2, p3)
+                                .clause(
+                                                new RelationViewAtom(personView, p1),
+                                                new RelationViewAtom(personView, p2),
+                                                new RelationViewAtom(friendMustView, p1, p3),
+                                                new RelationViewAtom(friendMustView, p2, p3),
+                                                new EquivalenceAtom(false, p1, p2)
+                                )
+                                .build();
 
                 QueryableModelStore store = new ViatraQueryableModelStore(Set.of(person, friend),
                                 Set.of(personView, friendMustView), Set.of(predicate));
@@ -294,24 +292,25 @@ class QueryTest {
 
                 Variable p1 = new Variable("p1");
                 Variable p2 = new Variable("p2");
-                List<Variable> parameters = Arrays.asList(p1, p2);
-
-                RelationAtom personRelationAtom1 = new RelationAtom(personView, List.of(p1));
-                RelationAtom personRelationAtom2 = new RelationAtom(personView, List.of(p2));
-                RelationAtom friendRelationAtom = new RelationAtom(friendMustView, Arrays.asList(p1, p2));
-                DNFAnd clause = new DNFAnd(Collections.emptySet(),
-                                Arrays.asList(personRelationAtom1, personRelationAtom2, friendRelationAtom));
-                DNFPredicate friendPredicate = new DNFPredicate("RelationConstraint", parameters, List.of(clause));
+                DNF friendPredicate = DNF.builder("RelationConstraint")
+                                .parameters(p1, p2)
+                                .clause(
+                                                new RelationViewAtom(personView, p1),
+                                                new RelationViewAtom(personView, p2),
+                                                new RelationViewAtom(friendMustView, p1, p2)
+                                )
+                                .build();
 
                 Variable p3 = new Variable("p3");
                 Variable p4 = new Variable("p4");
-                List<Variable> substitution = Arrays.asList(p3, p4);
-                RelationAtom personRelationAtom3 = new RelationAtom(personView, List.of(p3));
-                RelationAtom personRelationAtom4 = new RelationAtom(personView, List.of(p4));
-                PredicateAtom friendPredicateAtom = new PredicateAtom(true, false, friendPredicate, substitution);
-                DNFAnd patternCallClause = new DNFAnd(Collections.emptySet(),
-                                Arrays.asList(personRelationAtom3, personRelationAtom4, friendPredicateAtom));
-                DNFPredicate predicate = new DNFPredicate("PatternCall", substitution, List.of(patternCallClause));
+                DNF predicate = DNF.builder("PositivePatternCall")
+                                .parameters(p3, p4)
+                                .clause(
+                                                new RelationViewAtom(personView, p3),
+                                                new RelationViewAtom(personView, p4),
+                                                new DNFCallAtom(friendPredicate, p3, p4)
+                                )
+                                .build();
 
                 QueryableModelStore store = new ViatraQueryableModelStore(Set.of(person, friend),
                                 Set.of(personView, friendMustView), Set.of(friendPredicate, predicate));
@@ -339,25 +338,25 @@ class QueryTest {
 
                 Variable p1 = new Variable("p1");
                 Variable p2 = new Variable("p2");
-                List<Variable> parameters = Arrays.asList(p1, p2);
-
-                RelationAtom personRelationAtom1 = new RelationAtom(personView, List.of(p1));
-                RelationAtom personRelationAtom2 = new RelationAtom(personView, List.of(p2));
-                RelationAtom friendRelationAtom = new RelationAtom(friendMustView, Arrays.asList(p1, p2));
-                DNFAnd clause = new DNFAnd(Collections.emptySet(),
-                                Arrays.asList(personRelationAtom1, personRelationAtom2, friendRelationAtom));
-                DNFPredicate friendPredicate = new DNFPredicate("RelationConstraint", parameters, List.of(clause));
+                DNF friendPredicate = DNF.builder("RelationConstraint")
+                                .parameters(p1, p2)
+                                .clause(
+                                                new RelationViewAtom(personView, p1),
+                                                new RelationViewAtom(personView, p2),
+                                                new RelationViewAtom(friendMustView, p1, p2)
+                                )
+                                .build();
 
                 Variable p3 = new Variable("p3");
                 Variable p4 = new Variable("p4");
-                List<Variable> substitution = Arrays.asList(p3, p4);
-                RelationAtom personRelationAtom3 = new RelationAtom(personView, List.of(p3));
-                RelationAtom personRelationAtom4 = new RelationAtom(personView, List.of(p4));
-                PredicateAtom friendPredicateAtom = new PredicateAtom(false, false, friendPredicate, substitution);
-                DNFAnd negativePatternCallClause = new DNFAnd(Collections.emptySet(),
-                                Arrays.asList(personRelationAtom3, personRelationAtom4, friendPredicateAtom));
-                DNFPredicate predicate = new DNFPredicate("NegativePatternCall", substitution,
-                                List.of(negativePatternCallClause));
+                DNF predicate = DNF.builder("NegativePatternCall")
+                                .parameters(p3, p4)
+                                .clause(
+                                                new RelationViewAtom(personView, p3),
+                                                new RelationViewAtom(personView, p4),
+                                                new DNFCallAtom(CallKind.NEGATIVE, friendPredicate, p3, p4)
+                                )
+                                .build();
 
                 QueryableModelStore store = new ViatraQueryableModelStore(Set.of(person, friend),
                                 Set.of(personView, friendMustView), Set.of(friendPredicate, predicate));
@@ -384,25 +383,25 @@ class QueryTest {
 
                 Variable p1 = new Variable("p1");
                 Variable p2 = new Variable("p2");
-                List<Variable> parameters = Arrays.asList(p1, p2);
-
-                RelationAtom personRelationAtom1 = new RelationAtom(personView, List.of(p1));
-                RelationAtom personRelationAtom2 = new RelationAtom(personView, List.of(p2));
-                RelationAtom friendRelationAtom = new RelationAtom(friendMustView, Arrays.asList(p1, p2));
-                DNFAnd clause = new DNFAnd(Collections.emptySet(),
-                                Arrays.asList(personRelationAtom1, personRelationAtom2, friendRelationAtom));
-                DNFPredicate friendPredicate = new DNFPredicate("RelationConstraint", parameters, List.of(clause));
+                DNF friendPredicate = DNF.builder("RelationConstraint")
+                                .parameters(p1, p2)
+                                .clause(
+                                                new RelationViewAtom(personView, p1),
+                                                new RelationViewAtom(personView, p2),
+                                                new RelationViewAtom(friendMustView, p1, p2)
+                                )
+                                .build();
 
                 Variable p3 = new Variable("p3");
                 Variable p4 = new Variable("p4");
-                List<Variable> substitution = Arrays.asList(p3, p4);
-                RelationAtom personRelationAtom3 = new RelationAtom(personView, List.of(p3));
-                RelationAtom personRelationAtom4 = new RelationAtom(personView, List.of(p4));
-                PredicateAtom friendPredicateAtom = new PredicateAtom(true, true, friendPredicate, substitution);
-                DNFAnd patternCallClause = new DNFAnd(Collections.emptySet(),
-                                Arrays.asList(personRelationAtom3, personRelationAtom4, friendPredicateAtom));
-                DNFPredicate predicate = new DNFPredicate("TransitivePatternCall", substitution,
-                                List.of(patternCallClause));
+                DNF predicate = DNF.builder("TransitivePatternCall")
+                                .parameters(p3, p4)
+                                .clause(
+                                                new RelationViewAtom(personView, p3),
+                                                new RelationViewAtom(personView, p4),
+                                                new DNFCallAtom(CallKind.TRANSITIVE, friendPredicate, p3, p4)
+                                )
+                                .build();
 
                 QueryableModelStore store = new ViatraQueryableModelStore(Set.of(person, friend),
                                 Set.of(personView, friendMustView), Set.of(friendPredicate, predicate));
diff --git a/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/tests/QueryTransactionTest.java b/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/tests/QueryTransactionTest.java
index f57bca2b..a555f024 100644
--- a/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/tests/QueryTransactionTest.java
+++ b/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/tests/QueryTransactionTest.java
@@ -1,20 +1,17 @@
 package tools.refinery.store.query.viatra.tests;
 
 import org.junit.jupiter.api.Test;
-import tools.refinery.store.tuple.Tuple;
 import tools.refinery.store.model.representation.Relation;
+import tools.refinery.store.query.DNF;
 import tools.refinery.store.query.QueryableModel;
 import tools.refinery.store.query.QueryableModelStore;
-import tools.refinery.store.query.building.DNFAnd;
-import tools.refinery.store.query.building.DNFPredicate;
-import tools.refinery.store.query.building.RelationAtom;
-import tools.refinery.store.query.building.Variable;
+import tools.refinery.store.query.Variable;
+import tools.refinery.store.query.atom.RelationViewAtom;
 import tools.refinery.store.query.viatra.ViatraQueryableModelStore;
 import tools.refinery.store.query.view.KeyOnlyRelationView;
 import tools.refinery.store.query.view.RelationView;
+import tools.refinery.store.tuple.Tuple;
 
-import java.util.Collections;
-import java.util.List;
 import java.util.Set;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -26,10 +23,11 @@ class QueryTransactionTest {
                 Relation<Boolean> asset = new Relation<>("Asset", 1, false);
                 RelationView<Boolean> personView = new KeyOnlyRelationView(person);
 
-                List<Variable> parameters = List.of(new Variable("p1"));
-                RelationAtom personRelationAtom = new RelationAtom(personView, parameters);
-                DNFAnd clause = new DNFAnd(Collections.emptySet(), List.of(personRelationAtom));
-                DNFPredicate predicate = new DNFPredicate("TypeConstraint", parameters, List.of(clause));
+                var p1 = new Variable("p1");
+                DNF predicate = DNF.builder("TypeConstraint")
+                                .parameters(p1)
+                                .clause(new RelationViewAtom(personView, p1))
+                                .build();
 
                 QueryableModelStore store = new ViatraQueryableModelStore(Set.of(person, asset), Set.of(personView),
                                 Set.of(predicate));