Skip to content

Commit 4b4d1e8

Browse files
msridharmsridhar
authored
JSpecify: Apply annotations on type variables to lambdas in more cases (#1428)
Consider this example: ```java import org.jspecify.annotations.NullMarked; import org.jspecify.annotations.Nullable; import java.util.function.Function; @NullMarked class Test { static class Foo<T> { T doSomething(Function<@nullable T, T> f) { return f.apply(null); } } void test(Foo<String> foo) { String s = foo.doSomething( t -> { return "length: " + t.length(); }); } } ``` NullAway should produce an error at `t.length()` since `t` may be null. Before this PR, NullAway would not give an error, since for a `Foo<String> foo`, it treated the type of the formal parameter as `Function<String, String>`, without considering the `@Nullable` annotation in `Function<@nullable T, T>`. This PR closes that gap, both for passing lambdas as parameters and for assigning them to fields. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Improvements** * Enhanced nullability analysis for lambda expressions in generic type contexts. * Improved detection of null-dereference issues when lambdas use generic type variables. * Better nullability tracking for lambdas assigned to fields or used as method arguments and return values. <sub>✏️ Tip: You can customize this high-level summary in your review settings.</sub> <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent cdb7675 commit 4b4d1e8

2 files changed

Lines changed: 151 additions & 2 deletions

File tree

‎nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java‎

Lines changed: 45 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -419,6 +419,16 @@ private void reportInvalidOverridingMethodParamTypeError(
419419
*/
420420
private @Nullable Type getTreeType(Tree tree, VisitorState state) {
421421
tree = ASTHelpers.stripParentheses(tree);
422+
if (tree instanceof LambdaExpressionTree lambdaExpressionTree) {
423+
Type result = inferredLambdaTypes.get(lambdaExpressionTree);
424+
if (result == null) {
425+
result = ASTHelpers.getType(tree);
426+
}
427+
if (result != null && result.isRaw()) {
428+
return null;
429+
}
430+
return result;
431+
}
422432
if (tree instanceof NewClassTree
423433
&& ((NewClassTree) tree).getIdentifier() instanceof ParameterizedTypeTree paramTypedTree) {
424434
if (paramTypedTree.getTypeArguments().isEmpty()) {
@@ -586,6 +596,11 @@ public void checkTypeParameterNullnessForAssignability(Tree tree, VisitorState s
586596
if (rhsTree == null || rhsTree.getKind().equals(Tree.Kind.NULL_LITERAL)) {
587597
return;
588598
}
599+
if (!assignedToLocal
600+
&& rhsTree instanceof LambdaExpressionTree lambdaExpressionTree
601+
&& isAssignmentToField(tree)) {
602+
maybeStoreLambdaTypeFromTarget(lambdaExpressionTree, lhsType);
603+
}
589604
Type rhsType = getTreeType(rhsTree, state);
590605
if (rhsType != null) {
591606
if (isGenericCallNeedingInference(rhsTree)) {
@@ -601,6 +616,14 @@ public void checkTypeParameterNullnessForAssignability(Tree tree, VisitorState s
601616
}
602617

603618
private static boolean isAssignmentToLocalVariable(Tree tree) {
619+
return isAssignmentToKind(tree, ElementKind.LOCAL_VARIABLE);
620+
}
621+
622+
private static boolean isAssignmentToField(Tree tree) {
623+
return isAssignmentToKind(tree, ElementKind.FIELD);
624+
}
625+
626+
private static boolean isAssignmentToKind(Tree tree, ElementKind kind) {
604627
Symbol treeSymbol;
605628
if (tree instanceof VariableTree variableTree) {
606629
treeSymbol = ASTHelpers.getSymbol(variableTree);
@@ -609,7 +632,7 @@ private static boolean isAssignmentToLocalVariable(Tree tree) {
609632
} else {
610633
throw new RuntimeException("Unexpected tree type: " + tree.getKind());
611634
}
612-
return treeSymbol != null && treeSymbol.getKind().equals(ElementKind.LOCAL_VARIABLE);
635+
return treeSymbol != null && treeSymbol.getKind().equals(kind);
613636
}
614637

615638
private ConstraintSolver makeSolver(VisitorState state, NullAway analysis) {
@@ -1293,7 +1316,8 @@ public void compareGenericTypeParameterNullabilityForCall(
12931316
}
12941317

12951318
Type actualParameterType = null;
1296-
if ((currentActualParam instanceof LambdaExpressionTree)) {
1319+
if ((currentActualParam instanceof LambdaExpressionTree lambdaExpressionTree)) {
1320+
maybeStoreLambdaTypeFromTarget(lambdaExpressionTree, formalParameter);
12971321
Type lambdaInferredType = inferredLambdaTypes.get(currentActualParam);
12981322
if (lambdaInferredType != null) {
12991323
actualParameterType = lambdaInferredType;
@@ -2036,6 +2060,25 @@ public boolean isNullableAnnotated(Type type) {
20362060
return Nullness.hasNullableAnnotation(type.getAnnotationMirrors().stream(), config);
20372061
}
20382062

2063+
/**
2064+
* Store a lambda's target type with explicit type-variable nullability from the assignment
2065+
* context, when the lambda's type has not already been cached.
2066+
*
2067+
* <p>This is used to compensate for javac dropping annotations on type variables in lambda target
2068+
* types, so later checks use the correctly annotated functional interface type.
2069+
*/
2070+
private void maybeStoreLambdaTypeFromTarget(
2071+
LambdaExpressionTree lambdaExpressionTree, Type targetType) {
2072+
if (targetType.isRaw() || inferredLambdaTypes.containsKey(lambdaExpressionTree)) {
2073+
return;
2074+
}
2075+
Type lambdaTreeType = castToNonNull(ASTHelpers.getType(lambdaExpressionTree));
2076+
Type lambdaTypeWithTargetAnnotations =
2077+
TypeSubstitutionUtils.restoreExplicitNullabilityAnnotations(
2078+
targetType, lambdaTreeType, config, Collections.emptyMap());
2079+
inferredLambdaTypes.put(lambdaExpressionTree, lambdaTypeWithTargetAnnotations);
2080+
}
2081+
20392082
private static @Nullable Type syntheticNullableAnnotType;
20402083
private static @Nullable Type syntheticNonNullAnnotType;
20412084

‎nullaway/src/test/java/com/uber/nullaway/jspecify/GenericLambdaTests.java‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
package com.uber.nullaway.jspecify;
2+
3+
import com.google.errorprone.CompilationTestHelper;
4+
import com.uber.nullaway.NullAwayTestsBase;
5+
import com.uber.nullaway.generics.JSpecifyJavacConfig;
6+
import java.util.Arrays;
7+
import org.junit.Test;
8+
9+
public class GenericLambdaTests extends NullAwayTestsBase {
10+
11+
@Test
12+
public void lambdaArgumentUsesAnnotatedTypeVar() {
13+
makeHelper()
14+
.addSourceLines(
15+
"Test.java",
16+
"""
17+
package com.uber;
18+
import org.jspecify.annotations.NullMarked;
19+
import org.jspecify.annotations.Nullable;
20+
import java.util.function.Function;
21+
@NullMarked
22+
class Test {
23+
static class Foo<T extends @Nullable CharSequence> {
24+
T doSomething(Function<@Nullable T, T> f) {
25+
return f.apply(null);
26+
}
27+
}
28+
void test(Foo<String> foo) {
29+
foo.doSomething(
30+
t -> {
31+
// BUG: Diagnostic contains: dereferenced expression t is @Nullable
32+
return "length: " + t.length();
33+
});
34+
}
35+
}
36+
""")
37+
.doTest();
38+
}
39+
40+
@Test
41+
public void lambdaAssignedToFieldOrLocalUsesAnnotatedTypeVar() {
42+
makeHelper()
43+
.addSourceLines(
44+
"Test.java",
45+
"""
46+
package com.uber;
47+
import org.jspecify.annotations.NullMarked;
48+
import org.jspecify.annotations.Nullable;
49+
import java.util.function.Function;
50+
@NullMarked
51+
class Test {
52+
static class Box<T extends @Nullable CharSequence> {
53+
@Nullable Function<@Nullable T, T> field;
54+
void assignField(Box<String> box) {
55+
box.field = t -> {
56+
// BUG: Diagnostic contains: dereferenced expression t is @Nullable
57+
t.length();
58+
return t;
59+
};
60+
}
61+
void assignLocal() {
62+
Function<@Nullable T, T> local =
63+
t -> {
64+
// BUG: Diagnostic contains: dereferenced expression t is @Nullable
65+
t.length();
66+
return t;
67+
};
68+
}
69+
}
70+
}
71+
""")
72+
.doTest();
73+
}
74+
75+
@Test
76+
public void lambdaReturnedUsesAnnotatedTypeVar() {
77+
makeHelper()
78+
.addSourceLines(
79+
"Test.java",
80+
"""
81+
package com.uber;
82+
import org.jspecify.annotations.NullMarked;
83+
import org.jspecify.annotations.Nullable;
84+
import java.util.function.Function;
85+
@NullMarked
86+
class Test {
87+
static class Box<T extends @Nullable CharSequence> {
88+
Function<@Nullable T, T> make() {
89+
return t -> {
90+
// BUG: Diagnostic contains: dereferenced expression t is @Nullable
91+
t.length();
92+
return t;
93+
};
94+
}
95+
}
96+
}
97+
""")
98+
.doTest();
99+
}
100+
101+
private CompilationTestHelper makeHelper() {
102+
return makeTestHelperWithArgs(
103+
JSpecifyJavacConfig.withJSpecifyModeArgs(
104+
Arrays.asList("-XepOpt:NullAway:AnnotatedPackages=com.uber")));
105+
}
106+
}

0 commit comments

Comments
 (0)