[epilogue] Use more specific loggers at runtime, if available

epilogue-processor/src/main/java/edu/wpi/first/epilogue/processor/LoggableHandler.java

@@ -5,39 +5,151 @@
 package edu.wpi.first.epilogue.processor;
 
 import edu.wpi.first.epilogue.Logged;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
 import javax.annotation.processing.ProcessingEnvironment;
 import javax.lang.model.element.Element;
+import javax.lang.model.element.ElementKind;
+import javax.lang.model.element.TypeElement;
 import javax.lang.model.type.DeclaredType;
 import javax.lang.model.type.TypeMirror;
+import javax.lang.model.util.ElementScanner9;
+import javax.lang.model.util.Elements;
+import javax.lang.model.util.Types;
 
 /** Handles logging for types annotated with the {@link Logged @Logged} annotation. */
 public class LoggableHandler extends ElementHandler {
+  private final Types m_typeUtils;
+  private final Elements m_elementUtils;
+
   protected LoggableHandler(ProcessingEnvironment processingEnv) {
     super(processingEnv);
+    this.m_typeUtils = processingEnv.getTypeUtils();
+    this.m_elementUtils = processingEnv.getElementUtils();
   }
 
   @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;
+        || (dataType instanceof DeclaredType decl
+            && decl.asElement().getAnnotation(Logged.class) != null);
   }
 
   @Override
   public String logInvocation(Element element) {
     TypeMirror dataType = dataType(element);
-    var reflectedType =
-        m_processingEnv
-            .getElementUtils()
-            .getTypeElement(m_processingEnv.getTypeUtils().erasure(dataType).toString());
-
-    return "Epilogue."
-        + StringUtils.loggerFieldName(reflectedType)
-        + ".tryUpdate(dataLogger.getSubLogger(\""
-        + loggedName(element)
-        + "\"), "
-        + elementAccess(element)
-        + ", Epilogue.getConfig().errorHandler)";
+    TypeElement reflectedType =
+        (TypeElement) m_processingEnv.getTypeUtils().asElement(m_typeUtils.erasure(dataType));
+
+    List<TypeMirror> subtypes = getLoggedSubtypes(dataType);
+
+    return subtypes.isEmpty()
+        ? generateSimpleLogInvocation(reflectedType, element)
+        : generateConditionalLogInvocation(reflectedType, element, subtypes);
+  }
+
+  private List<TypeMirror> getLoggedSubtypes(TypeMirror type) {
+    List<TypeMirror> loggedSubtypes = new ArrayList<>();
+    for (TypeMirror subtype : getAllSubtypes(type)) {
+      Element subtypeElement = m_typeUtils.asElement(subtype);
+      if (subtypeElement != null && subtypeElement.getAnnotation(Logged.class) != null) {
+        System.out.println("Detected logged subtype: " + subtypeElement); // Debugging line
+        loggedSubtypes.add(subtype);
+      }
+    }
+    return loggedSubtypes;
+  }
+
+  private List<TypeMirror> getAllSubtypes(TypeMirror type) {
+    List<TypeMirror> allSubtypes = new ArrayList<>();
+    Set<TypeElement> allTypes = findAllTypes();
+
+    for (TypeElement potentialSubtype : allTypes) {
+      TypeMirror potentialSubtypeMirror = potentialSubtype.asType();
+      if (m_typeUtils.isSubtype(potentialSubtypeMirror, type)
+          && !potentialSubtypeMirror.equals(type)) {
+        allSubtypes.add(potentialSubtypeMirror);
+      }
+    }
+
+    return allSubtypes;
+  }
+
+  private Set<TypeElement> findAllTypes() {
+    Set<TypeElement> allTypes = new HashSet<>();
+
+    // Scanner to collect all TypeElements from root elements
+    ElementScanner9<Void, Void> scanner =
+        new ElementScanner9<>() {
+          @Override
+          public Void visitType(TypeElement e, Void p) {
+            allTypes.add(e);
+            return super.visitType(e, p);
+          }
+        };
+
+    // Iterate through all elements and scan
+    for (Element rootElement : m_elementUtils.getAllModuleElements()) {
+      if (rootElement.getKind() == ElementKind.PACKAGE) {
+        rootElement.accept(scanner, null);
+      }
+    }
+
+    return allTypes;
+  }
+
+  @SuppressWarnings("checkstyle:LineLength")
+  private String generateSimpleLogInvocation(TypeElement reflectedType, Element element) {
+    return String.format(
+        "Epilogue.%s.tryUpdate(dataLogger.getSubLogger(\"%s\"), %s, Epilogue.getConfig().errorHandler)",
+        StringUtils.loggerFieldName(reflectedType), loggedName(element), elementAccess(element));
+  }
+
+  @SuppressWarnings("PMD.ConsecutiveLiteralAppends")
+  private String generateConditionalLogInvocation(
+      TypeElement reflectedType, Element element, List<TypeMirror> subtypes) {
+    StringBuilder builder = new StringBuilder(256);
+
+    builder.append("if (Epilogue.shouldLog(Logged.Importance.DEBUG)) {\n");
+    builder.append("  var obj = ").append(elementAccess(element)).append(";\n");
+    builder
+        .append("  var logger = dataLogger.getSubLogger(\"")
+        .append(loggedName(element))
+        .append("\");\n");
+
+    for (int i = 0; i < subtypes.size(); i++) {
+      TypeMirror subtype = subtypes.get(i);
+      TypeElement subtypeElement = (TypeElement) m_typeUtils.asElement(subtype);
+      String typeName = subtypeElement.getQualifiedName().toString();
+
+      // Generate the conditional block for each subtype
+      if (i == 0) {
+        builder.append("  if (obj instanceof ").append(typeName).append(") {\n");
+      } else {
+        builder.append("  else if (obj instanceof ").append(typeName).append(") {\n");
+      }
+
+      builder
+          .append("    Epilogue.")
+          .append(StringUtils.loggerFieldName(subtypeElement))
+          .append(".tryUpdate(logger, (")
+          .append(typeName)
+          .append(") obj, Epilogue.getConfig().errorHandler);\n");
+      builder.append("  }\n");
+    }
+
+    // Fallback to the parent type handling if no subtype matches
+    builder.append("  else {\n");
+    builder
+        .append("    Epilogue.")
+        .append(StringUtils.loggerFieldName(reflectedType))
+        .append(".tryUpdate(logger, obj, Epilogue.getConfig().errorHandler);\n");
+    builder.append("  }\n");
+
+    builder.append("}");
+    return builder.toString();
   }
 }

