Add direct heading control on swerve input stream

src/main/java/swervelib/SwerveInputStream.java

@@ -71,6 +71,10 @@ public class SwerveInputStream implements Supplier<ChassisSpeeds>
    * Controller supplier as heading.
    */
   private Optional<DoubleSupplier>        controllerHeadingY                  = Optional.empty();
+  /**
+   * Direct heading supplier as Rotation2d.
+   */
+  private Optional<Supplier<Rotation2d>>  headingSupplier                     = Optional.empty();
   /**
    * Axis deadband for the controller.
    */
@@ -151,6 +155,10 @@ public class SwerveInputStream implements Supplier<ChassisSpeeds>
    * Current {@link SwerveInputMode} to use.
    */
   private SwerveInputMode                 currentMode                         = SwerveInputMode.ANGULAR_VELOCITY;
+  /**
+   * Output {@link ChassisSpeeds} based on direct heading while this is True.
+   */
+  private Optional<BooleanSupplier>       directHeadingEnabled                = Optional.empty();
 
 
   /**
@@ -222,6 +230,7 @@ public SwerveInputStream copy()
     newStream.controllerOmega = controllerOmega;
     newStream.controllerHeadingX = controllerHeadingX;
     newStream.controllerHeadingY = controllerHeadingY;
+    newStream.headingSupplier = headingSupplier;
     newStream.axisDeadband = axisDeadband;
     newStream.translationAxisScale = translationAxisScale;
     newStream.omegaAxisScale = omegaAxisScale;
@@ -242,6 +251,7 @@ public SwerveInputStream copy()
     newStream.allianceRelative = allianceRelative;
     newStream.translationHeadingOffsetEnabled = translationHeadingOffsetEnabled;
     newStream.translationHeadingOffset = translationHeadingOffset;
+    newStream.directHeadingEnabled = directHeadingEnabled;
     return newStream;
   }
 
@@ -451,6 +461,21 @@ public SwerveInputStream withControllerHeadingAxis(DoubleSupplier headingX, Doub
   {
     controllerHeadingX = Optional.of(headingX);
     controllerHeadingY = Optional.of(headingY);
+    headingSupplier = Optional.empty(); // Clear any direct heading supplier
+    return this;
+  }
+
+  /**
+   * Add direct heading supplier for Heading based control.
+   *
+   * @param heading Supplier that provides the desired heading as Rotation2d
+   * @return self
+   */
+  public SwerveInputStream withHeading(Supplier<Rotation2d> heading)
+  {
+    headingSupplier = Optional.of(heading);
+    controllerHeadingX = Optional.empty(); // Clear controller axes
+    controllerHeadingY = Optional.empty(); // Clear controller axes
     return this;
   }
 
@@ -520,6 +545,36 @@ public SwerveInputStream headingWhile(boolean headingState)
     return this;
   }
 
