Update swerve drivetrain position messages & logic

This makes it so that the swerve drivetrain positions are sent in two
separate position messages (for FPGA-based and CAN-based positions),
using message definitions from the frc971/ folder.

Also refactors the SwerveModule class/usage a bit.

Removes the follower wheels from the y2024_swerve code.

Change-Id: I36898c3337b5a1437ce0c2a0189fd317929f1986
Signed-off-by: James Kuszmaul <jabukuszmaul+collab@gmail.com>
diff --git a/frc971/control_loops/swerve/BUILD b/frc971/control_loops/swerve/BUILD
index 4e06932..ea6dd7c 100644
--- a/frc971/control_loops/swerve/BUILD
+++ b/frc971/control_loops/swerve/BUILD
@@ -8,6 +8,13 @@
 )
 
 static_flatbuffer(
+    name = "swerve_drivetrain_can_position_fbs",
+    srcs = ["swerve_drivetrain_can_position.fbs"],
+    visibility = ["//visibility:public"],
+    deps = ["//frc971/control_loops:can_talonfx_fbs"],
+)
+
+static_flatbuffer(
     name = "swerve_drivetrain_position_fbs",
     srcs = ["swerve_drivetrain_position.fbs"],
     visibility = ["//visibility:public"],
diff --git a/frc971/control_loops/swerve/swerve_drivetrain_can_position.fbs b/frc971/control_loops/swerve/swerve_drivetrain_can_position.fbs
new file mode 100644
index 0000000..957bce8
--- /dev/null
+++ b/frc971/control_loops/swerve/swerve_drivetrain_can_position.fbs
@@ -0,0 +1,18 @@
+include "frc971/control_loops/can_talonfx.fbs";
+
+namespace frc971.control_loops.swerve;
+
+table SwerveModuleCanPosition {
+  rotation: frc971.control_loops.CANTalonFX (id: 0);
+  translation: frc971.control_loops.CANTalonFX (id: 1);
+}
+
+// CAN Readings from the CAN sensor reader loop for each swerve module
+table CanPosition {
+  front_left: SwerveModuleCanPosition (id: 0);
+  front_right: SwerveModuleCanPosition (id: 1);
+  back_left: SwerveModuleCanPosition (id: 2);
+  back_right: SwerveModuleCanPosition (id: 3);
+}
+
+root_type CanPosition;
diff --git a/frc971/control_loops/swerve/swerve_drivetrain_output.fbs b/frc971/control_loops/swerve/swerve_drivetrain_output.fbs
index 43ba0ed..72517a6 100644
--- a/frc971/control_loops/swerve/swerve_drivetrain_output.fbs
+++ b/frc971/control_loops/swerve/swerve_drivetrain_output.fbs
@@ -1,4 +1,4 @@
-namespace frc971.control_loops.drivetrain.swerve;
+namespace frc971.control_loops.swerve;
 
 table SwerveModuleOutput {
   // Current in Amps.
diff --git a/frc971/control_loops/swerve/swerve_drivetrain_position.fbs b/frc971/control_loops/swerve/swerve_drivetrain_position.fbs
index bad8af1..a5c921a 100644
--- a/frc971/control_loops/swerve/swerve_drivetrain_position.fbs
+++ b/frc971/control_loops/swerve/swerve_drivetrain_position.fbs
@@ -1,17 +1,19 @@
-namespace frc971.control_loops.drivetrain.swerve;
+include "frc971/control_loops/control_loops.fbs";
+
+namespace frc971.control_loops.swerve;
 
 table SwerveModulePosition {
-  // Rotation in radians
-  rotation_encoder:double (id: 0);
-  // Translation in meters
-  translation_encoder:double (id: 1);
+  // Position of the mag encoder for the rotation of the module.
+  rotation_position: frc971.AbsolutePosition (id: 0);
 }
 
+// Captures all of the roborio-sourced position information for a
+// swerve drivetrain.
 table Position {
-  front_left_position:SwerveModulePosition (id: 0);
-  front_right_position:SwerveModulePosition (id: 1);
-  back_left_position:SwerveModulePosition (id: 2);
-  back_right_position:SwerveModulePosition (id: 3);
+  front_left:SwerveModulePosition (id: 0);
+  front_right:SwerveModulePosition (id: 1);
+  back_left:SwerveModulePosition (id: 2);
+  back_right:SwerveModulePosition (id: 3);
 }
 
 root_type Position;
diff --git a/frc971/wpilib/swerve/BUILD b/frc971/wpilib/swerve/BUILD
index 8eb5651..0b2cc1c 100644
--- a/frc971/wpilib/swerve/BUILD
+++ b/frc971/wpilib/swerve/BUILD
@@ -12,7 +12,9 @@
         ":swerve_module",
         "//aos/logging",
         "//frc971:can_configuration_fbs",
+        "//frc971/control_loops/swerve:swerve_drivetrain_can_position_fbs",
         "//frc971/control_loops/swerve:swerve_drivetrain_output_fbs",
+        "//frc971/control_loops/swerve:swerve_drivetrain_position_fbs",
         "//frc971/wpilib:loop_output_handler",
         "//frc971/wpilib:talonfx",
         "//third_party:phoenix6",
@@ -25,6 +27,7 @@
         "swerve_module.h",
     ],
     deps = [
+        "//frc971/wpilib:encoder_and_potentiometer",
         "//frc971/wpilib:talonfx",
     ],
 )
