Add mecanum support

src-core/src/generation/generate.rs

@@ -3,7 +3,7 @@
 use std::sync::OnceLock;
 use std::sync::mpsc::{Receiver, Sender, channel};
 
-use trajoptlib::{DifferentialTrajectory, SwerveTrajectory};
+use trajoptlib::{DifferentialTrajectory, MecanumTrajectory, SwerveTrajectory};
 
 use super::transformers::{
     CallbackSetter, ConstraintSetter, DrivetrainAndBumpersSetter, IntervalCountSetter,
@@ -30,6 +30,9 @@ pub enum LocalProgressUpdate {
         // Diff variant
         update: Vec<Sample>,
     },
+    MecanumTrajectory {
+        update: Vec<Sample>,
+    },
     DiagnosticText {
         update: String,
     },
@@ -68,6 +71,14 @@ impl From<Vec<usize>> for LocalProgressUpdate {
     }
 }
 
+impl From<MecanumTrajectory> for LocalProgressUpdate {
+    fn from(trajectory: MecanumTrajectory) -> Self {
+        LocalProgressUpdate::MecanumTrajectory {
+            update: trajectory.samples.iter().map(Sample::from).collect(),
+        }
+    }
+}
+
 pub struct HandledLocalProgressUpdate {
     pub handle: i64,
     pub update: LocalProgressUpdate,

src-core/src/generation/remote.rs

@@ -94,6 +94,9 @@ pub enum RemoteProgressUpdate {
     IncompleteTankTrajectory(Vec<Sample>),
     IntervalCounts(Vec<usize>),
     CompleteTrajectory(TrajectoryFile),
+    // Mecanum variant
+    IncompleteMecanumTrajectory(Vec<Sample>),
+    CompleteTrajectory(Trajectory),
     Error(ChoreoError),
 }
 
@@ -123,6 +126,11 @@ pub fn remote_generate_child(args: RemoteArgs) {
                         update: LocalProgressUpdate::IntervalCounts { update },
                         ..
                     } => serde_json::to_string(&RemoteProgressUpdate::IntervalCounts(update)),
+                        update: LocalProgressUpdate::MecanumTrajectory { update },
+                        ..
+                    } => serde_json::to_string(&RemoteProgressUpdate::IncompleteMecanumTrajectory(
+                        update,
+                    )),
                     _ => continue,
                 }
                 .expect("Failed to serialize progress update");
@@ -317,6 +325,13 @@ pub async fn remote_generate_parent(
                                     }.handled(handle)
                                 );
                             },
+                            Ok(RemoteProgressUpdate::IncompleteMecanumTrajectory(trajectory)) => {
+                                remote_resources.emit_progress(
+                                    LocalProgressUpdate::MecanumTrajectory {
+                                        update: trajectory
+                                    }.handled(handle)
+                                );
+                            },
                             Ok(RemoteProgressUpdate::CompleteTrajectory(trajectory)) => {
                                 break Ok(trajectory);
                             },

src-core/src/generation/transformers/callback.rs

@@ -1,4 +1,4 @@
-use trajoptlib::{DifferentialTrajectory, SwerveTrajectory};
+use trajoptlib::{DifferentialTrajectory, MecanumTrajectory, SwerveTrajectory};
 
 use crate::{
     ResultExt,
@@ -7,7 +7,7 @@ use crate::{
 
 use super::{
     DifferentialGenerationTransformer, FeatureLockedTransformer, GenerationContext,
-    SwerveGenerationTransformer,
+    MecanumGenerationTransformer, SwerveGenerationTransformer,
 };
 
 pub struct CallbackSetter;
@@ -30,6 +30,15 @@ fn differential_status_callback(trajectory: DifferentialTrajectory, handle: i64)
     };
 }
 
