1use std::collections::{hash_map::Entry, HashMap, HashSet};
2use std::hash::Hash;
3use std::marker::PhantomData;
4
5use solverforge_core::score::Score;
6use solverforge_core::{ConstraintRef, ImpactType};
7
8use crate::api::constraint_set::IncrementalConstraint;
9use crate::stream::collector::{Accumulator, UniCollector};
10use crate::stream::filter::UniFilter;
11use crate::stream::{ProjectedRowCoordinate, ProjectedRowOwner, ProjectedSource};
12
13struct GroupState<Acc> {
14 accumulator: Acc,
15 count: usize,
16}
17
18pub struct ProjectedGroupedConstraint<S, Out, K, Src, F, KF, C, W, Sc>
19where
20 Src: ProjectedSource<S, Out>,
21 C: UniCollector<Out>,
22 Sc: Score,
23{
24 constraint_ref: ConstraintRef,
25 impact_type: ImpactType,
26 source: Src,
27 filter: F,
28 key_fn: KF,
29 collector: C,
30 weight_fn: W,
31 is_hard: bool,
32 source_state: Option<Src::State>,
33 groups: HashMap<K, GroupState<C::Accumulator>>,
34 row_outputs: HashMap<ProjectedRowCoordinate, Out>,
35 rows_by_owner: HashMap<ProjectedRowOwner, Vec<ProjectedRowCoordinate>>,
36 _phantom: PhantomData<(fn() -> S, fn() -> Out, fn() -> Sc)>,
37}
38
39impl<S, Out, K, Src, F, KF, C, W, Sc> ProjectedGroupedConstraint<S, Out, K, Src, F, KF, C, W, Sc>
40where
41 S: Send + Sync + 'static,
42 Out: Send + Sync + 'static,
43 K: Eq + Hash + Send + Sync + 'static,
44 Src: ProjectedSource<S, Out>,
45 F: UniFilter<S, Out>,
46 KF: Fn(&Out) -> K + Send + Sync,
47 C: UniCollector<Out> + Send + Sync + 'static,
48 C::Accumulator: Send + Sync,
49 C::Result: Send + Sync,
50 C::Value: Send + Sync,
51 W: Fn(&C::Result) -> Sc + Send + Sync,
52 Sc: Score + 'static,
53{
54 #[allow(clippy::too_many_arguments)]
55 pub fn new(
56 constraint_ref: ConstraintRef,
57 impact_type: ImpactType,
58 source: Src,
59 filter: F,
60 key_fn: KF,
61 collector: C,
62 weight_fn: W,
63 is_hard: bool,
64 ) -> Self {
65 Self {
66 constraint_ref,
67 impact_type,
68 source,
69 filter,
70 key_fn,
71 collector,
72 weight_fn,
73 is_hard,
74 source_state: None,
75 groups: HashMap::new(),
76 row_outputs: HashMap::new(),
77 rows_by_owner: HashMap::new(),
78 _phantom: PhantomData,
79 }
80 }
81
82 fn compute_score(&self, result: &C::Result) -> Sc {
83 let base = (self.weight_fn)(result);
84 match self.impact_type {
85 ImpactType::Penalty => -base,
86 ImpactType::Reward => base,
87 }
88 }
89
90 fn ensure_source_state(&mut self, solution: &S) {
91 if self.source_state.is_none() {
92 self.source_state = Some(self.source.build_state(solution));
93 }
94 }
95
96 fn index_coordinate(&mut self, coordinate: ProjectedRowCoordinate) {
97 coordinate.for_each_owner(|owner| {
98 self.rows_by_owner
99 .entry(owner)
100 .or_default()
101 .push(coordinate);
102 });
103 }
104
105 fn unindex_coordinate(&mut self, coordinate: ProjectedRowCoordinate) {
106 coordinate.for_each_owner(|owner| {
107 let mut remove_bucket = false;
108 if let Some(rows) = self.rows_by_owner.get_mut(&owner) {
109 rows.retain(|candidate| *candidate != coordinate);
110 remove_bucket = rows.is_empty();
111 }
112 if remove_bucket {
113 self.rows_by_owner.remove(&owner);
114 }
115 });
116 }
117
118 fn insert_value(&mut self, key: K, value: &C::Value) -> Sc {
119 let impact = self.impact_type;
120 let weight_fn = &self.weight_fn;
121 let group = match self.groups.entry(key) {
122 Entry::Occupied(entry) => entry.into_mut(),
123 Entry::Vacant(entry) => entry.insert(GroupState {
124 accumulator: self.collector.create_accumulator(),
125 count: 0,
126 }),
127 };
128 let old = if group.count == 0 {
129 Sc::zero()
130 } else {
131 let old_base = weight_fn(&group.accumulator.finish());
132 match impact {
133 ImpactType::Penalty => -old_base,
134 ImpactType::Reward => old_base,
135 }
136 };
137 group.accumulator.accumulate(value);
138 group.count += 1;
139 let new_base = weight_fn(&group.accumulator.finish());
140 let new_score = match self.impact_type {
141 ImpactType::Penalty => -new_base,
142 ImpactType::Reward => new_base,
143 };
144 new_score - old
145 }
146
147 fn retract_value(&mut self, key: K, value: &C::Value) -> Sc {
148 let impact = self.impact_type;
149 let weight_fn = &self.weight_fn;
150 let Entry::Occupied(mut entry) = self.groups.entry(key) else {
151 return Sc::zero();
152 };
153 let group = entry.get_mut();
154 let old_base = weight_fn(&group.accumulator.finish());
155 let old = match impact {
156 ImpactType::Penalty => -old_base,
157 ImpactType::Reward => old_base,
158 };
159 group.accumulator.retract(value);
160 group.count = group.count.saturating_sub(1);
161 let new_score = if group.count == 0 {
162 entry.remove();
163 Sc::zero()
164 } else {
165 let new_base = weight_fn(&group.accumulator.finish());
166 match impact {
167 ImpactType::Penalty => -new_base,
168 ImpactType::Reward => new_base,
169 }
170 };
171
172 new_score - old
173 }
174
175 fn insert_row(&mut self, solution: &S, coordinate: ProjectedRowCoordinate, output: Out) -> Sc {
176 if self.row_outputs.contains_key(&coordinate) || !self.filter.test(solution, &output) {
177 return Sc::zero();
178 }
179 let key = (self.key_fn)(&output);
180 let value = self.collector.extract(&output);
181 let delta = self.insert_value(key, &value);
182 self.row_outputs.insert(coordinate, output);
183 self.index_coordinate(coordinate);
184 delta
185 }
186
187 fn retract_row(&mut self, coordinate: ProjectedRowCoordinate) -> Sc {
188 let Some(output) = self.row_outputs.remove(&coordinate) else {
189 return Sc::zero();
190 };
191 self.unindex_coordinate(coordinate);
192 let key = (self.key_fn)(&output);
193 let value = self.collector.extract(&output);
194 self.retract_value(key, &value)
195 }
196
197 fn localized_owners(
198 &self,
199 descriptor_index: usize,
200 entity_index: usize,
201 ) -> Vec<ProjectedRowOwner> {
202 let mut owners = Vec::new();
203 for slot in 0..self.source.source_count() {
204 if self
205 .source
206 .change_source(slot)
207 .assert_localizes(descriptor_index, &self.constraint_ref.name)
208 {
209 owners.push(ProjectedRowOwner {
210 source_slot: slot,
211 entity_index,
212 });
213 }
214 }
215 owners
216 }
217
218 fn coordinates_for_owners(&self, owners: &[ProjectedRowOwner]) -> Vec<ProjectedRowCoordinate> {
219 let mut seen = HashSet::new();
220 let mut coordinates = Vec::new();
221 for owner in owners {
222 let Some(rows) = self.rows_by_owner.get(owner) else {
223 continue;
224 };
225 for &coordinate in rows {
226 if seen.insert(coordinate) {
227 coordinates.push(coordinate);
228 }
229 }
230 }
231 coordinates
232 }
233}
234
235impl<S, Out, K, Src, F, KF, C, W, Sc> IncrementalConstraint<S, Sc>
236 for ProjectedGroupedConstraint<S, Out, K, Src, F, KF, C, W, Sc>
237where
238 S: Send + Sync + 'static,
239 Out: Send + Sync + 'static,
240 K: Eq + Hash + Send + Sync + 'static,
241 Src: ProjectedSource<S, Out>,
242 F: UniFilter<S, Out>,
243 KF: Fn(&Out) -> K + Send + Sync,
244 C: UniCollector<Out> + Send + Sync + 'static,
245 C::Accumulator: Send + Sync,
246 C::Result: Send + Sync,
247 C::Value: Send + Sync,
248 W: Fn(&C::Result) -> Sc + Send + Sync,
249 Sc: Score + 'static,
250{
251 fn evaluate(&self, solution: &S) -> Sc {
252 let state = self.source.build_state(solution);
253 let mut groups: HashMap<K, C::Accumulator> = HashMap::new();
254 self.source.collect_all(solution, &state, |_, output| {
255 if !self.filter.test(solution, &output) {
256 return;
257 }
258 let key = (self.key_fn)(&output);
259 let value = self.collector.extract(&output);
260 groups
261 .entry(key)
262 .or_insert_with(|| self.collector.create_accumulator())
263 .accumulate(&value);
264 });
265 groups.values().fold(Sc::zero(), |total, acc| {
266 total + self.compute_score(&acc.finish())
267 })
268 }
269
270 fn match_count(&self, solution: &S) -> usize {
271 let state = self.source.build_state(solution);
272 let mut keys = HashMap::<K, ()>::new();
273 self.source.collect_all(solution, &state, |_, output| {
274 if self.filter.test(solution, &output) {
275 keys.insert((self.key_fn)(&output), ());
276 }
277 });
278 keys.len()
279 }
280
281 fn initialize(&mut self, solution: &S) -> Sc {
282 self.reset();
283 let state = self.source.build_state(solution);
284 let mut total = Sc::zero();
285 let mut rows = Vec::new();
286 self.source
287 .collect_all(solution, &state, |coordinate, output| {
288 rows.push((coordinate, output));
289 });
290 self.source_state = Some(state);
291 for (coordinate, output) in rows {
292 total = total + self.insert_row(solution, coordinate, output);
293 }
294 total
295 }
296
297 fn on_insert(&mut self, solution: &S, entity_index: usize, descriptor_index: usize) -> Sc {
298 let owners = self.localized_owners(descriptor_index, entity_index);
299 self.ensure_source_state(solution);
300 {
301 let state = self.source_state.as_mut().expect("projected source state");
302 for owner in &owners {
303 self.source.insert_entity_state(
304 solution,
305 state,
306 owner.source_slot,
307 owner.entity_index,
308 );
309 }
310 }
311 let mut rows = Vec::new();
312 let state = self.source_state.as_ref().expect("projected source state");
313 for owner in &owners {
314 self.source.collect_entity(
315 solution,
316 state,
317 owner.source_slot,
318 owner.entity_index,
319 |coordinate, output| rows.push((coordinate, output)),
320 );
321 }
322 let mut total = Sc::zero();
323 for (coordinate, output) in rows {
324 total = total + self.insert_row(solution, coordinate, output);
325 }
326 total
327 }
328
329 fn on_retract(&mut self, solution: &S, entity_index: usize, descriptor_index: usize) -> Sc {
330 let owners = self.localized_owners(descriptor_index, entity_index);
331 let mut total = Sc::zero();
332 for coordinate in self.coordinates_for_owners(&owners) {
333 total = total + self.retract_row(coordinate);
334 }
335 if let Some(state) = self.source_state.as_mut() {
336 for owner in &owners {
337 self.source.retract_entity_state(
338 solution,
339 state,
340 owner.source_slot,
341 owner.entity_index,
342 );
343 }
344 }
345 total
346 }
347
348 fn reset(&mut self) {
349 self.source_state = None;
350 self.groups.clear();
351 self.row_outputs.clear();
352 self.rows_by_owner.clear();
353 }
354
355 fn name(&self) -> &str {
356 &self.constraint_ref.name
357 }
358
359 fn constraint_ref(&self) -> &ConstraintRef {
360 &self.constraint_ref
361 }
362
363 fn is_hard(&self) -> bool {
364 self.is_hard
365 }
366}