Add support for Mechanism2d

docs/CODE-STRUCTURE.md

@@ -137,6 +137,13 @@ Logger.getInstance().recordOutput("MySwerveModuleStates", stateA, stateB, stateC
 Logger.getInstance().recordOutput("MySwerveModuleStates", new SwerveModuleState[] { stateA, stateB, stateC, stateD });
 ```
 
+AdvantageKit can also log [`Mechanism2d`](https://docs.wpilib.org/en/stable/docs/software/dashboards/glass/mech2d-widget.html) objects as outputs, which can be viewed using AdvantageScope. Note that the call below only records the current state of the `Mechanism2d`, so it must be called every loop cycle after the object is updated.
+
+```java
+Mechanism2d mechanism = new Mechanism2d(3, 3);
+Logger.getInstance().recordOutput("MyMechanism", mechanism);
+```
+
 ## `@AutoLog` Annotation
 
 As of version 1.8, a new `@AutoLog` annotation was added. By adding this annotation to your inputs class, AdvantageKit will automatically generate implementations of `toLog` and `fromLog` for your inputs.

junction/core/src/org/littletonrobotics/junction/Logger.java

@@ -1,5 +1,6 @@
 package org.littletonrobotics.junction;
 
+import java.lang.reflect.InvocationTargetException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -25,6 +26,7 @@
 import edu.wpi.first.util.WPIUtilJNI;
 import edu.wpi.first.wpilibj.DriverStation;
 import edu.wpi.first.wpilibj.RobotBase;
+import edu.wpi.first.wpilibj.smartdashboard.Mechanism2d;
 
 /** Central class for recording and replaying log data. */
 public class Logger {
@@ -528,4 +530,27 @@ public void recordOutput(String key, SwerveModuleState... value) {
     }
     recordOutput(key, data);
   }
+
+  /**
+   * 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.
+   * 
+   * The current position of the Mechanism2d is logged once as a set of nested
+   * fields. If the position is updated, this method must be called again.
+   * 
+   * @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 void recordOutput(String key, Mechanism2d value) {
+    if (running) {
+      try {
+        // Use reflection because we don't explicitly depend on the shimmed classes
+        Mechanism2d.class.getMethod("akitLog", LogTable.class).invoke(value, outputTable.getSubtable(key));
+      } catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException | NoSuchMethodException
+          | SecurityException e) {
+        e.printStackTrace();
+      }
+    }
+  }
 }

junction/shims/wpilib/src/edu/wpi/first/wpilibj/smartdashboard/Mechanism2d.java

@@ -0,0 +1,159 @@
+// 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.wpilibj.smartdashboard;
+
+import edu.wpi.first.networktables.DoubleArrayPublisher;
+import edu.wpi.first.networktables.NTSendable;
+import edu.wpi.first.networktables.NTSendableBuilder;
+import edu.wpi.first.networktables.NetworkTable;
+import edu.wpi.first.networktables.StringPublisher;
+import edu.wpi.first.wpilibj.util.Color8Bit;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Map.Entry;
+import org.littletonrobotics.junction.LogTable;
+
+/**
+ * Visual 2D representation of arms, elevators, and general mechanisms through a
+ * node-based API.
+ *
+ * <p>
+ * A Mechanism2d object is published and contains at least one root node. A root
+ * is the anchor point of other nodes (such as ligaments). Other nodes are
+ * recursively appended based on other nodes.
+ *
+ * @see MechanismObject2d
+ * @see MechanismLigament2d
+ * @see MechanismRoot2d
+ */
+public final class Mechanism2d implements NTSendable, AutoCloseable {
+  private NetworkTable m_table;
+  private final Map<String, MechanismRoot2d> m_roots;
+  private final double[] m_dims = new double[2];
+  private String m_color;
+  private DoubleArrayPublisher m_dimsPub;
+  private StringPublisher m_colorPub;
+
+  /**
+   * Create a new Mechanism2d with the given dimensions and default color (dark
+   * blue).
+   *
+   * <p>
+   * The dimensions represent the canvas that all the nodes are drawn on.
+   *
+   * @param width  the width
+   * @param height the height
+   */
+  public Mechanism2d(double width, double height) {
+    this(width, height, new Color8Bit(0, 0, 32));
+  }
+
+  /**
+   * Create a new Mechanism2d with the given dimensions.
+   *
+   * <p>
+   * The dimensions represent the canvas that all the nodes are drawn on.
+   *
+   * @param width           the width
+   * @param height          the height
+   * @param backgroundColor the background color. Defaults to dark blue.
+   */
+  public Mechanism2d(double width, double height, Color8Bit backgroundColor) {
+    m_roots = new HashMap<>();
+    m_dims[0] = width;
+    m_dims[1] = height;
+    setBackgroundColor(backgroundColor);
+  }
+
+  @Override
+  public void close() {
+    if (m_dimsPub != null) {
+      m_dimsPub.close();
+    }
+    if (m_colorPub != null) {
+      m_colorPub.close();
+    }
+    for (MechanismRoot2d root : m_roots.values()) {
+      root.close();
+    }
+  }
+
+  /**
+   * Get or create a root in this Mechanism2d with the given name and position.
+   *
+   * <p>
+   * If a root with the given name already exists, the given x and y coordinates
+   * are not used.
+   *
+   * @param name the root name
+   * @param x    the root x coordinate
+   * @param y    the root y coordinate
+   * @return a new root joint object, or the existing one with the given name.
+   */
+  public synchronized MechanismRoot2d getRoot(String name, double x, double y) {
+    MechanismRoot2d existing = m_roots.get(name);
+    if (existing != null) {
+      return existing;
+    }
+
+    MechanismRoot2d root = new MechanismRoot2d(name, x, y);
+    m_roots.put(name, root);
+    if (m_table != null) {
+      root.update(m_table.getSubTable(name));
+    }
+    return root;
+  }
+
+  /**
+   * Set the Mechanism2d background color.
+   *
+   * @param color the new color
+   */
+  public synchronized void setBackgroundColor(Color8Bit color) {
+    m_color = color.toHexString();
+    if (m_colorPub != null) {
+      m_colorPub.set(m_color);
+    }
+  }
+
+  @Override
+  public void initSendable(NTSendableBuilder builder) {
+    builder.setSmartDashboardType("Mechanism2d");
+    synchronized (this) {
+      m_table = builder.getTable();
+      if (m_dimsPub != null) {
+        m_dimsPub.close();
+      }
+      m_dimsPub = m_table.getDoubleArrayTopic("dims").publish();
+      m_dimsPub.set(m_dims);
+      if (m_colorPub != null) {
+        m_colorPub.close();
+      }
+      m_colorPub = m_table.getStringTopic("backgroundColor").publish();
+      m_colorPub.set(m_color);
+      for (Entry<String, MechanismRoot2d> entry : m_roots.entrySet()) {
+        String name = entry.getKey();
+        MechanismRoot2d root = entry.getValue();
+        synchronized (root) {
+          root.update(m_table.getSubTable(name));
+        }
+      }
+    }
+  }
+
+  public synchronized void akitLog(LogTable table) {
+    table.put(".type", "Mechanism2d");
+    table.put(".controllable", false);
+    table.put("dims", m_dims);
+    table.put("backgroundColor", m_color);
+    for (Entry<String, MechanismRoot2d> entry : m_roots.entrySet()) {
+      String name = entry.getKey();
+      MechanismRoot2d root = entry.getValue();
+      synchronized (root) {
+        root.akitLog(table.getSubtable(name));
+      }
+    }
+  }
+}