diff --git a/src/main/java/org/openrewrite/staticanalysis/DeclarationSiteTypeVariance.java b/src/main/java/org/openrewrite/staticanalysis/DeclarationSiteTypeVariance.java index 5ab493f0e..94761d153 100644 --- a/src/main/java/org/openrewrite/staticanalysis/DeclarationSiteTypeVariance.java +++ b/src/main/java/org/openrewrite/staticanalysis/DeclarationSiteTypeVariance.java @@ -27,6 +27,7 @@ import java.util.Arrays; import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; import static java.util.stream.Collectors.toList; import static org.openrewrite.java.tree.J.Wildcard.Bound.Extends; @@ -95,7 +96,7 @@ public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, Ex if (varParam.getTypeExpression() instanceof J.ParameterizedType) { J.ParameterizedType pt = (J.ParameterizedType) varParam.getTypeExpression(); for (VariantTypeSpec variantTypeSpec : variantTypeSpecs) { - if (variantTypeSpec.hasType(pt)) { + if (variantTypeSpec.hasType(pt) && !isStoredInvariantly(m, varParam)) { return varParam.withTypeExpression(useDeclarationSiteVariance(pt, variantTypeSpec)); } } @@ -106,6 +107,72 @@ public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, Ex })); } + private boolean isStoredInvariantly(J.MethodDeclaration method, J.VariableDeclarations parameter) { + if (method.getBody() == null || parameter.getVariables().size() != 1) { + return false; + } + JavaType.Variable parameterType = parameter.getVariables().get(0).getVariableType(); + if (parameterType == null) { + return false; + } + return new JavaIsoVisitor() { + @Override + public J.Assignment visitAssignment(J.Assignment assignment, AtomicBoolean stored) { + J.Assignment visitedAssignment = super.visitAssignment(assignment, stored); + if (TypeUtils.isOfType(visitedAssignment.getVariable().getType(), parameterType.getType()) && + !references(visitedAssignment.getVariable(), parameterType) && + directlyReferences(visitedAssignment.getAssignment(), parameterType)) { + stored.set(true); + } + return visitedAssignment; + } + + @Override + public J.VariableDeclarations.NamedVariable visitVariable( + J.VariableDeclarations.NamedVariable variable, AtomicBoolean stored) { + J.VariableDeclarations.NamedVariable visitedVariable = super.visitVariable(variable, stored); + J.VariableDeclarations declarations = getCursor().firstEnclosing(J.VariableDeclarations.class); + boolean explicitlyTyped = declarations != null && + !(declarations.getTypeExpression() instanceof J.Identifier && + "var".equals(((J.Identifier) declarations.getTypeExpression()).getSimpleName())); + if (explicitlyTyped && visitedVariable.getInitializer() != null && + TypeUtils.isOfType(visitedVariable.getType(), parameterType.getType()) && + directlyReferences(visitedVariable.getInitializer(), parameterType)) { + stored.set(true); + } + return visitedVariable; + } + }.reduce(method.getBody(), new AtomicBoolean()).get(); + } + + private boolean references(Expression expression, JavaType.Variable parameterType) { + return new JavaIsoVisitor() { + @Override + public J.Lambda visitLambda(J.Lambda lambda, AtomicBoolean found) { + return lambda; + } + + @Override + public J.MemberReference visitMemberReference(J.MemberReference memberRef, AtomicBoolean found) { + return memberRef; + } + + @Override + public J.Identifier visitIdentifier(J.Identifier identifier, AtomicBoolean found) { + if (parameterType.equals(identifier.getFieldType())) { + found.set(true); + } + return identifier; + } + }.reduce(expression, new AtomicBoolean()).get(); + } + + private boolean directlyReferences(Expression expression, JavaType.Variable parameterType) { + Expression unwrapped = expression.unwrap(); + return unwrapped instanceof J.Identifier && + parameterType.equals(((J.Identifier) unwrapped).getFieldType()); + } + private J.ParameterizedType useDeclarationSiteVariance(J.ParameterizedType pt, VariantTypeSpec spec) { return pt.withTypeParameters(ListUtils.map(pt.getTypeParameters(), (i, tp) -> { VariantTypeSpec.Variance variance = spec.getVariances().get(i); diff --git a/src/test/java/org/openrewrite/staticanalysis/DeclarationSiteTypeVarianceTest.java b/src/test/java/org/openrewrite/staticanalysis/DeclarationSiteTypeVarianceTest.java index d633e6d0f..62f137ddb 100644 --- a/src/test/java/org/openrewrite/staticanalysis/DeclarationSiteTypeVarianceTest.java +++ b/src/test/java/org/openrewrite/staticanalysis/DeclarationSiteTypeVarianceTest.java @@ -235,4 +235,220 @@ public void test(Function f) { ) ); } + @Test + void doesNotAddVarianceToParameterStoredInvariantly() { + rewriteRun( + //language=java + java( + """ + import java.util.function.Function; + + class FieldTest { + private final Function mapper; + + FieldTest(Function mapper) { + this.mapper = mapper; + } + } + """ + ), + //language=java + java( + """ + import java.util.function.Function; + + class LocalTest { + void test(Function mapper) { + Function stored = mapper; + } + } + """ + ) + ); + } + + @Test + void addsVarianceWhenFieldAcceptsIt() { + rewriteRun( + //language=java + java( + """ + import java.util.function.Function; + + class Test { + private final Function mapper; + + Test(Function mapper) { + this.mapper = mapper; + } + } + """, + """ + import java.util.function.Function; + + class Test { + private final Function mapper; + + Test(Function mapper) { + this.mapper = mapper; + } + } + """ + ) + ); + } + + @Test + void addsVarianceWhenInvariantFieldStoresAdapter() { + rewriteRun( + //language=java + java( + """ + import java.util.function.Function; + + class Test { + private Function mapper; + + void set(Function mapper) { + this.mapper = input -> mapper.apply(input); + } + + void forward(Function mapper) { + set(mapper); + } + } + """, + """ + import java.util.function.Function; + + class Test { + private Function mapper; + + void set(Function mapper) { + this.mapper = input -> mapper.apply(input); + } + + void forward(Function mapper) { + set(mapper); + } + } + """ + ) + ); + } + + @Test + void addsVarianceWhenParameterIsNormalizedInPlace() { + rewriteRun( + //language=java + java( + """ + import java.util.Objects; + import java.util.function.Function; + + class Test { + void use(Function mapper) { + mapper = Objects.requireNonNull(mapper); + } + + void forward(Function mapper) { + use(mapper); + } + } + """, + """ + import java.util.Objects; + import java.util.function.Function; + + class Test { + void use(Function mapper) { + mapper = Objects.requireNonNull(mapper); + } + + void forward(Function mapper) { + use(mapper); + } + } + """ + ) + ); + } + + @Test + void addsVarianceWhenParameterIsStoredInVar() { + rewriteRun( + //language=java + java( + """ + import java.util.function.Function; + + class Test { + void use(Function mapper) { + var stored = mapper; + } + + void forward(Function mapper) { + use(mapper); + } + } + """, + """ + import java.util.function.Function; + + class Test { + void use(Function mapper) { + var stored = mapper; + } + + void forward(Function mapper) { + use(mapper); + } + } + """ + ) + ); + } + + @Test + void addsVarianceWhenInvariantLocalStoresAdapterResult() { + rewriteRun( + //language=java + java( + """ + import java.util.function.Function; + + class Test { + Function adapt(Function mapper) { + return input -> mapper.apply(input); + } + + void use(Function mapper) { + Function stored = adapt(mapper); + } + + void forward(Function mapper) { + use(mapper); + } + } + """, + """ + import java.util.function.Function; + + class Test { + Function adapt(Function mapper) { + return input -> mapper.apply(input); + } + + void use(Function mapper) { + Function stored = adapt(mapper); + } + + void forward(Function mapper) { + use(mapper); + } + } + """ + ) + ); + } }