Add a swerve drivetrain writer

Signed-off-by: Nathan Leong <100028864@mvla.net>
Change-Id: I4750a4acba814e774befb0de4b0ba547ab67efcf
diff --git a/frc971/BUILD b/frc971/BUILD
index d6a3af8..5613737 100644
--- a/frc971/BUILD
+++ b/frc971/BUILD
@@ -1,3 +1,5 @@
+load("@com_github_google_flatbuffers//:build_defs.bzl", "flatbuffer_cc_library")
+
 package(default_visibility = ["//visibility:public"])
 
 cc_library(
@@ -21,3 +23,12 @@
     target_compatible_with = ["@platforms//os:linux"],
     visibility = ["//visibility:public"],
 )
+
+flatbuffer_cc_library(
+    name = "can_configuration_fbs",
+    srcs = [
+        ":can_configuration.fbs",
+    ],
+    gen_reflections = 1,
+    visibility = ["//visibility:public"],
+)
diff --git a/y2023/can_configuration.fbs b/frc971/can_configuration.fbs
similarity index 93%
rename from y2023/can_configuration.fbs
rename to frc971/can_configuration.fbs
index 75e2691..4ce70ab 100644
--- a/y2023/can_configuration.fbs
+++ b/frc971/can_configuration.fbs
@@ -1,4 +1,4 @@
-namespace y2023;
+namespace frc971;
 
 // Message which triggers wpilib_interface to print out the current
 // configuration, and optionally re-apply it.
