1use std::{
12 any::{Any, TypeId},
13 hash::{Hash, Hasher},
14 ops::Deref,
15};
16
17use crate::numeric_id::{DenseIdMap, IdVec, NumericId, define_id};
18use crossbeam_queue::SegQueue;
19use dashmap::SharedValue;
20use rustc_hash::FxHasher;
21
22use crate::{
23 ColumnId, CounterId, ExecutionState, Offset, SubsetRef, TableId, TaggedRowBuffer, Value,
24 WrappedTable,
25 common::{DashMap, IndexSet, SubsetTracker},
26 parallel,
27 parallel_heuristics::{parallelize_inter_container_op, parallelize_intra_container_op},
28 table_spec::{Rebuilder, ValueRebuilder},
29};
30
31#[cfg(test)]
32mod tests;
33
34define_id!(pub ContainerValueId, u32, "an identifier for containers");
35
36pub trait MergeFn:
37 Fn(&mut ExecutionState, Value, Value) -> Value + dyn_clone::DynClone + Send + Sync
38{
39}
40impl<T: Fn(&mut ExecutionState, Value, Value) -> Value + Clone + Send + Sync> MergeFn for T {}
41
42dyn_clone::clone_trait_object!(MergeFn);
44
45#[derive(Clone, Default)]
46struct ContainerIds {
47 ids: IndexSet<TypeId>,
48}
49
50impl ContainerIds {
51 fn insert(&mut self, ty: TypeId) -> ContainerValueId {
52 if let Some(idx) = self.ids.get_index_of(&ty) {
53 ContainerValueId::from_usize(idx)
54 } else {
55 let idx = self.ids.len();
56 self.ids.insert(ty);
57 ContainerValueId::from_usize(idx)
58 }
59 }
60
61 fn get(&self, ty: &TypeId) -> Option<ContainerValueId> {
62 self.ids.get_index_of(ty).map(ContainerValueId::from_usize)
63 }
64}
65
66#[derive(Clone, Default)]
67pub struct ContainerValues {
68 subset_tracker: SubsetTracker,
69 container_ids: ContainerIds,
70 data: DenseIdMap<ContainerValueId, Box<dyn DynamicContainerEnv + Send + Sync>>,
71}
72
73#[derive(Clone, Default)]
88pub struct ContainerRebuildSummary {
89 changed: bool,
90 dirty_ids: IndexSet<Value>,
93}
94
95impl ContainerRebuildSummary {
96 pub fn changed(&self) -> bool {
98 self.changed
99 }
100
101 pub fn dirty_ids(&self) -> &IndexSet<Value> {
103 &self.dirty_ids
104 }
105
106 fn note_change(&mut self) {
107 self.changed = true;
108 }
109
110 fn note_dirty_id(&mut self, value: Value) {
111 self.changed = true;
112 self.dirty_ids.insert(value);
113 }
114
115 fn extend(&mut self, other: Self) {
116 self.changed |= other.changed;
117 self.dirty_ids.extend(other.dirty_ids);
118 }
119}
120
121impl ContainerValues {
122 pub fn new() -> Self {
123 Default::default()
124 }
125
126 fn get<C: ContainerValue>(&self) -> Option<&ContainerEnv<C>> {
127 let id = self.container_ids.get(&TypeId::of::<C>())?;
128 let res = self.data.get(id)?.as_any();
129 Some(res.downcast_ref::<ContainerEnv<C>>().unwrap())
130 }
131
132 pub fn for_each<C: ContainerValue>(&self, mut f: impl FnMut(&C, Value)) {
134 let Some(env) = self.get::<C>() else {
135 return;
136 };
137 for ent in env.to_id.iter() {
138 f(ent.key(), *ent.value());
139 }
140 }
141
142 pub fn get_val<C: ContainerValue>(&self, val: Value) -> Option<impl Deref<Target = C> + '_> {
148 self.get::<C>()?.get_container(val)
149 }
150
151 pub fn register_val<C: ContainerValue>(
152 &self,
153 container: C,
154 exec_state: &mut ExecutionState,
155 ) -> Value {
156 let env = self
157 .get::<C>()
158 .expect("must register container type before registering a value");
159 env.get_or_insert(&container, exec_state)
160 }
161
162 pub fn rebuild_val_with(
171 &self,
172 type_id: TypeId,
173 value: Value,
174 exec_state: &mut ExecutionState,
175 remap: &(dyn Fn(Value) -> Value + Send + Sync),
176 ) -> Value {
177 let Some(id) = self.container_ids.get(&type_id) else {
178 return value;
179 };
180 let Some(env) = self.data.get(id) else {
181 return value;
182 };
183 env.rebuild_val_with(value, exec_state, remap)
184 .unwrap_or(value)
185 }
186
187 pub fn rebuild_all(
189 &mut self,
190 table_id: TableId,
191 table: &WrappedTable,
192 exec_state: &mut ExecutionState,
193 ) -> ContainerRebuildSummary {
194 let Some(rebuilder) = table.rebuilder(&[]) else {
195 return Default::default();
196 };
197 let to_scan = rebuilder.hint_col().map(|_| {
198 self.subset_tracker.recent_updates(table_id, table)
200 });
201 let mut summary = if parallelize_inter_container_op(self.data.next_id().index()) {
202 parallel::map_dense_id_map_mut(&mut self.data, |_, env| {
203 let mut exec_state = exec_state.clone();
204 env.apply_rebuild(
205 table,
206 &*rebuilder,
207 to_scan.as_ref().map(|x| x.as_ref()),
208 &mut exec_state,
209 )
210 })
211 .into_iter()
212 .fold(ContainerRebuildSummary::default(), |mut acc, summary| {
213 acc.extend(summary);
214 acc
215 })
216 } else {
217 let mut summary = ContainerRebuildSummary::default();
218 for (_, env) in self.data.iter_mut() {
219 summary.extend(env.apply_rebuild(
220 table,
221 &*rebuilder,
222 to_scan.as_ref().map(|x| x.as_ref()),
223 exec_state,
224 ));
225 }
226 summary
227 };
228 self.expand_dirty_id_closure(&mut summary);
229 summary
230 }
231
232 fn expand_dirty_id_closure(&self, summary: &mut ContainerRebuildSummary) {
244 let mut frontier = summary.dirty_ids.clone();
245 let mut seen = frontier.iter().copied().collect::<IndexSet<_>>();
246
247 while !frontier.is_empty() {
248 let mut next = IndexSet::default();
249 for (_, env) in self.data.iter() {
250 env.extend_containers_containing(&frontier, &mut next);
251 }
252 frontier.clear();
253 for value in next {
254 if seen.insert(value) {
255 summary.note_dirty_id(value);
256 frontier.insert(value);
257 }
258 }
259 }
260 }
261
262 pub fn register_type<C: ContainerValue>(
267 &mut self,
268 id_counter: CounterId,
269 merge_fn: impl MergeFn + 'static,
270 ) -> ContainerValueId {
271 let id = self.container_ids.insert(TypeId::of::<C>());
272 self.data.get_or_insert(id, || {
273 Box::new(ContainerEnv::<C>::new(Box::new(merge_fn), id_counter))
274 });
275 id
276 }
277}
278
279pub trait ContainerValue: Hash + Eq + Clone + Send + Sync + 'static {
285 fn rebuild_contents(&mut self, rebuilder: &dyn ValueRebuilder) -> bool;
290
291 fn iter(&self) -> impl Iterator<Item = Value> + '_;
298}
299
300pub trait DynamicContainerEnv: Any + dyn_clone::DynClone + Send + Sync {
301 fn as_any(&self) -> &dyn Any;
302 fn apply_rebuild(
303 &mut self,
304 table: &WrappedTable,
305 rebuilder: &dyn Rebuilder,
306 subset: Option<SubsetRef>,
307 exec_state: &mut ExecutionState,
308 ) -> ContainerRebuildSummary;
309 fn extend_containers_containing(&self, values: &IndexSet<Value>, out: &mut IndexSet<Value>);
315 fn rebuild_val_with(
319 &self,
320 value: Value,
321 exec_state: &mut ExecutionState,
322 remap: &(dyn Fn(Value) -> Value + Send + Sync),
323 ) -> Option<Value>;
324}
325
326dyn_clone::clone_trait_object!(DynamicContainerEnv);
328
329fn hash_container(container: &impl ContainerValue) -> u64 {
330 let mut hasher = FxHasher::default();
331 container.hash(&mut hasher);
332 hasher.finish()
333}
334
335#[derive(Clone)]
336struct ContainerEnv<C: Eq + Hash> {
337 merge_fn: Box<dyn MergeFn>,
338 counter: CounterId,
339 to_id: DashMap<C, Value>,
340 to_container: DashMap<Value, (usize , usize )>,
341 val_index: DashMap<Value, IndexSet<Value>>,
343}
344
345impl<C: ContainerValue> DynamicContainerEnv for ContainerEnv<C> {
346 fn as_any(&self) -> &dyn Any {
347 self
348 }
349
350 fn apply_rebuild(
351 &mut self,
352 table: &WrappedTable,
353 rebuilder: &dyn Rebuilder,
354 subset: Option<SubsetRef>,
355 exec_state: &mut ExecutionState,
356 ) -> ContainerRebuildSummary {
357 if let Some(subset) = subset
358 && incremental_rebuild(
359 subset.size(),
360 self.to_id.len(),
361 parallelize_intra_container_op(self.to_id.len()),
362 )
363 {
364 return self.apply_rebuild_incremental(
365 table,
366 rebuilder,
367 exec_state,
368 subset,
369 rebuilder.hint_col().unwrap(),
370 );
371 }
372 self.apply_rebuild_nonincremental(rebuilder, exec_state)
373 }
374
375 fn extend_containers_containing(&self, values: &IndexSet<Value>, out: &mut IndexSet<Value>) {
376 for value in values {
377 if let Some(containers) = self.val_index.get(value) {
378 out.extend(containers.iter().copied());
379 }
380 }
381 }
382
383 fn rebuild_val_with(
384 &self,
385 value: Value,
386 exec_state: &mut ExecutionState,
387 remap: &(dyn Fn(Value) -> Value + Send + Sync),
388 ) -> Option<Value> {
389 let mut container = self.get_container(value)?.clone();
392 container.rebuild_contents(&ClosureRebuilder { remap });
393 Some(self.get_or_insert(&container, exec_state))
394 }
395}
396
397impl<C: ContainerValue> ContainerEnv<C> {
398 pub fn new(merge_fn: Box<dyn MergeFn>, counter: CounterId) -> Self {
399 Self {
400 merge_fn,
401 counter,
402 to_id: DashMap::default(),
403 to_container: DashMap::default(),
404 val_index: DashMap::default(),
405 }
406 }
407
408 fn get_or_insert(&self, container: &C, exec_state: &mut ExecutionState) -> Value {
409 if let Some(value) = self.to_id.get(container) {
410 return *value;
411 }
412
413 let value = Value::from_usize(exec_state.inc_counter(self.counter));
418 let target_map = self.to_id.determine_map(container);
419 debug_assert_eq!(
423 target_map,
424 self.to_container
425 .determine_shard(hash_container(container) as usize)
426 );
427 self.to_container
428 .insert(value, (hash_container(container) as usize, target_map));
429
430 match self.to_id.entry(container.clone()) {
433 dashmap::Entry::Vacant(vac) => {
434 vac.insert(value);
436 for val in container.iter() {
437 self.val_index.entry(val).or_default().insert(value);
438 }
439 value
440 }
441 dashmap::Entry::Occupied(occ) => {
442 let res = *occ.get();
446 std::mem::drop(occ); self.to_container.remove(&value);
448 res
449 }
450 }
451 }
452
453 fn insert_owned(&self, container: C, value: Value, exec_state: &mut ExecutionState) -> Value {
454 let hc = hash_container(&container);
455 let target_map = self.to_id.determine_map(&container);
456 match self.to_id.entry(container) {
457 dashmap::Entry::Occupied(mut occ) => {
458 let result = (self.merge_fn)(exec_state, *occ.get(), value);
459 let old_val = *occ.get();
460 if result != old_val {
461 self.to_container.remove(&old_val);
462 self.to_container.insert(result, (hc as usize, target_map));
463 *occ.get_mut() = result;
464 for val in occ.key().iter() {
465 let mut index = self.val_index.entry(val).or_default();
466 index.swap_remove(&old_val);
467 index.insert(result);
468 }
469 }
470 result
471 }
472 dashmap::Entry::Vacant(vacant_entry) => {
473 self.to_container.insert(value, (hc as usize, target_map));
474 for val in vacant_entry.key().iter() {
475 self.val_index.entry(val).or_default().insert(value);
476 }
477 vacant_entry.insert(value);
478 value
479 }
480 }
481 }
482
483 fn reinsert_incremental(
484 &self,
485 container: C,
486 old_id: Value,
487 rebuilt_id: Value,
488 container_changed: bool,
489 exec_state: &mut ExecutionState,
490 summary: &mut ContainerRebuildSummary,
491 ) {
492 if container_changed || rebuilt_id != old_id {
493 summary.note_change();
494 }
495 if rebuilt_id != old_id {
496 self.to_container.remove(&old_id);
499 }
500 let actual = self.insert_owned(container, rebuilt_id, exec_state);
501 if container_changed && rebuilt_id == old_id && actual == old_id {
502 summary.note_dirty_id(old_id);
503 }
504 }
505
506 fn apply_rebuild_incremental(
507 &mut self,
508 table: &WrappedTable,
509 rebuilder: &dyn Rebuilder,
510 exec_state: &mut ExecutionState,
511 to_scan: SubsetRef,
512 search_col: ColumnId,
513 ) -> ContainerRebuildSummary {
514 let mut summary = ContainerRebuildSummary::default();
521 let mut buf = TaggedRowBuffer::new(1);
522 table.scan_project(
523 to_scan,
524 &[search_col],
525 Offset::new(0),
526 usize::MAX,
527 &[],
528 &mut buf,
529 );
530 let mut to_rebuild = IndexSet::<Value>::default();
532 for (_, row) in buf.iter() {
533 to_rebuild.insert(row[0]);
534 let Some(ids) = self.val_index.get(&row[0]) else {
535 continue;
536 };
537 to_rebuild.extend(&*ids);
538 }
539 for id in to_rebuild {
540 let Some((hc, target_map)) = self.to_container.get(&id).map(|x| *x) else {
541 continue;
542 };
543 let shard_mut = self.to_id.shards_mut()[target_map].get_mut();
544 let Some((mut container, _)) =
545 shard_mut.remove_entry(hc as u64, |(_, v)| *v.get() == id)
546 else {
547 continue;
548 };
549 let rebuilt_id = rebuilder.rebuild_val(id);
550 let container_changed = container.rebuild_contents(rebuilder);
551 self.reinsert_incremental(
552 container,
553 id,
554 rebuilt_id,
555 container_changed,
556 exec_state,
557 &mut summary,
558 );
559 }
560 summary
561 }
562
563 fn apply_rebuild_nonincremental(
564 &mut self,
565 rebuilder: &dyn Rebuilder,
566 exec_state: &mut ExecutionState,
567 ) -> ContainerRebuildSummary {
568 if parallelize_inter_container_op(self.to_id.len()) {
569 return self.apply_rebuild_nonincremental_parallel(rebuilder, exec_state);
570 }
571 let mut summary = ContainerRebuildSummary::default();
572 let mut to_reinsert = Vec::new();
573 let shards = self.to_id.shards_mut();
574 for shard in shards.iter_mut() {
575 let shard = shard.get_mut();
576 for bucket in unsafe { shard.iter() } {
578 let (container, val) = unsafe { bucket.as_mut() };
580 let old_val = *val.get();
581 let new_val = rebuilder.rebuild_val(old_val);
582 let container_changed = container.rebuild_contents(rebuilder);
583 if !container_changed && new_val == old_val {
584 continue;
586 }
587 summary.note_change();
588 if container_changed {
589 let ((container, _), _) = unsafe { shard.remove(bucket) };
593 self.to_container.remove(&old_val);
594 to_reinsert.push((container, new_val, new_val == old_val));
595 } else {
596 *val.get_mut() = new_val;
598 let prev = self.to_container.remove(&old_val).unwrap().1;
599 self.to_container.insert(new_val, prev);
600 }
601 }
602 }
603 for (container, val, stable_id) in to_reinsert {
604 let actual = self.insert_owned(container, val, exec_state);
605 if stable_id && actual == val {
609 summary.note_dirty_id(val);
610 }
611 }
612 summary
613 }
614
615 fn apply_rebuild_nonincremental_parallel(
616 &mut self,
617 rebuilder: &dyn Rebuilder,
618 exec_state: &mut ExecutionState,
619 ) -> ContainerRebuildSummary {
620 let mut to_reinsert =
625 IdVec::<usize , SegQueue<(C, Value, bool)>>::default();
626 to_reinsert.resize_with(self.to_id.shards().len(), Default::default);
627
628 let shards = self.to_id.shards_mut();
629 let changed = parallel::map_mut(shards, |_, shard| {
630 let mut changed = false;
631 let shard = shard.get_mut();
632 for bucket in unsafe { shard.iter() } {
634 let (container, val) = unsafe { bucket.as_mut() };
636 let old_val = *val.get();
637 let new_val = rebuilder.rebuild_val(old_val);
638 let container_changed = container.rebuild_contents(rebuilder);
639 if !container_changed && new_val == old_val {
640 continue;
642 }
643 changed = true;
644 if container_changed {
645 let ((container, _), _) = unsafe { shard.remove(bucket) };
649 self.to_container.remove(&old_val);
650 let shard = self
655 .to_container
656 .determine_shard(hash_container(&container) as usize);
657 to_reinsert[shard].push((container, new_val, new_val == old_val));
658 } else {
659 *val.get_mut() = new_val;
661 let prev = self.to_container.remove(&old_val).unwrap().1;
662 self.to_container.insert(new_val, prev);
663 }
664 }
665 changed
666 })
667 .into_iter()
668 .any(|changed| changed);
669
670 let dirty_ids = SegQueue::new();
671 parallel::for_each_mut(shards, |shard_id, shard| {
672 let mut exec_state = exec_state.clone();
673 let shard = shard.get_mut();
679 let queue = &to_reinsert[shard_id];
680 while let Some((container, val, stable_id)) = queue.pop() {
681 let hc = hash_container(&container);
682 let target_map = self.to_container.determine_shard(hc as usize);
683 match shard.find_or_find_insert_slot(
684 hc,
685 |(c, _)| c == &container,
686 |(c, _)| hash_container(c),
687 ) {
688 Ok(bucket) => {
689 let (container, val_slot) = unsafe { bucket.as_mut() };
692 let old_val = *val_slot.get();
693 let result = (self.merge_fn)(&mut exec_state, old_val, val);
694 if result != old_val {
695 self.to_container.remove(&old_val);
696 self.to_container.insert(result, (hc as usize, target_map));
697 *val_slot.get_mut() = result;
698 for val in container.iter() {
699 let mut index = self.val_index.entry(val).or_default();
700 index.swap_remove(&old_val);
701 index.insert(result);
702 }
703 }
704 if stable_id && result == val {
707 dirty_ids.push(val);
708 }
709 }
710 Err(slot) => {
711 self.to_container.insert(val, (hc as usize, target_map));
712 for v in container.iter() {
713 self.val_index.entry(v).or_default().insert(val);
714 }
715 unsafe {
718 shard.insert_in_slot(hc, slot, (container, SharedValue::new(val)));
719 }
720 if stable_id {
721 dirty_ids.push(val);
722 }
723 }
724 }
725 }
726 });
727 let mut summary = ContainerRebuildSummary::default();
728 if changed {
729 summary.note_change();
730 }
731 while let Some(value) = dirty_ids.pop() {
732 summary.note_dirty_id(value);
733 }
734 summary
735 }
736
737 fn get_container(&self, value: Value) -> Option<impl Deref<Target = C> + '_> {
738 let (hc, target_map) = *self.to_container.get(&value)?;
739 let shard = &self.to_id.shards()[target_map];
740 let read_guard = shard.read();
741 let val_ptr: *const (C, _) = shard
742 .read()
743 .find(hc as u64, |(_, v)| *v.get() == value)?
744 .as_ptr();
745 struct ValueDeref<'a, T, Guard> {
746 _guard: Guard,
747 data: &'a T,
748 }
749
750 impl<T, Guard> Deref for ValueDeref<'_, T, Guard> {
751 type Target = T;
752
753 fn deref(&self) -> &T {
754 self.data
755 }
756 }
757
758 Some(ValueDeref {
759 _guard: read_guard,
760 data: unsafe {
762 let unwrapped: &(C, _) = &*val_ptr;
763 &unwrapped.0
764 },
765 })
766 }
767}
768
769fn incremental_rebuild(uf_size: usize, table_size: usize, parallel: bool) -> bool {
770 if parallel {
771 table_size > 1000 && uf_size * 512 <= table_size
772 } else {
773 table_size > 1000 && uf_size * 8 <= table_size
774 }
775}
776
777struct ClosureRebuilder<'a> {
781 remap: &'a (dyn Fn(Value) -> Value + Send + Sync),
782}
783
784impl ValueRebuilder for ClosureRebuilder<'_> {
785 fn rebuild_val(&self, val: Value) -> Value {
786 (self.remap)(val)
787 }
788}