private static @Nullable Problem doCheckLambda(@NotNull PsiLambdaExpression expression) {
    final PsiElement body = expression.getBody();
    if (body == null) return null;

    Predicate<PsiMethodCallExpression> checkBody = call -> {
      final PsiElement callParent = call.getParent();

      if (!(body instanceof PsiCodeBlock)) {
        return callParent instanceof PsiStatement || callParent instanceof PsiLocalVariable || expression.isValueCompatible();
      }

      if (callParent instanceof PsiReturnStatement) {
        return true;
      }
      
      PsiStatement[] statements = ((PsiCodeBlock)body).getStatements();
      if (statements.length == 1) {
        return statements[0] instanceof PsiReturnStatement && expression.isValueCompatible();
      }

      final PsiReturnStatement[] returnStatements = PsiUtil.findReturnStatements((PsiCodeBlock)body);
      if (returnStatements.length > 1) {
        return false;
      }

      if (returnStatements.length == 1) {
        if (!(ArrayUtil.getLastElement(statements) instanceof PsiReturnStatement)) {
          return false;
        }
      }

      return CodeBlockSurrounder.canSurround(call);
    };
    Predicate<PsiMethodCallExpression> checkWrites = call ->
      !ContainerUtil.exists(expression.getParameterList().getParameters(), parameter -> VariableAccessUtils.variableIsAssigned(parameter, body));

    return doCheckMethodCallOnFunctionalExpression(expression, checkBody.and(checkWrites));
  }
private static @Nullable Problem doCheckLambda(@NotNull PsiLambdaExpression expression)