+fn mecanum_status_callback(trajectory: MecanumTrajectory, handle: i64) {
+    let tx_opt = PROGRESS_SENDER_LOCK.get();
+    if let Some(tx) = tx_opt {
+        let _ = tx
+            .send(LocalProgressUpdate::from(trajectory).handled(handle))
+            .trace_warn();
+    };
+}
+
 impl SwerveGenerationTransformer for CallbackSetter {
     fn initialize(_: &GenerationContext) -> FeatureLockedTransformer<Self> {
         FeatureLockedTransformer::always(Self)
@@ -49,3 +58,13 @@ impl DifferentialGenerationTransformer for CallbackSetter {
         generator.add_callback(differential_status_callback);
     }
 }
+
+impl MecanumGenerationTransformer for CallbackSetter {
+    fn initialize(_: &GenerationContext) -> FeatureLockedTransformer<Self> {
+        FeatureLockedTransformer::always(Self)
+    }
+
+    fn transform(&self, generator: &mut trajoptlib::MecanumTrajectoryGenerator) {
+        generator.add_callback(mecanum_status_callback);
+    }
+}

src-core/src/generation/transformers/constraints.rs

@@ -2,7 +2,7 @@ use crate::spec::trajectory::{ConstraintData, ConstraintIDX, ConstraintScope, Wa
 
 use super::{
     DifferentialGenerationTransformer, FeatureLockedTransformer, GenerationContext,
-    SwerveGenerationTransformer,
+    MecanumGenerationTransformer, SwerveGenerationTransformer,
 };
 
 fn fix_scope(idx: usize, removed_idxs: &[usize]) -> usize {
@@ -228,3 +228,71 @@ impl DifferentialGenerationTransformer for ConstraintSetter {
         }
     }
 }
+
+impl MecanumGenerationTransformer for ConstraintSetter {
+    fn initialize(context: &GenerationContext) -> FeatureLockedTransformer<Self> {
+        Self::initialize(context)
+    }
+
+    fn transform(&self, generator: &mut trajoptlib::MecanumTrajectoryGenerator) {
+        for constraint in &self.constraint_idx {
+            let from = fix_scope(constraint.from, &self.guess_points);
+            let to_opt = constraint.to.map(|idx| fix_scope(idx, &self.guess_points));
+            match constraint.data {
+                ConstraintData::PointAt {
+                    x,
+                    y,
+                    tolerance,
+                    flip,
+                } => match to_opt {
+                    None => generator.wpt_point_at(from, x, y, tolerance, flip),
+                    Some(to) => generator.sgmt_point_at(from, to, x, y, tolerance, flip),
+                },
+                ConstraintData::MaxVelocity { max } => match to_opt {
+                    None => generator.wpt_linear_velocity_max_magnitude(from, max),
+                    Some(to) => generator.sgmt_linear_velocity_max_magnitude(from, to, max),
+                },
+                ConstraintData::MaxAcceleration { max } => match to_opt {
+                    None => generator.wpt_linear_acceleration_max_magnitude(from, max),
+                    Some(to) => generator.sgmt_linear_acceleration_max_magnitude(from, to, max),
+                },
+                ConstraintData::MaxAngularVelocity { max } => match to_opt {
+                    None => generator.wpt_angular_velocity_max_magnitude(from, max),
+                    Some(to) => generator.sgmt_angular_velocity_max_magnitude(from, to, max),
+                },
+                ConstraintData::StopPoint {} => {
+                    if to_opt.is_none() {
+                        generator.wpt_linear_velocity_max_magnitude(from, 0.0f64);
+                        generator.wpt_angular_velocity_max_magnitude(from, 0.0f64);
+                    }
+                }
+                ConstraintData::KeepInCircle { x, y, r } => match to_opt {
+                    None => generator.wpt_keep_in_circle(from, x, y, r),
+                    Some(to) => generator.sgmt_keep_in_circle(from, to, x, y, r),
+                },
+                ConstraintData::KeepInRectangle { x, y, w, h } => {
+                    let xs = vec![x, x + w, x + w, x];
+                    let ys = vec![y, y, y + h, y + h];
+                    match to_opt {
+                        None => generator.wpt_keep_in_polygon(from, xs, ys),
+                        Some(to) => generator.sgmt_keep_in_polygon(from, to, xs, ys),
+                    }
+                }
+                ConstraintData::KeepInLane { tolerance } => {
+                    if let Some(idx_to) = to_opt
+                        && let Some(wpt_from) = self.waypoint_idx.get(from)
+                        && let Some(wpt_to) = self.waypoint_idx.get(idx_to)
+                    {
+                        generator.sgmt_keep_in_lane(
+                            from, idx_to, wpt_from.x, wpt_from.y, wpt_to.x, wpt_to.y, tolerance,
+                        );
+                    }
+                }
+                ConstraintData::KeepOutCircle { x, y, r } => match to_opt {
+                    None => generator.wpt_keep_out_circle(from, x, y, r),
+                    Some(to) => generator.sgmt_keep_out_circle(from, to, x, y, r),
+                },
+            };
+        }
+    }
+}

src-core/src/generation/transformers/drivetrain_and_bumpers.rs

@@ -1,10 +1,10 @@
-use trajoptlib::{DifferentialDrivetrain, SwerveDrivetrain};
+use trajoptlib::{DifferentialDrivetrain, MecanumDrivetrain, SwerveDrivetrain};
 
 use crate::spec::project::RobotConfig;
 
 use super::{
     DifferentialGenerationTransformer, FeatureLockedTransformer, GenerationContext,
-    SwerveGenerationTransformer,
+    MecanumGenerationTransformer, SwerveGenerationTransformer,
 };
 
 pub struct DrivetrainAndBumpersSetter {
@@ -70,3 +70,35 @@ impl DifferentialGenerationTransformer for DrivetrainAndBumpersSetter {
         );
     }
 }
