Add swerve control loops

Add the control loops for swerve in frc971/control_loops/swerve
Sets the current for each swerve module to 0
Does not use status and position flatbuffers

Change-Id: Id706459fabc1326ee1b78d1f35b90cf8cd3f7df9
Signed-off-by: Yash Maheshwari <yashmahe2018@gmail.com>
diff --git a/frc971/control_loops/swerve/BUILD b/frc971/control_loops/swerve/BUILD
index 3bd3453..ccba681 100644
--- a/frc971/control_loops/swerve/BUILD
+++ b/frc971/control_loops/swerve/BUILD
@@ -1,3 +1,4 @@
+load("//aos:config.bzl", "aos_config")
 load("//aos/flatbuffers:generate.bzl", "static_flatbuffer")
 
 package(default_visibility = ["//visibility:public"])
@@ -16,21 +17,18 @@
 static_flatbuffer(
     name = "swerve_drivetrain_output_fbs",
     srcs = ["swerve_drivetrain_output.fbs"],
-    visibility = ["//visibility:public"],
     deps = ["//frc971/control_loops:can_talonfx_fbs"],
 )
 
 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"],
     deps = ["//frc971/control_loops:control_loops_fbs"],
 )
 
@@ -47,3 +45,54 @@
         "@pip//sympy",
     ],
 )
