Add support for rotated rectangle constraints

choreolib/py/choreo/test/choreolib_test.py

@@ -5,7 +5,7 @@
 TRAJECTORY = """
 {
  "name":"New Path",
- "version":1,
+ "version":2,
  "snapshot":{
   "waypoints":[
     {"x":0.0, "y":0.0, "heading":0.0, "intervals":9, "split":false, "fixTranslation":true, "fixHeading":true, "overrideIntervals":false},

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

@@ -1,14 +1,14 @@
 {
  "name":"test",
- "version":1,
+ "version":2,
  "snapshot":{
   "waypoints":[
     {"x":2.6185336112976074, "y":6.034867286682129, "heading":0.0, "intervals":16, "split":false, "fixTranslation":true, "fixHeading":true, "overrideIntervals":false},
     {"x":6.458227634429932, "y":4.456109523773193, "heading":0.0, "intervals":40, "split":false, "fixTranslation":true, "fixHeading":true, "overrideIntervals":false}],
   "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":0.0, "y":0.0, "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":"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}, "rotation":{"exp":"0 deg", "val":0.0}}}, "enabled":true}],
   "targetDt":{
    "exp":"0.05 s",
    "val":0.05

choreolib/py/choreo/util/traj_schema_version.py

@@ -1,2 +1,2 @@
 # Auto-generated by update_traj_schema.py
-TRAJ_SCHEMA_VERSION = 1
+TRAJ_SCHEMA_VERSION = 2

choreolib/src/main/java/choreo/util/TrajSchemaVersion.java

@@ -7,7 +7,7 @@
 /** Internal autogenerated class for storing the current trajectory schema version. */
 public class TrajSchemaVersion {
   /** The current trajectory schema version. */
-  public static final int TRAJ_SCHEMA_VERSION = 1;
+  public static final int TRAJ_SCHEMA_VERSION = 2;
 
   /** Utility class. */
   private TrajSchemaVersion() {}

choreolib/src/main/native/include/choreo/util/TrajSchemaVersion.h

@@ -9,7 +9,7 @@
 namespace choreo {
 
 [[deprecated("Use kTrajSchemaVersion.")]]
-inline constexpr uint32_t kTrajSpecVersion = 1;
-inline constexpr uint32_t kTrajSchemaVersion = 1;
+inline constexpr uint32_t kTrajSpecVersion = 2;
+inline constexpr uint32_t kTrajSchemaVersion = 2;
 
 }  // namespace choreo

choreolib/src/test/java/choreo/ChoreoTests.java

@@ -15,7 +15,7 @@ public class ChoreoTests {
       """
 {
  "name":"New Path",
- "version":1,
+ "version":2,
  "snapshot":{
   "waypoints":[
     {"x":0.0, "y":0.0, "heading":0.0, "intervals":9, "split":false, "fixTranslation":true, "fixHeading":true, "overrideIntervals":false},

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

@@ -131,9 +131,40 @@ 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 } => {
+                    // Calculate the center of the rectangle
+                    let center_x = x + w / 2.0;
+                    let center_y = y + h / 2.0;
+
+                    // Original corner points relative to bottom-left origin
+                    let corners = vec![
+                        (x, y),         // bottom-left
+                        (x + w, y),     // bottom-right
+                        (x + w, y + h), // top-right
+                        (x, y + h),     // top-left
+                    ];
+
+                    // Apply rotation around center
+                    let cos_r = rotation.cos();
+                    let sin_r = rotation.sin();
+
+                    let mut xs = Vec::new();
+                    let mut ys = Vec::new();
+
+                    for (corner_x, corner_y) in corners {
+                        // Translate to origin (center of rectangle)
+                        let rel_x = corner_x - center_x;
+                        let rel_y = corner_y - center_y;
+
+                        // 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
+                        xs.push(center_x + rotated_x);
+                        ys.push(center_y + 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 +240,40 @@ 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 } => {
+                    // Calculate the center of the rectangle
+                    let center_x = x + w / 2.0;
+                    let center_y = y + h / 2.0;
+
+                    // Original corner points relative to bottom-left origin
+                    let corners = vec![
+                        (x, y),         // bottom-left
+                        (x + w, y),     // bottom-right
+                        (x + w, y + h), // top-right
+                        (x, y + h),     // top-left
+                    ];
+
+                    // Apply rotation around center
+                    let cos_r = rotation.cos();
+                    let sin_r = rotation.sin();
+
+                    let mut xs = Vec::new();
+                    let mut ys = Vec::new();
+
+                    for (corner_x, corner_y) in corners {
+                        // Translate to origin (center of rectangle)
+                        let rel_x = corner_x - center_x;
+                        let rel_y = corner_y - center_y;
+
+                        // 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
+                        xs.push(center_x + rotated_x);
+                        ys.push(center_y + 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 = 1;
+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(),

src-core/src/spec/upgraders.rs

@@ -17,6 +17,7 @@ mod traj_file {
     fn make_upgrader() -> Upgrader {
         let mut upgrader = Upgrader::new(TRAJ_SCHEMA_VERSION);
         upgrader.add_version_action(up_0_1);
+        upgrader.add_version_action(up_1_2);
         // Ensure the new upgrader is added here
         upgrader
     }
@@ -34,6 +35,61 @@ mod traj_file {
         )
     }
 
+    fn up_1_2(editor: &mut Editor) -> ChoreoResult<()> {
+        use crate::spec::Expr;
+        use serde_json::Value as JsonValue;
+
+        // Add rotation field to all KeepInRectangle constraints in both snapshot and params
+
+        // Handle snapshot constraints
+        if editor.has_path("snapshot.constraints") {
+            let snapshot_constraints: Vec<JsonValue> = editor.get_path("snapshot.constraints")?;
+            let mut updated_constraints = Vec::new();
+
+            for mut constraint in snapshot_constraints {
+                if let Some(data_type) = constraint["data"]["type"].as_str() {
+                    if data_type == "KeepInRectangle" {
+                        // Add rotation field to props
+                        if let Some(props) = constraint["data"]["props"].as_object_mut() {
+                            props.insert(
+                                "rotation".to_string(),
+                                JsonValue::Number(serde_json::Number::from_f64(0.0).unwrap()),
+                            );
+                        }
+                    }
+                }
+                updated_constraints.push(constraint);
+            }
+
+            editor.set_path_serialize("snapshot.constraints", updated_constraints)?;
+        }
+
+        // Handle params constraints
+        if editor.has_path("params.constraints") {
+            let params_constraints: Vec<JsonValue> = editor.get_path("params.constraints")?;
+            let mut updated_constraints = Vec::new();
+
+            for mut constraint in params_constraints {
+                if let Some(data_type) = constraint["data"]["type"].as_str() {
+                    if data_type == "KeepInRectangle" {
+                        // Add rotation field to props
+                        if let Some(props) = constraint["data"]["props"].as_object_mut() {
+                            props.insert(
+                                "rotation".to_string(),
+                                serde_json::to_value(Expr::new("0 deg", 0.0))?,
+                            );
+                        }
+                    }
+                }
+                updated_constraints.push(constraint);
+            }
+
+            editor.set_path_serialize("params.constraints", updated_constraints)?;
+        }
+
+        Ok(())
+    }
+
     #[cfg(test)]
     mod tests {
         use crate::spec::upgraders::testing_shared::{get_contents, FileType};