Repository: openrewrite/rewrite (Java)

Issue:
`JavaTemplate` is unaware of `instanceof` pattern variables
When `JavaTemplate` is used to insert a piece of code which references an `instanceof` pattern variable in the code where the template is embedded, the reference in the result is not properly type attributed.

So if in the example:
```java
              class A {
                  void foo(Object o) {
                      int i = 0;
                      if (o instanceof String s) {
                          i = 1;
                      }
                  }
              }
```
the `i = 1` assignment is replaced by `i = s.length()` using a template, the type attribution on `s` in this assignment is missing in the result.

The `BlockStatementTemplateGenerator` would have to include any `instanceof` expressions which are in `if` conditions and ternary operator conditions as placeholders. In the example for example: `if (null instanceof String s)`.


Discussion: The big challenge here is the flow scoping of the `instanceof` pattern variable. Depending on what an `if` condition with an `instanceof` expression exactly looks like, this pattern variable may be accessible in the following scopes (in addition to the any additional expressions in the `if` condition joined with the logical `&&` operator):
- None. This is for example if the `if` condition contains additional boolean expressions joined via `||`.
- The `if` block only. If the `instanceof` is not negated.
- The `else` block only. If the `instanceof` expression is negated and the `else` block contains an exit node (`return`, `break` or `continue`).
- The `else` block **and** all statements following the `if` statement: If the `if` block contains an exit node but at least one of the `else` and `else if` blocks doesn't.

This is essentially what the control flow analysis of the `instanceof` pattern matching boils down to.

Instead of introducing synthetic local variables with the same names in the different scopes, the `JavaTemplate` engine could leave the `instanceof` expression more or less intact. Assuming it isn't the `instanceof` expression itself or anything in the LHS of it which is being substituted, the template engine could emit a corresponding `o() instanceof SomeType someVar`. And since a pattern variable can be visible both after an `if` statement and inside the `if` and `else` blocks, the template engine must basically for every `if` statement include some parts of it, if it contains an `instanceof` pattern matching expression in the `if` condition.



Fix (patch):
diff --git a/rewrite-java/build.gradle.kts b/rewrite-java/build.gradle.kts
index 41c6720849..a9bc325f73 100644
--- a/rewrite-java/build.gradle.kts
+++ b/rewrite-java/build.gradle.kts
@@ -46,6 +46,7 @@ dependencies {
     }
     testImplementation(project(":rewrite-test"))
     testImplementation(project(":rewrite-java-test"))
+    testRuntimeOnly(project(":rewrite-java-17"))
     testImplementation("com.tngtech.archunit:archunit:1.0.1")
     testImplementation("com.tngtech.archunit:archunit-junit5:1.0.1")
 
diff --git a/rewrite-java/src/main/java/org/openrewrite/java/internal/template/BlockStatementTemplateGenerator.java b/rewrite-java/src/main/java/org/openrewrite/java/internal/template/BlockStatementTemplateGenerator.java
index 51f60fe105..bae5311ff1 100644
--- a/rewrite-java/src/main/java/org/openrewrite/java/internal/template/BlockStatementTemplateGenerator.java
+++ b/rewrite-java/src/main/java/org/openrewrite/java/internal/template/BlockStatementTemplateGenerator.java
@@ -43,8 +43,7 @@
 public class BlockStatementTemplateGenerator {
     private static final String TEMPLATE_COMMENT = "__TEMPLATE__";
     private static final String STOP_COMMENT = "__TEMPLATE_STOP__";
-    static final String EXPR_STATEMENT_PARAM = "" +
-                                               "class __P__ {" +
+    static final String EXPR_STATEMENT_PARAM = "class __P__ {" +
                                                "  static native <T> T p();" +
                                                "  static native <T> T[] arrp();" +
                                                "  static native boolean booleanp();" +
@@ -56,8 +55,7 @@ public class BlockStatementTemplateGenerator {
                                                "  static native short shortp();" +
                                                "  static native float floatp();" +
                                                "}";
-    private static final String METHOD_INVOCATION_STUBS = "" +
-                                                          "class __M__ {" +
+    private static final String METHOD_INVOCATION_STUBS = "class __M__ {" +
                                                           "  static native Object any(Object o);" +
                                                           "  static native <T> Object anyT();" +
                                                           "}";
@@ -81,7 +79,7 @@ public String template(Cursor cursor, String template, Space.Location location)
                         after.append('}');
                     }
 
-                    template(next(cursor), cursor.getValue(), before, after, newSetFromMap(new IdentityHashMap<>()));
+                    template(next(cursor), cursor.getValue(), before, after, newSetFromMap(new IdentityHashMap<>()), cursor.getValue());
 
                     return before.toString().trim() + "\n/*" + TEMPLATE_COMMENT + "*/" + template + "/*" + STOP_COMMENT + "*/" + "\n" + after;
                 });