+
+cc_library(
+    name = "swerve_control_loops",
+    srcs = ["swerve_control_loops.cc"],
+    hdrs = ["swerve_control_loops.h"],
+    target_compatible_with = ["@platforms//os:linux"],
+    deps = [
+        ":swerve_drivetrain_can_position_fbs",
+        ":swerve_drivetrain_goal_fbs",
+        ":swerve_drivetrain_output_fbs",
+        ":swerve_drivetrain_position_fbs",
+        ":swerve_drivetrain_status_fbs",
+        "//frc971/control_loops:control_loop",
+    ],
+)
+
+cc_test(
+    name = "swerve_control_test",
+    srcs = ["swerve_control_test.cc"],
+    data = [
+        ":aos_config",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    deps = [
+        ":swerve_control_loops",
+        ":swerve_drivetrain_can_position_fbs",
+        ":swerve_drivetrain_goal_fbs",
+        ":swerve_drivetrain_output_fbs",
+        ":swerve_drivetrain_position_fbs",
+        ":swerve_drivetrain_status_fbs",
+        "//aos/events:shm_event_loop",
+        "//aos/testing:googletest",
+        "//frc971/control_loops:control_loop_test",
+        "//frc971/control_loops:state_feedback_loop",
+        "//frc971/control_loops:team_number_test_environment",
+    ],
+)
+
+aos_config(
+    name = "aos_config",
+    src = "swerve.json",
+    flatbuffers = [
+        ":swerve_drivetrain_goal_fbs",
+        ":swerve_drivetrain_output_fbs",
+        ":swerve_drivetrain_position_fbs",
+        ":swerve_drivetrain_can_position_fbs",
+        ":swerve_drivetrain_status_fbs",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    deps = ["//frc971/input:aos_config"],
+)
diff --git a/frc971/control_loops/swerve/swerve.json b/frc971/control_loops/swerve/swerve.json
new file mode 100644
index 0000000..1282b54
--- /dev/null
+++ b/frc971/control_loops/swerve/swerve.json
@@ -0,0 +1,27 @@
+{
+    "channels": [
+        {
+            "name": "/swerve",
+            "type": "frc971.control_loops.swerve.Goal"
+        },
+        {
+            "name": "/swerve",
+            "type": "frc971.control_loops.swerve.Output"
+        },
+        {
+            "name": "/swerve",
+            "type": "frc971.control_loops.swerve.CanPosition"
+        },
+        {
+            "name": "/swerve",
+            "type": "frc971.control_loops.swerve.Position"
+        },
+        {
+            "name": "/swerve",
+            "type": "frc971.control_loops.swerve.Status"
+        }
+    ],
+    "imports": [
+        "../../../frc971/input/robot_state_config.json"
+    ]
+}
diff --git a/frc971/control_loops/swerve/swerve_control_loops.cc b/frc971/control_loops/swerve/swerve_control_loops.cc
new file mode 100644
index 0000000..4689e55
--- /dev/null
+++ b/frc971/control_loops/swerve/swerve_control_loops.cc
@@ -0,0 +1,71 @@
+#include "frc971/control_loops/swerve/swerve_control_loops.h"
+
+namespace frc971::control_loops::swerve {
+
+SwerveControlLoops::SwerveControlLoops(::aos::EventLoop *event_loop,
+                                       const ::std::string &name)
+    : frc971::controls::ControlLoop<Goal, Position, StatusStatic, OutputStatic>(
+          event_loop, name) {}
+
+void SwerveControlLoops::RunIteration(
+    const Goal *goal, const Position *position,
+    aos::Sender<OutputStatic>::StaticBuilder *output_builder,
+    aos::Sender<StatusStatic>::StaticBuilder *status_builder) {
+  (void)goal, (void)position;
+
+  if (output_builder != nullptr && goal != nullptr) {
+    OutputStatic *output = output_builder->get();
+
+    auto front_left_output = output->add_front_left_output();
+    front_left_output->set_rotation_current(0);
+    front_left_output->set_translation_current(
+        goal->front_left_goal()->translation_current());
+
+    auto front_right_output = output->add_front_right_output();
+    front_right_output->set_rotation_current(0);
+    front_right_output->set_translation_current(
+        goal->front_right_goal()->translation_current());
+
+    auto back_left_output = output->add_back_left_output();
+    back_left_output->set_rotation_current(0);
+    back_left_output->set_translation_current(
+        goal->back_left_goal()->translation_current());
+
+    auto back_right_output = output->add_back_right_output();
+    back_right_output->set_rotation_current(0);
+    back_right_output->set_translation_current(
+        goal->back_right_goal()->translation_current());
+
+    // Ignore the return value of Send
+    output_builder->CheckOk(output_builder->Send());
+  }
+
+  if (status_builder != nullptr) {
+    StatusStatic *status = status_builder->get();
+
+    auto front_left_status = status->add_front_left_status();
+    front_left_status->set_goal_translation_speed(0);
+    front_left_status->set_translation_speed(0);
+    front_left_status->add_rotation();
+
+    auto front_right_status = status->add_front_right_status();
+    front_right_status->set_goal_translation_speed(0);
+    front_right_status->set_translation_speed(0);
+    front_right_status->add_rotation();
+
+    auto back_left_status = status->add_back_left_status();
+    back_left_status->set_goal_translation_speed(0);
+    back_left_status->set_translation_speed(0);
+    back_left_status->add_rotation();
+
+    auto back_right_status = status->add_back_right_status();
+    back_right_status->set_goal_translation_speed(0);
+    back_right_status->set_translation_speed(0);
+    back_right_status->add_rotation();
+
+    // Ignore the return value of Send
+    status_builder->CheckOk(status_builder->Send());
+  }
+}
+
+}  // namespace frc971::control_loops::swerve
diff --git a/frc971/control_loops/swerve/swerve_control_loops.h b/frc971/control_loops/swerve/swerve_control_loops.h
new file mode 100644
index 0000000..78f6ba0
--- /dev/null
+++ b/frc971/control_loops/swerve/swerve_control_loops.h
@@ -0,0 +1,39 @@
+#ifndef FRC971_CONTROL_LOOPS_SWERVE_SWERVE_CONTROL_LOOPS_H_
+#define FRC971_CONTROL_LOOPS_SWERVE_SWERVE_CONTROL_LOOPS_H_
+
+#include "aos/time/time.h"
+#include "frc971/control_loops/control_loop.h"
+#include "frc971/control_loops/profiled_subsystem_generated.h"
+#include "frc971/control_loops/swerve/swerve_drivetrain_can_position_generated.h"
+#include "frc971/control_loops/swerve/swerve_drivetrain_goal_generated.h"
+#include "frc971/control_loops/swerve/swerve_drivetrain_output_static.h"
+#include "frc971/control_loops/swerve/swerve_drivetrain_position_generated.h"
+#include "frc971/control_loops/swerve/swerve_drivetrain_status_static.h"
+
+namespace frc971::control_loops::swerve {
+
+inline void PopulateSwerveModuleRotation(
+    SwerveModuleStatusStatic *swerve_module_table) {
+  auto rotation = swerve_module_table->add_rotation();
+  auto estimator_state = rotation->add_estimator_state();
+  (void)estimator_state;
+}
+
+// Handles the translation and rotation current for each swerve module
+class SwerveControlLoops
+    : public ::frc971::controls::ControlLoop<Goal, Position, StatusStatic,
+                                             OutputStatic> {
+ public:
+  explicit SwerveControlLoops(::aos::EventLoop *event_loop,
+                              const ::std::string &name = "/swerve");
+
+ protected:
+  void RunIteration(
+      const Goal *goal, const Position *position,
+      aos::Sender<OutputStatic>::StaticBuilder *output_builder,
+      aos::Sender<StatusStatic>::StaticBuilder *status_builder) override;
+};
+
+}  // namespace frc971::control_loops::swerve
+
+#endif  // FRC971_CONTROL_LOOPS_SWERVE_SWERVE_CONTROL_LOOPS_H_
diff --git a/frc971/control_loops/swerve/swerve_control_test.cc b/frc971/control_loops/swerve/swerve_control_test.cc
new file mode 100644
index 0000000..9378f1c
--- /dev/null
+++ b/frc971/control_loops/swerve/swerve_control_test.cc
@@ -0,0 +1,161 @@
+#include <unistd.h>
+
+#include <chrono>
+#include <memory>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+#include "aos/events/shm_event_loop.h"
+#include "frc971/control_loops/control_loop_test.h"
+#include "frc971/control_loops/swerve/swerve_control_loops.h"
+#include "frc971/control_loops/swerve/swerve_drivetrain_can_position_generated.h"
+#include "frc971/control_loops/swerve/swerve_drivetrain_goal_generated.h"
+#include "frc971/control_loops/swerve/swerve_drivetrain_output_generated.h"
+#include "frc971/control_loops/swerve/swerve_drivetrain_position_generated.h"
+#include "frc971/control_loops/swerve/swerve_drivetrain_status_generated.h"
+#include "frc971/control_loops/team_number_test_environment.h"
+
+using namespace std;
+
+namespace frc971::control_loops::swerve::testing {
+namespace chrono = ::std::chrono;
+using aos::monotonic_clock;
+
+// Class which simulates stuff and sends out queue messages with the position.
+class SwerveControlSimulation {
+ public:
+  SwerveControlSimulation(::aos::EventLoop *event_loop, chrono::nanoseconds dt)
+      : event_loop_(event_loop),
+        position_sender_(event_loop_->MakeSender<Position>("/swerve")),
+        can_position_sender_(event_loop_->MakeSender<CanPosition>("/swerve")),
+        goal_fetcher_(event_loop_->MakeFetcher<Goal>("/swerve")),
+        status_fetcher_(event_loop_->MakeFetcher<Status>("/swerve")),
+        output_fetcher_(event_loop_->MakeFetcher<Output>("/swerve")) {
+    event_loop_->AddPhasedLoop(
+        [this](int) {
+          if (!first_) {
+            Simulate();
+          }
+          first_ = false;
+          SendPositionMessage();
+        },
+        dt);
+  }
+
+  // Sends a queue message with the position data.
+  void SendPositionMessage() {
+    auto builder = position_sender_.MakeBuilder();
+
+    Position::Builder position_builder = builder.MakeBuilder<Position>();
+
+    EXPECT_EQ(builder.Send(position_builder.Finish()),
+              aos::RawSender::Error::kOk);
+  }
+
+  void VerifyNearGoal() {
+    goal_fetcher_.Fetch();
+    status_fetcher_.Fetch();
+
+    ASSERT_TRUE(goal_fetcher_.get() != nullptr) << ": No goal";
+    ASSERT_TRUE(status_fetcher_.get() != nullptr) << ": No status";
+
+    constexpr double kEpsRotationAngle = 0.03;
+    constexpr double kEpsRotationVelocity = 0.03;
+    constexpr double kEpsTranslationSpeed = 0.03;
+
+    std::vector<const SwerveModuleStatus *> modules_status{
+        status_fetcher_->front_left_status(),
+        status_fetcher_->front_right_status(),
+        status_fetcher_->back_left_status(),
+        status_fetcher_->back_right_status()};
+
+    for (auto &module_status : modules_status) {
+      EXPECT_NEAR(module_status->rotation()->position(),
+                  module_status->rotation()->goal_position(),
+                  kEpsRotationAngle);
+      EXPECT_NEAR(module_status->rotation()->velocity(),
+                  module_status->rotation()->goal_velocity(),
+                  kEpsRotationVelocity);
+      EXPECT_NEAR(module_status->goal_translation_speed(),
+                  module_status->translation_speed(), kEpsTranslationSpeed);
+    }
+  }
+
+  // Simulates basic control loop for a single timestep.
+  void Simulate() { EXPECT_TRUE(output_fetcher_.Fetch()); }
+
+ private:
+  ::aos::EventLoop *event_loop_;
+
+  ::aos::Sender<Position> position_sender_;
+  ::aos::Sender<CanPosition> can_position_sender_;
+  ::aos::Sender<Goal> goal_sender_;
+
+  ::aos::Fetcher<CanPosition> can_position_fetcher_;
+  ::aos::Fetcher<Position> position_fetcher_;
+  ::aos::Fetcher<Goal> goal_fetcher_;
+  ::aos::Fetcher<Status> status_fetcher_;
+  ::aos::Fetcher<Output> output_fetcher_;
+
+  bool first_ = true;
+};
+
+class SwerveControlLoopTest : public ::frc971::testing::ControlLoopTest {
+ public:
+  SwerveControlLoopTest()
+      : ::frc971::testing::ControlLoopTest(
+            aos::configuration::ReadConfig(
+                "frc971/control_loops/swerve/aos_config.json"),
+            chrono::microseconds(5050)),
+        swerve_test_event_loop_(MakeEventLoop("test")),
+        goal_sender_(swerve_test_event_loop_->MakeSender<Goal>("/swerve")),
+
+        swerve_control_event_loop_(MakeEventLoop("swerve_control")),
+        swerve_control_loops_(swerve_control_event_loop_.get(), "/swerve"),
+
+        swerve_control_simulation_event_loop_(MakeEventLoop("simulation")),
+        swerve_control_simulation_(swerve_control_simulation_event_loop_.get(),
+                                   dt()) {
+    set_team_id(control_loops::testing::kTeamNumber);
+    SetEnabled(true);
+  }
+
+  ::std::unique_ptr<::aos::EventLoop> swerve_test_event_loop_;
+  ::aos::Sender<Goal> goal_sender_;
+
+  ::std::unique_ptr<::aos::EventLoop> swerve_control_event_loop_;
+  SwerveControlLoops swerve_control_loops_;
+
+  ::std::unique_ptr<::aos::EventLoop> swerve_control_simulation_event_loop_;
+  SwerveControlSimulation swerve_control_simulation_;
+};
+
+// Tests that the swerve modules' speeds are all set to 0.
+TEST_F(SwerveControlLoopTest, SwerveModulesDontMove) {
+  {
+    auto builder = goal_sender_.MakeBuilder();
+    SwerveModuleGoal::Builder swerve_module_builder =
+        builder.MakeBuilder<SwerveModuleGoal>();
+    swerve_module_builder.add_translation_control_type_goal(
+        TranslationControlTypeGoal::CURRENT);
+    swerve_module_builder.add_rotation_angle(0.0);
+    swerve_module_builder.add_translation_current(0.0);
+    swerve_module_builder.add_translation_speed(0.0);
+    auto empty_module_offset = swerve_module_builder.Finish();
+
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_front_left_goal(empty_module_offset);
+    goal_builder.add_front_right_goal(empty_module_offset);
+    goal_builder.add_back_left_goal(empty_module_offset);
+    goal_builder.add_back_right_goal(empty_module_offset);
+
+    ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
+  }
+
+  RunFor(dt() * 2);
+
+  swerve_control_simulation_.VerifyNearGoal();
+}
+
+}  // namespace frc971::control_loops::swerve::testing
\ No newline at end of file
diff --git a/frc971/wpilib/swerve/swerve_drivetrain_writer.cc b/frc971/wpilib/swerve/swerve_drivetrain_writer.cc
index ccadb99..44aff53 100644
--- a/frc971/wpilib/swerve/swerve_drivetrain_writer.cc
+++ b/frc971/wpilib/swerve/swerve_drivetrain_writer.cc
@@ -14,7 +14,7 @@
 }
 
 void DrivetrainWriter::set_talonfxs(SwerveModules modules) {
-  modules = std::move(modules);
+  modules_ = std::move(modules);
 }
 
 void DrivetrainWriter::HandleCANConfiguration(
diff --git a/y2024_swerve/y2024_swerve_roborio.json b/y2024_swerve/y2024_swerve_roborio.json
index bd8ddff..1330e43 100644
--- a/y2024_swerve/y2024_swerve_roborio.json
+++ b/y2024_swerve/y2024_swerve_roborio.json
@@ -27,6 +27,14 @@
         },
         {
             "name": "/roborio/aos",
+            "type": "aos.RobotState",
+            "source_node": "roborio",
+            "frequency": 50,
+            "num_senders": 20,
+            "max_size": 192
+        },
+        {
+            "name": "/roborio/aos",
             "type": "aos.logging.LogMessageFbs",
             "source_node": "roborio",
             "frequency": 500,
@@ -39,7 +47,7 @@
             "source_node": "roborio",
             "frequency": 50,
             "num_senders": 20,
-            "max_size": 2000
+            "max_size": 3000
         },
         {
             "name": "/roborio/aos",
@@ -115,7 +123,7 @@
             "source_node": "roborio",
             "frequency": 250,
             "num_senders": 1,
-            "max_size": 480
+            "max_size": 656
         },
         {
             "name": "/can",