From baa413b1ed86ac4c3f906730668ab46c6db60421 Mon Sep 17 00:00:00 2001
From: Knut Wannheden <knut@moderne.io>
Date: Wed, 8 Nov 2023 22:35:37 +0100
Subject: [PATCH] Don't stop iteration on constructor calls, only anonymous
 subclasses

The lambda support is basically already there. What was missing for the existing `LambdaTest` test case was that the iteration in `LocalVariableDecl#findNearestParentCallable()` stopped on any constructor invocation, rather than only when iterating out of the body of an anonymous subclasses.

Fixes: #31
---
 .../java/org/openrewrite/analysis/InvocationMatcher.java  | 2 +-
 .../org/openrewrite/analysis/trait/variable/Field.java    | 3 ++-
 .../analysis/trait/variable/LocalVariableDecl.java        | 8 +++++++-
 .../openrewrite/analysis/trait/variable/LambdaTest.java   | 1 -
 4 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/main/java/org/openrewrite/analysis/InvocationMatcher.java b/src/main/java/org/openrewrite/analysis/InvocationMatcher.java
index 4174898e9..e5446b4d5 100644
--- a/src/main/java/org/openrewrite/analysis/InvocationMatcher.java
+++ b/src/main/java/org/openrewrite/analysis/InvocationMatcher.java
@@ -164,7 +164,7 @@ private static Optional<MethodCall> nearestMethodCall(Cursor cursor) {
         }
 
         private static boolean asExpression(Cursor cursor, Predicate<Expression> expressionPredicate) {
-            return cursor.getValue() instanceof Expression && expressionPredicate.test((Expression) cursor.getValue());
+            return cursor.getValue() instanceof Expression && expressionPredicate.test(cursor.getValue());
         }
     }
 }
diff --git a/src/main/java/org/openrewrite/analysis/trait/variable/Field.java b/src/main/java/org/openrewrite/analysis/trait/variable/Field.java
index 4699593f5..295576a54 100644
--- a/src/main/java/org/openrewrite/analysis/trait/variable/Field.java
+++ b/src/main/java/org/openrewrite/analysis/trait/variable/Field.java
@@ -55,7 +55,8 @@ public Validation<TraitErrors, Field> viewOf(Cursor cursor) {
                 Cursor maybeVariableDecl = cursor.getParentTreeCursor();
                 Cursor maybeBlock = maybeVariableDecl.getParentTreeCursor();
                 Cursor maybeClassDecl = maybeBlock.getParentTreeCursor();
-                if (maybeClassDecl.getValue() instanceof J.ClassDeclaration || maybeClassDecl.getValue() instanceof J.NewClass) {
+                if (maybeClassDecl.getValue() instanceof J.ClassDeclaration ||
+                    maybeClassDecl.getValue() instanceof J.NewClass && ((J.NewClass) maybeClassDecl.getValue()).getBody() != null) {
                     return Validation.success(new FieldFromCursor(cursor, cursor.getValue(), maybeVariableDecl.getValue(), maybeBlock));
                 }
                 return TraitErrors.invalidTraitCreationError("Field must be declared in a class, interface, or anonymous class");
diff --git a/src/main/java/org/openrewrite/analysis/trait/variable/LocalVariableDecl.java b/src/main/java/org/openrewrite/analysis/trait/variable/LocalVariableDecl.java
index 97d898a78..30a20b712 100644
--- a/src/main/java/org/openrewrite/analysis/trait/variable/LocalVariableDecl.java
+++ b/src/main/java/org/openrewrite/analysis/trait/variable/LocalVariableDecl.java
@@ -105,7 +105,7 @@ static Validation<TraitErrors, Callable> findNearestParentCallable(Cursor cursor
         while (path.hasNext()) {
             Cursor c = path.next();
             Tree t = c.getValue();
-            if (t instanceof J.ClassDeclaration || t instanceof J.NewClass) {
+            if (t instanceof J.ClassDeclaration || t instanceof J.NewClass && previous != null && ((J.NewClass) t).getBody() == previous.getValue()) {
                 break;
             }
             if (t instanceof J.Block && J.Block.isStaticOrInitBlock(c)) {
@@ -122,6 +122,12 @@ static Validation<TraitErrors, Callable> findNearestParentCallable(Cursor cursor
                 if (m.getParameters().contains(previous.<Statement>getValue())) {
                     break;
                 }
+            } else if (t instanceof J.Lambda) {
+                J.Lambda l = (J.Lambda) t;
+                assert previous != null : "previous should not be null";
+                if (l.getParameters().getParameters().contains(previous.<Statement>getValue())) {
+                    break;
+                }
             }
             Validation<TraitErrors, Method> v = Method.viewOf(c);
             if (v.isSuccess()) {
diff --git a/src/test/java/org/openrewrite/analysis/trait/variable/LambdaTest.java b/src/test/java/org/openrewrite/analysis/trait/variable/LambdaTest.java
index ce23a52ae..95549a985 100644
--- a/src/test/java/org/openrewrite/analysis/trait/variable/LambdaTest.java
+++ b/src/test/java/org/openrewrite/analysis/trait/variable/LambdaTest.java
@@ -46,7 +46,6 @@ public J preVisit(J tree, ExecutionContext executionContext) {
     }
 
     @Test
-    @Disabled("https://github.com/openrewrite/rewrite-analysis/issues/31")
     void lambdaException() {
         //language=java
         rewriteRun(