Improve preconditions check

src/main/java/org/openrewrite/java/template/processor/Precondition.java

@@ -0,0 +1,193 @@
+/*
+ * Copyright 2024 the original author or authors.
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.openrewrite.java.template.processor;
+
+import lombok.EqualsAndHashCode;
+import lombok.Value;
+
+import java.util.*;
+
+import static java.util.stream.Collectors.toCollection;
+
+public abstract class Precondition {
+    private static final Comparator<String> BY_USES_TYPE_FIRST = Comparator
+            .comparing((String s) -> !s.startsWith("new UsesType"))
+            .thenComparing(Comparator.naturalOrder());
+
+    abstract boolean fitsInto(Precondition p);
+
+    public Precondition prune() {
+        return this;
+    }
+
+    @Value
+    @EqualsAndHashCode(callSuper = false)
+    public static class Rule extends Precondition {
+        String rule;
+
+        @Override
+        boolean fitsInto(Precondition p) {
+            if (p instanceof Rule) {
+                return this.equals(p);
+            } else if (p instanceof Or) {
+                return ((Or) p).preconditions.stream().anyMatch(this::fitsInto);
+            } else if (p instanceof And) {
+                return ((And) p).preconditions.stream().anyMatch(this::fitsInto);
+            }
+            return false; // unreachable code
+        }
+
+        @Override
+        public String toString() {
+            return rule;
+        }
+    }
+
+    @Value
+    @EqualsAndHashCode(callSuper = false, of = "preconditions")
+    public static class Or extends Precondition {
+        Set<Precondition> preconditions;
+        int indent;
+
+        @Override
+        boolean fitsInto(Precondition p) {
+            if (p instanceof Or) {
+                return this.equals(p);
+            }
+            return false;
+        }
+
+        @Override
+        public Precondition prune() {
+            Precondition pruned = takeElementIfItFitsInAllOtherElements();
+            return pruned == null ? extractCommonElements() : pruned;
+        }
+
+        /**
+         * If element fits in all others, take element as precondition. Eg:
+         * <pre>
+         * or(
+         *    and(new UsesType<>("Map"), new UsesMethod<>("PrintStream println(..)")),
+         *    new UsesMethod<>("PrintStream println(..)")
+         * )
+         * </pre>
+         * <p>
+         * will return:
+         * <pre>
+         * new UsesMethod<>("PrintStream println()")
+         * </pre>
+         */
+        private Precondition takeElementIfItFitsInAllOtherElements() {
+            for (Precondition p : preconditions) {
+                int matches = 0;
+                for (Precondition p2 : preconditions) {
+                    if (p == p2 || p.fitsInto(p2)) {
+                        matches++;
+                    }
+                    if (matches == preconditions.size()) {
+                        return p;
+                    }
+                }
+            }
+            return null;
+        }
+
+        /**
+         * If child element of an element exist as child element in all others, move child element up. Eg:
+         * <pre>
+         * or(
+         *    and(new UsesType<>("Map"), new UsesType<>("HashMap"), new UsesMethod<>("PrintStream println(..)")),
+         *    and(new UsesType<>("List"), new UsesType<>("ArrayList"), new UsesMethod<>("PrintStream println(..)"))
+         * )
+         * </pre>
+         * <p>
+         * will return:
+         * <pre>
+         * and(
+         *    new UsesMethod<>("PrintStream println()"),
+         *    or(
+         *      and(new UsesType<>("Map"), new UsesType<>("HashMap")),
+         *      and(new UsesType<>("List"), new UsesType<>("ArrayList")),
+         *    )
+         * )
+         * </pre>
+         */
+        private Precondition extractCommonElements() {
+            boolean first = true;
+            Set<Precondition> commons = new HashSet<>();
+            for (Precondition p : preconditions) {
+                if (!(p instanceof And)) {
+                    return this;
+                }
+                if (first) {
+                    commons.addAll(((And) p).preconditions);
+                    first = false;
+                } else {
+                    commons.retainAll(((And) p).preconditions);
+                }
+            }
+
+            if (!commons.isEmpty()) {
+                preconditions.forEach(it -> ((And) it).preconditions.removeAll(commons));
+                commons.add(new Or(preconditions, indent));
+                return new And(commons, indent);
+            }
+
+            return this;
+        }
+
+        @Override
+        public String toString() {
+            return joinPreconditions(preconditions, "or", indent);
+        }
+    }
+
+    @Value
+    @EqualsAndHashCode(callSuper = false, of = "preconditions")
+    public static class And extends Precondition {
+        Set<Precondition> preconditions;
+        int indent;
+
+        @Override
+        boolean fitsInto(Precondition p) {
+            if (p instanceof And) {
+                if (preconditions.size() > ((And) p).preconditions.size()) {
+                    return false;
+                }
+                return preconditions.stream().allMatch(it -> it.fitsInto(p));
+            }
+            return false;
+        }
+
+        @Override
+        public String toString() {
+            return joinPreconditions(preconditions, "and", indent);
+        }
+    }
+
+    private static String joinPreconditions(Collection<Precondition> rules, String op, int indent) {
+        if (rules.isEmpty()) {
+            return "";
+        } else if (rules.size() == 1) {
+            return rules.iterator().next().toString();
+        }
+        String whitespace = String.format("%" + indent + "s", " ");
+        Set<String> preconditions = rules.stream().map(Object::toString).sorted(BY_USES_TYPE_FIRST).collect(toCollection(LinkedHashSet::new));
+        return "Preconditions." + op + "(\n" +
+                whitespace + String.join(",\n" + whitespace, preconditions) + "\n" +
+                whitespace.substring(0, indent - 4) + ')';
+    }
+}