+
+impl MecanumGenerationTransformer for DrivetrainAndBumpersSetter {
+    fn initialize(context: &GenerationContext) -> FeatureLockedTransformer<Self> {
+        FeatureLockedTransformer::always(Self {
+            config: context.project.config.snapshot(),
+        })
+    }
+
+    fn transform(&self, generator: &mut trajoptlib::MecanumTrajectoryGenerator) {
+        let config = &self.config;
+        let drivetrain = MecanumDrivetrain {
+            mass: config.mass,
+            moi: config.inertia,
+            wheel_radius: config.radius,
+            // rad per sec
+            wheel_max_angular_velocity: config.vmax / config.gearing,
+            wheel_max_torque: config.tmax * config.gearing,
+            wheel_cof: config.cof,
+            static_friction_coefficient: 0.0,
+            strafe_efficiency: 0.75,
+            modules: config.module_translations(),
+        };
+
+        generator.set_drivetrain(&drivetrain);
+        generator.set_bumpers(
+            config.bumper.front,
+            config.bumper.side,
+            config.bumper.side,
+            config.bumper.back,
+        );
+    }
+}

src-core/src/generation/transformers/interval_count.rs

@@ -4,7 +4,7 @@ use crate::spec::trajectory::Waypoint;
 
 use super::{
     DifferentialGenerationTransformer, FeatureLockedTransformer, GenerationContext,
-    SwerveGenerationTransformer,
+    MecanumGenerationTransformer, SwerveGenerationTransformer,
 };
 
 pub struct IntervalCountSetter {
@@ -108,3 +108,56 @@ impl DifferentialGenerationTransformer for IntervalCountSetter {
         generator.set_control_interval_counts(control_interval_counts);
     }
 }
+
+impl MecanumGenerationTransformer for IntervalCountSetter {
+    fn initialize(context: &GenerationContext) -> FeatureLockedTransformer<Self> {
+        FeatureLockedTransformer::always(Self {
+            counts: guess_control_interval_counts(
+                &context.project.config.snapshot(),
+                &context.params,
+            )
+            .unwrap_or_default(),
+            waypoints: context.params.waypoints.clone(),
+        })
+    }
+
+    fn transform(&self, generator: &mut trajoptlib::MecanumTrajectoryGenerator) {
+        let waypoints = &self.waypoints;
+        let mut guess_points_after_waypoint = Vec::new();
+        let mut control_interval_counts = Vec::new();
+        let mut wpt_cnt = 0;
+        for i in 0..waypoints.len() {
+            let wpt = &waypoints[i];
+            // add initial guess points (actually unconstrained empty wpts in Choreo terms)
+            if wpt.is_initial_guess && !wpt.fix_heading && !wpt.fix_translation {
+                let guess_point = Pose2d {
+                    x: wpt.x,
+                    y: wpt.y,
+                    heading: wpt.heading,
+                };
+                guess_points_after_waypoint.push(guess_point);
+                if let Some(last) = control_interval_counts.last_mut() {
+                    *last += self.counts[i];
+                }
+            } else {
+                if wpt_cnt > 0 {
+                    generator.sgmt_initial_guess_points(wpt_cnt - 1, &guess_points_after_waypoint);
+                }
+                guess_points_after_waypoint.clear();
+                if wpt.fix_heading && wpt.fix_translation {
+                    generator.pose_wpt(wpt_cnt, wpt.x, wpt.y, wpt.heading);
+                } else if wpt.fix_translation {
+                    generator.translation_wpt(wpt_cnt, wpt.x, wpt.y, wpt.heading);
+                } else {
+                    generator.empty_wpt(wpt_cnt, wpt.x, wpt.y, wpt.heading);
+                }
+                wpt_cnt += 1;
+                if i != waypoints.len() - 1 {
+                    control_interval_counts.push(self.counts[i]);
+                }
+            }
+        }
+
+        generator.set_control_interval_counts(control_interval_counts);
+    }
+}