Add support for rotated rectangle constraints

choreolib/py/choreo/test/resources/swerve_test.traj

@@ -8,7 +8,7 @@
   "constraints":[
     {"from":"first", "to":null, "data":{"type":"StopPoint", "props":{}}, "enabled":true},
     {"from":"last", "to":null, "data":{"type":"StopPoint", "props":{}}, "enabled":true},
-    {"from":"first", "to":"last", "data":{"type":"KeepInRectangle", "props":{"x":0.0, "y":0.0, "w":16.54, "h":8.21}}, "enabled":true}],
+    {"from":"first", "to":"last", "data":{"type":"KeepInRectangle", "props":{"x":8.27, "y":4.105, "w":16.54, "h":8.21, "rotation":0.0}}, "enabled":true}],
   "targetDt":0.05
  },
  "params":{
@@ -18,7 +18,7 @@
   "constraints":[
     {"from":"first", "to":null, "data":{"type":"StopPoint", "props":{}}, "enabled":true},
     {"from":"last", "to":null, "data":{"type":"StopPoint", "props":{}}, "enabled":true},
-    {"from":"first", "to":"last", "data":{"type":"KeepInRectangle", "props":{"x":{"exp":"0 m", "val":0.0}, "y":{"exp":"0 m", "val":0.0}, "w":{"exp":"16.54 m", "val":16.54}, "h":{"exp":"8.21 m", "val":8.21}}}, "enabled":true}],
+    {"from":"first", "to":"last", "data":{"type":"KeepInRectangle", "props":{"x":{"exp":"8.27 m", "val":8.27}, "y":{"exp":"4.105 m", "val":4.105}, "w":{"exp":"16.54 m", "val":16.54}, "h":{"exp":"8.21 m", "val":8.21}, "rotation":{"exp":"0 deg", "val":0.0}}}, "enabled":true}],
   "targetDt":{
    "exp":"0.05 s",
    "val":0.05

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

@@ -15,6 +15,36 @@ fn fix_scope(idx: usize, removed_idxs: &[usize]) -> usize {
     idx - to_subtract
 }
 
+/// Rotates a point around another point in 2D space.
+///
+/// ```text
+/// [x_new]   [rot.cos, -rot.sin][x - other.x]   [other.x]
+/// [y_new] = [rot.sin,  rot.cos][y - other.y] + [other.y]
+/// ```
+///
+/// # Arguments
+/// * `point` - The point to rotate as (x, y)
+/// * `center` - The center point to rotate around as (x, y)
+/// * `rotation` - The rotation angle in radians
+///
+/// # Returns
+/// The new rotated point as (x, y)
+fn rotate_around(point: (f64, f64), center: (f64, f64), rotation: f64) -> (f64, f64) {
+    let cos_r = rotation.cos();
+    let sin_r = rotation.sin();
+
+    // Translate to origin (center of rotation)
+    let rel_x = point.0 - center.0;
+    let rel_y = point.1 - center.1;
+
+    // Apply rotation
+    let rotated_x = rel_x * cos_r - rel_y * sin_r;
+    let rotated_y = rel_x * sin_r + rel_y * cos_r;
+
+    // Translate back
+    (center.0 + rotated_x, center.1 + rotated_y)
+}
+
 pub struct ConstraintSetter {
     guess_points: Vec<usize>,
     constraint_idx: Vec<ConstraintIDX<f64>>,
@@ -131,9 +161,27 @@ impl SwerveGenerationTransformer for ConstraintSetter {
                     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];
+                ConstraintData::KeepInRectangle { x, y, w, h, rotation } => {
+                    // x, y now represent the center of the rectangle
+                    let center = (x, y);
+
+                    // Original corner points relative to center
+                    let corners = vec![
+                        (x - w / 2.0, y - h / 2.0), // bottom-left
+                        (x + w / 2.0, y - h / 2.0), // bottom-right
+                        (x + w / 2.0, y + h / 2.0), // top-right
+                        (x - w / 2.0, y + h / 2.0), // top-left
+                    ];
+
+                    let mut xs = Vec::new();
+                    let mut ys = Vec::new();
+
+                    for corner in corners {
+                        let (rotated_x, rotated_y) = rotate_around(corner, center, rotation);
+                        xs.push(rotated_x);
+                        ys.push(rotated_y);
+                    }
+
                     match to_opt {
                         None => generator.wpt_keep_in_polygon(from, xs, ys),
                         Some(to) => generator.sgmt_keep_in_polygon(from, to, xs, ys),
@@ -209,9 +257,27 @@ impl DifferentialGenerationTransformer for ConstraintSetter {
                     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];
+                ConstraintData::KeepInRectangle { x, y, w, h, rotation } => {
+                    // x, y now represent the center of the rectangle
+                    let center = (x, y);
+
+                    // Original corner points relative to center
+                    let corners = vec![
+                        (x - w / 2.0, y - h / 2.0), // bottom-left
+                        (x + w / 2.0, y - h / 2.0), // bottom-right
+                        (x + w / 2.0, y + h / 2.0), // top-right
+                        (x - w / 2.0, y + h / 2.0), // top-left
+                    ];
+
+                    let mut xs = Vec::new();
+                    let mut ys = Vec::new();
+
+                    for corner in corners {
+                        let (rotated_x, rotated_y) = rotate_around(corner, center, rotation);
+                        xs.push(rotated_x);
+                        ys.push(rotated_y);
+                    }
+
                     match to_opt {
                         None => generator.wpt_keep_in_polygon(from, xs, ys),
                         Some(to) => generator.sgmt_keep_in_polygon(from, to, xs, ys),

src-core/src/spec/traj_schema_version.rs

@@ -1,2 +1,2 @@
-// Auto-generated by update_traj_schema.py
-pub const TRAJ_SCHEMA_VERSION: u32 = 2;
+// Auto-generated by update_traj_schema.py
+pub const TRAJ_SCHEMA_VERSION: u32 = 2;

src-core/src/spec/trajectory.rs

@@ -146,7 +146,7 @@ pub enum ConstraintData<T: SnapshottableType> {
     /// A constraint to contain the bumpers within a circlular region of the field
     KeepInCircle { x: T, y: T, r: T },
     /// A constraint to contain the bumpers within a rectangular region of the field
-    KeepInRectangle { x: T, y: T, w: T, h: T },
+    KeepInRectangle { x: T, y: T, w: T, h: T, rotation: T },
     /// A constraint to contain the bumpers within two line
     KeepInLane { tolerance: T },
     /// A constraint to contain the bumpers outside a circlular region of the field
@@ -192,11 +192,18 @@ impl<T: SnapshottableType> ConstraintData<T> {
                 y: y.snapshot(),
                 r: r.snapshot(),
             },
-            ConstraintData::KeepInRectangle { x, y, w, h } => ConstraintData::KeepInRectangle {
+            ConstraintData::KeepInRectangle {
+                x,
+                y,
+                w,
+                h,
+                rotation,
+            } => ConstraintData::KeepInRectangle {
                 x: x.snapshot(),
                 y: y.snapshot(),
                 w: w.snapshot(),
                 h: h.snapshot(),
+                rotation: rotation.snapshot(),
             },
             ConstraintData::KeepInLane { tolerance } => ConstraintData::KeepInLane {
                 tolerance: tolerance.snapshot(),