diff --git a/frc971/wpilib/swerve/swerve_drivetrain_writer.cc b/frc971/wpilib/swerve/swerve_drivetrain_writer.cc
index e4b4ee9..ccadb99 100644
--- a/frc971/wpilib/swerve/swerve_drivetrain_writer.cc
+++ b/frc971/wpilib/swerve/swerve_drivetrain_writer.cc
@@ -6,27 +6,21 @@
                                    int drivetrain_writer_priority,
                                    double max_voltage)
     : ::frc971::wpilib::LoopOutputHandler<
-          ::frc971::control_loops::drivetrain::swerve::Output>(event_loop,
-                                                               "/drivetrain"),
+          ::frc971::control_loops::swerve::Output>(event_loop, "/drivetrain"),
       max_voltage_(max_voltage) {
   event_loop->SetRuntimeRealtimePriority(drivetrain_writer_priority);
 
   event_loop->OnRun([this]() { WriteConfigs(); });
 }
 
-void DrivetrainWriter::set_talonfxs(std::shared_ptr<SwerveModule> front_left,
-                                    std::shared_ptr<SwerveModule> front_right,
-                                    std::shared_ptr<SwerveModule> back_left,
-                                    std::shared_ptr<SwerveModule> back_right) {
-  front_left_ = std::move(front_left);
-  front_right_ = std::move(front_right);
-  back_left_ = std::move(back_left);
-  back_right_ = std::move(back_right);
+void DrivetrainWriter::set_talonfxs(SwerveModules modules) {
+  modules = std::move(modules);
 }
 
 void DrivetrainWriter::HandleCANConfiguration(
     const CANConfiguration &configuration) {
-  for (auto module : {front_left_, front_right_, back_left_, back_right_}) {
+  for (auto module : {modules_.front_left, modules_.front_right,
+                      modules_.back_left, modules_.back_right}) {
     module->rotation->PrintConfigs();
     module->translation->PrintConfigs();
   }
@@ -36,24 +30,26 @@
 }
 
 void DrivetrainWriter::WriteConfigs() {
-  for (auto module : {front_left_, front_right_, back_left_, back_right_}) {
+  for (auto module : {modules_.front_left, modules_.front_right,
+                      modules_.back_left, modules_.back_right}) {
     module->rotation->WriteConfigs();
     module->translation->WriteConfigs();
   }
 }
 
 void DrivetrainWriter::Write(
-    const ::frc971::control_loops::drivetrain::swerve::Output &output) {
-  front_left_->WriteModule(output.front_left_output(), max_voltage_);
-  front_right_->WriteModule(output.front_right_output(), max_voltage_);
-  back_left_->WriteModule(output.back_left_output(), max_voltage_);
-  back_right_->WriteModule(output.back_right_output(), max_voltage_);
+    const ::frc971::control_loops::swerve::Output &output) {
+  modules_.front_left->WriteModule(output.front_left_output(), max_voltage_);
+  modules_.front_right->WriteModule(output.front_right_output(), max_voltage_);
+  modules_.back_left->WriteModule(output.back_left_output(), max_voltage_);
+  modules_.back_right->WriteModule(output.back_right_output(), max_voltage_);
 }
 
 void DrivetrainWriter::Stop() {
   AOS_LOG(WARNING, "drivetrain output too old\n");
 
-  for (auto module : {front_left_, front_right_, back_left_, back_right_}) {
+  for (auto module : {modules_.front_left, modules_.front_right,
+                      modules_.back_left, modules_.back_right}) {
     module->rotation->WriteCurrent(0, 0);
     module->translation->WriteCurrent(0, 0);
   }
diff --git a/frc971/wpilib/swerve/swerve_drivetrain_writer.h b/frc971/wpilib/swerve/swerve_drivetrain_writer.h
index 16db158..3b79186 100644
--- a/frc971/wpilib/swerve/swerve_drivetrain_writer.h
+++ b/frc971/wpilib/swerve/swerve_drivetrain_writer.h
@@ -13,32 +13,26 @@
 
 // Reads from the swerve output flatbuffer and uses wpilib to set the current
 // for each motor.
-class DrivetrainWriter
-    : public ::frc971::wpilib::LoopOutputHandler<
-          ::frc971::control_loops::drivetrain::swerve::Output> {
+class DrivetrainWriter : public ::frc971::wpilib::LoopOutputHandler<
+                             ::frc971::control_loops::swerve::Output> {
  public:
   DrivetrainWriter(::aos::EventLoop *event_loop, int drivetrain_writer_priority,
                    double max_voltage);
 
-  void set_talonfxs(std::shared_ptr<SwerveModule> front_left,
-                    std::shared_ptr<SwerveModule> front_right,
-                    std::shared_ptr<SwerveModule> back_left,
-                    std::shared_ptr<SwerveModule> back_right);
+  void set_talonfxs(SwerveModules modules);
 
   void HandleCANConfiguration(const CANConfiguration &configuration);
 
  private:
   void WriteConfigs();
 
-  void Write(const ::frc971::control_loops::drivetrain::swerve::Output &output)
-      override;
+  void Write(const ::frc971::control_loops::swerve::Output &output) override;
 
   void Stop() override;
 
   double SafeSpeed(double voltage);
 
-  std::shared_ptr<SwerveModule> front_left_, front_right_, back_left_,
-      back_right_;
+  SwerveModules modules_;
 
   double max_voltage_;
 };
diff --git a/frc971/wpilib/swerve/swerve_module.h b/frc971/wpilib/swerve/swerve_module.h
index 729ebab..c01df24 100644
--- a/frc971/wpilib/swerve/swerve_module.h
+++ b/frc971/wpilib/swerve/swerve_module.h
@@ -1,11 +1,17 @@
 #ifndef FRC971_WPILIB_SWERVE_SWERVE_MODULE_H_
 #define FRC971_WPILIB_SWERVE_SWERVE_MODULE_H_
 
+#include "frc971/control_loops/swerve/swerve_drivetrain_can_position_static.h"
 #include "frc971/control_loops/swerve/swerve_drivetrain_output_generated.h"
+#include "frc971/control_loops/swerve/swerve_drivetrain_position_static.h"
+#include "frc971/wpilib/encoder_and_potentiometer.h"
 #include "frc971/wpilib/talonfx.h"
 
 namespace frc971::wpilib::swerve {
 
+// Contains the objects for interacting with the hardware for a given swerve
+// module, assuming that the module uses two TalonFX-based motor controllers and
+// has a CTRE mag encoder on the rotation of the module.
 struct SwerveModule {
   SwerveModule(TalonFXParams rotation_params, TalonFXParams translation_params,
                std::string canbus,
@@ -18,9 +24,10 @@
                                               signals, stator_current_limit,
                                               supply_current_limit)) {}
 
+  // Writes the requested torque currents from the module_output to the motors,
+  // setting the maximum voltage of the motor outputs to the requested value.
   void WriteModule(
-      const frc971::control_loops::drivetrain::swerve::SwerveModuleOutput
-          *module_output,
+      const frc971::control_loops::swerve::SwerveModuleOutput *module_output,
       double max_voltage) {
     double rotation_current = 0.0;
     double translation_current = 0.0;
@@ -34,8 +41,58 @@
     translation->WriteCurrent(translation_current, max_voltage);
   }
 
+  // Used during initialization to set the WPILib objects used by the mag
+  // encoder on the rotation joint.
+  void set_rotation_encoder(std::unique_ptr<frc::Encoder> encoder,
+                            std::unique_ptr<frc::DigitalInput> absolute_pwm) {
+    rotation_encoder.set_encoder(std::move(encoder));
+    rotation_encoder.set_absolute_pwm(std::move(absolute_pwm));
+  }
+
+  // Populates the Position message with the mag encoder values.
+  void PopulatePosition(
+      frc971::control_loops::swerve::SwerveModulePositionStatic *fbs) {
+    auto rotation_position = fbs->add_rotation_position();
+    rotation_position->set_encoder(rotation_encoder.ReadRelativeEncoder());
+    rotation_position->set_absolute_encoder(
+        rotation_encoder.ReadAbsoluteEncoder());
+  }
+
+  // Populates a CAN-position message with the CAN-based devices (currently,
+  // just the motors themselves).
+  void PopulateCanPosition(
+      frc971::control_loops::swerve::SwerveModuleCanPositionStatic
+          *can_position) {
+    // TODO(james): Source these numbers from the constants.json.
+    rotation->SerializePosition(can_position->add_rotation(), 1.0);
+    translation->SerializePosition(can_position->add_translation(), 1.0);
+  }
+
   std::shared_ptr<TalonFX> rotation;
   std::shared_ptr<TalonFX> translation;
+  frc971::wpilib::AbsoluteEncoder rotation_encoder;
+};
+
+// Represents all the modules in a swerve drivetrain.
+struct SwerveModules {
+  void PopulateFalconsVector(std::vector<std::shared_ptr<TalonFX>> *falcons) {
+    CHECK_NOTNULL(falcons)->push_back(front_left->rotation);
+    falcons->push_back(front_left->translation);
+
+    falcons->push_back(front_right->rotation);
+    falcons->push_back(front_right->translation);
+
+    falcons->push_back(back_left->rotation);
+    falcons->push_back(back_left->translation);
+
+    falcons->push_back(back_right->rotation);
+    falcons->push_back(back_right->translation);
+  }
+
+  std::shared_ptr<SwerveModule> front_left;
+  std::shared_ptr<SwerveModule> front_right;
+  std::shared_ptr<SwerveModule> back_left;
+  std::shared_ptr<SwerveModule> back_right;
 };
 
 }  // namespace frc971::wpilib::swerve