epilogue-processor/src/test/java/edu/wpi/first/epilogue/processor/LoggableHandlerTest.java

@@ -0,0 +1,141 @@
+package edu.wpi.first.epilogue.processor;
+
+import static com.google.testing.compile.CompilationSubject.assertThat;
+import static com.google.testing.compile.Compiler.javac;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import com.google.testing.compile.Compilation;
+import com.google.testing.compile.JavaFileObjects;
+import java.io.IOException;
+import org.junit.jupiter.api.Test;
+
+@SuppressWarnings("PMD.ConsecutiveLiteralAppends")
+class LoggableHandlerTest {
+
+  @Test
+  void noSubtypes() {
+    String source =
+        """
+                package edu.wpi.first.epilogue;
+
+                @Logged
+                class Example {}
+                """;
+
+    String expected =
+        """
+                if (Epilogue.shouldLog(Logged.Importance.DEBUG)) {
+                  Epilogue.exampleLogger.tryUpdate(dataLogger.getSubLogger("example"), object, Epilogue.getConfig().errorHandler);
+                }
+                """;
+
+    assertLogInvocation(source, expected);
+  }
+
+  @Test
+  void withSubtypes() {
+    String source =
+        """
+                package edu.wpi.first.epilogue;
+
+                @Logged
+                class Parent {}
+
+                @Logged
+                class Child extends Parent {}
+
+                @Logged
+                class Example {
+                  Parent exampleField;
+                }
+                """;
+
+    String expected =
+        """
+                if (Epilogue.shouldLog(Logged.Importance.DEBUG)) {
+                  var obj = object.exampleField;
+                  var logger = dataLogger.getSubLogger("exampleField");
+                  if (obj instanceof edu.wpi.first.epilogue.Child) {
+                    Epilogue.childLogger.tryUpdate(logger, (edu.wpi.first.epilogue.Child) obj, Epilogue.getConfig().errorHandler);
+                  } else {
+                    Epilogue.parentLogger.tryUpdate(logger, obj, Epilogue.getConfig().errorHandler);
+                  }
+                }
+                """;
+
+    // Additional debug information
+    System.out.println("Expected Log Invocation:\n" + expected);
+
+    assertLogInvocation(source, expected);
+  }
+
+  private void assertLogInvocation(String classContent, String expectedLogInvocation) {
+    Compilation compilation =
+        javac()
+            .withProcessors(new AnnotationProcessor())
+            .compile(
+                JavaFileObjects.forSourceString("edu.wpi.first.epilogue.Example", classContent));
+
+    assertThat(compilation).succeededWithoutWarnings();
+
+    var generatedFile =
+        compilation.generatedSourceFiles().stream()
+            .filter(jfo -> jfo.getName().contains("ExampleLogger"))
+            .findFirst()
+            .orElseThrow(() -> new IllegalStateException("ExampleLogger file was not generated!"));
+
+    try {
+      var content = generatedFile.getCharContent(false).toString();
+      System.out.println("Full generated content:\n" + content);
+
+      String generatedLogInvocation = extractLogInvocation(content);
+      assertEquals(
+          expectedLogInvocation.strip().replace("\r\n", "\n"),
+          generatedLogInvocation.strip().replace("\r\n", "\n"));
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private String extractLogInvocation(String loggerClassContent) {
+    System.out.println("Generated Logger Class Content:\n" + loggerClassContent);
+
+    // Find the start index of the main 'if' condition
+    int startIndex = loggerClassContent.indexOf("if (Epilogue.shouldLog(Logged.Importance.DEBUG))");
+    if (startIndex == -1) {
+      throw new IllegalStateException(
+          "Failed to locate log invocation in the generated logger class.");
+    }
+
+    // Initialize brace tracking from the start of the block
+    int braceCount = 0;
+    int endIndex = startIndex;
+
+    // Traverse the content starting from the found 'if' block
+    boolean startedTracking = false;
+
+    while (endIndex < loggerClassContent.length()) {
+      char currentChar = loggerClassContent.charAt(endIndex);
+
+      if (currentChar == '{') {
+        braceCount++;
+        startedTracking = true; // Start tracking when the first '{' is encountered
+      } else if (currentChar == '}') {
+        braceCount--;
+        // Stop when the block fully closes
+        if (braceCount == 0 && startedTracking) {
+          endIndex++; // Move one step past the closing brace
+          break;
+        }
+      }
+      endIndex++;
+    }
+
+    if (braceCount != 0) {
+      throw new IllegalStateException("Failed to identify the complete log invocation block.");
+    }
+
+    // Extract the substring containing the complete log invocation block
+    return loggerClassContent.substring(startIndex, endIndex).strip();
+  }
+}