+  /**
+   * Enable direct heading control while the supplier is True.
+   *
+   * @param trigger Supplier to use.
+   * @return this.
+   */
+  public SwerveInputStream directHeadingWhile(BooleanSupplier trigger)
+  {
+    directHeadingEnabled = Optional.of(trigger);
+    return this;
+  }
+
+  /**
+   * Set the direct heading enable state.
+   *
+   * @param headingState Direct heading enabled state.
+   * @return this
+   */
+  public SwerveInputStream directHeadingWhile(boolean headingState)
+  {
+    if (headingState)
+    {
+      directHeadingEnabled = Optional.of(() -> true);
+    } else
+    {
+      directHeadingEnabled = Optional.empty();
+    }
+    return this;
+  }
+
   /**
    * Aim the {@link SwerveDrive} at this pose while driving.
    *
@@ -625,15 +680,26 @@ private SwerveInputMode findMode()
             "Attempting to enter AIM mode without target, please use SwerveInputStream.aim() to select a target first!",
             false);
       }
+    } else if (directHeadingEnabled.isPresent() && directHeadingEnabled.get().getAsBoolean())
+    {
+      if (headingSupplier.isPresent())
+      {
+        return SwerveInputMode.DIRECT_HEADING;
+      } else
+      {
+        DriverStation.reportError(
+            "Attempting to enter DIRECT_HEADING mode without heading information, please use SwerveInputStream.withHeading to add heading supplier!",
+            false);
+      }
     } else if (headingEnabled.isPresent() && headingEnabled.get().getAsBoolean())
     {
-      if (controllerHeadingX.isPresent() && controllerHeadingY.isPresent())
+      if (headingSupplier.isPresent() || (controllerHeadingX.isPresent() && controllerHeadingY.isPresent()))
       {
         return SwerveInputMode.HEADING;
       } else
       {
         DriverStation.reportError(
-            "Attempting to enter HEADING mode without heading axis, please use SwerveInputStream.withControllerHeadingAxis to add heading axis!",
+            "Attempting to enter HEADING mode without heading information, please use SwerveInputStream.withHeading or SwerveInputStream.withControllerHeadingAxis to add heading axis!",
             false);
       }
     } else if (controllerOmega.isEmpty())
@@ -822,6 +888,32 @@ private Translation2d applyAllianceAwareTranslation(Translation2d fieldRelativeT
     return fieldRelativeTranslation;
   }
 
+  /**
+   * Apply alliance aware translation which flips the {@link Rotation2d} if the robot is on the Blue alliance.
+   *
+   * @param fieldRelativeRotation Field-relative {@link Rotation2d} to flip.
+   * @return Alliance-oriented {@link Rotation2d}
+   */
+  private Rotation2d applyAllianceAwareRotation(Rotation2d fieldRelativeRotation)
+  {
+    if (allianceRelative.isPresent() && allianceRelative.get().getAsBoolean())
+    {
+      if (robotRelative.isPresent() && robotRelative.get().getAsBoolean())
+      {
+        if (driveToPoseEnabled.isPresent() && driveToPoseEnabled.get().getAsBoolean())
+        {
+          return fieldRelativeRotation;
+        }
+        throw new RuntimeException("Cannot use robot oriented control with Alliance aware movement!");
+      }
+      if (DriverStation.getAlliance().isPresent() && DriverStation.getAlliance().get() == Alliance.Red)
+      {
+        return fieldRelativeRotation.rotateBy(Rotation2d.k180deg);
+      }
+    }
+    return fieldRelativeRotation;
+  }
+
   /**
    * Adds offset to translation if one is set.
    *
@@ -841,6 +933,24 @@ private ChassisSpeeds applyTranslationHeadingOffset(ChassisSpeeds speeds)
     }
     return speeds;
   }
+
+  /**
+   * Adds offset to heading if one is set.
+   *
+   * @param fieldRelativeRotation Field-relative {@link Rotation2d} to offset
+   * @return Offsetted {@link Rotation2d}
+   */
+  private Rotation2d applyHeadingOffset(Rotation2d fieldRelativeRotation)
+  {
+    if (translationHeadingOffsetEnabled.isPresent() && translationHeadingOffsetEnabled.get().getAsBoolean())
+    {
+      if (translationHeadingOffset.isPresent())
+      {
+        return fieldRelativeRotation.rotateBy(translationHeadingOffset.get());
+      }
+    }
+    return fieldRelativeRotation;
+  }
 
   /**
    * When the {@link SwerveInputStream} is in {@link SwerveInputMode#DRIVE_TO_POSE} this function will return if the
@@ -916,14 +1026,15 @@ public ChassisSpeeds get()
       }
       case HEADING ->
       {
+        // Use controller joystick inputs
         omegaRadiansPerSecond = swerveController.headingCalculate(swerveDrive.getOdometryHeading().getRadians(),
-                                                                  Rotation2d.fromRadians(
-                                                                                swerveController.getJoystickAngle(
-                                                                                    controllerHeadingX.get()
-                                                                                                      .getAsDouble(),
-                                                                                    controllerHeadingY.get()
-                                                                                                      .getAsDouble()))
-                                                                            .getRadians());
+                                                                Rotation2d.fromRadians(
+                                                                              swerveController.getJoystickAngle(
+                                                                                  controllerHeadingX.get()
+                                                                                                    .getAsDouble(),
+                                                                                  controllerHeadingY.get()
+                                                                                                    .getAsDouble()))
+                                                                          .getRadians());
 
         // Prevent rotation if controller heading inputs are not past axisDeadband
         if (Math.abs(controllerHeadingX.get().getAsDouble()) + Math.abs(controllerHeadingY.get().getAsDouble()) <
@@ -934,6 +1045,20 @@ public ChassisSpeeds get()
         speeds = new ChassisSpeeds(vxMetersPerSecond, vyMetersPerSecond, omegaRadiansPerSecond);
         break;
       }
+      case DIRECT_HEADING ->
+      {
+        // Use direct heading supplier
+        Rotation2d targetHeading = applyHeadingOffset(
+                                      applyAllianceAwareRotation(
+                                        headingSupplier.get().get()));
+
+        omegaRadiansPerSecond = swerveController.headingCalculate(
+            swerveDrive.getOdometryHeading().getRadians(),
+            targetHeading.getRadians());
+
+        speeds = new ChassisSpeeds(vxMetersPerSecond, vyMetersPerSecond, omegaRadiansPerSecond);
+        break;
+      }
       case AIM ->
       {
         Rotation2d    currentHeading = swerveDrive.getOdometryHeading();
@@ -1005,9 +1130,13 @@ enum SwerveInputMode
      */
     ANGULAR_VELOCITY,
     /**
-     * Output based off of heading.
+     * Output based off of heading from controller axes.
      */
     HEADING,
+    /**
+     * Output based off of direct heading supplier.
+     */
+    DIRECT_HEADING,
     /**
      * Output based off of targeting.
      */