[WIP] optimization support

Cargo.toml

@@ -9,7 +9,7 @@ edition = "2021"
 argmin = { version = "0.8.1", features = [] }
 argmin-math = { version = "0.3.0", features = ["nalgebra_v0_32"] }
 bevy_ecs = "0.10.1"
-eframe = { version = "0.21.3" }
+eframe = { version = "0.21.3", default-features = false, features = ["default_fonts", "glow"] }
 indexmap = "1.8.1"
 levenberg-marquardt = "0.13.0"
 nalgebra = { version = "0.32.0", features = ["rand"] }
@@ -17,13 +17,18 @@ nalgebra-sparse = "0.9.0"
 nohash-hasher = "0.2.0"
 rand = "0.8.5"
 rand_distr = "0.4.3"
-tracing = { version = "0.1.37", optional = true }
+slotmap = "1.0.6"
+tracing = "0.1.37"
 tracing-chrome = { version = "0.7.1", optional = true }
 tracing-error = { version = "0.2.0", optional = true }
-tracing-subscriber = { version = "0.3.17", features = ["env-filter"], optional = true }
+tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
 
 [dev-dependencies]
 criterion = "0.4.0"
 
 [features]
-trace = ["tracing", "tracing-subscriber", "tracing-chrome", "tracing-error", "bevy_ecs/trace"]
+trace = ["tracing-chrome", "tracing-error", "bevy_ecs/trace"]
+
+[[bench]]
+name = "my_benchmark"
+harness = false

src/lib.rs

@@ -0,0 +1,5 @@
+#![feature(generic_const_exprs)]
+
+pub mod geometry;
+pub mod optimization;
+// mod relations;

src/main.rs

@@ -1,18 +1,40 @@
+#![feature(generic_const_exprs)]
+
 use bevy_ecs::prelude::*;
 use eframe::egui;
+use rand::seq::SliceRandom;
+use rand_distr::Distribution;
+
+use sketchrs::geometry;
 
-mod geometry;
-pub mod optimization;
-mod relations;
 #[cfg(feature = "trace")]
