From 726c2cfbb96650890017535a3f5b94237cb4eef7 Mon Sep 17 00:00:00 2001
From: see-quick <maros.orsak159@gmail.com>
Date: Wed, 6 Nov 2024 15:52:24 +0100
Subject: [PATCH] update

Signed-off-by: see-quick <maros.orsak159@gmail.com>
---
 .../ExcludedAnnotationInterceptor.java        | 48 +++++++++++++++++--
 .../ExcludedAnnotationInterceptorTest.java    | 28 ++++++-----
 2 files changed, 61 insertions(+), 15 deletions(-)

diff --git a/pitest-entry/src/main/java/org/pitest/mutationtest/build/intercept/annotations/ExcludedAnnotationInterceptor.java b/pitest-entry/src/main/java/org/pitest/mutationtest/build/intercept/annotations/ExcludedAnnotationInterceptor.java
index 2d33ea16e..2992cba58 100644
--- a/pitest-entry/src/main/java/org/pitest/mutationtest/build/intercept/annotations/ExcludedAnnotationInterceptor.java
+++ b/pitest-entry/src/main/java/org/pitest/mutationtest/build/intercept/annotations/ExcludedAnnotationInterceptor.java
@@ -14,8 +14,10 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 public class ExcludedAnnotationInterceptor implements MutationInterceptor {
 
@@ -38,18 +40,56 @@ public void begin(ClassTree clazz) {
     this.skipClass = clazz.annotations().stream()
         .anyMatch(avoidedAnnotation());
     if (!this.skipClass) {
-      final List<Predicate<MutationDetails>> methods = clazz.methods().stream()
+      // 1. Collect methods with avoided annotations
+      final List<MethodTree> avoidedMethods = clazz.methods().stream()
           .filter(hasAvoidedAnnotation())
+          .collect(Collectors.toList());
+
+      final Set<String> avoidedMethodNames = avoidedMethods.stream()
+          .map(method -> method.rawNode().name)
+          .collect(Collectors.toSet());
+
+      // 2. Collect lambda methods with being inside avoided methods
+      final List<MethodTree> lambdaMethods = clazz.methods().stream()
+          .filter(MethodTree::isGeneratedLambdaMethod)
+          .filter(lambdaMethod -> {
+            String lambdaName = lambdaMethod.rawNode().name; // e.g., lambda$fooWithLambdas$0
+            String enclosingMethodName = extractEnclosingMethodName(lambdaName);
+
+            return avoidedMethodNames.contains(enclosingMethodName);
+          })
+          .collect(Collectors.toList());
+
+      // 3. Merge the two lists into a single list and cast MethodTree to Predicate<MutationDetails>
+      final List<Predicate<MutationDetails>> mutationPredicates = Stream.concat(avoidedMethods.stream(), lambdaMethods.stream())
           .map(AnalysisFunctions.matchMutationsInMethod())
           .collect(Collectors.toList());
-      this.annotatedMethodMatcher = Prelude.or(methods);
+
+      this.annotatedMethodMatcher = Prelude.or(mutationPredicates);
+    }
+  }
+
+  /**
+   * TODO: maybe move to MethodTree class?? WDYT?
+   * Extracts the enclosing method name from a lambda method's name.
+   * Assumes lambda methods follow the naming convention: lambda$enclosingMethodName$number
+   *
+   * @param lambdaName The name of the lambda method (e.g., "lambda$fooWithLambdas$0")
+   * @return The name of the enclosing method (e.g., "fooWithLambdas")
+   */
+  private String extractEnclosingMethodName(String lambdaName) {
+    int firstDollar = lambdaName.indexOf('$');
+    int secondDollar = lambdaName.indexOf('$', firstDollar + 1);
+
+    if (firstDollar != -1 && secondDollar != -1) {
+      return lambdaName.substring(firstDollar + 1, secondDollar);
     }
+    return lambdaName;
   }
 
   private Predicate<MethodTree> hasAvoidedAnnotation() {
     return methodTree ->
-        // count also lambda generated methods
-        methodTree.isGeneratedLambdaMethod() || methodTree.annotations().stream().anyMatch(avoidedAnnotation());
+        methodTree.annotations().stream().anyMatch(avoidedAnnotation());
   }
 
   private Predicate<AnnotationNode> avoidedAnnotation() {
diff --git a/pitest-entry/src/test/java/org/pitest/mutationtest/build/intercept/annotations/ExcludedAnnotationInterceptorTest.java b/pitest-entry/src/test/java/org/pitest/mutationtest/build/intercept/annotations/ExcludedAnnotationInterceptorTest.java
index 1dbdbdd11..3d449fecb 100644
--- a/pitest-entry/src/test/java/org/pitest/mutationtest/build/intercept/annotations/ExcludedAnnotationInterceptorTest.java
+++ b/pitest-entry/src/test/java/org/pitest/mutationtest/build/intercept/annotations/ExcludedAnnotationInterceptorTest.java
@@ -71,8 +71,11 @@ public void shouldFilterMethodsWithGeneratedAnnotation() {
   public void shouldFilterMethodsWithGeneratedAnnotationAndLambdasInside() {
     final List<MutationDetails> mutations = this.mutator.findMutations(ClassName.fromClass(ClassAnnotatedWithGeneratedWithLambdas.class));
     final Collection<MutationDetails> actual = runWithTestee(mutations, ClassAnnotatedWithGeneratedWithLambdas.class);
-    assertThat(actual).hasSize(1);
-    assertThat(actual.iterator().next().getId().getLocation().getMethodName()).isEqualTo("bar");
+    assertThat(actual).hasSize(3);
+
+    for (MutationDetails mutationDetails : actual) {
+       assertThat(mutationDetails.getId().getLocation().getMethodName()).isIn("barWithLambdas", "lambda$barWithLambdas$2", "lambda$barWithLambdas$3");
+    }
   }
 
   @Test
@@ -153,15 +156,6 @@ public void bar() {
 
 class ClassAnnotatedWithGeneratedWithLambdas {
 
-  @TestGeneratedAnnotation
-  public void foo() {
-      System.out.println("don't mutate me");
-  }
-
-  public void bar() {
-    System.out.println("mutate me");
-  }
-
   @TestGeneratedAnnotation
   public void fooWithLambdas() {
     System.out.println("don't mutate me");
@@ -174,6 +168,18 @@ public void fooWithLambdas() {
         };
     };
   }
+
+  public void barWithLambdas() {
+    System.out.println("mutate me");
+
+    Runnable runnable = () -> {
+      System.out.println("mutate me also in lambdas");
+
+      Runnable anotherOne = () -> {
+        System.out.println("mutate me also recursive lambdas");
+      };
+    };
+  }
 }