How to make IntelliJ IDEA warn when null might be returned from lambda?

There are very useful code inspections in IntelliJ IDEA Constant conditions & exceptions and Return of 'null'. They show warnings when a method without @Nullable annotation returns null, a result of @Nullable method, @Nullable field or variable which might evaluate to null.

Unfortunately, the same not work with lambda expressions. IntelliJ IDEA just ignores possible returning of null from them. But I am pretty sure that there is a way to make Search Template for Structural Search Inspection which will find all possible returnings null from lambda expressions.

By this moment I've managed to make a template which finds references to @Nullable methods and template that wraps single-line lambdas with a method accepting @NotNull generic argument and returning it back, it doesn't change code semantics and shows warnings when lambda might return null.

Here are my templates:

Nullable methods:

@Nullable
$MethodType$ $Method$($ParameterType$ $Parameter$);

Nullable method references

$Qualifier$::$Method$

//Filters for $Method$:
reference = Nullable methods

Lambdas with unchecked result:

($Parameter$) -> $Statement$

//Filters for $Statement$:
text = !warnIfNull
type = !void|Optional|int|boolean|long|char|short|byte

//Replace template for $Statement$:
WarnUtils.warnIfNull($Statement$)

WarnUtils:

@UtilityClass
public class WarnUtils {
    public <T> @NotNull T warnIfNull(@NotNull T t) {
        return t;
    }
}

I don't know how to make a template for multiline lambdas and I don't like wrapping all single-line lambdas with an additional method.

I think script filters may help, but I have no idea what to do with __context__ to check whether lambda probably returns null.

Solution

This problem can be solved without additional code inspections just by using functional interface which has @NotNull annotation on its abstract method.

For example, if you need a supplier which mustn't return null you can use this interface:

@FunctionalInterface
public interface NotNullSupplier<T> {
    @NotNull
    T get();
}

If you try to return null from lambda expression of this type, Constant conditions & exceptions or Return of 'null' code inspection will detect it:

NotNullSupplier<Integer> supplier = () -> null; // warning - 'null' is returned by the method declared as @NotNull

I want to notice that using Supplier<@NotNull T> may lead to incorrect behavior of mentioned code inspections (it must be a bug which probably will be fixed later). Here is an example of such behavior:

Supplier<@NonNull Integer> supplier = () -> null; // no warnings
Integer result = supplier.get();
if(result != null) // warning - Condition 'result != null' is always 'true'
    result++;

Constant conditions & exceptions code inspection suggests to "Unwrap 'if' statement", following this advice in this case will lead to NullPointerException.

To avoid such mistakes you can add the following Search Templates to Structural Search Inspection:

//Search template:
Supplier<@NotNull $T$>
//Replace template:
NotNullSupplier<$T$>

//Search template:
Function<$T$, @NotNull $R$>
//Replace template:
NotNullFunction<$T$, $R$>

//Search template:
BiFunction<$T$, $U$, @NotNull $R$>
//Replace template:
NotNullBiFunction<$T$, $U$, $R$>