-mod tracing;
+mod trace;
 mod ui;
 
 fn init(mut commands: Commands) {
-    let p1 = geometry::insert_point_at(&mut commands, (10., 30.));
+    // let p1 = geometry::insert_point_at(&mut commands, (10., 30.));
-    let p2 = geometry::insert_point_at(&mut commands, (-20., 15.));
+    // let p2 = geometry::insert_point_at(&mut commands, (-20., 15.));
-    geometry::insert_point_at(&mut commands, (0., -10.));
+    // geometry::insert_point_at(&mut commands, (0., -10.));
+    // commands.spawn(geometry::LineBundle::new(p1, p2));
+
+    // Insert 1000 random points
+    let dist = rand_distr::Uniform::new(-100., 100.);
+    let mut rng = rand::thread_rng();
+    let mut points = Vec::new();
+    for _ in 0..1000 {
+        let p = geometry::insert_point_at(
+            &mut commands,
+            (dist.sample(&mut rng), dist.sample(&mut rng)),
+        );
+        points.push(p);
+    }
+
+    // Insert 1000 random lines
+    for _ in 0..1000 {
+        let p1 = *points.choose(&mut rng).unwrap();
+        let p2 = *points.choose(&mut rng).unwrap();
         commands.spawn(geometry::LineBundle::new(p1, p2));
+    }
 }
 
 struct MyApp {
@@ -31,9 +53,9 @@ impl MyApp {
         let mut world = World::default();
 
         #[cfg(feature = "trace")]
-        tracing::setup(&mut world);
+        trace::setup(&mut world);
 
-        world.init_resource::<ui::Tool>();
+        ui::init(&mut world);
 
         let mut init_sched = Schedule::new();
         init_sched.add_system(init);
@@ -51,6 +73,9 @@ impl MyApp {
 
 impl eframe::App for MyApp {
     fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
+        let span = tracing::info_span!("update");
+        let _enter = span.enter();
+
         self.world.insert_resource(ui::ContextRes(ctx.clone()));
 
         self.update_schedule.run(&mut self.world);
@@ -60,6 +85,11 @@ impl eframe::App for MyApp {
 }
 
 fn main() {
+    tracing_subscriber::fmt()
+        .with_max_level(tracing::Level::INFO)
+        .with_target(false)
+        .init();
+
     let options = eframe::NativeOptions::default();
     eframe::run_native("sketchrs", options, Box::new(|_cc| Box::<MyApp>::default())).unwrap();
 }

src/optimization.rs

@@ -1,4 +1,9 @@
-use std::{borrow::Borrow, fmt::Debug};
+use std::{borrow::Borrow, fmt::Debug, ops::Index};
+
+mod params;
+mod residuals;
+
+use params::ReadParams;
 
 use argmin_math;
 use levenberg_marquardt::{self, TerminationReason};
@@ -7,11 +12,154 @@ use nalgebra_sparse::CooMatrix;
 
 use crate::geometry::Scalar;
 
+pub type Idx = u32;
+
+use self::{params::{Parameters, Param}, residuals::WriteResiduals};
+
 type SVector<const D: usize> = nalgebra::SVector<Scalar, D>;
 type SMatrix<const R: usize, const C: usize> = nalgebra::SMatrix<Scalar, R, C>;
 type DVector = nalgebra::DVector<Scalar>;
 type DMatrix = nalgebra::DMatrix<Scalar>;
 
+
+trait CostFunction {
+    type Params;
+    // const NUM_RESIDUALS: usize;
+    type Residuals;
+    type Jacobian;
+
+    fn evaluate(
+        &self,
+        params: &Self::Params,
+        residual: Option<&mut Self::Residuals>,
+        jacobian: Option<&mut Self::Jacobian>,
+    );
+
+    fn residual(&self, x: &Self::Params) -> Self::Residuals
+    where
+        Self::Residuals: Default,
+    {
+        let mut residual = Self::Residuals::default();
+        self.evaluate(x, Some(&mut residual), None);
+        residual
+    }
+
+    fn jacobian(&self, x: &Self::Params) -> Self::Jacobian
+    where
+        Self::Jacobian: Default,
+    {
+        let mut jacobian = Self::Jacobian::default();
+        self.evaluate(x, None, Some(&mut jacobian));
+        jacobian
+    }
+
+    fn with_indices<PI, RI, JI>(
+        self,
+        param_indices: PI,
+        residual_indices: RI,
+        jacobian_indices: JI,
+    ) -> IndexedCostFunctionImpl<Self, PI, RI, JI>
+    where
+        Self: Sized,
+    {
+        IndexedCostFunctionImpl {
+            residual: self,
+            param_indices,
+            residual_indices,
+            jacobian_indices,
+        }
+    }
+}
+
+trait IndexedCostFunction {
+    fn evaluate(
+        &self,
+        params: &Parameters,
+        residuals: Option<&mut [Scalar]>,
+        jacobian: Option<&mut [Scalar]>,
+    );
+}
+
+struct IndexedCostFunctionImpl<R, PI, RI, JI> {
+    residual: R,
+    param_indices: PI,
+    residual_indices: RI,
+    jacobian_indices: JI,
+}
+
+impl<R, PI, RI, JI> IndexedCostFunction for IndexedCostFunctionImpl<R, PI, RI, JI>
+where
+    R: CostFunction,
+    R::Params: ReadParams<PI>,
+    R::Residuals: WriteResiduals<RI> + Default,
+    R::Jacobian: WriteResiduals<JI> + Default,
+{
+    fn evaluate(
+        &self,
+        params: &Parameters,
+        residuals: Option<&mut [Scalar]>,
+        jacobian: Option<&mut [Scalar]>,
+    ) {
+        // let p = R::Params::read(params, self.param_indices);
+
+        // let mut res = residuals.and_then(|_| Some(R::Residuals::default()));
+        // let mut jac = jacobian.and_then(|_| Some(R::Jacobian::default()));
+        // self.residual.evaluate(&p, res.as_mut(), jac.as_mut());
+
+        // if let (Some(res), Some(residuals)) = (res, residuals) {
+        //     res.write(self.residual_indices, residuals);
+        // }
+        // if let (Some(jac), Some(jacobian)) = (jac, jacobian) {
+        //     jac.write(self.jacobian_indices, jacobian);
+        // }
+    }
+}
+
+slotmap::new_key_type!{ struct ParamKey; }
+
+struct ParamIndex(Idx);
+
+#[derive(Default)]
+struct CostFunctions {
+    functions: Vec<Box<dyn IndexedCostFunction>>,
+    param_indices: slotmap::SlotMap<ParamKey, Idx>,
+    num_residuals: Idx,
+}
+
+impl CostFunctions {
+    fn new() -> Self {
+        Self::default()
+    }
+
+    // fn add_param<P: Param>(&mut self) -> ParamIndex {
+    //     // let index = self.params_size;
+    //     // self.params_size += P::SIZE;
+    //     ParamIndex(index)
+    // }
+
+    // fn add<R, PI, RI, JI>(&mut self, residual: R, param_indices: PI, residual_indices: RI, jacobian_indices: JI)
+    // where
+    //     R: CostFunction,
+    //     R::Params: ReadParams<PI>,
+    //     R::Residuals: WriteResiduals<RI> + Default,
+    //     R::Jacobian: WriteResiduals<JI> + Default,
+    // {
+    //     self.functions.push(Box::new(
+    //         residual
+    //             .with_indices(param_indices, residual_indices, jacobian_indices),
+    //     ));
+    //     self.param_indices.push(param_indices.into());
+    //     self.residual_indices.push(residual_indices.into());
+    //     self.jacobian_indices.push(jacobian_indices.into());
+    // }
+
+    fn evaluate(&self, params: &Parameters, residuals: Option<&mut [Scalar]>, jacobian: Option<&mut [Scalar]>) {
+        // for (i, f) in self.functions.iter().enumerate() {
+        //     f.evaluate(params, residuals, jacobian);
+        // }
+    }
+}
+
 trait SResBlock<const NIN: usize, const NOUT: usize = 1>: Clone + Debug {
     fn evaluate(
         &self,
@@ -39,20 +187,20 @@ trait SResBlock<const NIN: usize, const NOUT: usize = 1>: Clone + Debug {
     }
 
     // Probably should move to own trait
-    fn into_indexed(
+    // fn into_indexed(
-        self,
+    //     self,
-        in_indices: [Index; NIN],
+    //     in_indices: [Index; NIN],
-        out_indices: [Index; NOUT],
+    //     out_indices: [Index; NOUT],
-    ) -> Box<dyn IndexedResBlock>
+    // ) -> Box<dyn IndexedResBlock>
-    where
+    // where
-        Self: Sized + 'static,
+    //     Self: Sized + 'static,
-    {
+    // {
-        Box::new(IndexedResBlockImpl {
+    //     Box::new(IndexedResBlockImpl {
-            residual: self,
+    //         residual: self,
-            in_indices,
+    //         in_indices,
-            out_indices,
+    //         out_indices,
-        })
+    //     })
-    }
+    // }
 }
 
 fn square(x: Scalar) -> Scalar {
@@ -84,6 +232,36 @@ impl SResBlock<2> for Equal1D {
 #[derive(Clone, Debug)]
 pub struct AxisDistance(Scalar);
 
+/*
+impl ResBlock for AxisDistance {
+    fn parameter_sizes(&self) -> &[usize] {
+        &[2]
+    }
+
+    fn n_residuals(&self) -> usize {
+        1
+    }
+
+    fn evaluate(
+        &self,
+        x: &[&[Scalar]],
+        residual: Option<&mut [Scalar]>,
+        jacobian: Option<&mut [&mut [Scalar]]>,
+    ) {
+        let [x1, x2] = x[0] else { panic!() };
+        let dx = x2 - x1;
+        let r = self.0;
+        if let Some(residual) = residual {
+            residual[0] = square(dx) - square(r);
+        }
+        if let Some(jacobian) = jacobian {
+            jacobian[0][0] = -2. * dx;
+            jacobian[0][1] = 2. * dx;
+        }
+    }
+}
+*/
+
 impl SResBlock<2> for AxisDistance {
     fn evaluate(
         &self,
@@ -107,6 +285,34 @@ impl SResBlock<2> for AxisDistance {
 #[derive(Clone, Debug)]
 pub struct PointPointDistance(Scalar);
 
+impl CostFunction for PointPointDistance {
+    type Params = (SVector<2>, SVector<2>);
+    type Residuals = Scalar;
+    type Jacobian = (SMatrix<2, 1>, SMatrix<2, 1>);
+
+    fn evaluate(
+        &self,
+        (p1, p2): &Self::Params,
+        residual: Option<&mut Self::Residuals>,
+        jacobian: Option<&mut Self::Jacobian>,
+    ) {
+        let d = p2 - p1;
+        if let Some(residual) = residual {
+            *residual = square(d.x) + square(d.y) - square(self.0);
+        }
+        if let Some(jacobian) = jacobian {
+            jacobian.0.copy_from(&(2. * d));
+            jacobian.1.copy_from(&(-2. * d));
+            // jacobian.row_mut(0).copy_from(&(2. * d).transpose());
+            // jacobian.row_mut(1).copy_from(&(-2. * d).transpose());
+            // jacobian[0] = -2. * d.x;
+            // jacobian[1] = -2. * d.y;
+            // jacobian[2] = 2. * d.x;
+            // jacobian[3] = 2. * d.y;
+        }
+    }
+}
+
 impl SResBlock<4> for PointPointDistance {
     fn evaluate(
         &self,
@@ -116,9 +322,8 @@ impl SResBlock<4> for PointPointDistance {
     ) {
         let [x1, y1, x2, y2] = *x.as_ref();
         let (dx, dy) = (x2 - x1, y2 - y1);
-        let r = self.0;
         if let Some(residual) = residual {
-            residual[0] = square(dx) + square(dy) - square(r);
+            residual[0] = square(dx) + square(dy) - square(self.0);
         }
         if let Some(jacobian) = jacobian {
             jacobian[0] = -2. * dx;
@@ -153,6 +358,33 @@ impl SResBlock<2> for PointDistance {
 #[derive(Clone, Debug)]
 pub struct Collinear;
 
+impl CostFunction for Collinear {
+    type Params = (SVector<2>, SVector<2>, SVector<2>);
+    type Residuals = Scalar;
+    type Jacobian = (SVector<2>, SVector<2>, SVector<2>);
+
+    fn evaluate(
+        &self,
+        (p1, p2, p3): &Self::Params,
+        residual: Option<&mut Self::Residuals>,
+        jacobian: Option<&mut Self::Jacobian>,
+    ) {
+        let d21 = p2 - p1;
+        let d32 = p3 - p2;
+        if let Some(residual) = residual {
+            *residual = d21.y * d32.x - d32.y * d21.x;
+        }
+        if let Some(jacobian) = jacobian {
+            jacobian.0.x = d32.y;
+            jacobian.0.y = -d32.x;
+            jacobian.1.x = p1.y - p3.y;
+            jacobian.1.y = p3.x - p1.x;
+            jacobian.2.x = d21.y;
+            jacobian.2.y = -d21.x;
+        }
+    }
+}
+
 impl SResBlock<6> for Collinear {
     fn evaluate(
         &self,
@@ -171,15 +403,14 @@ impl SResBlock<6> for Collinear {
         if let Some(jacobian) = jacobian {
             jacobian[0] = dy32;
             jacobian[1] = -dx32;
-            jacobian[2] = -dy32;
+            jacobian[2] = y1 - y3;
-            jacobian[3] = -dx21;
+            jacobian[3] = x3 - x1;
             jacobian[4] = dy21;
             jacobian[5] = -dx21;
         }
     }
 }
 
-type Index = u32;
 
 trait ScatterResBlock: Debug {
     fn shape(&self) -> (usize, usize);
@@ -194,9 +425,9 @@ trait ScatterResBlock: Debug {
 #[derive(Clone, Debug)]
 struct ScatterResBlockImpl<const NIN: usize, const NOUT: usize, R> {
     residual: R,
-    in_indices: [Index; NIN],
+    in_indices: [Idx; NIN],
-    residual_indices: [Index; NOUT],
+    residual_indices: [Idx; NOUT],
-    jacobian_indices: [[Index; NIN]; NOUT],
+    jacobian_indices: [[Idx; NIN]; NOUT],
 }
 
 impl<const NIN: usize, const NOUT: usize, R> ScatterResBlockImpl<NIN, NOUT, R> {
@@ -278,6 +509,7 @@ where
     }
 }
 
+/*
 trait IndexedResBlock {
     fn shape(&self) -> (usize, usize);
     fn in_indices(&self) -> &[Index];
@@ -323,20 +555,22 @@ where
         })
     }
 }
+*/
 
 #[derive(Debug)]
-struct Residuals {
+struct ResidualBlocks {
     res_blocks: Vec<Box<dyn ScatterResBlock>>,
     shape: (usize, usize),
 }
 
-impl Residuals {
+impl ResidualBlocks {
-    fn new(res_blocks: Vec<Box<dyn IndexedResBlock>>, shape: (usize, usize)) -> Self {
+    fn new(res_blocks: Vec<Box<dyn IndexedCostFunction>>, shape: (usize, usize)) -> Self {
-        let res_blocks = res_blocks
+        // let res_blocks = res_blocks
-            .into_iter()
+        //     .into_iter()
-            .map(|r| r.into_scatter(shape))
+        //     .map(|r| r.into_scatter(shape))
-            .collect();
+        //     .collect();
-        Self { res_blocks, shape }
+        // Self { res_blocks, shape }
+        todo!()
     }
 
     fn residuals_into<X: AsRef<[Scalar]>>(&self, x: X, residuals: &mut [Scalar]) {
@@ -379,7 +613,7 @@ impl Residuals {
 
 #[derive(Debug)]
 struct Problem {
-    residuals: Residuals,
+    residuals: ResidualBlocks,
     x: DVector,
 }
 
@@ -412,18 +646,18 @@ pub fn test() {
 
     let nvar = 4;
     let nres = 3;
-    let residuals = Residuals::new(
+    // let residuals = CostFunction::new(
-        vec![
+    //     vec![
-            PointDistance(SVector::<2>::new(1., 0.), 0.75).into_indexed([0, 1], [0]),
+    //         PointDistance(SVector::<2>::new(1., 0.), 0.75).into_indexed([0, 1], [0]),
-            PointPointDistance(1.0).into_indexed([0, 1, 2, 3], [2]),
+    //         PointPointDistance(1.0).into_indexed([0, 1, 2, 3], [2]),
-            PointDistance(SVector::<2>::new(0., 0.), 1.25).into_indexed([2, 3], [1]),
+    //         PointDistance(SVector::<2>::new(0., 0.), 1.25).into_indexed([2, 3], [1]),
-        ],
+    //     ],
-        (nres, nvar),
+    //     (nres, nvar),
-    );
+    // );
 
     let dist = rand_distr::Normal::new(0., 1.0).unwrap();
     let prob = Problem {
-        residuals,
+        residuals: todo!(),
         x: DVector::from_distribution_generic(
             Dyn(nvar),
             Const::<1>,

src/optimization/params.rs

@@ -0,0 +1,173 @@
+use nalgebra::{SMatrix, SVector};
+
+use crate::geometry::Scalar;
+
+use super::Idx;
+
+pub struct Parameters {
+    values: Vec<Scalar>,
+}
+
+pub trait Param {
+    const SIZE: usize;
+    type Jacobian<const N: usize>;
+
+    fn from_array(array: [Scalar; Self::SIZE]) -> Self;
+
+    fn from_slice(slice: &[Scalar]) -> Self
+    where
+        [Scalar; Self::SIZE]: Sized,
+        Self: Sized,
+    {
+        let array =
+            TryInto::<[Scalar; Self::SIZE]>::try_into(slice).expect("Slice has wrong length");
+        Self::from_array(array)
+    }
+
+    fn copy_to_array(&self, array: &mut [Scalar; Self::SIZE]);
+
+    fn copy_to_slice(&self, slice: &mut [Scalar])
+    where
+        [Scalar; Self::SIZE]: Sized,
+    {
+        let array =
+            TryInto::<&mut [Scalar; Self::SIZE]>::try_into(slice).expect("Slice has wrong length");
+        self.copy_to_array(array)
+    }
+}
+
+trait WithJacobian<const N: usize> {
+    type Jacobian<const M: usize>;
+}
+
+trait Jacobian {}
+
+impl Param for Scalar {
+    const SIZE: usize = 1;
+
+    type Jacobian<const N: usize> = SVector<Scalar, N>;
+
+    fn from_slice(slice: &[Scalar]) -> Self {
+        let [x] = slice else {
+            panic!("Slice has wrong length");
+        };
+        *x
+    }
+
+    fn copy_to_slice(&self, slice: &mut [Scalar]) {
+        let [x] = slice else {
+            panic!("Slice has wrong length");
+        };
+        *x = *self;
+    }
+
+    fn copy_to_array(&self, array: &mut [Scalar; Self::SIZE]) {
+        array[0] = *self;
+    }
+
+    fn from_array(array: [Scalar; Self::SIZE]) -> Self {
+        array[0]
+    }
+}
+
+pub trait ReadParams<I> {
+    fn read(params: &Parameters, index: I) -> Self;
+}
+
+// impl ReadParams<Idx> for Scalar {
+//     fn read(params: &Parameters, index: Idx) -> Self {
+//         params.values[index as usize]
+//     }
+// }
+
+impl<const M: usize> Param for [Scalar; M] {
+    const SIZE: usize = M;
+    type Jacobian<const N: usize> = SMatrix<Scalar, M, N>;
+
+    fn from_array(array: [Scalar; Self::SIZE]) -> Self {
+        unsafe { std::mem::transmute_copy(&array) }
+    }
+
+    fn copy_to_array(&self, slice: &mut [Scalar; Self::SIZE]) {
+        slice.copy_from_slice(self)
+    }
+}
+
+impl<const M: usize> Param for SMatrix<Scalar, M, 1> {
+    const SIZE: usize = M;
+    type Jacobian<const N: usize> = SMatrix<Scalar, M, N>;
+
+    fn from_slice(slice: &[Scalar]) -> Self {
+        Self::from_column_slice(slice)
+    }
+
+    fn from_array(array: [Scalar; Self::SIZE]) -> Self {
+        Self::from_slice(&array)
+    }
+
+    fn copy_to_slice(&self, slice: &mut [Scalar]) {
+        slice.copy_from_slice(self.as_slice())
+    }
+
+    fn copy_to_array(&self, array: &mut [Scalar; Self::SIZE]) {
+        array.copy_from_slice(self.as_slice())
+    }
+}
+
+impl<P: Param> ReadParams<Idx> for P
+where
+    [Scalar; P::SIZE]: Sized,
+{
+    fn read(params: &Parameters, index: Idx) -> Self {
+        Self::from_slice(&params.values[index as usize..][..Self::SIZE])
+    }
+}
+
+// impl<const R: usize, const C: usize> ReadParams<Idx> for SMatrix<Scalar, R, C> {
+//     fn read(params: &Parameters, index: Idx) -> Self {
+//         SMatrix::from_column_slice(&params.values[index as usize..][..R * C])
+//     }
+// }
+
+// TODO: Implement for tuples of arbitrary length using a macro
+impl<T1, T2, I1, I2> ReadParams<(I1, I2)> for (T1, T2)
+where
+    T1: ReadParams<I1>,
+    T2: ReadParams<I2>,
+{
+    fn read(params: &Parameters, index: (I1, I2)) -> Self {
+        (T1::read(params, index.0), T2::read(params, index.1))
+    }
+}
+
+impl<T1, T2, T3, I1, I2, I3> ReadParams<(I1, I2, I3)> for (T1, T2, T3)
+where
+    T1: ReadParams<I1>,
+    T2: ReadParams<I2>,
+    T3: ReadParams<I3>,
+{
+    fn read(params: &Parameters, index: (I1, I2, I3)) -> Self {
+        (
+            T1::read(params, index.0),
+            T2::read(params, index.1),
+            T3::read(params, index.2),
+        )
+    }
+}
+
+impl<T1, T2, T3, T4, I1, I2, I3, I4> ReadParams<(I1, I2, I3, I4)> for (T1, T2, T3, T4)
+where
+    T1: ReadParams<I1>,
+    T2: ReadParams<I2>,
+    T3: ReadParams<I3>,
+    T4: ReadParams<I4>,
+{
+    fn read(params: &Parameters, index: (I1, I2, I3, I4)) -> Self {
+        (
+            T1::read(params, index.0),
+            T2::read(params, index.1),
+            T3::read(params, index.2),
+            T4::read(params, index.3),
+        )
+    }
+}

src/optimization/residuals.rs

@@ -0,0 +1,64 @@
+use nalgebra::SMatrix;
+
+use crate::geometry::Scalar;
+
+pub trait WriteResiduals<Index> {
+    fn write(&self, index: Index, residuals: &mut [Scalar]);
+}
+
+impl<I> WriteResiduals<I> for Scalar
+where
+    usize: From<I>,
+{
+    fn write(&self, index: I, residuals: &mut [Scalar]) {
+        residuals[usize::from(index)] = *self;
+    }
+}
+
+impl<I, const R: usize, const C: usize> WriteResiduals<I> for SMatrix<Scalar, R, C>
+where
+    usize: From<I>,
+{
+    fn write(&self, index: I, residuals: &mut [Scalar]) {
+        residuals[usize::from(index)..][..R * C].copy_from_slice(self.as_slice());
+    }
+}
+
+impl<T1, T2, I1, I2> WriteResiduals<(I1, I2)> for (T1, T2)
+where
+    T1: WriteResiduals<I1>,
+    T2: WriteResiduals<I2>,
+{
+    fn write(&self, index: (I1, I2), residuals: &mut [Scalar]) {
+        self.0.write(index.0, residuals);
+        self.1.write(index.1, residuals);
+    }
+}
+
+impl<T1, T2, T3, I1, I2, I3> WriteResiduals<(I1, I2, I3)> for (T1, T2, T3)
+where
+    T1: WriteResiduals<I1>,
+    T2: WriteResiduals<I2>,
+    T3: WriteResiduals<I3>,
+{
+    fn write(&self, index: (I1, I2, I3), residuals: &mut [Scalar]) {
+        self.0.write(index.0, residuals);
+        self.1.write(index.1, residuals);
+        self.2.write(index.2, residuals);
+    }
+}
+
+impl<T1, T2, T3, T4, I1, I2, I3, I4> WriteResiduals<(I1, I2, I3, I4)> for (T1, T2, T3, T4)
+where
+    T1: WriteResiduals<I1>,
+    T2: WriteResiduals<I2>,
+    T3: WriteResiduals<I3>,
+    T4: WriteResiduals<I4>,
+{
+    fn write(&self, index: (I1, I2, I3, I4), residuals: &mut [Scalar]) {
+        self.0.write(index.0, residuals);
+        self.1.write(index.1, residuals);
+        self.2.write(index.2, residuals);
+        self.3.write(index.3, residuals);
+    }
+}

src/ui.rs

@@ -1,11 +1,12 @@
-use std::ops::Deref;
+use std::{cmp::Ordering, ops::Deref};
 
 use bevy_ecs::{prelude::*, query::WorldQuery, system::ReadOnlySystem};
 use eframe::{
-    egui::{self, CursorIcon, Painter, Response, Sense},
+    egui::{self, plot::Plot, CursorIcon, Painter, Response, ScrollArea, Sense},
     emath::RectTransform,
-    epaint::{Color32, Hsva, Pos2, Rect, Stroke, Vec2},
+    epaint::{Color32, Hsva, Pos2, Rect, Rounding, Stroke, Vec2},
 };
+use tracing::info;
 
 use crate::geometry::{self};
 
@@ -25,12 +26,46 @@ impl Default for Tool {
     }
 }
 
+#[derive(Resource, Clone)]
+struct Viewport {
+    scale: f32,
+    center: Pos2,
+}
+
+impl Default for Viewport {
+    fn default() -> Self {
+        Self {
+            scale: 1.0,
+            center: Pos2::ZERO,
+        }
+    }
+}
+
+fn update_viewport(mut viewport: ResMut<Viewport>, ui: Res<UiRes>) {
+    // if ui.response.hovered() && ui.response.ctx.input(|i| i.modifiers.command_only()) {
+    if ui.response.hovered() {
+        let delta = ui.response.ctx.input(|i| i.scroll_delta);
+        if delta != Vec2::ZERO {
+            // viewport.scale *= 1.0 + delta.y * 0.05;
+            viewport.scale *= (1.0f32 - 0.005).powf(-delta.y);
+        }
+    }
+}
+
 #[derive(Resource)]
 struct UiRes {
     response: Response,
     painter: Painter,
     to_screen: ToScreen,
 }
+
+impl UiRes {
+    fn to_screen(&self) -> ToScreen {
+        self.to_screen
+    }
+}
+
+#[derive(Clone, Copy)]
 struct ToScreen(RectTransform);
 
 impl ToScreen {
@@ -64,15 +99,36 @@ fn color_for_var_status(status: geometry::VarStatus) -> Hsva {
     }
 }
 
+#[derive(Default, Resource)]
+struct SelectionState {
+    start_point: Option<Pos2>,
+    end_point: Option<Pos2>,
+}
+
+impl SelectionState {
+    fn selection_rect(&self) -> Option<Rect> {
+        if let (Some(start), Some(end)) = (self.start_point, self.end_point) {
+            Some(Rect::from_two_pos(start, end))
+        } else {
+            None
+        }
+    }
+}
+
 const POINT_RADIUS: f32 = 3.0;
 
 fn update_hover_point(
     ui: Res<UiRes>,
     points: Query<(geometry::PointId, &geometry::ComputedPointPos)>,
+    sel: Res<SelectionState>,
     mut commands: Commands,
 ) {
     points.for_each(|(id, pos)| {
-        let hovered = if let Some(hover_pos) = ui.response.hover_pos() {
+        let hovered = if let Some(sel_rect) = sel.selection_rect() {
+            let center = ui.to_screen.transform_pos(pos);
+
+            sel_rect.contains(center)
+        } else if let Some(hover_pos) = ui.response.hover_pos() {
             let center = ui.to_screen.transform_pos(pos);
 
             (hover_pos - center).length() < (POINT_RADIUS * 3.)
@@ -91,15 +147,18 @@ fn update_hover_point(
 fn update_hover_line(
     ui: Res<UiRes>,
     lines: Query<(Entity, &geometry::ComputedLinePos)>,
+    sel: Res<SelectionState>,
     mut commands: Commands,
 ) {
     lines.for_each(|(id, pos)| {
-        let hovered = if let Some(hover_pos) = ui.response.hover_pos() {
         let points = [
             ui.to_screen.transform_pos(&pos.start),
             ui.to_screen.transform_pos(&pos.end),
         ];
-
+        let hovered = if let Some(sel_rect) = sel.selection_rect() {
+            // TODO: check if line intersects rect
+            sel_rect.contains(points[0]) || sel_rect.contains(points[1])
+        } else if let Some(hover_pos) = ui.response.hover_pos() {
             let b = points[1] - points[0];
             let a = hover_pos - points[0];
             let p = a.dot(b) / b.dot(b);
@@ -121,27 +180,44 @@ fn select_tool(
     ui: Res<UiRes>,
     hovered: Query<Entity, With<Hovered>>,
     selected: Query<Entity, With<Selected>>,
+    mut sel: ResMut<SelectionState>,
     mut commands: Commands,
 ) {
     ui.response.ctx.output_mut(|output| {
-        output.cursor_icon = if !hovered.is_empty() {
+        output.cursor_icon = if sel.selection_rect().is_some() {
+            CursorIcon::Grabbing
+        } else if !hovered.is_empty() {
             CursorIcon::PointingHand
         } else {
             CursorIcon::Default
         }
     });
 
-    if ui.response.clicked() {
+    if ui.response.drag_released() {
         if !ui.response.ctx.input(|input| input.modifiers.shift) {
             selected.for_each(|selected| {
                 commands.entity(selected).remove::<Selected>();
             });
         }
-        // TODO: choose which to select
+        for entity in hovered.iter() {
-        if let Some(hovered) = hovered.iter().next() {
+            commands.entity(entity).insert(Selected);
-            commands.entity(hovered).insert(Selected);
+            // Only select one if not dragging
+            // if sel.selection_rect().is_none() {
+            //     break;
+            // }
         }
     }
+    if ui.response.dragged() {
+        // TODO: choose which to select
+        if sel.start_point.is_none() {
+            sel.start_point = ui.response.hover_pos();
+        } else if ui.response.hover_pos() != sel.start_point {
+            sel.end_point = ui.response.hover_pos();
+        }
+    } else {
+        sel.start_point = None;
+        sel.end_point = None;
+    }
 
     if ui
         .response
@@ -162,6 +238,17 @@ fn select_tool(
     }
 }
 
+fn paint_selection(ui: Res<UiRes>, sel: Res<SelectionState>) {
+    if let Some(rect) = sel.selection_rect() {
+        ui.painter.rect(
+            rect,
+            Rounding::none(),
+            Hsva::new(0.0, 0.0, 0.5, 0.5),
+            Stroke::new(1.0, Color32::WHITE),
+        )
+    }
+}
+
 #[derive(Default)]
 struct DragDelta(Vec2);
 
@@ -199,7 +286,7 @@ fn move_tool(
         selected.for_each(|selected| {
             commands.entity(selected.0).remove::<Selected>();
         });
-        drag_delta.0 = Vec2::ZERO;
+        (*drag_delta).0 = Vec2::ZERO;
         None
     } else {
         selected.iter().next()
@@ -208,7 +295,7 @@ fn move_tool(
     if let Some((_, point)) = selected {
         if response.drag_started() {
             let drag_point_pos = point_pos.get(point).unwrap();
-            drag_delta.0 = hover_pos - to_screen.transform_pos(&drag_point_pos);
+            (*drag_delta).0 = hover_pos - to_screen.transform_pos(&drag_point_pos);
         }
         let move_to = to_screen.inverse_transform_to_point(hover_pos - drag_delta.0);
         point_pos.set(point, move_to).unwrap();
@@ -330,8 +417,11 @@ fn paint_lines(
 
         let mut color = color_for_var_status(pos.status);
         color.v -= 0.6;
-        if hovered.contains(id) || selected.contains(id) {
+        if selected.contains(id) {
             color.s -= 0.8;
+            color.v += 0.2;
+        } else if hovered.contains(id) {
+            color.s -= 0.4;
         }
 
         let stroke = Stroke::new(2.0, color);
@@ -350,9 +440,11 @@ fn paint_points(
         let center = ui.to_screen.transform_pos(pos);
 
         let color = color_for_var_status(pos.status);
-        let stroke = if selected.contains(id) || hovered.contains(id) {
+        let stroke = if selected.contains(id) {
-            // color.s -= 0.8;
             Stroke::new(1.0, Color32::WHITE)
+        } else if hovered.contains(id) {
+            // color.s -= 0.8;
+            Stroke::new(1.0, Color32::WHITE.linear_multiply(0.1))
         } else {
             Stroke::default()
         };
@@ -414,6 +506,7 @@ pub enum ShowEntitiesStage {
     PostTools,
     ToolsFlush,
     Paint,
+    PostPaint,
 }
 
 impl Default for ShowEntitiesSchedule {
@@ -430,11 +523,14 @@ impl Default for ShowEntitiesSchedule {
                     ShowEntitiesStage::ToolsFlush,
                     ShowEntitiesStage::PostTools,
                     ShowEntitiesStage::Paint,
+                    ShowEntitiesStage::PostPaint,
                 )
                     .chain(),
             )
+            .add_system(update_viewport.in_base_set(ShowEntitiesStage::Update))
             .add_systems(
                 (geometry::update_point_pos, geometry::update_line_pos)
+                    .chain()
                     .in_base_set(ShowEntitiesStage::Update),
             )
             .add_system(apply_system_buffers.in_base_set(ShowEntitiesStage::UpdateFlush))
@@ -456,6 +552,7 @@ impl Default for ShowEntitiesSchedule {
             .add_system(geometry::remove_dangling_lines.in_base_set(ShowEntitiesStage::PostTools))
             .add_systems(
                 (geometry::update_point_pos, geometry::update_line_pos)
+                    .chain()
                     .in_base_set(ShowEntitiesStage::PostTools),
             )
             .add_system(apply_system_buffers.in_base_set(ShowEntitiesStage::ToolsFlush))
@@ -463,7 +560,8 @@ impl Default for ShowEntitiesSchedule {
                 (paint_lines, paint_points)
                     .chain()
                     .in_base_set(ShowEntitiesStage::Paint),
-            );
+            )
+            .add_system((paint_selection).in_base_set(ShowEntitiesStage::PostPaint));
         Self(schedule)
     }
 }
@@ -474,13 +572,18 @@ fn central_panel(world: &mut World, mut schedule: Local<ShowEntitiesSchedule>) {
         .frame(egui::Frame::none())
         .show(&ctx, |ui| {
             let sense = match *world.resource::<Tool>() {
-                Tool::Move => Sense::drag(),
+                Tool::Select | Tool::Move => Sense::drag(),
-                Tool::Select | Tool::AddPoint | Tool::AddLine | Tool::AddRelation => Sense::click(),
+                Tool::AddPoint | Tool::AddLine | Tool::AddRelation => Sense::click(),
             };
+            info!("render");
 
             let (response, painter) = ui.allocate_painter(ui.available_size(), sense);
+            let viewport = world.get_resource::<Viewport>().unwrap().clone();
             let to_screen = ToScreen(RectTransform::from_to(
-                Rect::from_center_size(Pos2::ZERO, response.rect.size() / 2.0),
+                Rect::from_center_size(
+                    viewport.center,
+                    response.rect.size() * (viewport.scale / 2.0),
+                ),
                 response.rect,
             ));
             world.insert_resource(UiRes {
@@ -496,10 +599,10 @@ fn central_panel(world: &mut World, mut schedule: Local<ShowEntitiesSchedule>) {
 }
 
 #[derive(WorldQuery)]
-struct SelectableEntity<'a> {
+struct SelectableEntity {
     id: Entity,
-    point: Option<&'a geometry::Point>,
+    point: Option<&'static geometry::Point>,
-    line: Option<&'a geometry::Line>,
+    line: Option<&'static geometry::Line>,
 }
 
 fn side_panel_ui(
@@ -508,10 +611,14 @@ fn side_panel_ui(
     tool: Res<Tool>,
 ) {
     let tool = *tool;
-    ui.vertical(|ui| match tool {
+    ui.heading("Selection:");
+    ScrollArea::vertical().show(ui, |ui| {
+        match tool {
             Tool::Select => {
                 let mut count = 0;
-            selected.for_each(|sel| {
+                let mut selection: Vec<_> = selected.iter().collect();
+                selection.sort_by(|a, b| a.id.index().cmp(&b.id.index()));
+                selection.iter().for_each(|sel| {
                     count += 1;
                     if sel.point.is_some() {
                         ui.label(format!("Selected point {}", sel.id.index()));
@@ -546,7 +653,9 @@ fn side_panel_ui(
             Tool::AddRelation => {
                 ui.label("Click to add a relation");
             }
+        };
     });
+    ui.separator();
 }
 
 fn side_panel(
@@ -569,6 +678,12 @@ fn bottom_panel(ctx: Res<ContextRes>) {
     });
 }
 
+pub fn init(world: &mut World) {
+    world.init_resource::<Tool>();
+    world.init_resource::<SelectionState>();
+    world.init_resource::<Viewport>();
+}
+
 pub fn add_to_schedule(schedule: &mut Schedule) {
     schedule.add_systems((prepare, toolbar, side_panel, bottom_panel, central_panel).chain());
 }