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(
172 &self,
173 type_id: TypeId,
174 value: Value,
175 exec_state: &mut ExecutionState,
176 remap: &(dyn Fn(Value) -> Value + Send + Sync),
177 ) -> Option<Value> {
178 let id = self.container_ids.get(&type_id)?;
179 let env = self.data.get(id)?;
180 env.rebuild_val_with(value, exec_state, remap)
181 }
182
183 pub fn rebuild_all(
185 &mut self,
186 table_id: TableId,
187 table: &WrappedTable,
188 exec_state: &mut ExecutionState,
189 ) -> ContainerRebuildSummary {
190 let Some(rebuilder) = table.rebuilder(&[]) else {
191 return Default::default();
192 };
193 let to_scan = rebuilder.hint_col().map(|_| {
194 self.subset_tracker.recent_updates(table_id, table)
196 });
197 let mut summary = if parallelize_inter_container_op(self.data.next_id().index()) {
198 parallel::map_dense_id_map_mut(&mut self.data, |_, env| {
199 let mut exec_state = exec_state.clone();
200 env.apply_rebuild(
201 table,
202 &*rebuilder,
203 to_scan.as_ref().map(|x| x.as_ref()),
204 &mut exec_state,
205 )
206 })
207 .into_iter()
208 .fold(ContainerRebuildSummary::default(), |mut acc, summary| {
209 acc.extend(summary);
210 acc
211 })
212 } else {
213 let mut summary = ContainerRebuildSummary::default();
214 for (_, env) in self.data.iter_mut() {
215 summary.extend(env.apply_rebuild(
216 table,
217 &*rebuilder,
218 to_scan.as_ref().map(|x| x.as_ref()),
219 exec_state,
220 ));
221 }
222 summary
223 };
224 self.expand_dirty_id_closure(&mut summary);
225 summary
226 }
227
228 fn expand_dirty_id_closure(&self, summary: &mut ContainerRebuildSummary) {
240 let mut frontier = summary.dirty_ids.clone();
241 let mut seen = frontier.iter().copied().collect::<IndexSet<_>>();
242
243 while !frontier.is_empty() {
244 let mut next = IndexSet::default();
245 for (_, env) in self.data.iter() {
246 env.extend_containers_containing(&frontier, &mut next);
247 }
248 frontier.clear();
249 for value in next {
250 if seen.insert(value) {
251 summary.note_dirty_id(value);
252 frontier.insert(value);
253 }
254 }
255 }
256 }
257
258 pub fn register_type<C: ContainerValue>(
263 &mut self,
264 id_counter: CounterId,
265 merge_fn: impl MergeFn + 'static,
266 ) -> ContainerValueId {
267 let id = self.container_ids.insert(TypeId::of::<C>());
268 self.data.get_or_insert(id, || {
269 Box::new(ContainerEnv::<C>::new(Box::new(merge_fn), id_counter))
270 });
271 id
272 }
273}
274
275pub trait ContainerValue: Hash + Eq + Clone + Send + Sync + 'static {
281 fn rebuild_contents(&mut self, rebuilder: &dyn ValueRebuilder) -> bool;
286
287 fn iter(&self) -> impl Iterator<Item = Value> + '_;
294}
295
296pub trait DynamicContainerEnv: Any + dyn_clone::DynClone + Send + Sync {
297 fn as_any(&self) -> &dyn Any;
298 fn apply_rebuild(
299 &mut self,
300 table: &WrappedTable,
301 rebuilder: &dyn Rebuilder,
302 subset: Option<SubsetRef>,
303 exec_state: &mut ExecutionState,
304 ) -> ContainerRebuildSummary;
305 fn extend_containers_containing(&self, values: &IndexSet<Value>, out: &mut IndexSet<Value>);
311 fn rebuild_val_with(
315 &self,
316 value: Value,
317 exec_state: &mut ExecutionState,
318 remap: &(dyn Fn(Value) -> Value + Send + Sync),
319 ) -> Option<Value>;
320}
321
322dyn_clone::clone_trait_object!(DynamicContainerEnv);
324
325fn hash_container(container: &impl ContainerValue) -> u64 {
326 let mut hasher = FxHasher::default();
327 container.hash(&mut hasher);
328 hasher.finish()
329}
330
331#[derive(Clone)]
332struct ContainerEnv<C: Eq + Hash> {
333 merge_fn: Box<dyn MergeFn>,
334 counter: CounterId,
335 to_id: DashMap<C, Value>,
336 to_container: DashMap<Value, (usize , usize )>,
337 val_index: DashMap<Value, IndexSet<Value>>,
339}
340
341impl<C: ContainerValue> DynamicContainerEnv for ContainerEnv<C> {
342 fn as_any(&self) -> &dyn Any {
343 self
344 }
345
346 fn apply_rebuild(
347 &mut self,
348 table: &WrappedTable,
349 rebuilder: &dyn Rebuilder,
350 subset: Option<SubsetRef>,
351 exec_state: &mut ExecutionState,
352 ) -> ContainerRebuildSummary {
353 if let Some(subset) = subset
354 && incremental_rebuild(
355 subset.size(),
356 self.to_id.len(),
357 parallelize_intra_container_op(self.to_id.len()),
358 )
359 {
360 return self.apply_rebuild_incremental(
361 table,
362 rebuilder,
363 exec_state,
364 subset,
365 rebuilder.hint_col().unwrap(),
366 );
367 }
368 self.apply_rebuild_nonincremental(rebuilder, exec_state)
369 }
370
371 fn extend_containers_containing(&self, values: &IndexSet<Value>, out: &mut IndexSet<Value>) {
372 for value in values {
373 if let Some(containers) = self.val_index.get(value) {
374 out.extend(containers.iter().copied());
375 }
376 }
377 }
378
379 fn rebuild_val_with(
380 &self,
381 value: Value,
382 exec_state: &mut ExecutionState,
383 remap: &(dyn Fn(Value) -> Value + Send + Sync),
384 ) -> Option<Value> {
385 let mut container = self.get_container(value)?.clone();
388 container.rebuild_contents(&ClosureRebuilder { remap });
389 Some(self.get_or_insert(&container, exec_state))
390 }
391}
392
393impl<C: ContainerValue> ContainerEnv<C> {
394 pub fn new(merge_fn: Box<dyn MergeFn>, counter: CounterId) -> Self {
395 Self {
396 merge_fn,
397 counter,
398 to_id: DashMap::default(),
399 to_container: DashMap::default(),
400 val_index: DashMap::default(),
401 }
402 }
403
404 fn get_or_insert(&self, container: &C, exec_state: &mut ExecutionState) -> Value {
405 if let Some(value) = self.to_id.get(container) {
406 return *value;
407 }
408
409 let value = Value::from_usize(exec_state.inc_counter(self.counter));
414 let target_map = self.to_id.determine_map(container);
415 debug_assert_eq!(
419 target_map,
420 self.to_container
421 .determine_shard(hash_container(container) as usize)
422 );
423 self.to_container
424 .insert(value, (hash_container(container) as usize, target_map));
425
426 match self.to_id.entry(container.clone()) {
429 dashmap::Entry::Vacant(vac) => {
430 vac.insert(value);
432 for val in container.iter() {
433 self.val_index.entry(val).or_default().insert(value);
434 }
435 value
436 }
437 dashmap::Entry::Occupied(occ) => {
438 let res = *occ.get();
442 std::mem::drop(occ); self.to_container.remove(&value);
444 res
445 }
446 }
447 }
448
449 fn insert_owned(&self, container: C, value: Value, exec_state: &mut ExecutionState) -> Value {
450 let hc = hash_container(&container);
451 let target_map = self.to_id.determine_map(&container);
452 match self.to_id.entry(container) {
453 dashmap::Entry::Occupied(mut occ) => {
454 let result = (self.merge_fn)(exec_state, *occ.get(), value);
455 let old_val = *occ.get();
456 if result != old_val {
457 self.to_container.remove(&old_val);
458 self.to_container.insert(result, (hc as usize, target_map));
459 *occ.get_mut() = result;
460 for val in occ.key().iter() {
461 let mut index = self.val_index.entry(val).or_default();
462 index.swap_remove(&old_val);
463 index.insert(result);
464 }
465 }
466 result
467 }
468 dashmap::Entry::Vacant(vacant_entry) => {
469 self.to_container.insert(value, (hc as usize, target_map));
470 for val in vacant_entry.key().iter() {
471 self.val_index.entry(val).or_default().insert(value);
472 }
473 vacant_entry.insert(value);
474 value
475 }
476 }
477 }
478
479 fn reinsert_incremental(
480 &self,
481 container: C,
482 old_id: Value,
483 rebuilt_id: Value,
484 container_changed: bool,
485 exec_state: &mut ExecutionState,
486 summary: &mut ContainerRebuildSummary,
487 ) {
488 if container_changed || rebuilt_id != old_id {
489 summary.note_change();
490 }
491 if rebuilt_id != old_id {
492 self.to_container.remove(&old_id);
495 }
496 let actual = self.insert_owned(container, rebuilt_id, exec_state);
497 if container_changed && rebuilt_id == old_id && actual == old_id {
498 summary.note_dirty_id(old_id);
499 }
500 }
501
502 fn apply_rebuild_incremental(
503 &mut self,
504 table: &WrappedTable,
505 rebuilder: &dyn Rebuilder,
506 exec_state: &mut ExecutionState,
507 to_scan: SubsetRef,
508 search_col: ColumnId,
509 ) -> ContainerRebuildSummary {
510 let mut summary = ContainerRebuildSummary::default();
517 let mut buf = TaggedRowBuffer::new(1);
518 table.scan_project(
519 to_scan,
520 &[search_col],
521 Offset::new(0),
522 usize::MAX,
523 &[],
524 &mut buf,
525 );
526 let mut to_rebuild = IndexSet::<Value>::default();
528 for (_, row) in buf.iter() {
529 to_rebuild.insert(row[0]);
530 let Some(ids) = self.val_index.get(&row[0]) else {
531 continue;
532 };
533 to_rebuild.extend(&*ids);
534 }
535 for id in to_rebuild {
536 let Some((hc, target_map)) = self.to_container.get(&id).map(|x| *x) else {
537 continue;
538 };
539 let shard_mut = self.to_id.shards_mut()[target_map].get_mut();
540 let Some((mut container, _)) =
541 shard_mut.remove_entry(hc as u64, |(_, v)| *v.get() == id)
542 else {
543 continue;
544 };
545 let rebuilt_id = rebuilder.rebuild_val(id);
546 let container_changed = container.rebuild_contents(rebuilder);
547 self.reinsert_incremental(
548 container,
549 id,
550 rebuilt_id,
551 container_changed,
552 exec_state,
553 &mut summary,
554 );
555 }
556 summary
557 }
558
559 fn apply_rebuild_nonincremental(
560 &mut self,
561 rebuilder: &dyn Rebuilder,
562 exec_state: &mut ExecutionState,
563 ) -> ContainerRebuildSummary {
564 if parallelize_inter_container_op(self.to_id.len()) {
565 return self.apply_rebuild_nonincremental_parallel(rebuilder, exec_state);
566 }
567 let mut summary = ContainerRebuildSummary::default();
568 let mut to_reinsert = Vec::new();
569 let shards = self.to_id.shards_mut();
570 for shard in shards.iter_mut() {
571 let shard = shard.get_mut();
572 for bucket in unsafe { shard.iter() } {
574 let (container, val) = unsafe { bucket.as_mut() };
576 let old_val = *val.get();
577 let new_val = rebuilder.rebuild_val(old_val);
578 let container_changed = container.rebuild_contents(rebuilder);
579 if !container_changed && new_val == old_val {
580 continue;
582 }
583 summary.note_change();
584 if container_changed {
585 let ((container, _), _) = unsafe { shard.remove(bucket) };
589 self.to_container.remove(&old_val);
590 to_reinsert.push((container, new_val, new_val == old_val));
591 } else {
592 *val.get_mut() = new_val;
594 let prev = self.to_container.remove(&old_val).unwrap().1;
595 self.to_container.insert(new_val, prev);
596 }
597 }
598 }
599 for (container, val, stable_id) in to_reinsert {
600 let actual = self.insert_owned(container, val, exec_state);
601 if stable_id && actual == val {
605 summary.note_dirty_id(val);
606 }
607 }
608 summary
609 }
610
611 fn apply_rebuild_nonincremental_parallel(
612 &mut self,
613 rebuilder: &dyn Rebuilder,
614 exec_state: &mut ExecutionState,
615 ) -> ContainerRebuildSummary {
616 let mut to_reinsert =
621 IdVec::<usize , SegQueue<(C, Value, bool)>>::default();
622 to_reinsert.resize_with(self.to_id.shards().len(), Default::default);
623
624 let shards = self.to_id.shards_mut();
625 let changed = parallel::map_mut(shards, |_, shard| {
626 let mut changed = false;
627 let shard = shard.get_mut();
628 for bucket in unsafe { shard.iter() } {
630 let (container, val) = unsafe { bucket.as_mut() };
632 let old_val = *val.get();
633 let new_val = rebuilder.rebuild_val(old_val);
634 let container_changed = container.rebuild_contents(rebuilder);
635 if !container_changed && new_val == old_val {
636 continue;
638 }
639 changed = true;
640 if container_changed {
641 let ((container, _), _) = unsafe { shard.remove(bucket) };
645 self.to_container.remove(&old_val);
646 let shard = self
651 .to_container
652 .determine_shard(hash_container(&container) as usize);
653 to_reinsert[shard].push((container, new_val, new_val == old_val));
654 } else {
655 *val.get_mut() = new_val;
657 let prev = self.to_container.remove(&old_val).unwrap().1;
658 self.to_container.insert(new_val, prev);
659 }
660 }
661 changed
662 })
663 .into_iter()
664 .any(|changed| changed);
665
666 let dirty_ids = SegQueue::new();
667 parallel::for_each_mut(shards, |shard_id, shard| {
668 let mut exec_state = exec_state.clone();
669 let shard = shard.get_mut();
675 let queue = &to_reinsert[shard_id];
676 while let Some((container, val, stable_id)) = queue.pop() {
677 let hc = hash_container(&container);
678 let target_map = self.to_container.determine_shard(hc as usize);
679 match shard.find_or_find_insert_slot(
680 hc,
681 |(c, _)| c == &container,
682 |(c, _)| hash_container(c),
683 ) {
684 Ok(bucket) => {
685 let (container, val_slot) = unsafe { bucket.as_mut() };
688 let old_val = *val_slot.get();
689 let result = (self.merge_fn)(&mut exec_state, old_val, val);
690 if result != old_val {
691 self.to_container.remove(&old_val);
692 self.to_container.insert(result, (hc as usize, target_map));
693 *val_slot.get_mut() = result;
694 for val in container.iter() {
695 let mut index = self.val_index.entry(val).or_default();
696 index.swap_remove(&old_val);
697 index.insert(result);
698 }
699 }
700 if stable_id && result == val {
703 dirty_ids.push(val);
704 }
705 }
706 Err(slot) => {
707 self.to_container.insert(val, (hc as usize, target_map));
708 for v in container.iter() {
709 self.val_index.entry(v).or_default().insert(val);
710 }
711 unsafe {
714 shard.insert_in_slot(hc, slot, (container, SharedValue::new(val)));
715 }
716 if stable_id {
717 dirty_ids.push(val);
718 }
719 }
720 }
721 }
722 });
723 let mut summary = ContainerRebuildSummary::default();
724 if changed {
725 summary.note_change();
726 }
727 while let Some(value) = dirty_ids.pop() {
728 summary.note_dirty_id(value);
729 }
730 summary
731 }
732
733 fn get_container(&self, value: Value) -> Option<impl Deref<Target = C> + '_> {
734 let (hc, target_map) = *self.to_container.get(&value)?;
735 let shard = &self.to_id.shards()[target_map];
736 let read_guard = shard.read();
737 let val_ptr: *const (C, _) = shard
738 .read()
739 .find(hc as u64, |(_, v)| *v.get() == value)?
740 .as_ptr();
741 struct ValueDeref<'a, T, Guard> {
742 _guard: Guard,
743 data: &'a T,
744 }
745
746 impl<T, Guard> Deref for ValueDeref<'_, T, Guard> {
747 type Target = T;
748
749 fn deref(&self) -> &T {
750 self.data
751 }
752 }
753
754 Some(ValueDeref {
755 _guard: read_guard,
756 data: unsafe {
758 let unwrapped: &(C, _) = &*val_ptr;
759 &unwrapped.0
760 },
761 })
762 }
763}
764
765fn incremental_rebuild(uf_size: usize, table_size: usize, parallel: bool) -> bool {
766 if parallel {
767 table_size > 1000 && uf_size * 512 <= table_size
768 } else {
769 table_size > 1000 && uf_size * 8 <= table_size
770 }
771}
772
773struct ClosureRebuilder<'a> {
777 remap: &'a (dyn Fn(Value) -> Value + Send + Sync),
778}
779
780impl ValueRebuilder for ClosureRebuilder<'_> {
781 fn rebuild_val(&self, val: Value) -> Value {
782 (self.remap)(val)
783 }
784}