[choreo] Add Lane Constraint to UI

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

@@ -1,6 +1,9 @@
-use crate::spec::trajectory::{ConstraintData, ConstraintIDX, ConstraintScope};
+use crate::spec::trajectory::{ConstraintData, ConstraintIDX, ConstraintScope, Waypoint};
 
-use super::{DifferentialGenerationTransformer, FeatureLockedTransformer, GenerationContext, SwerveGenerationTransformer};
+use super::{
+    DifferentialGenerationTransformer, FeatureLockedTransformer, GenerationContext,
+    SwerveGenerationTransformer,
+};
 
 fn fix_scope(idx: usize, removed_idxs: &[usize]) -> usize {
     let mut to_subtract: usize = 0;
@@ -14,19 +17,30 @@ fn fix_scope(idx: usize, removed_idxs: &[usize]) -> usize {
 
 pub struct ConstraintSetter {
     guess_points: Vec<usize>,
-    constraint_idx: Vec<ConstraintIDX<f64>>
+    constraint_idx: Vec<ConstraintIDX<f64>>,
+    /// A vector of remaining waypoints matching the indexing scheme of constraint_idx
+    waypoint_idx: Vec<Waypoint<f64>>
 }
 
 impl ConstraintSetter {
     fn initialize(context: &GenerationContext) -> FeatureLockedTransformer<Self> {
         let mut guess_points: Vec<usize> = Vec::new();
         let mut constraint_idx = Vec::<ConstraintIDX<f64>>::new();
+        let mut waypoint_idx = Vec::<Waypoint<f64>>::new();
         let num_wpts = context.params.waypoints.len();
 
-        context.params.waypoints.iter()
+        context.params
+            .waypoints
+            .iter()
             .enumerate()
-            .filter(|(_, w)| w.is_initial_guess && !w.fix_heading && !w.fix_translation)
-            .for_each(|(idx, _)| guess_points.push(idx));
+            .for_each(|(idx, w)| {
+                if w.is_initial_guess && !w.fix_heading && !w.fix_translation {
+                    // filtered out, save the index
+                    guess_points.push(idx);
+                } else {
+                    waypoint_idx.push(*w);
+                }
+            });
 
         for constraint in context.params.get_enabled_constraints() {
             let from = constraint.from.get_idx(num_wpts);
@@ -61,7 +75,7 @@ impl ConstraintSetter {
                             from: fixed_from,
                             to: fixed_to,
                             data: constraint.data,
-                            enabled: constraint.enabled
+                            enabled: constraint.enabled,
                         });
                     }
                 }
@@ -70,7 +84,8 @@ impl ConstraintSetter {
 
         FeatureLockedTransformer::always(Self {
             guess_points,
-            constraint_idx
+            constraint_idx,
+            waypoint_idx
         })
     }
 }
@@ -92,7 +107,7 @@ impl SwerveGenerationTransformer for ConstraintSetter {
                     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)
+                    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),
@@ -123,13 +138,30 @@ impl SwerveGenerationTransformer for ConstraintSetter {
                         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 {
+                      if let Some(wpt_to) = self.waypoint_idx.get(idx_to) {
+                      if let Some(wpt_from) = self.waypoint_idx.get(from) {
+                        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),
                     }
                 },
+                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),
+                },
             };
         }
     }
