From 790571b6cfaf8e0adc20bc940fdf5e2c019bfcef Mon Sep 17 00:00:00 2001
From: Kevin Clark <kevin.clark@gmail.com>
Date: Sat, 13 Jan 2024 09:23:52 -0800
Subject: [PATCH] Refactor LogTable struct writes

---
 .../littletonrobotics/junction/LogTable.java  | 45 +++++++++++--------
 .../junction/LogTableTest.java                | 38 +++++++++-------
 2 files changed, 50 insertions(+), 33 deletions(-)

diff --git a/junction/core/src/main/org/littletonrobotics/junction/LogTable.java b/junction/core/src/main/org/littletonrobotics/junction/LogTable.java
index 65c05296..7b5ba1d9 100644
--- a/junction/core/src/main/org/littletonrobotics/junction/LogTable.java
+++ b/junction/core/src/main/org/littletonrobotics/junction/LogTable.java
@@ -338,15 +338,7 @@ private void addStructSchema(Struct<?> struct, Set<String> seen) {
   @SuppressWarnings("unchecked")
   public <T> void put(String key, Struct<T> struct, T value) {
     if (writeAllowed(key, LoggableType.Raw)) {
-      addStructSchema(struct, new HashSet<>());
-      if (!structBuffers.containsKey(struct.getTypeString())) {
-        structBuffers.put(struct.getTypeString(), StructBuffer.create(struct));
-      }
-      StructBuffer<T> buffer = (StructBuffer<T>) structBuffers.get(struct.getTypeString());
-      ByteBuffer bb = buffer.write(value);
-      byte[] array = new byte[bb.position()];
-      bb.position(0);
-      bb.get(array);
+      byte[] array = pack(struct, value);
       put(key, new LogValue(array, struct.getTypeString()));
     }
   }
@@ -358,19 +350,36 @@ public <T> void put(String key, Struct<T> struct, T value) {
   @SuppressWarnings("unchecked")
   public <T> void put(String key, Struct<T> struct, T... value) {
     if (writeAllowed(key, LoggableType.Raw)) {
-      addStructSchema(struct, new HashSet<>());
-      if (!structBuffers.containsKey(struct.getTypeString())) {
-        structBuffers.put(struct.getTypeString(), StructBuffer.create(struct));
-      }
-      StructBuffer<T> buffer = (StructBuffer<T>) structBuffers.get(struct.getTypeString());
-      ByteBuffer bb = buffer.writeArray(value);
-      byte[] array = new byte[bb.position()];
-      bb.position(0);
-      bb.get(array);
+      byte[] array = pack(struct, value);
       put(key, new LogValue(array, struct.getTypeString() + "[]"));
     }
   }
 
+  private <T> StructBuffer<T> bufferForStruct(Struct<T> struct) {
+    addStructSchema(struct, new HashSet<>());
+    if (!structBuffers.containsKey(struct.getTypeString())) {
+      structBuffers.put(struct.getTypeString(), StructBuffer.create(struct));
+    }
+   return (StructBuffer<T>) structBuffers.get(struct.getTypeString());
+  }
+
+  private <T> byte[] pack(Struct<T> struct, T[] value) {
+    StructBuffer<T> buffer = bufferForStruct(struct);
+    return copyBytes(buffer.writeArray(value));
+  }
+
+  private <T> byte[] pack(Struct<T> struct, T value) {
+    StructBuffer<T> buffer = bufferForStruct(struct);
+    return copyBytes(buffer.write(value));
+  }
+
+  private byte[] copyBytes(ByteBuffer bb) {
+    byte[] array = new byte[bb.position()];
+    bb.position(0);
+    bb.get(array);
+    return array;
+  }
+
   /**
    * Writes a new protobuf value to the table. Skipped if the key already exists
    * as a different type.
diff --git a/junction/core/src/test/org/littletonrobotics/junction/LogTableTest.java b/junction/core/src/test/org/littletonrobotics/junction/LogTableTest.java
index c3e2d053..95c64aba 100644
--- a/junction/core/src/test/org/littletonrobotics/junction/LogTableTest.java
+++ b/junction/core/src/test/org/littletonrobotics/junction/LogTableTest.java
@@ -1,53 +1,62 @@
 package org.littletonrobotics.junction;
 
 import edu.wpi.first.math.geometry.Rotation2d;
-import edu.wpi.first.math.geometry.proto.Rotation2dProto;
 import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
-import org.littletonrobotics.junction.LogTable.LogValue;
 
 public class LogTableTest {
+    private LogTable table;
+
+    @BeforeEach
+    void setup() {
+        table = new LogTable(0);
+    }
+
     @Test
     void supportsIntegers() {
-        LogTable table = new LogTable(0);
         table.put("int", 5);
-        LogValue val = table.get("int");
-        Assertions.assertEquals(5, val.getInteger());
+        Assertions.assertEquals(5, table.get("int").getInteger());
     }
 
     @Test
     void supportsIntArrays() {
         long[] expected = { 1, 2, 3 };
 
-        LogTable table = new LogTable(0);
         table.put("int-array", expected);
 
-        LogValue val = table.get("int-array");
-        Assertions.assertArrayEquals(expected, val.getIntegerArray());
+        Assertions.assertArrayEquals(expected, table.get("int-array").getIntegerArray());
     }
 
     @Test
     void supportsProtobufSerialization() {
         Rotation2d expected = new Rotation2d(1, 2);
 
-        LogTable table = new LogTable(0);
         // We're forcing protobuf based serialization so Struct based doesn't run
         table.put("rot", Rotation2d.proto, expected);
 
-        Rotation2d actual = table.get("rot", Rotation2d.proto, new Rotation2d());
-        Assertions.assertEquals(expected, actual);
+        Assertions.assertEquals(expected, table.get("rot", Rotation2d.proto, new Rotation2d()));
+    }
+
+    @Test
+    void protobufIsntWrittenIfKeyAlreadyInUse() {
+        Rotation2d rot = new Rotation2d(1, 2);
+
+        table.put("rot", 5);
+        table.put("rot", Rotation2d.proto, rot); // This should be skipped
+
+        Assertions.assertNotEquals(rot, table.get("rot", Rotation2d.proto, new Rotation2d()));
+        Assertions.assertEquals(5, table.get("rot").getInteger());
     }
 
     @Test
     void supportsStructSerialization() {
         Rotation2d expected = new Rotation2d(1, 2);
 
-        LogTable table = new LogTable(0);
         // We're forcing Struct based serialization so Protobuf based doesn't run
         table.put("rot", Rotation2d.struct, expected);
 
-        Rotation2d actual = table.get("rot", Rotation2d.struct, new Rotation2d());
-        Assertions.assertEquals(expected, actual);
+        Assertions.assertEquals(expected, table.get("rot", Rotation2d.struct, new Rotation2d()));
     }
 
     @Test
@@ -55,7 +64,6 @@ void supportsStructArraySerialization() {
         Rotation2d first = new Rotation2d(1, 2);
         Rotation2d second = new Rotation2d(3, 4);
 
-        LogTable table = new LogTable(0);
         // We're forcing Struct serialization
         table.put("rot", Rotation2d.struct, first, second);