Broaden criteria for a loggable type

wpilibsuite · Nov 28, 2024 · b7d58a9 · b7d58a9
1 parent 2253477
commit b7d58a9
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 14 deletions.
diff --git a/epilogue-processor/src/main/java/edu/wpi/first/epilogue/processor/AnnotationProcessor.java b/epilogue-processor/src/main/java/edu/wpi/first/epilogue/processor/AnnotationProcessor.java
@@ -14,6 +14,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import javax.annotation.processing.AbstractProcessor;
 import javax.annotation.processing.RoundEnvironment;
 import javax.annotation.processing.SupportedAnnotationTypes;
@@ -97,7 +99,7 @@ public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment
         List.of(
             new LoggableHandler(
                 processingEnv,
-                roundEnv.getRootElements()), // prioritize epilogue logging over Sendable
+                getLoggedTypes(roundEnv)), // prioritize epilogue logging over Sendable
             new ConfiguredLoggerHandler(
                 processingEnv, customLoggers), // then customized logging configs
             new ArrayHandler(processingEnv),
@@ -117,11 +119,38 @@ public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment
         .findAny()
         .ifPresent(
             epilogue -> {
-              processEpilogue(roundEnv);
+              processEpilogue(roundEnv, epilogue);
             });
 
     return false;
   }
+
+  /**
+   * Gets the set of all loggable types in the compilation unit. A type is considered loggable if
+   * it is directly annotated with {@code @Logged} or contains a field or method with a
+   * {@code @Logged} annotation.
+   *
+   * @param roundEnv the compilation round environment
+   *
+   * @return the set of all loggable types
+   */
+  private Set<TypeElement> getLoggedTypes(RoundEnvironment roundEnv) {
+    return Stream.concat(
+        // 1. All type elements (classes, interfaces, or enums) with the @Logged annotation
+        roundEnv.getRootElements()
+            .stream()
+            .filter(e -> e instanceof TypeElement)
+            .map(e -> (TypeElement) e)
+            .filter(t -> t.getAnnotation(Logged.class) != null),
+        // 2. All type elements containing a field or method with the @Logged annotation
+        roundEnv.getElementsAnnotatedWith(Logged.class)
+            .stream()
+            .filter(e -> e instanceof VariableElement || e instanceof ExecutableElement)
+            .map(e -> e.getEnclosingElement())
+            .filter(e -> e instanceof TypeElement)
+            .map(e -> (TypeElement) e)
+    ).collect(Collectors.toSet()); // Collect to a set to avoid duplicates
+  }
 
   private boolean validateFields(Set<? extends Element> annotatedElements) {
     var fields =
@@ -329,8 +358,8 @@ private Map<DeclaredType, DeclaredType> processCustomLoggers(
     return customLoggers;
   }
 
-  private void processEpilogue(RoundEnvironment roundEnv) {
-    var annotatedElements = roundEnv.getRootElements();
+  private void processEpilogue(RoundEnvironment roundEnv, TypeElement epilogueAnnotation) {
+    var annotatedElements = roundEnv.getElementsAnnotatedWith(epilogueAnnotation);
 
     List<String> loggerClassNames = new ArrayList<>();
     var mainRobotClasses = new ArrayList<TypeElement>();
@@ -347,11 +376,8 @@ private void processEpilogue(RoundEnvironment roundEnv) {
       return;
     }
 
-    var classes =
-        annotatedElements.stream()
-            .filter(e -> e instanceof TypeElement)
-            .map(e -> (TypeElement) e)
-            .toList();
+    var classes = getLoggedTypes(roundEnv);
+
     for (TypeElement clazz : classes) {
       try {
         warnOfNonLoggableElements(clazz);

diff --git a/epilogue-processor/src/main/java/edu/wpi/first/epilogue/processor/LoggableHandler.java b/epilogue-processor/src/main/java/edu/wpi/first/epilogue/processor/LoggableHandler.java
@@ -15,7 +15,6 @@
 import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.TypeElement;
 import javax.lang.model.element.VariableElement;
-import javax.lang.model.type.DeclaredType;
 import javax.lang.model.type.TypeKind;
 import javax.lang.model.type.TypeMirror;
 
@@ -35,10 +34,8 @@ protected LoggableHandler(
 
   @Override
   public boolean isLoggable(Element element) {
-    var dataType = dataType(element);
-    return dataType.getAnnotation(Logged.class) != null
-        || dataType instanceof DeclaredType decl
-            && decl.asElement().getAnnotation(Logged.class) != null;
+    return m_processingEnv.getTypeUtils().asElement(dataType(element)) instanceof TypeElement t
+        && m_loggedTypes.contains(t);
   }
 
   @Override

diff --git a/...gue-processor/src/test/java/edu/wpi/first/epilogue/processor/AnnotationProcessorTest.java b/...gue-processor/src/test/java/edu/wpi/first/epilogue/processor/AnnotationProcessorTest.java
@@ -1329,6 +1329,47 @@ public void update(DataLogger dataLogger, Example object) {
     assertLoggerGenerates(source, expectedRootLogger);
   }
 
+  @Test
+  void nestedImplicit() {
+    String source =
+        """
+        package edu.wpi.first.epilogue;
+        
+        class Implicit {
+          @Logged double x;
+        }
+        
+        class Example {
+          @Logged Implicit i;
+        }
+        """;
+
+    String expectedRootLogger =
+        """
+        package edu.wpi.first.epilogue;
+
+        import edu.wpi.first.epilogue.Logged;
+        import edu.wpi.first.epilogue.Epilogue;
+        import edu.wpi.first.epilogue.logging.ClassSpecificLogger;
+        import edu.wpi.first.epilogue.logging.DataLogger;
+
+        public class ExampleLogger extends ClassSpecificLogger<Example> {
+          public ExampleLogger() {
+            super(Example.class);
+          }
+
+          @Override
+          public void update(DataLogger dataLogger, Example object) {
+            if (Epilogue.shouldLog(Logged.Importance.DEBUG)) {
+              Epilogue.implicitLogger.tryUpdate(dataLogger.getSubLogger("i"), object.i, Epilogue.getConfig().errorHandler);
+            }
+          }
+        }
+        """;
+
+    assertLoggerGenerates(source, expectedRootLogger);
+  }
+
   @Test
   void customLogger() {
     String source =