diff --git a/frc971/wpilib/BUILD b/frc971/wpilib/BUILD
index 6c59d35..33a3cd1 100644
--- a/frc971/wpilib/BUILD
+++ b/frc971/wpilib/BUILD
@@ -506,6 +506,7 @@
     target_compatible_with = ["//tools/platforms/hardware:roborio"],
     deps = [
         "//aos:init",
+        "//aos:math",
         "//aos/events:shm_event_loop",
         "//aos/logging",
         "//frc971/control_loops/drivetrain:drivetrain_can_position_fbs",
diff --git a/frc971/wpilib/falcon.cc b/frc971/wpilib/falcon.cc
index 3f551eb..6be83aa 100644
--- a/frc971/wpilib/falcon.cc
+++ b/frc971/wpilib/falcon.cc
@@ -1,6 +1,7 @@
 #include "frc971/wpilib/falcon.h"
 
 using frc971::wpilib::Falcon;
+using frc971::wpilib::kMaxBringupPower;
 
 Falcon::Falcon(int device_id, std::string canbus,
                std::vector<ctre::phoenix6::BaseStatusSignal *> *signals,
@@ -76,6 +77,23 @@
   PrintConfigs();
 }
 
+ctre::phoenix::StatusCode Falcon::WriteCurrent(double current,
+                                               double max_voltage) {
+  ctre::phoenix6::controls::TorqueCurrentFOC control(
+      static_cast<units::current::ampere_t>(current));
+  // Using 0_Hz here makes it a one-shot update.
+  control.UpdateFreqHz = 0_Hz;
+  control.MaxAbsDutyCycle =
+      ::aos::Clip(max_voltage, -kMaxBringupPower, kMaxBringupPower) / 12.0;
+  ctre::phoenix::StatusCode status = talon()->SetControl(control);
+  if (!status.IsOK()) {
+    AOS_LOG(ERROR, "Failed to write control to falcon %d: %s: %s", device_id(),
+            status.GetName(), status.GetDescription());
+  }
+
+  return status;
+}
+
 void Falcon::SerializePosition(flatbuffers::FlatBufferBuilder *fbb) {
   control_loops::CANFalcon::Builder builder(*fbb);
   builder.add_id(device_id_);
diff --git a/frc971/wpilib/falcon.h b/frc971/wpilib/falcon.h
index f6858fd..8f0f1f0 100644
--- a/frc971/wpilib/falcon.h
+++ b/frc971/wpilib/falcon.h
@@ -8,6 +8,7 @@
 #include "ctre/phoenix6/TalonFX.hpp"
 #include "glog/logging.h"
 
+#include "aos/commonmath.h"
 #include "aos/init.h"
 #include "aos/logging/logging.h"
 #include "frc971/control_loops/drivetrain/drivetrain_can_position_generated.h"
@@ -18,6 +19,7 @@
 namespace wpilib {
 
 static constexpr units::frequency::hertz_t kCANUpdateFreqHz = 200_Hz;
+static constexpr double kMaxBringupPower = 12.0;
 
 // Gets info from and writes to falcon motors using the TalonFX controller.
 class Falcon {
@@ -29,6 +31,7 @@
   void PrintConfigs();
 
   void WriteConfigs(ctre::phoenix6::signals::InvertedValue invert);
+  ctre::phoenix::StatusCode WriteCurrent(double current, double max_voltage);
 
   ctre::phoenix6::hardware::TalonFX *talon() { return &talon_; }
 
diff --git a/frc971/wpilib/swerve/BUILD b/frc971/wpilib/swerve/BUILD
new file mode 100644
index 0000000..9f559f6
--- /dev/null
+++ b/frc971/wpilib/swerve/BUILD
@@ -0,0 +1,30 @@
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "swerve_drivetrain_writer",
+    srcs = [
+        "swerve_drivetrain_writer.cc",
+    ],
+    hdrs = [
+        "swerve_drivetrain_writer.h",
+    ],
+    deps = [
+        ":swerve_module",
+        "//aos/logging",
+        "//frc971:can_configuration_fbs",
+        "//frc971/control_loops/drivetrain/swerve:swerve_drivetrain_output_fbs",
+        "//frc971/wpilib:falcon",
+        "//frc971/wpilib:loop_output_handler",
+        "//third_party:phoenix6",
+    ],
+)
+
+cc_library(
+    name = "swerve_module",
+    hdrs = [
+        "swerve_module.h",
+    ],
+    deps = [
+        "//frc971/wpilib:falcon",
+    ],
+)
diff --git a/frc971/wpilib/swerve/swerve_drivetrain_writer.cc b/frc971/wpilib/swerve/swerve_drivetrain_writer.cc
new file mode 100644
index 0000000..2b6ef9e
--- /dev/null
+++ b/frc971/wpilib/swerve/swerve_drivetrain_writer.cc
@@ -0,0 +1,60 @@
+#include "frc971/wpilib/swerve/swerve_drivetrain_writer.h"
+
+using frc971::wpilib::swerve::DrivetrainWriter;
+
+DrivetrainWriter::DrivetrainWriter(::aos::EventLoop *event_loop,
+                                   int drivetrain_writer_priority,
+                                   double max_voltage)
+    : ::frc971::wpilib::LoopOutputHandler<
+          ::frc971::control_loops::drivetrain::swerve::Output>(event_loop,
+                                                               "/drivetrain"),
+      max_voltage_(max_voltage) {
+  event_loop->SetRuntimeRealtimePriority(drivetrain_writer_priority);
+
+  event_loop->OnRun([this]() { WriteConfigs(); });
+}
+
+void DrivetrainWriter::set_falcons(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::HandleCANConfiguration(
+    const CANConfiguration &configuration) {
+  for (auto module : {front_left_, front_right_, back_left_, back_right_}) {
+    module->rotation->PrintConfigs();
+    module->translation->PrintConfigs();
+  }
+  if (configuration.reapply()) {
+    WriteConfigs();
+  }
+}
+
+void DrivetrainWriter::WriteConfigs() {
+  for (auto module : {front_left_, front_right_, back_left_, back_right_}) {
+    module->rotation->WriteConfigs(false);
+    module->translation->WriteConfigs(false);
+  }
+}
+
+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_);
+}
+
+void DrivetrainWriter::Stop() {
+  AOS_LOG(WARNING, "drivetrain output too old\n");
+
+  for (auto module : {front_left_, front_right_, back_left_, 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
new file mode 100644
index 0000000..4bd6639
--- /dev/null
+++ b/frc971/wpilib/swerve/swerve_drivetrain_writer.h
@@ -0,0 +1,52 @@
+#ifndef FRC971_WPILIB_SWERVE_DRIVETRAIN_WRITER_H_
+#define FRC971_WPILIB_SWERVE_DRIVETRAIN_WRITER_H_
+
+#include "ctre/phoenix6/TalonFX.hpp"
+
+#include "frc971/can_configuration_generated.h"
+#include "frc971/control_loops/drivetrain/swerve/swerve_drivetrain_output_generated.h"
+#include "frc971/wpilib/falcon.h"
+#include "frc971/wpilib/loop_output_handler.h"
+#include "frc971/wpilib/swerve/swerve_module.h"
+
+namespace frc971 {
+namespace wpilib {
+namespace swerve {
+
+// 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> {
+ public:
+  DrivetrainWriter(::aos::EventLoop *event_loop, int drivetrain_writer_priority,
+                   double max_voltage);
+
+  void set_falcons(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 HandleCANConfiguration(const CANConfiguration &configuration);
+
+ private:
+  void WriteConfigs();
+
+  void Write(const ::frc971::control_loops::drivetrain::swerve::Output &output)
+      override;
+
+  void Stop() override;
+
+  double SafeSpeed(double voltage);
+
+  std::shared_ptr<SwerveModule> front_left_, front_right_, back_left_,
+      back_right_;
+
+  double max_voltage_;
+};
+
+}  // namespace swerve
+}  // namespace wpilib
+}  // namespace frc971
+
+#endif  // FRC971_WPILIB_SWERVE_DRIVETRAIN_WRITER_H_
diff --git a/frc971/wpilib/swerve/swerve_module.h b/frc971/wpilib/swerve/swerve_module.h
new file mode 100644
index 0000000..534f0ce
--- /dev/null
+++ b/frc971/wpilib/swerve/swerve_module.h
@@ -0,0 +1,44 @@
+#ifndef FRC971_WPILIB_SWERVE_SWERVE_MODULE_H_
+#define FRC971_WPILIB_SWERVE_SWERVE_MODULE_H_
+
+#include "frc971/wpilib/falcon.h"
+
+namespace frc971 {
+namespace wpilib {
+namespace swerve {
+
+struct SwerveModule {
+  SwerveModule(int rotation_id, int translation_id, std::string canbus,
+               std::vector<ctre::phoenix6::BaseStatusSignal *> *signals,
+               double stator_current_limit, double supply_current_limit)
+      : rotation(std::make_shared<Falcon>(rotation_id, canbus, signals,
+                                          stator_current_limit,
+                                          supply_current_limit)),
+        translation(std::make_shared<Falcon>(translation_id, canbus, signals,
+                                             stator_current_limit,
+                                             supply_current_limit)) {}
+
+  void WriteModule(
+      const frc971::control_loops::drivetrain::swerve::SwerveModuleOutput
+          *module_output,
+      double max_voltage) {
+    double rotation_current = 0.0;
+    double translation_current = 0.0;
+
+    if (module_output != nullptr) {
+      rotation_current = module_output->rotation_current();
+      translation_current = module_output->translation_current();
+    }
+
+    rotation->WriteCurrent(rotation_current, max_voltage);
+    translation->WriteCurrent(translation_current, max_voltage);
+  }
+
+  std::shared_ptr<Falcon> rotation;
+  std::shared_ptr<Falcon> translation;
+};
+
+}  // namespace swerve
+}  // namespace wpilib
+}  // namespace frc971
+#endif  // FRC971_WPILIB_SWERVE_SWERVE_MODULE_H_
diff --git a/y2023/BUILD b/y2023/BUILD
index 992ba5d..e200f87 100644
--- a/y2023/BUILD
+++ b/y2023/BUILD
@@ -1,7 +1,6 @@
 load("//frc971:downloader.bzl", "robot_downloader")
 load("//aos:config.bzl", "aos_config")
 load("//tools/build_rules:template.bzl", "jinja2_template")
-load("@com_github_google_flatbuffers//:build_defs.bzl", "flatbuffer_cc_library")
 load("//aos/util:config_validator_macro.bzl", "config_validator_test")
 
 config_validator_test(
@@ -203,7 +202,7 @@
     name = "config_roborio",
     src = "y2023_roborio.json",
     flatbuffers = [
-        ":can_configuration_fbs",
+        "//frc971:can_configuration_fbs",
         "//aos/network:remote_message_fbs",
         "//aos/network:message_bridge_client_fbs",
         "//aos/network:message_bridge_server_fbs",
@@ -271,7 +270,6 @@
     ],
     target_compatible_with = ["//tools/platforms/hardware:roborio"],
     deps = [
-        ":can_configuration_fbs",
         ":constants",
         "//aos:init",
         "//aos:math",
@@ -283,6 +281,7 @@
         "//aos/util:log_interval",
         "//aos/util:phased_loop",
         "//aos/util:wrapping_counter",
+        "//frc971:can_configuration_fbs",
         "//frc971/autonomous:auto_mode_fbs",
         "//frc971/control_loops:control_loop",
         "//frc971/control_loops:control_loops_fbs",
@@ -384,12 +383,3 @@
         "@com_github_google_glog//:glog",
     ],
 )
-
-flatbuffer_cc_library(
-    name = "can_configuration_fbs",
-    srcs = [
-        ":can_configuration.fbs",
-    ],
-    gen_reflections = 1,
-    visibility = ["//visibility:public"],
-)
diff --git a/y2023/wpilib_interface.cc b/y2023/wpilib_interface.cc
index 7975163..abf11da 100644
--- a/y2023/wpilib_interface.cc
+++ b/y2023/wpilib_interface.cc
@@ -40,6 +40,7 @@
 #include "aos/util/phased_loop.h"
 #include "aos/util/wrapping_counter.h"
 #include "frc971/autonomous/auto_mode_generated.h"
+#include "frc971/can_configuration_generated.h"
 #include "frc971/control_loops/drivetrain/drivetrain_can_position_generated.h"
 #include "frc971/control_loops/drivetrain/drivetrain_position_generated.h"
 #include "frc971/input/robot_state_generated.h"
@@ -56,7 +57,6 @@
 #include "frc971/wpilib/pdp_fetcher.h"
 #include "frc971/wpilib/sensor_reader.h"
 #include "frc971/wpilib/wpilib_robot_base.h"
-#include "y2023/can_configuration_generated.h"
 #include "y2023/constants.h"
 #include "y2023/control_loops/superstructure/led_indicator.h"
 #include "y2023/control_loops/superstructure/superstructure_output_generated.h"
@@ -67,6 +67,7 @@
             "devices on the CAN bus using Phoenix Tuner");
 
 using ::aos::monotonic_clock;
+using ::frc971::CANConfiguration;
 using ::y2023::constants::Values;
 namespace superstructure = ::y2023::control_loops::superstructure;
 namespace drivetrain = ::y2023::control_loops::drivetrain;
diff --git a/y2023/y2023_roborio.json b/y2023/y2023_roborio.json
index 33d9904..c1e42e6 100644
--- a/y2023/y2023_roborio.json
+++ b/y2023/y2023_roborio.json
@@ -300,7 +300,7 @@
     },
     {
       "name": "/roborio",
-      "type": "y2023.CANConfiguration",
+      "type": "frc971.CANConfiguration",
       "source_node": "roborio",
       "frequency": 2
     },