1use geop_core_math::scalars::Scalar;
16use std::{any::Any, cell::RefCell, collections::BTreeMap, rc::Rc};
17
18use geop_core_math::{
19 geop_error::{GeopError, GeopResult},
20 solvers::system::Variable,
21};
22use geop_core_program::{
23 element::{ElementReference, HasElements, Kind},
24 part::{Make, StepContext},
25 placed::key_behind,
26 reader::{Reader, Reads},
27 target::{Target, TargetReference},
28};
29
30use crate::Generated;
31
32pub use geop_core_math::solvers::system::Bound;
33
34#[derive(Debug, Clone, PartialEq)]
38pub struct Constraint {
39 pub residuals: Vec<Scalar>,
40 pub jacobian: BTreeMap<String, Vec<Vec<Scalar>>>,
45 pub bound: Bound,
46}
47
48impl Constraint {
49 pub(crate) fn check(
54 &self,
55 name: &str,
56 dimension: impl Fn(&str) -> Option<usize>,
57 ) -> GeopResult<()> {
58 for (state, rows) in &self.jacobian {
59 let Some(columns) = dimension(state) else {
60 return Err(GeopError::new(format!(
61 "constraint `{name}` gives a derivative by state `{state}`, which it does not read"
62 )));
63 };
64 if rows.len() != self.residuals.len() {
65 return Err(GeopError::new(format!(
66 "constraint `{name}` has {} residuals but {} rows of derivatives by state `{state}`",
67 self.residuals.len(),
68 rows.len()
69 )));
70 }
71 if let Some(row) = rows.iter().find(|row| row.len() != columns) {
72 return Err(GeopError::new(format!(
73 "constraint `{name}` gives a derivative of {} numbers by state `{state}`, whose step takes {columns}",
74 row.len()
75 )));
76 }
77 }
78 Ok(())
79 }
80
81 pub fn holds(&self, tolerance: Scalar) -> bool {
84 (self.residuals.iter()).all(|&r| self.bound.holds(r, tolerance))
85 }
86}
87
88const CONSTRAINTS: &str = "constraints";
90
91#[derive(Debug, Clone, Default, PartialEq)]
94pub struct Constraints(BTreeMap<String, Defined>);
95
96#[derive(Debug, Clone, PartialEq)]
98struct Defined {
99 step: String,
100 constraint: Constraint,
101 at: DefinedAt,
102}
103
104#[derive(Debug, Clone, PartialEq)]
108pub struct DefinedAt {
109 path: Vec<String>,
110 step: String,
111 name: String,
112}
113
114impl DefinedAt {
115 fn inside<'a>(&self, reader: &Reader<'a>) -> GeopResult<Reader<'a>> {
118 let mut inside = reader.clone();
119 for instance in &self.path {
120 inside = inside.placed_part(instance)?;
121 }
122 Ok(inside)
123 }
124
125 fn make(&self, after: &Reader) -> Made {
128 let own = TargetReference::<StepConstraints>::new(format!("{CONSTRAINTS}:{}", self.step));
129 let own = match after.retrieve_element(&ElementReference::whole(own)) {
130 Ok(own) => own,
131 Err(e) => return (Reads::default(), Err(e)),
132 };
133 MADE_AGAIN.set(MADE_AGAIN.get() + 1);
134 let (made, reads) = after.recorded(|reader| (own.again.make)(reader, &*own.again.defining));
135 (reads, made)
136 }
137}
138
139thread_local! {
140 static MADE_AGAIN: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
141}
142
143pub fn made_again() -> usize {
148 MADE_AGAIN.get()
149}
150
151#[derive(Default)]
157pub struct Remade(RefCell<BTreeMap<(Vec<String>, String), Made>>);
158
159type Made = (Reads, GeopResult<StepConstraints>);
161
162impl Remade {
163 pub fn constraints(
178 &self,
179 reader: &Reader,
180 which: &[&DefinedAt],
181 ) -> Vec<GeopResult<Constraint>> {
182 let mut remade = self.0.borrow_mut();
183 (which.iter())
184 .map(|at| {
185 let after = (at.inside(reader)).map(|inside| inside.after_step(&at.step))?;
186 let key = (at.path.clone(), at.step.clone());
187 let kept =
188 (remade.get(&key)).is_some_and(|(reads, _)| after.reads_same_state(reads));
189 if !kept {
190 remade.insert(key.clone(), at.make(&after));
191 }
192 let (_, step) = remade.get(&key).ok_or_else(|| GeopError::new("made"))?;
193 let step = (step.as_ref()).map_err(|e| GeopError::new(e.to_string()))?;
194 let constraint = step.constraints.get(&at.name).ok_or_else(|| {
195 GeopError::new(format!(
196 "step `{}` defines no constraint `{}`",
197 at.step, at.name
198 ))
199 })?;
200 let mut constraint = constraint.clone();
201 for instance in at.path.iter().rev() {
202 constraint.jacobian = (constraint.jacobian.into_iter())
203 .map(|(state, rows)| (key_behind(instance, &state), rows))
204 .collect();
205 }
206 Ok(constraint)
207 })
208 .collect()
209 }
210}
211
212impl Constraints {
213 pub fn of(reader: &Reader) -> GeopResult<Self> {
224 let mut all = Self::own(reader)?;
225 for name in reader.instances() {
226 let inside = (reader
227 .placed_part(name)
228 .and_then(|inside| Self::of(&inside)))
229 .map_err(|e| e.with_context(format!("the constraints of the placed part `{name}`")))?;
230 for (constraint, defined) in inside.0 {
231 let jacobian = (defined.constraint.jacobian.into_iter())
232 .map(|(state, rows)| (key_behind(name, &state), rows))
233 .collect();
234 let mut at = defined.at;
235 at.path.insert(0, name.to_string());
236 let defined = Defined {
237 step: key_behind(name, &defined.step),
238 constraint: Constraint {
239 jacobian,
240 ..defined.constraint
241 },
242 at,
243 };
244 all.0.insert(key_behind(name, &constraint), defined);
245 }
246 }
247 Ok(all)
248 }
249
250 fn own(reader: &Reader) -> GeopResult<Self> {
252 if let Some(failed) = reader.failed_targets_of::<StepConstraints>().pop() {
253 return Err(failed);
254 }
255 let mut all = Self::default();
256 for own in reader.targets_of::<StepConstraints>() {
257 let own = reader.retrieve_element(&ElementReference::whole(own))?;
258 for (name, constraint) in &own.constraints {
259 if let Some(earlier) = all.0.get(name) {
260 return Err(GeopError::new(format!(
261 "constraint `{name}` is defined twice: by step `{}` and by step `{}`",
262 earlier.step, own.step
263 )));
264 }
265 let defined = Defined {
266 step: own.step.clone(),
267 constraint: constraint.clone(),
268 at: DefinedAt {
269 path: Vec::new(),
270 step: own.step.clone(),
271 name: name.clone(),
272 },
273 };
274 all.0.insert(name.clone(), defined);
275 }
276 }
277 Ok(all)
278 }
279
280 pub fn get(&self, name: &str) -> Option<&Constraint> {
282 self.0.get(name).map(|defined| &defined.constraint)
283 }
284
285 pub fn iter(&self) -> impl Iterator<Item = (&str, &Constraint)> {
287 (self.0.iter()).map(|(name, defined)| (name.as_str(), &defined.constraint))
288 }
289
290 pub fn defining(&self) -> impl Iterator<Item = (&str, &DefinedAt)> {
292 (self.0.iter()).map(|(name, defined)| (name.as_str(), &defined.at))
293 }
294
295 pub fn len(&self) -> usize {
296 self.0.len()
297 }
298
299 pub fn is_empty(&self) -> bool {
300 self.0.is_empty()
301 }
302}
303
304#[derive(Clone, PartialEq)]
308struct StepConstraints {
309 step: String,
310 constraints: BTreeMap<String, Constraint>,
311 again: Again,
312}
313
314#[derive(Clone)]
317struct Again {
318 defining: Rc<dyn Any>,
319 make: fn(&Reader, &dyn Any) -> GeopResult<StepConstraints>,
320 equal: fn(&dyn Any, &dyn Any) -> bool,
321}
322
323impl PartialEq for Again {
325 fn eq(&self, other: &Self) -> bool {
326 std::ptr::fn_addr_eq(self.equal, other.equal)
327 && (self.equal)(&*self.defining, &*other.defining)
328 }
329}
330
331fn make_any<Args: PartialEq + Clone + std::fmt::Debug + 'static>(
333 reader: &Reader,
334 defining: &dyn Any,
335) -> GeopResult<StepConstraints> {
336 let defining = (defining.downcast_ref::<Defining<Args>>())
337 .ok_or_else(|| GeopError::new("a step's constraints made of another's arguments"))?;
338 make(reader, defining)
339}
340
341fn equal_any<Args: PartialEq + 'static>(a: &dyn Any, b: &dyn Any) -> bool {
343 (a.downcast_ref::<Defining<Args>>())
344 .is_some_and(|a| b.downcast_ref::<Defining<Args>>() == Some(a))
345}
346
347impl Target for StepConstraints {}
349
350impl Kind for StepConstraints {
352 const NAME: &'static str = "step constraints";
353}
354
355impl HasElements<StepConstraints> for StepConstraints {
356 type Element<'a> = &'a StepConstraints;
357
358 fn element(&self, name: &str) -> GeopResult<&Self> {
359 match name.is_empty() {
360 true => Ok(self),
361 false => Err(GeopError::new(format!(
362 "the constraints of step `{}` have no element `{name}`: they are read whole",
363 self.step
364 ))),
365 }
366 }
367}
368
369pub type ConstraintGenerator<Args> = Make<Args, BTreeMap<String, Constraint>>;
371
372pub fn define_constraints<Args: PartialEq + Clone + std::fmt::Debug + 'static>(
388 ctx: &mut StepContext,
389 args: &Args,
390 generator: ConstraintGenerator<Args>,
391) -> GeopResult<()> {
392 let step = ctx.step();
393 let own = TargetReference::<StepConstraints>::new(format!("{CONSTRAINTS}:{step}"));
394 let generated = Generated {
395 args: args.clone(),
396 generator,
397 };
398 ctx.produce(&own, &(step.to_string(), generated), make)
399}
400
401type Defining<Args> = (String, Generated<Args, BTreeMap<String, Constraint>>);
404
405fn make<Args: PartialEq + Clone + std::fmt::Debug + 'static>(
407 reader: &Reader,
408 defining: &Defining<Args>,
409) -> GeopResult<StepConstraints> {
410 let (step, generated) = defining;
411 let constraints = (generated.generator)(reader, &generated.args)?;
412 if let Some((target, state)) = reader.targets_reading_state().into_iter().next() {
413 let names: Vec<&str> = constraints.keys().map(String::as_str).collect();
414 return Err(GeopError::new(format!(
415 "constraint `{}` reads target `{target}`, which itself reads the state `{state}`: a \
416 solve could not see how it moves with the state — a constraint may read the state \
417 only itself",
418 names.join("`, `")
419 )));
420 }
421 for (name, constraint) in &constraints {
422 constraint.check(name, |state| {
423 (reader.has_read_state(state)).then(|| reader.retrieve_state(state).dimension())
424 })?;
425 }
426 Ok(StepConstraints {
427 step: step.clone(),
428 constraints,
429 again: Again {
430 defining: Rc::new(defining.clone()),
431 make: make_any::<Args>,
432 equal: equal_any::<Args>,
433 },
434 })
435}