KEYCLOAK-19749 Optimize DefaultModelCriteria creation

tnorimat · Nov 3, 2021 · 6966e0c · 6966e0c
1 parent b37f2d5
commit 6966e0c
Show file tree

Hide file tree

Showing 7 changed files with 294 additions and 52 deletions.
diff --git a/model/map/src/main/java/org/keycloak/models/map/storage/ModelCriteriaBuilder.java b/model/map/src/main/java/org/keycloak/models/map/storage/ModelCriteriaBuilder.java
@@ -117,7 +117,7 @@ public enum Operator {
      * @return
      * @throws CriterionNotSupported If the operator is not supported for the given field.
      */
-    ModelCriteriaBuilder<M> compare(SearchableModelField<M> modelField, Operator op, Object... value);
+    ModelCriteriaBuilder<M> compare(SearchableModelField<? super M> modelField, Operator op, Object... value);
 
     /**
      * Creates and returns a new instance of {@code ModelCriteriaBuilder} that

diff --git a/model/map/src/main/java/org/keycloak/models/map/storage/chm/ConcurrentHashMapStorage.java b/model/map/src/main/java/org/keycloak/models/map/storage/chm/ConcurrentHashMapStorage.java
@@ -53,7 +53,7 @@ public class ConcurrentHashMapStorage<K, V extends AbstractEntity & UpdatableEnt
 
     protected final ConcurrentMap<K, V> store = new ConcurrentHashMap<>();
 
-    protected final Map<SearchableModelField<M>, UpdatePredicatesFunc<K, V, M>> fieldPredicates;
+    protected final Map<SearchableModelField<? super M>, UpdatePredicatesFunc<K, V, M>> fieldPredicates;
     protected final StringKeyConvertor<K> keyConvertor;
     protected final DeepCloner cloner;
 

diff --git a/model/map/src/main/java/org/keycloak/models/map/storage/chm/MapFieldPredicates.java b/model/map/src/main/java/org/keycloak/models/map/storage/chm/MapFieldPredicates.java
@@ -527,7 +527,7 @@ public static <K, V extends AbstractEntity, M> Comparator<V> getComparator(Strea
     }
 
     @SuppressWarnings("unchecked")
-    public static <K, V extends AbstractEntity, M> Map<SearchableModelField<M>, UpdatePredicatesFunc<K, V, M>> getPredicates(Class<M> clazz) {
+    public static <K, V extends AbstractEntity, M> Map<SearchableModelField<? super M>, UpdatePredicatesFunc<K, V, M>> getPredicates(Class<M> clazz) {
         return PREDICATES.get(clazz);
     }
 }
diff --git a/model/map/src/main/java/org/keycloak/models/map/storage/chm/MapModelCriteriaBuilder.java b/model/map/src/main/java/org/keycloak/models/map/storage/chm/MapModelCriteriaBuilder.java
@@ -43,22 +43,22 @@ public static interface UpdatePredicatesFunc<K, V extends AbstractEntity, M> {
     private static final Predicate<Object> ALWAYS_FALSE = (e) -> false;
     private final Predicate<? super K> keyFilter;
     private final Predicate<? super V> entityFilter;
-    private final Map<SearchableModelField<M>, UpdatePredicatesFunc<K, V, M>> fieldPredicates;
+    private final Map<SearchableModelField<? super M>, UpdatePredicatesFunc<K, V, M>> fieldPredicates;
     private final StringKeyConvertor<K> keyConvertor;
 
-    public MapModelCriteriaBuilder(StringKeyConvertor<K> keyConvertor, Map<SearchableModelField<M>, UpdatePredicatesFunc<K, V, M>> fieldPredicates) {
+    public MapModelCriteriaBuilder(StringKeyConvertor<K> keyConvertor, Map<SearchableModelField<? super M>, UpdatePredicatesFunc<K, V, M>> fieldPredicates) {
         this(keyConvertor, fieldPredicates, ALWAYS_TRUE, ALWAYS_TRUE);
     }
 
-    private MapModelCriteriaBuilder(StringKeyConvertor<K> keyConvertor, Map<SearchableModelField<M>, UpdatePredicatesFunc<K, V, M>> fieldPredicates, Predicate<? super K> indexReadFilter, Predicate<? super V> sequentialReadFilter) {
+    private MapModelCriteriaBuilder(StringKeyConvertor<K> keyConvertor, Map<SearchableModelField<? super M>, UpdatePredicatesFunc<K, V, M>> fieldPredicates, Predicate<? super K> indexReadFilter, Predicate<? super V> sequentialReadFilter) {
         this.keyConvertor = keyConvertor;
         this.fieldPredicates = fieldPredicates;
         this.keyFilter = indexReadFilter;
         this.entityFilter = sequentialReadFilter;
     }
 
     @Override
-    public MapModelCriteriaBuilder<K, V, M> compare(SearchableModelField<M> modelField, Operator op, Object... values) {
+    public MapModelCriteriaBuilder<K, V, M> compare(SearchableModelField<? super M> modelField, Operator op, Object... values) {
         UpdatePredicatesFunc<K, V, M> method = fieldPredicates.get(modelField);
         if (method == null) {
             throw new IllegalArgumentException("Filter not implemented for field " + modelField);

diff --git a/model/map/src/main/java/org/keycloak/models/map/storage/criteria/DefaultModelCriteria.java b/model/map/src/main/java/org/keycloak/models/map/storage/criteria/DefaultModelCriteria.java
@@ -19,62 +19,122 @@
 import org.keycloak.models.map.storage.ModelCriteriaBuilder;
 import org.keycloak.models.map.storage.criteria.ModelCriteriaNode.ExtOperator;
 import org.keycloak.storage.SearchableModelField;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.Stream;
 
 /**
  * Descriptive model criteria implementation which in other words represents a Boolean formula on searchable fields.
  * @author hmlnarik
  */
 public class DefaultModelCriteria<M> implements ModelCriteriaBuilder<M> {
 
-    private final ModelCriteriaNode<M> node;
+    private static final DefaultModelCriteria<?> INSTANCE = new DefaultModelCriteria<>(null);
 
-    public DefaultModelCriteria() {
-        this.node = null;
-    }
+    private final ModelCriteriaNode<M> node;
 
     private DefaultModelCriteria(ModelCriteriaNode<M> node) {
         this.node = node;
     }
 
+    public static <M> DefaultModelCriteria<M> criteria() {
+        return (DefaultModelCriteria<M>) INSTANCE;
+    }
+
     @Override
-    public DefaultModelCriteria<M> compare(SearchableModelField<M> modelField, Operator op, Object... value) {
-        final ModelCriteriaNode<M> targetNode;
+    public DefaultModelCriteria<M> compare(SearchableModelField<? super M> modelField, Operator op, Object... value) {
+        return compare(new ModelCriteriaNode<>(modelField, op, value));
+    }
+
+    private DefaultModelCriteria<M> compare(final ModelCriteriaNode<M> nodeToAdd) {
+        ModelCriteriaNode<M> targetNode;
+
         if (isEmpty()) {
-            targetNode = new ModelCriteriaNode<>(modelField, op, value);
+            targetNode = nodeToAdd;
         } else if (node.getNodeOperator() == ExtOperator.AND) {
             targetNode = node.cloneTree();
-            targetNode.addChild(new ModelCriteriaNode<>(modelField, op, value));
+            targetNode.addChild(nodeToAdd);
         } else {
             targetNode = new ModelCriteriaNode<>(ExtOperator.AND);
             targetNode.addChild(node.cloneTree());
-            targetNode.addChild(new ModelCriteriaNode<>(modelField, op, value));
+            targetNode.addChild(nodeToAdd);
         }
+
         return new DefaultModelCriteria<>(targetNode);
     }
 
     @Override
     public DefaultModelCriteria<M> and(ModelCriteriaBuilder<M>... mcbs) {
+        if (mcbs.length == 1) {
+            ModelCriteriaNode<M> toBeChild = ((DefaultModelCriteria<M>) mcbs[0].unwrap(DefaultModelCriteria.class)).node;
+            if (toBeChild.getNodeOperator() == ExtOperator.AND || toBeChild.getNodeOperator() == ExtOperator.OR) {
+                return ((DefaultModelCriteria<M>) mcbs[0].unwrap(DefaultModelCriteria.class));
+            }
+        }
+
         final ModelCriteriaNode<M> targetNode = new ModelCriteriaNode<>(ExtOperator.AND);
+        AtomicBoolean hasFalseNode = new AtomicBoolean(false);
         for (ModelCriteriaBuilder<M> mcb : mcbs) {
-            targetNode.addChild(((DefaultModelCriteria<M>) mcb.unwrap(DefaultModelCriteria.class)).node);
+            final ModelCriteriaNode<M> nodeToAdd = ((DefaultModelCriteria<M>) mcb.unwrap(DefaultModelCriteria.class)).node;
+            getNodesToAddForAndOr(nodeToAdd, ExtOperator.AND)
+              .filter(ModelCriteriaNode::isNotTrueNode)
+              .peek(n -> { if (n.isFalseNode()) hasFalseNode.lazySet(true); })
+              .map(ModelCriteriaNode::cloneTree)
+              .forEach(targetNode::addChild);
+
+            if (hasFalseNode.get()) {
+                return compare(new ModelCriteriaNode<>(ExtOperator.__FALSE__));
+            }
         }
-        return new DefaultModelCriteria<>(targetNode);
+
+        if (targetNode.getChildren().isEmpty()) {
+            // AND on empty set of formulae is TRUE: It does hold that there all formulae are satisfied
+            return compare(new ModelCriteriaNode<>(ExtOperator.__TRUE__));
+        }
+
+        return compare(targetNode);
     }
 
     @Override
     public DefaultModelCriteria<M> or(ModelCriteriaBuilder<M>... mcbs) {
+        if (mcbs.length == 1) {
+            ModelCriteriaNode<M> toBeChild = ((DefaultModelCriteria<M>) mcbs[0].unwrap(DefaultModelCriteria.class)).node;
+            if (toBeChild.getNodeOperator() == ExtOperator.AND || toBeChild.getNodeOperator() == ExtOperator.OR) {
+                return ((DefaultModelCriteria<M>) mcbs[0].unwrap(DefaultModelCriteria.class));
+            }
+        }
+
         final ModelCriteriaNode<M> targetNode = new ModelCriteriaNode<>(ExtOperator.OR);
+        AtomicBoolean hasTrueNode = new AtomicBoolean(false);
         for (ModelCriteriaBuilder<M> mcb : mcbs) {
-            targetNode.addChild(((DefaultModelCriteria<M>) mcb.unwrap(DefaultModelCriteria.class)).node);
+            final ModelCriteriaNode<M> nodeToAdd = ((DefaultModelCriteria<M>) mcb.unwrap(DefaultModelCriteria.class)).node;
+            getNodesToAddForAndOr(nodeToAdd, ExtOperator.OR)
+              .filter(ModelCriteriaNode::isNotFalseNode)
+              .peek(n -> { if (n.isTrueNode()) hasTrueNode.lazySet(true); })
+              .map(ModelCriteriaNode::cloneTree)
+              .forEach(targetNode::addChild);
+
+            if (hasTrueNode.get()) {
+                return compare(new ModelCriteriaNode<>(ExtOperator.__TRUE__));
+            }
         }
-        return new DefaultModelCriteria<>(targetNode);
+
+        if (targetNode.getChildren().isEmpty()) {
+            // OR on empty set of formulae is FALSE: It does not hold that there is at least one satisfied formula
+            return compare(new ModelCriteriaNode<>(ExtOperator.__FALSE__));
+        }
+
+        return compare(targetNode);
     }
 
     @Override
     public DefaultModelCriteria<M> not(ModelCriteriaBuilder<M> mcb) {
         final ModelCriteriaNode<M> targetNode = new ModelCriteriaNode<>(ExtOperator.NOT);
-        targetNode.addChild(((DefaultModelCriteria<M>) mcb.unwrap(DefaultModelCriteria.class)).node);
-        return new DefaultModelCriteria<>(targetNode);
+        ModelCriteriaNode<M> toBeChild = ((DefaultModelCriteria<M>) mcb.unwrap(DefaultModelCriteria.class)).node;
+        if (toBeChild.getNodeOperator() == ExtOperator.NOT) {
+            return compare(toBeChild.getChildren().get(0).cloneTree());
+        }
+        targetNode.addChild(toBeChild.cloneTree());
+        return compare(targetNode);
     }
 
     /**
@@ -83,10 +143,35 @@ public DefaultModelCriteria<M> not(ModelCriteriaBuilder<M> mcb) {
      * @param mcb {@code ModelCriteriaBuilder} to copy the contents onto
      * @return Updated {@code ModelCriteriaBuilder}
      */
-    public ModelCriteriaBuilder<M> flashToModelCriteriaBuilder(ModelCriteriaBuilder<M> mcb) {
+    public <C extends ModelCriteriaBuilder<M>> C flashToModelCriteriaBuilder(C mcb) {
         return mcb == null ? null : node.flashToModelCriteriaBuilder(mcb);
     }
 
+    /**
+     * Optimizes this formula into another {@code ModelCriteriaBuilder}, using the values of
+     * {@link ExtOperator#__TRUE__} and {@link ExtOperator#__FALSE__} accordingly.
+     * @return New instance of {@code }
+     */
+    public DefaultModelCriteria<M> optimize() {
+        return flashToModelCriteriaBuilder(criteria());
+    }
+
+    @FunctionalInterface
+    public interface AtomicFormulaTester<M> {
+        public Boolean test(SearchableModelField<? super M> field, Operator operator, Object[] operatorArguments);
+    }
+
+    public DefaultModelCriteria<M> partiallyEvaluate(AtomicFormulaTester<M> tester) {
+        return new DefaultModelCriteria<>(node.cloneTree((field, operator, operatorArguments) -> {
+            Boolean res = tester.test(field, operator, operatorArguments);
+            if (res == null) {
+                return new ModelCriteriaNode<>(field, operator, operatorArguments);
+            } else {
+                return new ModelCriteriaNode<>(res ? ExtOperator.__TRUE__ : ExtOperator.__FALSE__);
+            }
+        }, ModelCriteriaNode::new));
+    }
+
     public boolean isEmpty() {
         return node == null;
     }
@@ -96,4 +181,14 @@ public String toString() {
         return isEmpty() ? "" : node.toString();
     }
 
+    private Stream<ModelCriteriaNode<M>> getNodesToAddForAndOr(ModelCriteriaNode<M> nodeToAdd, ExtOperator operatorBeingAdded) {
+        final ExtOperator op = nodeToAdd.getNodeOperator();
+
+        if (op == operatorBeingAdded) {
+            return nodeToAdd.getChildren().stream();
+        }
+
+        return Stream.of(nodeToAdd);
+    }
+
 }