Add enum array logging

Mechanical-Advantage · Nov 1, 2024 · e004975 · e004975
1 parent 2cf3242
commit e004975
Show file tree

Hide file tree

Showing 2 changed files with 100 additions and 2 deletions.
diff --git a/junction/core/src/org/littletonrobotics/junction/LogTable.java b/junction/core/src/org/littletonrobotics/junction/LogTable.java
@@ -431,6 +431,33 @@ public <E extends Enum<E>> void put(String key, E value) {
     put(key, new LogValue(value.name(), null));
   }
 
+  /**
+   * Writes a new enum array value to the table. Skipped if the key already exists
+   * as a different type.
+   */
+  public <E extends Enum<E>> void put(String key, E[] value) {
+    if (value == null)
+      return;
+    String[] stringValues = new String[value.length];
+    for (int i = 0; i < value.length; i++) {
+      stringValues[i] = value[i].name();
+    }
+    put(key, new LogValue(stringValues, null));
+  }
+
+  /**
+   * Writes a new 2D enum array value to the table. Skipped if the key already
+   * exists as a different type.
+   */
+  public <E extends Enum<E>> void put(String key, E[][] value) {
+    if (value == null)
+      return;
+    put(key + "/length", value.length);
+    for (int i = 0; i < value.length; i++) {
+      put(key + "/" + Integer.toString(i), value[i]);
+    }
+  }
+
   /**
    * Writes a new Measure value to the table. Skipped if the key already exists as
    * a different type.
@@ -886,6 +913,39 @@ public <E extends Enum<E>> E get(String key, E defaultValue) {
     }
   }
 
+  /** Reads an enum array value from the table. */
+  public <E extends Enum<E>> E[] get(String key, E[] defaultValue) {
+    if (data.containsKey(prefix + key)) {
+      String[] names = get(key).getStringArray(null);
+      if (names == null)
+        return defaultValue;
+      Class<? extends Enum> enumClass = (Class<? extends Enum>) defaultValue.getClass().getComponentType();
+      E[] values = (E[]) Array.newInstance(enumClass, names.length);
+      for (int i = 0; i < names.length; i++) {
+        values[i] = (E) Enum.valueOf(enumClass, names[i]);
+      }
+      return values;
+    } else {
+      return defaultValue;
+    }
+  }
+
+  /** Reads a 2D enum array value from the table. */
+  public <E extends Enum<E>> E[][] get(String key, E[][] defaultValue) {
+    if (data.containsKey(prefix + key + "/length")) {
+      int length = get(key + "/length", 0);
+      E[][] value = (E[][]) Array.newInstance(defaultValue.getClass().getComponentType(), length);
+      for (int i = 0; i < length; i++) {
+        E[] defaultItemValue = (E[]) Array.newInstance(defaultValue.getClass().getComponentType().getComponentType(),
+            0);
+        value[i] = get(key + "/" + Integer.toString(i), defaultItemValue);
+      }
+      return value;
+    } else {
+      return defaultValue;
+    }
+  }
+
   /** Reads a Measure value from the table. */
   public <U extends Unit, M extends Measure<U>> M get(String key, M defaultValue) {
     if (data.containsKey(prefix + key)) {
@@ -942,7 +1002,7 @@ public <T> T[][] get(String key, Struct<T> struct, T[][] defaultValue) {
     if (data.containsKey(prefix + key + "/length")) {
       int length = get(key + "/length", 0);
       T[][] value = (T[][]) Array.newInstance(defaultValue.getClass().getComponentType(), length);
-      for (int i = 0; i < value.length; i++) {
+      for (int i = 0; i < length; i++) {
         T[] defaultItemValue = (T[]) Array.newInstance(defaultValue.getClass().getComponentType().getComponentType(),
             0);
         value[i] = get(key + "/" + Integer.toString(i), struct, defaultItemValue);
@@ -1014,7 +1074,7 @@ public <T extends StructSerializable> T[][] get(String key, T[][] defaultValue)
     if (data.containsKey(prefix + key + "/length")) {
       int length = get(key + "/length", 0);
       T[][] value = (T[][]) Array.newInstance(defaultValue.getClass().getComponentType(), length);
-      for (int i = 0; i < value.length; i++) {
+      for (int i = 0; i < length; i++) {
         T[] defaultItemValue = (T[]) Array.newInstance(defaultValue.getClass().getComponentType().getComponentType(),
             0);
         value[i] = get(key + "/" + Integer.toString(i), defaultItemValue);

diff --git a/junction/core/src/org/littletonrobotics/junction/Logger.java b/junction/core/src/org/littletonrobotics/junction/Logger.java
@@ -886,6 +886,44 @@ public static <E extends Enum<E>> void recordOutput(String key, E value) {
     }
   }
 
+  /**
+   * Records a single output field for easy access when viewing the log. On the
+   * simulator, use this method to record extra data based on the original inputs.
+   * 
+   * <p>
+   * This method is <b>not thread-safe</b> and should only be called from the
+   * main thread. See the "Common Issues" page in the documentation for more
+   * details.
+   * 
+   * @param key   The name of the field to record. It will be stored under
+   *              "/RealOutputs" or "/ReplayOutputs"
+   * @param value The value of the field.
+   */
+  public static <E extends Enum<E>> void recordOutput(String key, E[] value) {
+    if (running) {
+      outputTable.put(key, value);
+    }
+  }
+
+  /**
+   * Records a single output field for easy access when viewing the log. On the
+   * simulator, use this method to record extra data based on the original inputs.
+   * 
+   * <p>
+   * This method is <b>not thread-safe</b> and should only be called from the
+   * main thread. See the "Common Issues" page in the documentation for more
+   * details.
+   * 
+   * @param key   The name of the field to record. It will be stored under
+   *              "/RealOutputs" or "/ReplayOutputs"
+   * @param value The value of the field.
+   */
+  public static <E extends Enum<E>> void recordOutput(String key, E[][] value) {
+    if (running) {
+      outputTable.put(key, value);
+    }
+  }
+
   /**
    * Records a single output field for easy access when viewing the log. On the
    * simulator, use this method to record extra data based on the original inputs.