Add a swerve drivetrain writer
Signed-off-by: Nathan Leong <100028864@mvla.net>
Change-Id: I4750a4acba814e774befb0de4b0ba547ab67efcf
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_