@@ -150,7 +182,11 @@ impl DifferentialGenerationTransformer for ConstraintSetter {
                     y,
                     tolerance,
                     flip,
-                } => if to_opt.is_none() { generator.wpt_point_at(from, x, y, tolerance, flip) },
+                } => {
+                    if to_opt.is_none() {
+                        generator.wpt_point_at(from, 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),
@@ -168,7 +204,7 @@ impl DifferentialGenerationTransformer for ConstraintSetter {
                         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),
@@ -180,13 +216,30 @@ impl DifferentialGenerationTransformer for ConstraintSetter {
                         None => generator.wpt_keep_in_polygon(from, xs, ys),
                         Some(to) => generator.sgmt_keep_in_polygon(from, to, xs, ys),
                     }
-                }
-                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),
+                },
+                ConstraintData::KeepInLane {
+                    tolerance
+                } => {
+                    if let Some(idx_to) = to_opt {
+                      if let Some(wpt_to) = self.waypoint_idx.get(idx_to) {
+                      if let Some(wpt_from) = self.waypoint_idx.get(from) {
+                        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/spec/trajectory.rs

@@ -147,6 +147,8 @@ pub enum ConstraintData<T: SnapshottableType> {
     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 },
+    /// 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
     KeepOutCircle { x: T, y: T, r: T },
 }
@@ -156,6 +158,7 @@ impl<T: SnapshottableType> ConstraintData<T> {
     pub fn scope(&self) -> ConstraintScope {
         match self {
             ConstraintData::StopPoint {} => ConstraintScope::Waypoint,
+            ConstraintData::KeepInLane { tolerance: _ } => ConstraintScope::Segment,
             _ => ConstraintScope::Both,
         }
     }
@@ -195,6 +198,9 @@ impl<T: SnapshottableType> ConstraintData<T> {
                 w: w.snapshot(),
                 h: h.snapshot(),
             },
+            ConstraintData::KeepInLane { tolerance } => ConstraintData::KeepInLane {
+                tolerance: tolerance.snapshot(),
+            },
             ConstraintData::KeepOutCircle { x, y, r } => ConstraintData::KeepOutCircle {
                 x: x.snapshot(),
                 y: y.snapshot(),

src/assets/KeepInLane.tsx

@@ -0,0 +1,23 @@
+import * as React from "react";
+import { SvgIcon as MuiSvgIcon, SvgIconProps, styled } from "@mui/material";
+const SvgIcon = styled(MuiSvgIcon, {
+  name: "KeepInLaneIcon",
+  shouldForwardProp: (prop) => prop !== "fill"
+})<SvgIconProps>(() => ({
+  fill: "currentColor",
+  stroke: "none"
+}));
+
+const KeepInLane: React.FunctionComponent<SvgIconProps> = (props) => {
+  return (
+    <SvgIcon
+      viewBox="0 -960 960 960"
+      focusable="false"
+      aria-hidden="true"
+      {...props}
+    >
+      <path d="M160-160v-640h80v640h-80Zm280 0v-160h80v160h-80Zm280 0v-640h80v640h-80ZM440-400v-160h80v160h-80Zm0-240v-160h80v160h-80Z" />
+    </SvgIcon>
+  );
+};
+export default KeepInLane;

src/components/field/svg/constraintDisplay/FieldConstraintDisplayLayer.tsx

@@ -7,6 +7,7 @@ import KeepInCircleOverlay from "./KeepInCircleOverlay";
 import KeepInRectangleOverlay from "./KeepInRectangleOverlay";
 import PointAtOverlay from "./PointAtOverlay";
 import KeepOutCircleOverlay from "./KeepOutCircleOverlay";
+import KeepInLaneOverlay from "./KeepInLaneOverlay";
 
 const overlays = {
   PointAt: (
@@ -42,6 +43,13 @@ const overlays = {
       lineColor={lineColor}
     ></KeepInRectangleOverlay>
   ),
+  KeepInLane: (constraint: IConstraintStoreKeyed<"KeepInLane">) => (
+    <KeepInLaneOverlay
+      data={constraint.data}
+      start={constraint.getStartWaypoint()}
+      end={constraint.getEndWaypoint()}
+    ></KeepInLaneOverlay>
+  ),
   KeepOutCircle: (
     constraint: IConstraintStoreKeyed<"KeepOutCircle">,
     lineColor: string

src/components/field/svg/constraintDisplay/KeepInLaneOverlay.tsx

@@ -0,0 +1,143 @@
+import React, { Component } from "react";
+import { IConstraintDataStore } from "../../../../document/ConstraintDataStore";
+import {
+  ConstraintKey,
+  DataMap
+} from "../../../../document/ConstraintDefinitions";
+import { IHolonomicWaypointStore } from "../../../../document/HolonomicWaypointStore";
+import * as d3 from "d3";
+import { observer } from "mobx-react";
+import { doc } from "../../../../document/DocumentManager";
+
+const STROKE = 0.02;
+const DOT = 0.1;
+
+type Props<K extends ConstraintKey> = {
+  data: IConstraintDataStore<K>;
+  start?: IHolonomicWaypointStore;
+  end?: IHolonomicWaypointStore;
+};
+class KeepInLaneOverlay extends Component<Props<"KeepInLane">, object> {
+  rootRef: React.RefObject<SVGGElement> = React.createRef<SVGGElement>();
+  componentDidMount() {
+    if (this.rootRef.current) {
+      const dragHandleDrag = d3
+        .drag<SVGCircleElement, undefined>()
+        .on("drag", (event) => this.dragPointTolerance(event))
+        .on("start", () => {
+          doc.history.startGroup(() => {});
+        })
+        .on("end", (_event) => {
+          doc.history.stopGroup();
+        })
+        .container(this.rootRef.current);
+      d3.select<SVGCircleElement, undefined>(
+        `#dragTarget-keepInLaneAbove`
+      ).call(dragHandleDrag);
+      d3.select<SVGCircleElement, undefined>(
+        `#dragTarget-keepInLaneBelow`
+      ).call(dragHandleDrag);
+    }
+  }
+  dragPointTolerance(event: DragEvent) {
+    const data = this.props.data;
+    const { x, y } = event;
+    const start = this.props.start;
+    const end = this.props.end;
+    if (start === undefined || end === undefined || start.uuid === end.uuid) {
+      return;
+    }
+    const startX = start.x.value;
+    const startY = start.y.value;
+    const endX = end.x.value;
+    const endY = end.y.value;
+    // https://en.wikipedia.org/wiki/Distance_from_a_point_to_a_line#Line_defined_by_two_points
+    const dx = endX - startX;
+    const dy = endY - startY;
+    const dist =
+      Math.abs(dy * x - dx * y + endX * startY - endY * startX) /
+      Math.hypot(dx, dy);
+    data.tolerance.set(dist);
+  }
+  render() {
+    const data = this.props.data.serialize as DataMap["KeepInLane"];
+    const tolerance = data.props.tolerance.val + STROKE / 2;
+    const start = this.props.start;
+    const end = this.props.end;
+    if (start === undefined || end === undefined || start.uuid === end.uuid) {
+      return <></>;
+    }
+    const startX = start.x.value;
+    const startY = start.y.value;
+    const endX = end.x.value;
+    const endY = end.y.value;
+
+    const dx = endX - startX;
+    const dy = endY - startY;
+    const dist = Math.hypot(dy, dx);
+    if (dist === 0) {
+      return <></>;
+    }
+    const [offsetX, offsetY] = [
+      -tolerance * (dy / dist),
+      tolerance * (dx / dist)
+    ];
+    const [startAboveX, startAboveY] = [startX + offsetX, startY + offsetY];
+    const [endAboveX, endAboveY] = [endX + offsetX, endY + offsetY];
+    const [midAboveX, midAboveY] = [
+      (endAboveX + startAboveX) / 2,
+      (endAboveY + startAboveY) / 2
+    ];
+    const [startBelowX, startBelowY] = [startX - offsetX, startY - offsetY];
+    const [endBelowX, endBelowY] = [endX - offsetX, endY - offsetY];
+    const [midBelowX, midBelowY] = [
+      (endBelowX + startBelowX) / 2,
+      (endBelowY + startBelowY) / 2
+    ];
+    return (
+      <g ref={this.rootRef}>
+        {/* Lines */}
+
+        <line
+          x1={startAboveX}
+          x2={endAboveX}
+          y1={startAboveY}
+          y2={endAboveY}
+          stroke="green"
+          strokeWidth={STROKE}
+          strokeOpacity={1.0}
+          id="line-keepInLaneAbove"
+        ></line>
+        <line
+          x1={startBelowX}
+          x2={endBelowX}
+          y1={startBelowY}
+          y2={endBelowY}
+          stroke="green"
+          strokeWidth={STROKE}
+          strokeOpacity={1.0}
+          id="line-keepInLaneBelow"
+        ></line>
+        <circle
+          cx={midAboveX}
+          cy={midAboveY}
+          r={DOT}
+          fill={"green"}
+          fillOpacity={1.0}
+          pointerEvents={"visible"}
+          id="dragTarget-keepInLaneAbove"
+        ></circle>
+        <circle
+          cx={midBelowX}
+          cy={midBelowY}
+          r={DOT}
+          fill={"green"}
+          fillOpacity={1.0}
+          pointerEvents={"visible"}
+          id="dragTarget-keepInLaneBelow"
+        ></circle>
+      </g>
+    );
+  }
+}
+export default observer(KeepInLaneOverlay);

src/components/field/svg/constraintDisplay/KeepInRectangleOverlay.tsx

@@ -117,7 +117,7 @@ class KeepInRectangleOverlay extends Component<
     const data = this.props.data;
 
     data.x.set(data.serialize.props.x.val + event.dx);
-    data.y.set(this.props.data.serialize.props.y.val + event.dy);
+    data.y.set(data.serialize.props.y.val + event.dy);
   }
 
   fixWidthHeight() {