[wpimath] Add support for swerve joystick normalization (#4516)

wpimath/src/main/native/include/frc/kinematics/SwerveDriveKinematics.h

@@ -215,6 +215,34 @@ class SwerveDriveKinematics {
       wpi::array<SwerveModuleState, NumModules>* moduleStates,
       units::meters_per_second_t attainableMaxSpeed);
 
+  /**
+   * Renormalizes the wheel speeds if any individual speed is above the
+   * specified maximum, as well as getting rid of joystick saturation at edges
+   * of joystick.
+   *
+   * Sometimes, after inverse kinematics, the requested speed
+   * from one or more modules may be above the max attainable speed for the
+   * driving motor on that module. To fix this issue, one can reduce all the
+   * wheel speeds to make sure that all requested module speeds are at-or-below
+   * the absolute threshold, while maintaining the ratio of speeds between
+   * modules.
+   *
+   * @param moduleStates Reference to array of module states. The array will be
+   * mutated with the normalized speeds!
+   * @param currentChassisSpeed Current speed of the robot
+   * @param attainableMaxModuleSpeed The absolute max speed a module can reach
+   * @param attainableMaxRobotTranslationSpeed The absolute max speed the robot
+   * can reach while translating
+   * @param attainableMaxRobotRotationSpeed The absolute max speed the robot can
+   * reach while rotating
+   */
+  static void DesaturateWheelSpeeds(
+      wpi::array<SwerveModuleState, NumModules>* moduleStates,
+      ChassisSpeeds currentChassisSpeed,
+      units::meters_per_second_t attainableMaxModuleSpeed,
+      units::meters_per_second_t attainableMaxRobotTranslationSpeed,
+      units::radians_per_second_t attainableMaxRobotRotationSpeed);
+
  private:
   mutable Matrixd<NumModules * 2, 3> m_inverseKinematics;
   Eigen::HouseholderQR<Matrixd<NumModules * 2, 3>> m_forwardKinematics;

wpimath/src/main/native/include/frc/kinematics/SwerveDriveKinematics.inc

@@ -7,6 +7,7 @@
 #include <algorithm>
 #include <utility>
 
+#include "frc/kinematics/ChassisSpeeds.h"
 #include "frc/kinematics/SwerveDriveKinematics.h"
 #include "units/math.h"
 
@@ -147,4 +148,41 @@ void SwerveDriveKinematics<NumModules>::DesaturateWheelSpeeds(
   }
 }
 
+template <size_t NumModules>
+void SwerveDriveKinematics<NumModules>::DesaturateWheelSpeeds(
+    wpi::array<SwerveModuleState, NumModules>* moduleStates,
+    ChassisSpeeds currentChassisSpeed,
+    units::meters_per_second_t attainableMaxModuleSpeed,
+    units::meters_per_second_t attainableMaxRobotTranslationSpeed,
+    units::radians_per_second_t attainableMaxRobotRotationSpeed) {
+  auto& states = *moduleStates;
+
+  auto realMaxSpeed = std::max_element(states.begin(), states.end(),
+                                       [](const auto& a, const auto& b) {
+                                         return units::math::abs(a.speed) <
+                                                units::math::abs(b.speed);
+                                       })
+                          ->speed;
+
+  if (attainableMaxRobotTranslationSpeed == 0_mps ||
+      attainableMaxRobotRotationSpeed == 0_rad_per_s || realMaxSpeed == 0_mps) {
+    return;
+  }
+
+  auto translationalK =
+      units::math::hypot(currentChassisSpeed.vx, currentChassisSpeed.vy) /
+      attainableMaxRobotTranslationSpeed;
+
+  auto rotationalK = units::math::abs(currentChassisSpeed.omega) /
+                     attainableMaxRobotRotationSpeed;
+
+  auto k = units::math::max(translationalK, rotationalK);
+
+  auto scale = units::math::min(k * attainableMaxModuleSpeed / realMaxSpeed,
+                                units::scalar_t(1));
+  for (auto& module : states) {
+    module.speed = module.speed * scale;
+  }
+}
+
 }  // namespace frc