src/main/java/org/openrewrite/java/template/processor/RefasterTemplateProcessor.java

@@ -226,7 +226,7 @@ public void visitClassDef(JCTree.JCClassDecl classDecl) {
 
                     String javaVisitor = newAbstractRefasterJavaVisitor(beforeTemplates, after, descriptor);
 
-                    String preconditions = generatePreconditions(descriptor.beforeTemplates, 16);
+                    Precondition preconditions = generatePreconditions(descriptor.beforeTemplates, 16);
                     if (preconditions == null) {
                         recipe.append(String.format("        return %s;\n", javaVisitor));
                     } else {
@@ -587,18 +587,17 @@ private Set<String> getImportsAsStrings(Map<TemplateDescriptor, Set<String>> imp
             }
 
             /* Generate the minimal precondition that would allow to match each before template individually. */
-            @SuppressWarnings("SameParameterValue")
-            private @Nullable String generatePreconditions(List<TemplateDescriptor> beforeTemplates, int indent) {
-                Map<String, Set<String>> preconditions = new LinkedHashMap<>();
+            private @Nullable Precondition generatePreconditions(List<TemplateDescriptor> beforeTemplates, int indent) {
+                Map<String, Set<Precondition>> preconditions = new LinkedHashMap<>();
                 for (TemplateDescriptor beforeTemplate : beforeTemplates) {
                     int arity = beforeTemplate.getArity();
                     for (int i = 0; i < arity; i++) {
-                        Set<String> usesVisitors = new LinkedHashSet<>();
+                        Set<Precondition> usesVisitors = new LinkedHashSet<>();
 
                         for (Symbol.ClassSymbol usedType : beforeTemplate.usedTypes(i)) {
                             String name = usedType.getQualifiedName().toString().replace('$', '.');
                             if (!name.startsWith("java.lang.") && !name.startsWith("com.google.errorprone.refaster.")) {
-                                usesVisitors.add("new UsesType<>(\"" + name + "\", true)");
+                                usesVisitors.add(new Precondition.Rule("new UsesType<>(\"" + name + "\", true)"));
                             }
                         }
                         for (Symbol.MethodSymbol method : beforeTemplate.usedMethods(i)) {
@@ -607,48 +606,28 @@ private Set<String> getImportsAsStrings(Map<TemplateDescriptor, Set<String>> imp
                             }
                             String methodName = method.name.toString();
                             methodName = methodName.equals("<init>") ? "<constructor>" : methodName;
-                            usesVisitors.add("new UsesMethod<>(\"" + method.owner.getQualifiedName().toString() + ' ' + methodName + "(..)\", true)");
+                            usesVisitors.add(new Precondition.Rule(String.format("new UsesMethod<>(\"%s %s(..)\", true)",
+                                    method.owner.getQualifiedName().toString(), methodName)));
                         }
 
-                        preconditions.put(beforeTemplate.method.name.toString() + (arity == 1 ? "" : "$" + i), usesVisitors);
-                    }
-                }
-
-                if (preconditions.size() == 1) {
-                    return joinPreconditions(preconditions.values().iterator().next(), "and", indent + 4);
-                } else if (preconditions.size() > 1) {
-                    Set<String> common = new LinkedHashSet<>();
-                    for (String dep : preconditions.values().iterator().next()) {
-                        if (preconditions.values().stream().allMatch(v -> v.contains(dep))) {
-                            common.add(dep);
-                        }
-                    }
-                    common.forEach(dep -> preconditions.values().forEach(v -> v.remove(dep)));
-                    preconditions.values().removeIf(Collection::isEmpty);
-
-                    if (common.isEmpty()) {
-                        return joinPreconditions(preconditions.values().stream().map(v -> joinPreconditions(v, "and", indent + 4)).collect(toList()), "or", indent + 4);
-                    } else {
-                        if (!preconditions.isEmpty()) {
-                            String uniqueConditions = joinPreconditions(preconditions.values().stream().map(v -> joinPreconditions(v, "and", indent + 12)).collect(toList()), "or", indent + 8);
-                            common.add(uniqueConditions);
+                        if (!usesVisitors.isEmpty()) {
+                            preconditions.put(beforeTemplate.method.name.toString() + (arity == 1 ? "" : "$" + i), usesVisitors);
+                        } else {
+                            return null; // At least one of the before templates has no preconditions, so we can not use any preconditions
                         }
-                        return joinPreconditions(common, "and", indent + 4);
                     }
                 }
-                return null;
-            }
 
-            private String joinPreconditions(Collection<String> preconditions, String op, int indent) {
                 if (preconditions.isEmpty()) {
                     return null;
-                } else if (preconditions.size() == 1) {
-                    return preconditions.iterator().next();
                 }
-                char[] indentChars = new char[indent];
-                Arrays.fill(indentChars, ' ');
-                String indentStr = new String(indentChars);
-                return "Preconditions." + op + "(\n" + indentStr + String.join(",\n" + indentStr, preconditions) + "\n" + indentStr.substring(0, indent - 4) + ')';
+
+                return new Precondition.Or(
+                        preconditions.values().stream()
+                                .map(it -> new Precondition.And(it, indent + 4))
+                                .collect(toSet())
+                        , indent + 4)
+                        .prune();
             }
         }.scan(cu);
     }
@@ -751,7 +730,7 @@ public RuleDescriptor(JCTree.JCClassDecl classDecl, JCCompilationUnit cu, Contex
             this.context = context;
         }
 
-        private RefasterTemplateProcessor.@Nullable RuleDescriptor validate() {
+        private @Nullable RuleDescriptor validate() {
             if (beforeTemplates.isEmpty()) {
                 return null;
             }
@@ -1089,7 +1068,7 @@ private static List<JCTree.JCAnnotation> getTemplateAnnotations(MethodTree metho
                 result.add((JCTree.JCAnnotation) annotation);
             } else if (type.getKind() == Tree.Kind.MEMBER_SELECT && type instanceof JCTree.JCFieldAccess &&
                        ((JCTree.JCFieldAccess) type).sym != null &&
-                       typePredicate.test(((JCTree.JCFieldAccess) type).sym.getQualifiedName().toString())) {
+                typePredicate.test(((JCTree.JCFieldAccess) type).sym.getQualifiedName().toString())) {
                 result.add((JCTree.JCAnnotation) annotation);
             }
         }

src/test/java/org/openrewrite/java/template/RefasterTemplateProcessorTest.java

@@ -96,6 +96,7 @@ void skipRecipeGeneration(String recipeName) {
       "SimplifyTernary",
       "RefasterAnyOf",
       "Parameters",
+      "PreconditionsVerifier",
     })
     void nestedRecipes(String recipeName) {
         Compilation compilation = compileResource("refaster/" + recipeName + ".java");