@@ -160,7 +158,7 @@ else if (tree != null && !js.isEmpty()) {
     }
 
     @SuppressWarnings("ConstantConditions")
-    private void template(Cursor cursor, J prior, StringBuilder before, StringBuilder after, Set<J> templated) {
+    private void template(Cursor cursor, J prior, StringBuilder before, StringBuilder after, Set<J> templated, J toReplace) {
         templated.add(cursor.getValue());
         J j = cursor.getValue();
         if (j instanceof JavaSourceFile) {
@@ -192,16 +190,7 @@ private void template(Cursor cursor, J prior, StringBuilder before, StringBuilde
                 J.MethodDeclaration m = (J.MethodDeclaration) parent;
 
                 // variable declarations up to the point of insertion
-                assert m.getBody() != null;
-                for (Statement statement : m.getBody().getStatements()) {
-                    if (referToSameElement(prior, statement)) {
-                        break;
-                    } else if (statement instanceof J.VariableDeclarations) {
-                        before.insert(0, "\n" +
-                                         variable((J.VariableDeclarations) statement, true, cursor) +
-                                         ";\n");
-                    }
-                }
+                addLeadingVariableDeclarations(cursor, prior, m.getBody(), before, toReplace);
 
                 if (m.getReturnTypeExpression() != null && !JavaType.Primitive.Void
                         .equals(m.getReturnTypeExpression().getType())) {
@@ -215,19 +204,11 @@ private void template(Cursor cursor, J prior, StringBuilder before, StringBuilde
                                          .withLeadingAnnotations(emptyList())
                                          .withPrefix(Space.EMPTY)
                                          .printTrimmed(cursor).trim() + '{');
-            } else if (parent instanceof J.Block || parent instanceof J.Lambda) {
+            } else if (parent instanceof J.Block || parent instanceof J.Lambda || parent instanceof J.Label || parent instanceof Loop) {
                 J.Block b = (J.Block) j;
 
                 // variable declarations up to the point of insertion
-                for (Statement statement : b.getStatements()) {
-                    if (referToSameElement(prior, statement)) {
-                        break;
-                    } else if (statement instanceof J.VariableDeclarations) {
-                        before.insert(0, "\n" +
-                                         variable((J.VariableDeclarations) statement, true, cursor) +
-                                         ";\n");
-                    }
-                }
+                addLeadingVariableDeclarations(cursor, prior, b, before, toReplace);
 
                 before.insert(0, "{\n");
                 if (b.isStatic()) {
@@ -314,10 +295,12 @@ private void template(Cursor cursor, J prior, StringBuilder before, StringBuilde
             }
         } else if (j instanceof J.ForLoop) {
             J.ForLoop f = (J.ForLoop) j;
-            insertControlWithBlock(f.getBody(), before, after, () -> before.insert(0,
-                    f.withBody(null).withPrefix(Space.EMPTY)
-                            .withControl(f.getControl().withCondition(null).withUpdate(emptyList()))
-                            .printTrimmed(cursor).trim()));
+            if (referToSameElement(prior, f.getBody())) {
+                insertControlWithBlock(f.getBody(), before, after, () -> before.insert(0,
+                        f.withBody(null).withPrefix(Space.EMPTY)
+                                .withControl(f.getControl().withCondition(null).withUpdate(emptyList()))
+                                .printTrimmed(cursor).trim()));
+            }
         } else if (j instanceof J.ForEachLoop.Control) {
             J.ForEachLoop.Control c = (J.ForEachLoop.Control) j;
             if (c.getVariable() == prior) {
@@ -394,10 +377,47 @@ private void template(Cursor cursor, J prior, StringBuilder before, StringBuilde
         } else if (j instanceof J.If) {
             J.If iff = (J.If) j;
             if (referToSameElement(prior, iff.getIfCondition())) {
-                insertControlWithBlock(iff.getThenPart(), before, after, () -> {
-                    before.insert(0, "Object __b" + cursor.getPathAsStream().count() + "__ =");
-                    after.append(";");
-                });
+                String condition = PatternVariables.simplifiedPatternVariableCondition(iff.getIfCondition().getTree(), toReplace);
+                if (condition != null) {
+                    int splitIdx = condition.indexOf('§');
+                    before.insert(0, "if (" + condition.substring(0, splitIdx) + '(');
+                    after.append(')').append(condition.substring(splitIdx + 1)).append(") {}");
+                } else {
+                    insertControlWithBlock(iff.getThenPart(), before, after, () -> {
+                        before.insert(0, "Object __b" + cursor.getPathAsStream().count() + "__ =");
+                        after.append(";");
+                    });
+                }
+            } else {
+                String condition = PatternVariables.simplifiedPatternVariableCondition(iff.getIfCondition().getTree(), toReplace);
+                if (condition != null) {
+                    if (referToSameElement(prior, iff.getThenPart())) {
+                        insertControlWithBlock(iff.getThenPart(), before, after, () ->
+                                before.insert(0, "if (" + condition + ") "));
+                    } else if (referToSameElement(prior, iff.getElsePart())) {
+                        insertControlWithBlock(iff.getElsePart().getBody(), before, after, () ->
+                                before.insert(0, "if (" + condition + ") {} else "));
+                    }
+                }
+            }
+        } else if (j instanceof J.Ternary) {
+            J.Ternary ternary = (J.Ternary) j;
+            if (referToSameElement(prior, ternary.getCondition())) {
+                String condition = PatternVariables.simplifiedPatternVariableCondition(ternary.getCondition(), toReplace);
+                if (condition != null) {
+                    int splitIdx = condition.indexOf('§');
+                    before.insert(0, condition.substring(0, splitIdx) + '(');
+                    after.append(')').append(condition.substring(splitIdx + 1))
+                            .append(" ? ").append(ternary.getTruePart().printTrimmed(cursor).trim())
+                            .append(" : ").append(ternary.getFalsePart().printTrimmed(cursor).trim());
+                }
+            } else if (referToSameElement(prior, ternary.getTruePart())) {
+                String condition = PatternVariables.simplifiedPatternVariableCondition(ternary.getCondition(), toReplace);
+                before.insert(0, condition + " ? ");
+                after.append(" : ").append(ternary.getFalsePart().printTrimmed(cursor).trim());
+            } else if (referToSameElement(prior, ternary.getFalsePart())) {
+                String condition = PatternVariables.simplifiedPatternVariableCondition(ternary.getCondition(), toReplace);
+                before.insert(0, condition + " ? " + ternary.getTruePart().printTrimmed(cursor).trim() + " : ");
             }
         } else if (j instanceof J.WhileLoop) {
             J.WhileLoop wl = (J.WhileLoop) j;
@@ -419,7 +439,36 @@ private void template(Cursor cursor, J prior, StringBuilder before, StringBuilde
         } else if (j instanceof J.EnumValueSet) {
             after.append(";");
         }
-        template(next(cursor), j, before, after, templated);
+        template(next(cursor), j, before, after, templated, toReplace);
+    }
+
+    private void addLeadingVariableDeclarations(Cursor cursor, J current, J.Block containingBlock, StringBuilder before, J toReplace) {
+        for (Statement statement : containingBlock.getStatements()) {
+            if (referToSameElement(current, statement)) {
+                break;
+            }
+            if (statement instanceof J.Label) {
+                statement = ((J.Label) statement).getStatement();
+            }
+            if (statement instanceof J.VariableDeclarations) {
+                before.insert(0, "\n" +
+                                 variable((J.VariableDeclarations) statement, true, cursor) +
+                                 ";\n");
+            } else if (statement instanceof J.If) {
+                J.If iff = (J.If) statement;
+                String condition = PatternVariables.simplifiedPatternVariableCondition(iff.getIfCondition().getTree(), toReplace);
+                if (condition != null) {
+                    boolean thenNeverCompletesNormally = PatternVariables.neverCompletesNormally(iff.getThenPart());
+                    boolean elseNeverCompletesNormally = iff.getElsePart() != null && PatternVariables.neverCompletesNormally(iff.getElsePart().getBody());
+                    if (thenNeverCompletesNormally || elseNeverCompletesNormally) {
+                        StringBuilder ifStatement = new StringBuilder("if (").append(condition).append(") {");
+                        ifStatement.append(thenNeverCompletesNormally ? " throw new RuntimeException(); }" : " }");
+                        ifStatement.append(elseNeverCompletesNormally ? " else { throw new RuntimeException(); }" : " else { }");
+                        before.insert(0, ifStatement);
+                    }
+                }
+            }
+        }
     }
 
     private void insertControlWithBlock(J body, StringBuilder before, StringBuilder after, Runnable insertion) {
diff --git a/rewrite-java/src/main/java/org/openrewrite/java/internal/template/PatternVariables.java b/rewrite-java/src/main/java/org/openrewrite/java/internal/template/PatternVariables.java
new file mode 100644
index 0000000000..a799089703
--- /dev/null
+++ b/rewrite-java/src/main/java/org/openrewrite/java/internal/template/PatternVariables.java
@@ -0,0 +1,293 @@
+/*
+ * Copyright 2023 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.internal.template;
+
+import org.openrewrite.internal.lang.Nullable;
+import org.openrewrite.java.tree.Expression;
+import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.Loop;
+import org.openrewrite.java.tree.Statement;
+
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+final class PatternVariables {
+
+    private static final String DEFAULT_LABEL = "!";
+
+    private static class ResultCollector {
+        final StringBuilder builder = new StringBuilder();
+        boolean instanceOfFound;
+    }
+
+    @Nullable
+    static String simplifiedPatternVariableCondition(Expression condition, @Nullable J toReplace) {
+        ResultCollector resultCollector = new ResultCollector();
+        simplifiedPatternVariableCondition0(condition, toReplace, resultCollector);
+        return resultCollector.instanceOfFound ? resultCollector.builder.toString() : null;
+    }
+
+    private static boolean simplifiedPatternVariableCondition0(J expr, @Nullable J toReplace, ResultCollector collector) {
+        if (expr == toReplace) {
+            collector.builder.append('§');
+            return true;
+        }
+
+        if (expr instanceof J.Parentheses) {
+            J.Parentheses<?> parens = (J.Parentheses<?>) expr;
+            collector.builder.append('(');
+            try {
+                return simplifiedPatternVariableCondition0(parens.getTree(), toReplace, collector);
+            } finally {
+                collector.builder.append(')');
+            }
+        } else if (expr instanceof J.Unary) {
+            J.Unary unary = (J.Unary) expr;
+            switch (unary.getOperator()) {
+                case PostIncrement: {
+                    boolean found = simplifiedPatternVariableCondition0(unary.getExpression(), toReplace, collector);
+                    collector.builder.append("++");
+                    return found;
+                }
+                case PostDecrement: {
+                    boolean found = simplifiedPatternVariableCondition0(unary.getExpression(), toReplace, collector);
+                    collector.builder.append("--");
+                    return found;
+                }
+                case PreIncrement:
+                    collector.builder.append("++");
+                    break;
+                case PreDecrement:
+                    collector.builder.append("--");
+                    break;
+                case Positive:
+                    collector.builder.append('+');
+                    break;
+                case Negative:
+                    collector.builder.append('-');
+                    break;
+                case Complement:
+                    collector.builder.append('~');
+                    break;
+                case Not:
+                    collector.builder.append('!');
+                    break;
+                default:
+                    throw new IllegalStateException("Unexpected unary operator: " + unary.getOperator());
+            }
+            return simplifiedPatternVariableCondition0(unary.getExpression(), toReplace, collector);
+        } else if (expr instanceof J.Binary) {
+            J.Binary binary = (J.Binary) expr;
+            int length = collector.builder.length();
+            boolean result = simplifiedPatternVariableCondition0(binary.getLeft(), toReplace, collector);
+            switch (binary.getOperator()) {
+                case Addition:
+                    collector.builder.append('+');
+                    break;
+                case Subtraction:
+                    collector.builder.append('-');
+                    break;
+                case Multiplication:
+                    collector.builder.append('*');
+                    break;
+                case Division:
+                    collector.builder.append('/');
+                    break;
+                case Modulo:
+                    collector.builder.append('%');
+                    break;
+                case LessThan:
+                    collector.builder.append('<');
+                    break;
+                case GreaterThan:
+                    collector.builder.append('>');
+                    break;
+                case LessThanOrEqual:
+                    collector.builder.append("<=");
+                    break;
+                case GreaterThanOrEqual:
+                    collector.builder.append(">=");
+                    break;
+                case Equal:
+                    collector.builder.append("==");
+                    break;
+                case NotEqual:
+                    collector.builder.append("!=");
+                    break;
+                case BitAnd:
+                    collector.builder.append('&');
+                    break;
+                case BitOr:
+                    collector.builder.append('|');
+                    break;
+                case BitXor:
+                    collector.builder.append('^');
+                    break;
+                case LeftShift:
+                    collector.builder.append("<<");
+                    break;
+                case RightShift:
+                    collector.builder.append(">>");
+                    break;
+                case UnsignedRightShift:
+                    collector.builder.append(">>>");
+                    break;
+                case Or:
+                    collector.builder.append("||");
+                    break;
+                case And:
+                    collector.builder.append("&&");
+                    break;
+                default:
+                    throw new IllegalStateException("Unexpected binary operator: " + binary.getOperator());
+            }
+            result |= simplifiedPatternVariableCondition0(binary.getRight(), toReplace, collector);
+            if (!result) {
+                switch (binary.getOperator()) {
+                    case LessThan:
+                    case GreaterThan:
+                    case LessThanOrEqual:
+                    case GreaterThanOrEqual:
+                    case Equal:
+                    case NotEqual:
+                    case Or:
+                    case And:
+                        collector.builder.setLength(length);
+                        collector.builder.append("true");
+                        return false;
+                }
+            }
+            return result;
+        } else if (expr instanceof J.InstanceOf) {
+            J.InstanceOf instanceOf = (J.InstanceOf) expr;
+            if (instanceOf.getPattern() != null) {
+                collector.builder.append("((Object)null) instanceof ").append(instanceOf.getClazz()).append(' ').append(instanceOf.getPattern());
+                collector.instanceOfFound = true;
+                return true;
+            }
+            collector.builder.append("true");
+        } else if (expr instanceof J.Literal) {
+            J.Literal literal = (J.Literal) expr;
+            collector.builder.append(literal.getValue());
+        } else if (expr instanceof Expression) {
+            collector.builder.append("null");
+        }
+        return false;
+    }
+
+    static boolean neverCompletesNormally(Statement statement) {
+        return neverCompletesNormally0(statement, new HashSet<>());
+    }
+
+    private static boolean neverCompletesNormally0(@Nullable Statement statement, Set<String> labelsToIgnore) {
+        if (statement instanceof J.Return || statement instanceof J.Throw) {
+            return true;
+        } else if (statement instanceof J.Break) {
+            J.Break breakStatement = (J.Break) statement;
+            return breakStatement.getLabel() != null && !labelsToIgnore.contains(breakStatement.getLabel().getSimpleName())
+                    || breakStatement.getLabel() == null && !labelsToIgnore.contains(DEFAULT_LABEL);
+        } else if (statement instanceof J.Continue) {
+            J.Continue continueStatement = (J.Continue) statement;
+            return continueStatement.getLabel() != null && !labelsToIgnore.contains(continueStatement.getLabel().getSimpleName())
+                    || continueStatement.getLabel() == null && !labelsToIgnore.contains(DEFAULT_LABEL);
+        } else if (statement instanceof J.Block) {
+            return neverCompletesNormally0(getLastStatement(statement), labelsToIgnore);
+        } else if (statement instanceof Loop) {
+            Loop loop = (Loop) statement;
+            return neverCompletesNormallyIgnoringLabel(loop.getBody(), DEFAULT_LABEL, labelsToIgnore);
+        } else if (statement instanceof J.If) {
+            J.If if_ = (J.If) statement;
+            return if_.getElsePart() != null
+                    && neverCompletesNormally0(if_.getThenPart(), labelsToIgnore)
+                    && neverCompletesNormally0(if_.getElsePart().getBody(), labelsToIgnore);
+        } else if (statement instanceof J.Switch) {
+            J.Switch switch_ = (J.Switch) statement;
+            if (switch_.getCases().getStatements().isEmpty()) {
+                return false;
+            }
+            Statement defaultCase = null;
+            for (Statement case_ : switch_.getCases().getStatements()) {
+                if (!neverCompletesNormallyIgnoringLabel(case_, DEFAULT_LABEL, labelsToIgnore)) {
+                    return false;
+                }
+                if (case_ instanceof J.Case) {
+                    Expression elem = ((J.Case) case_).getPattern();
+                    if (elem instanceof J.Identifier && ((J.Identifier) elem).getSimpleName().equals("default")) {
+                        defaultCase = case_;
+                    }
+                }
+            }
+            return neverCompletesNormallyIgnoringLabel(defaultCase, DEFAULT_LABEL, labelsToIgnore);
+        } else if (statement instanceof J.Case) {
+            J.Case case_ = (J.Case) statement;
+            if (case_.getStatements().isEmpty()) {
+                // fallthrough to next case
+                return true;
+            }
+            return neverCompletesNormally0(getLastStatement(case_), labelsToIgnore);
+        } else if (statement instanceof J.Try) {
+            J.Try try_ = (J.Try) statement;
+            if (try_.getFinally() != null && !try_.getFinally().getStatements().isEmpty()
+                    && neverCompletesNormally0(try_.getFinally(), labelsToIgnore)) {
+                return true;
+            }
+            boolean bodyHasExit = false;
+            if (!try_.getBody().getStatements().isEmpty()
+                    && !(bodyHasExit = neverCompletesNormally0(try_.getBody(), labelsToIgnore))) {
+                return false;
+            }
+            for (J.Try.Catch catch_ : try_.getCatches()) {
+                if (!neverCompletesNormally0(catch_.getBody(), labelsToIgnore)) {
+                    return false;
+                }
+            }
+            return bodyHasExit;
+        } else if (statement instanceof J.Synchronized) {
+            return neverCompletesNormally0(((J.Synchronized) statement).getBody(), labelsToIgnore);
+        } else if (statement instanceof J.Label) {
+            String label = ((J.Label) statement).getLabel().getSimpleName();
+            Statement labeledStatement = ((J.Label) statement).getStatement();
+            return neverCompletesNormallyIgnoringLabel(labeledStatement, label, labelsToIgnore);
+        }
+        return false;
+    }
+
+    private static boolean neverCompletesNormallyIgnoringLabel(@Nullable Statement statement, String label, Set<String> labelsToIgnore) {
+        boolean added = labelsToIgnore.add(label);
+        try {
+            return neverCompletesNormally0(statement, labelsToIgnore);
+        } finally {
+            if (added) {
+                labelsToIgnore.remove(label);
+            }
+        }
+    }
+
+    @Nullable
+    private static Statement getLastStatement(Statement statement) {
+        if (statement instanceof J.Block) {
+            List<Statement> statements = ((J.Block) statement).getStatements();
+            return statements.isEmpty() ? null : getLastStatement(statements.get(statements.size() - 1));
+        } else if (statement instanceof J.Case) {
+            List<Statement> statements = ((J.Case) statement).getStatements();
+            return statements.isEmpty() ? null : getLastStatement(statements.get(statements.size() - 1));
+        } else if (statement instanceof Loop) {
+            return getLastStatement(((Loop) statement).getBody());
+        }
+        return statement;
+    }
+}