[epilogue] Support logging of protobuf-serializable types (#8229)

For parity with struct-serializable types.

Change struct serialization to only apply to types with a public static final <type> struct field, instead of relying only on the marker interface (which is not always followed). Doing this allows fallthrough to the protobuf handler for types with dynamic structs but static protobuf serializers.

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

@@ -112,7 +112,8 @@ public class AnnotationProcessor extends AbstractProcessor {
             new MeasureHandler(processingEnv),
             new PrimitiveHandler(processingEnv),
             new SupplierHandler(processingEnv),
-            new StructHandler(processingEnv), // prioritize struct over sendable
+            new StructHandler(processingEnv), // prioritize struct over sendable and protobuf
+            new ProtobufHandler(processingEnv), // then protobuf
             new SendableHandler(processingEnv));
 
     m_epiloguerGenerator = new EpilogueGenerator(processingEnv, customLoggers);

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

@@ -0,0 +1,102 @@
+// Copyright (c) FIRST and other WPILib contributors.
+// Open Source Software; you can modify and/or share it under the terms of
+// the WPILib BSD license file in the root directory of this project.
+
+package edu.wpi.first.epilogue.processor;
+
+import java.util.Set;
+import javax.annotation.processing.ProcessingEnvironment;
+import javax.lang.model.element.Element;
+import javax.lang.model.element.Modifier;
+import javax.lang.model.element.TypeElement;
+import javax.lang.model.element.VariableElement;
+import javax.lang.model.type.TypeMirror;
+import javax.lang.model.util.Elements;
+import javax.lang.model.util.Types;
+
+/**
+ * Supports protobuf serializable types. Protobuf-serializable types are loggable if they have a
+ * public static final {@code proto} field of a type that inherits from {@code Protobuf}.
+ */
+public class ProtobufHandler extends ElementHandler {
+  private final TypeMirror m_serializable;
+  private final TypeElement m_protobufType;
+  private final Types m_typeUtils;
+  private final Elements m_elementUtils;
+
+  protected ProtobufHandler(ProcessingEnvironment processingEnv) {
+    super(processingEnv);
+
+    m_serializable =
+        processingEnv
+            .getElementUtils()
+            .getTypeElement("edu.wpi.first.util.protobuf.ProtobufSerializable")
+            .asType();
+    m_protobufType =
+        processingEnv.getElementUtils().getTypeElement("edu.wpi.first.util.protobuf.Protobuf");
+    m_typeUtils = processingEnv.getTypeUtils();
+    m_elementUtils = processingEnv.getElementUtils();
+  }
+
+  @Override
+  public boolean isLoggable(Element element) {
+    return isLoggableType(dataType(element));
+  }
+
+  /**
+   * Checks if a type is protobuf-serializable: implements the ProtobufSerializable marker interface
+   * and has a `public static final proto` field of a type that inherits from Protobuf with a
+   * compatible generic type bound.
+   *
+   * @param type The type to check
+   * @return true if the type is protobuf-serializable, false otherwise
+   */
+  public boolean isLoggableType(TypeMirror type) {
+    var serializableType = m_typeUtils.erasure(type);
+    var typeElement = m_elementUtils.getTypeElement(serializableType.toString());
+    if (typeElement == null) {
+      return false;
+    }
+
+    // eg `Protobuf<Rotation2d, ?>` instead of the raw `Protobuf` type. The message type doesn't
+    // really matter here; we can leave it as a wildcard.
+    var sharpProtobufType =
+        m_typeUtils.getDeclaredType(
+            m_protobufType,
+            typeElement.asType(), // the serializable type
+            m_typeUtils.getWildcardType(
+                m_elementUtils.getTypeElement("us.hebi.quickbuf.ProtoMessage").asType(), null));
+
+    boolean hasProto =
+        typeElement.getEnclosedElements().stream()
+            .filter(e -> e instanceof VariableElement)
+            .map(e -> (VariableElement) e)
+            .anyMatch(
+                field -> {
+                  var nameMatch = field.getSimpleName().contentEquals("proto");
+                  var modifiersMatch =
+                      field
+                          .getModifiers()
+                          .containsAll(Set.of(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL));
+                  var typeMatch =
+                      m_typeUtils.isAssignable(
+                          m_typeUtils.erasure(field.asType()), sharpProtobufType);
+                  return nameMatch && modifiersMatch && typeMatch;
+                });
+    return m_typeUtils.isAssignable(type, m_serializable) && hasProto;
+  }
+
+  public String protoAccess(TypeMirror serializableType) {
+    var className = m_typeUtils.erasure(serializableType).toString();
+    return className + ".proto";
+  }
+
+  @Override
+  public String logInvocation(Element element, TypeElement loggedClass) {
+    return "backend.log(\"%s\", %s, %s)"
+        .formatted(
+            loggedName(element),
+            elementAccess(element, loggedClass),
+            protoAccess(dataType(element)));
+  }
+}

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

@@ -4,15 +4,26 @@
 
 package edu.wpi.first.epilogue.processor;
 
+import java.util.Set;
 import javax.annotation.processing.ProcessingEnvironment;
 import javax.lang.model.element.Element;
+import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
+import javax.lang.model.element.VariableElement;
+import javax.lang.model.type.ArrayType;
 import javax.lang.model.type.TypeMirror;
+import javax.lang.model.util.Elements;
 import javax.lang.model.util.Types;
 
+/**
+ * Supports struct serializable types. Struct-serializable types are loggable if they have a public
+ * static final {@code struct} field of a type that inherits from {@code Struct}.
+ */
 public class StructHandler extends ElementHandler {
   private final TypeMirror m_serializable;
+  private final TypeElement m_structType;
   private final Types m_typeUtils;
+  private final Elements m_elementUtils;
 
   protected StructHandler(ProcessingEnvironment processingEnv) {
     super(processingEnv);
@@ -21,16 +32,57 @@ public class StructHandler extends ElementHandler {
             .getElementUtils()
             .getTypeElement("edu.wpi.first.util.struct.StructSerializable")
             .asType();
+    m_structType =
+        processingEnv.getElementUtils().getTypeElement("edu.wpi.first.util.struct.Struct");
     m_typeUtils = processingEnv.getTypeUtils();
+    m_elementUtils = processingEnv.getElementUtils();
   }
 
   @Override
   public boolean isLoggable(Element element) {
-    return m_typeUtils.isAssignable(dataType(element), m_serializable);
+    return isLoggableType(dataType(element));
   }
 
+  /**
+   * Checks if a type is struct-serializable: implements the StructSerializable marker interface and
+   * has a `public static final struct` field of a type that inherits from Struct with a compatible
+   * generic type bound.
+   *
+   * @param type The type to check
+   * @return true if the type is struct-serializable, false otherwise
+   */
   public boolean isLoggableType(TypeMirror type) {
-    return m_typeUtils.isAssignable(type, m_serializable);
+    TypeMirror serializableType;
+    if (type instanceof ArrayType arr) {
+      serializableType = arr.getComponentType();
+    } else {
+      serializableType = m_typeUtils.erasure(type);
+    }
+    var typeElement = m_elementUtils.getTypeElement(serializableType.toString());
+    if (typeElement == null) {
+      return false;
+    }
+
+    // eg `Struct<Rotation2d>` instead of the raw `Struct` type
+    var sharpStructType = m_typeUtils.getDeclaredType(m_structType, typeElement.asType());
+
+    boolean hasStruct =
+        typeElement.getEnclosedElements().stream()
+            .filter(e -> e instanceof VariableElement)
+            .map(e -> (VariableElement) e)
+            .anyMatch(
+                field -> {
+                  var nameMatch = field.getSimpleName().contentEquals("struct");
+                  var modifiersMatch =
+                      field
+                          .getModifiers()
+                          .containsAll(Set.of(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL));
+                  var typeMatch =
+                      m_typeUtils.isAssignable(
+                          m_typeUtils.erasure(field.asType()), sharpStructType);
+                  return nameMatch && modifiersMatch && typeMatch;
+                });
+    return m_typeUtils.isAssignable(type, m_serializable) && hasStruct;
   }
 
   public String structAccess(TypeMirror serializableType) {

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

@@ -1141,6 +1141,76 @@ class AnnotationProcessorTest {
     assertLoggerGenerates(source, expectedGeneratedSource);
   }
 
+  @Test
+  void protobuf() {
+    String source =
+        """
+      package edu.wpi.first.epilogue;
+
+      import edu.wpi.first.util.protobuf.Protobuf;
+      import edu.wpi.first.util.protobuf.ProtobufSerializable;
+      import java.util.List;
+      import us.hebi.quickbuf.*;
+
+      class ProtobufType implements ProtobufSerializable {
+        // Message type is necessary - Epilogue can't log with a wildcard message type in the proto
+        public static final Protobuf<ProtobufType, Message> proto = null; // value doesn't matter
+
+        static class Message extends ProtoMessage<Message> {
+          // Implement stubs for the abstract base class.
+          // This code never runs so actual implementations are unnecessary.
+          @Override
+          public Message copyFrom(Message other) { return null; }
+          @Override
+          public Message clear() { return null; }
+          @Override
+          public int computeSerializedSize() { return 0; }
+          @Override
+          public void writeTo(ProtoSink output) {}
+          @Override
+          public Message mergeFrom(ProtoSource input) { return null; }
+          @Override
+          public boolean equals(Object obj) { return false; }
+          @Override
+          public Message clone() { return null; }
+        }
+      }
+
+      @Logged
+      class Example {
+        ProtobufType x;          // Should be logged
+        ProtobufType[] arr1;     // Should not be logged
+        ProtobufType[][] arr2;   // Should not be logged
+        List<ProtobufType> list; // Should not be logged
+      }
+      """;
+
+    String expectedGeneratedSource =
+        """
+      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.EpilogueBackend;
+
+      public class ExampleLogger extends ClassSpecificLogger<Example> {
+        public ExampleLogger() {
+          super(Example.class);
+        }
+
+        @Override
+        public void update(EpilogueBackend backend, Example object) {
+          if (Epilogue.shouldLog(Logged.Importance.DEBUG)) {
+            backend.log("x", object.x, edu.wpi.first.epilogue.ProtobufType.proto);
+          }
+        }
+      }
+      """;
+
+    assertLoggerGenerates(source, expectedGeneratedSource);
+  }
+
   @Test
   void lists() {
     String source =