1use std::alloc::Layout;
23use std::sync::atomic::{fence, AtomicPtr, AtomicU64, AtomicUsize, Ordering};
24use std::sync::{LazyLock, Once};
25
26use catallaxy_gc::{RootSource, Tracer};
27use parking_lot::Mutex;
28use rustc_hash::FxHashMap;
29
30use crate::buf::{release, WordBuf};
31use crate::object::{mark_value, str_of};
32use crate::ownership::Access;
33use crate::value::Value;
34
35pub const MAX_SHAPE_KEYS: usize = 32;
37const MAX_SHAPES: usize = 1 << 16;
40
41pub struct Shape {
44 pub id: u32,
45 keys: Vec<(u64, Value)>,
46 transitions: Mutex<FxHashMap<Box<str>, &'static Shape>>,
47}
48
49struct ShapeRegistry {
50 count: usize,
51 keys: Vec<Value>,
53}
54
55static REGISTRY: Mutex<ShapeRegistry> = Mutex::new(ShapeRegistry { count: 1, keys: Vec::new() });
56static ROOT_SHAPE: LazyLock<&'static Shape> =
57 LazyLock::new(|| Box::leak(Box::new(Shape { id: 0, keys: Vec::new(), transitions: Mutex::default() })));
58static SHAPE_ROOTS: Once = Once::new();
59
60struct ShapeKeyRoots;
61impl RootSource for ShapeKeyRoots {
62 fn trace_roots(&self, t: &mut Tracer) {
63 for &k in REGISTRY.lock().keys.iter() {
65 mark_value(t, k);
66 }
67 }
68}
69
70impl Shape {
71 pub fn len(&self) -> usize {
72 self.keys.len()
73 }
74 pub fn is_empty(&self) -> bool {
75 self.keys.is_empty()
76 }
77
78 fn add(&'static self, hash: u64, key: Value) -> Option<&'static Shape> {
81 let k = str_of(key)?;
82 let mut tr = self.transitions.lock();
83 if let Some(&next) = tr.get(k) {
84 return Some(next);
85 }
86 let mut reg = REGISTRY.lock();
87 if reg.count >= MAX_SHAPES {
88 return None;
89 }
90 SHAPE_ROOTS.call_once(|| catallaxy_gc::add_global_roots(Box::new(ShapeKeyRoots)));
91 if let Some(h) = key.header() {
93 h.tenure();
94 }
95 let id = reg.count as u32;
96 reg.count += 1;
97 reg.keys.push(key);
98 let mut keys = self.keys.clone();
99 keys.push((hash, key));
100 let next: &'static Shape = Box::leak(Box::new(Shape { id, keys, transitions: Mutex::default() }));
101 tr.insert(k.into(), next);
102 Some(next)
103 }
104}
105
106#[derive(Copy, Clone)]
107pub struct Entry {
108 pub hash: u64,
109 pub key: Value,
110 pub value: Value,
111}
112
113const TOMB: Entry = Entry { hash: 0, key: Value::TOMBSTONE, value: Value::NIL };
114
115const EMPTY_IX: u64 = u64::MAX;
117const W: usize = 3;
119
120pub struct DictTable {
121 seq: AtomicU64,
124 live: AtomicUsize,
125 used: AtomicUsize,
127 entries: AtomicPtr<WordBuf>,
129 indices: AtomicPtr<WordBuf>,
131 shape: AtomicPtr<Shape>,
133}
134
135pub enum Probe {
137 Found(usize),
139 Absent,
141 Check(Vec<usize>),
143}
144
145impl Default for DictTable {
146 fn default() -> Self {
147 DictTable::new()
148 }
149}
150
151#[inline(always)]
154fn young(v: Value) -> bool {
155 v.header().is_some_and(|h| h.is_young())
156}
157
158fn index_cap(n: usize) -> usize {
159 (n * 3 / 2 + 1).next_power_of_two().max(8)
160}
161
162impl DictTable {
163 pub fn new() -> Self {
164 DictTable {
165 seq: AtomicU64::new(0),
166 live: AtomicUsize::new(0),
167 used: AtomicUsize::new(0),
168 entries: AtomicPtr::new(WordBuf::empty()),
169 indices: AtomicPtr::new(WordBuf::empty()),
170 shape: AtomicPtr::new(std::ptr::null_mut()),
171 }
172 }
173
174 pub fn with_capacity(n: usize) -> Self {
175 Self::try_with_capacity(n).unwrap_or_else(|layout| std::alloc::handle_alloc_error(layout))
176 }
177
178 pub fn try_with_capacity(n: usize) -> Result<Self, Layout> {
181 let t = DictTable::new();
182 if n > 0 {
183 let ents = n.checked_mul(W).ok_or_else(|| Layout::from_size_align(isize::MAX as usize & !7, 8).unwrap())?;
185 t.entries.store(WordBuf::try_alloc(ents, Value::NIL.0)?, Ordering::Relaxed);
186 t.indices.store(WordBuf::try_alloc(index_cap(n), EMPTY_IX)?, Ordering::Relaxed);
188 }
189 Ok(t)
190 }
191
192 pub fn record(n: usize) -> Self {
194 Self::try_record(n).unwrap_or_else(|layout| std::alloc::handle_alloc_error(layout))
195 }
196
197 pub fn try_record(n: usize) -> Result<Self, Layout> {
200 let t = DictTable::new();
201 t.entries.store(WordBuf::try_alloc(n, Value::NIL.0)?, Ordering::Relaxed);
202 t.shape.store(*ROOT_SHAPE as *const Shape as *mut Shape, Ordering::Relaxed);
203 Ok(t)
204 }
205
206 #[inline]
212 fn ents(&self) -> &[AtomicU64] {
213 unsafe { WordBuf::slots(self.entries.load(Ordering::Acquire)) }
214 }
215
216 #[inline]
218 pub fn shape(&self) -> Option<&'static Shape> {
219 let p = self.shape.load(Ordering::Acquire);
220 (!p.is_null()).then(|| unsafe { &*p })
221 }
222
223 #[inline]
225 pub fn slot_value(&self, slot: usize) -> Option<Value> {
226 if slot >= self.used.load(Ordering::Relaxed) {
227 return None;
228 }
229 self.ents().get(slot).map(|w| Value(w.load(Ordering::Relaxed)))
230 }
231
232 pub fn len(&self) -> usize {
233 self.live.load(Ordering::Relaxed)
234 }
235 pub fn is_empty(&self) -> bool {
236 self.len() == 0
237 }
238 pub fn version(&self) -> u64 {
239 self.seq.load(Ordering::Acquire)
240 }
241
242 pub fn heap_bytes(&self) -> usize {
243 unsafe {
244 8 * (WordBuf::cap(self.entries.load(Ordering::Relaxed))
245 + WordBuf::cap(self.indices.load(Ordering::Relaxed)))
246 }
247 }
248
249 pub fn trace(&self, t: &mut Tracer) {
251 for e in self.iter() {
253 mark_value(t, e.key);
254 mark_value(t, e.value);
255 }
256 }
257
258 pub fn trace_written(&self, t: &mut Tracer) {
262 let r = unsafe { WordBuf::take_written(self.entries.load(Ordering::Relaxed)) };
265 let w = if self.shape().is_some() { 1 } else { W };
266 for i in r.start / w..r.end.div_ceil(w).min(self.raw_len()) {
267 if let Some(e) = self.raw_entry(i) {
268 mark_value(t, e.key);
269 mark_value(t, e.value);
270 }
271 }
272 }
273
274 #[doc(hidden)]
277 pub fn written_range(&self) -> std::ops::Range<usize> {
278 unsafe { WordBuf::written(self.entries.load(Ordering::Acquire)) }
280 }
281
282 pub fn iter(&self) -> impl Iterator<Item = Entry> + '_ {
284 (0..self.raw_len()).filter_map(|i| self.raw_entry(i))
285 }
286
287 pub fn entry(&self, idx: usize) -> Entry {
289 let ents = self.ents();
290 match self.shape() {
291 Some(shape) => match (shape.keys.get(idx), ents.get(idx)) {
292 (Some(&(hash, key)), Some(v)) => Entry { hash, key, value: Value(v.load(Ordering::Relaxed)) },
293 _ => TOMB,
294 },
295 None => match ents.get(idx * W..idx * W + W) {
296 Some(w) => Entry {
297 hash: w[0].load(Ordering::Relaxed),
298 key: Value(w[1].load(Ordering::Relaxed)),
299 value: Value(w[2].load(Ordering::Relaxed)),
300 },
301 None => TOMB,
302 },
303 }
304 }
305
306 pub fn raw_len(&self) -> usize {
309 self.used.load(Ordering::Relaxed)
310 }
311 pub fn raw_entry(&self, i: usize) -> Option<Entry> {
312 if i >= self.raw_len() {
313 return None;
314 }
315 Some(self.entry(i)).filter(|e| e.key != Value::TOMBSTONE)
316 }
317
318 pub fn first_live(&self) -> Option<usize> {
319 (0..self.raw_len()).find(|&i| self.raw_entry(i).is_some())
320 }
321
322 pub fn probe(&self, hash: u64, key: Value) -> Probe {
323 let mut check = Vec::new();
324 if let Some(shape) = self.shape() {
325 let n = self.raw_len().min(shape.keys.len());
326 for (i, &(h, k)) in shape.keys[..n].iter().enumerate() {
327 if h == hash {
328 if k == key {
329 return Probe::Found(i);
330 }
331 check.push(i);
332 }
333 }
334 return if check.is_empty() { Probe::Absent } else { Probe::Check(check) };
335 }
336 let idx = unsafe { WordBuf::slots(self.indices.load(Ordering::Acquire)) };
337 if idx.is_empty() {
338 return Probe::Absent;
339 }
340 let mask = idx.len() - 1;
341 let mut i = (hash as usize) & mask;
342 let mut perturb = hash;
343 for _ in 0..idx.len() {
345 let ix = idx[i].load(Ordering::Relaxed);
346 if ix == EMPTY_IX {
347 break;
348 }
349 let e = self.entry(ix as usize);
350 if e.key != Value::TOMBSTONE && e.hash == hash {
351 if e.key == key {
352 return Probe::Found(ix as usize);
353 }
354 check.push(ix as usize);
355 }
356 perturb >>= 5;
357 i = (i.wrapping_mul(5).wrapping_add(1).wrapping_add(perturb as usize)) & mask;
358 }
359 if check.is_empty() {
360 Probe::Absent
361 } else {
362 Probe::Check(check)
363 }
364 }
365
366 pub fn read_optimistic<R>(&self, f: &impl Fn(&DictTable) -> R) -> Option<R> {
369 for _ in 0..4 {
370 let s1 = self.seq.load(Ordering::Acquire);
371 if s1 & 1 == 1 {
372 std::hint::spin_loop();
373 continue;
374 }
375 let r = f(self);
376 fence(Ordering::Acquire);
377 if self.seq.load(Ordering::Relaxed) == s1 {
378 return Some(r);
379 }
380 }
381 None
382 }
383
384 fn begin(&self, access: Access) -> u64 {
389 let s = self.seq.load(Ordering::Relaxed);
390 self.seq.store(s + 1, Ordering::Relaxed);
391 if access == Access::Shared {
392 fence(Ordering::Release);
393 }
394 s
395 }
396
397 fn end(&self, s: u64) {
398 self.seq.store(s + 2, Ordering::Release);
399 }
400
401 pub unsafe fn set_value_at(&self, idx: usize, v: Value, access: Access) {
404 let s = self.begin(access);
405 let slot = if self.shape().is_some() { idx } else { idx * W + 2 };
406 let ents = self.entries.load(Ordering::Relaxed);
407 if let Some(w) = WordBuf::slots(ents).get(slot) {
408 w.store(v.0, Ordering::Relaxed);
409 if young(v) {
410 WordBuf::note_written(ents, slot, slot + 1);
411 }
412 }
413 self.end(s);
414 }
415
416 pub unsafe fn insert_new(&self, hash: u64, key: Value, value: Value, access: Access) {
422 if let Err(layout) = self.try_insert_new(hash, key, value, access) {
423 std::alloc::handle_alloc_error(layout);
424 }
425 }
426
427 pub unsafe fn try_insert_new(&self, hash: u64, key: Value, value: Value, access: Access) -> Result<(), Layout> {
433 let s = self.begin(access);
434 let r = self.insert_in_section(hash, key, value, access);
435 self.end(s);
436 r
437 }
438
439 unsafe fn insert_in_section(&self, hash: u64, key: Value, value: Value, access: Access) -> Result<(), Layout> {
441 if let Some(shape) = self.shape() {
442 if shape.len() < MAX_SHAPE_KEYS {
443 if let Some(next) = shape.add(hash, key) {
444 let n = self.raw_len();
445 let cap = WordBuf::cap(self.entries.load(Ordering::Relaxed));
446 if n == cap {
447 let buf = WordBuf::try_alloc((cap * 2).max(4), Value::NIL.0)?;
448 let old = self.entries.load(Ordering::Relaxed);
449 for (d, w) in WordBuf::slots(buf).iter().zip(WordBuf::slots(old)) {
450 d.store(w.load(Ordering::Relaxed), Ordering::Relaxed);
451 }
452 WordBuf::copy_written(old, buf);
454 release(self.entries.swap(buf, Ordering::AcqRel), access);
455 }
456 let ents = self.entries.load(Ordering::Relaxed);
457 WordBuf::slots(ents)[n].store(value.0, Ordering::Relaxed);
458 if young(value) {
459 WordBuf::note_written(ents, n, n + 1);
460 }
461 self.shape.store(next as *const Shape as *mut Shape, Ordering::Release);
462 self.used.store(n + 1, Ordering::Relaxed);
463 self.live.fetch_add(1, Ordering::Relaxed);
464 return Ok(());
465 }
466 }
467 self.try_leave_shape_mode(access)?;
468 }
469 let n = self.raw_len();
470 let icap = WordBuf::cap(self.indices.load(Ordering::Relaxed));
471 let ecap = WordBuf::cap(self.entries.load(Ordering::Relaxed)) / W;
472 if (n + 1) * 3 >= icap * 2 || n == ecap {
473 self.rebuild(self.len() + 1, access)?;
474 }
475 let n = self.raw_len();
476 let ents = self.entries.load(Ordering::Relaxed);
477 let w = &WordBuf::slots(ents)[n * W..n * W + W];
478 w[0].store(hash, Ordering::Relaxed);
479 w[1].store(key.0, Ordering::Relaxed);
480 w[2].store(value.0, Ordering::Relaxed);
481 if young(key) || young(value) {
482 WordBuf::note_written(ents, n * W, n * W + W);
483 }
484 Self::place(WordBuf::slots(self.indices.load(Ordering::Relaxed)), hash, n);
485 self.used.store(n + 1, Ordering::Relaxed);
486 self.live.fetch_add(1, Ordering::Relaxed);
487 Ok(())
488 }
489
490 fn place(idx: &[AtomicU64], hash: u64, at: usize) {
491 let mask = idx.len() - 1;
492 let mut i = (hash as usize) & mask;
493 let mut perturb = hash;
494 while idx[i].load(Ordering::Relaxed) != EMPTY_IX {
495 perturb >>= 5;
496 i = (i.wrapping_mul(5).wrapping_add(1).wrapping_add(perturb as usize)) & mask;
497 }
498 idx[i].store(at as u64, Ordering::Relaxed);
499 }
500
501 unsafe fn rebuild(&self, want: usize, access: Access) -> Result<(), Layout> {
504 let live = self.len();
505 let (ents, idx) = Self::alloc_pair((want.max(live) * 2).max(4) * W, index_cap(want.max(live) * 2))?;
506 let (ew, iw) = (WordBuf::slots(ents), WordBuf::slots(idx));
507 let mut n = 0;
510 for (i, e) in self.iter().enumerate() {
511 ew[i * W].store(e.hash, Ordering::Relaxed);
512 ew[i * W + 1].store(e.key.0, Ordering::Relaxed);
513 ew[i * W + 2].store(e.value.0, Ordering::Relaxed);
514 Self::place(iw, e.hash, i);
515 n = i + 1;
516 }
517 release(self.entries.swap(ents, Ordering::AcqRel), access);
518 release(self.indices.swap(idx, Ordering::AcqRel), access);
519 self.used.store(n, Ordering::Relaxed);
520 Ok(())
521 }
522
523 fn alloc_pair(ecap: usize, icap: usize) -> Result<(*mut WordBuf, *mut WordBuf), Layout> {
525 let ents = WordBuf::try_alloc(ecap, Value::NIL.0)?;
526 match WordBuf::try_alloc(icap, EMPTY_IX) {
527 Ok(idx) => Ok((ents, idx)),
528 Err(layout) => {
529 unsafe { WordBuf::free(ents as *mut u8) };
531 Err(layout)
532 }
533 }
534 }
535
536 unsafe fn leave_shape_mode(&self, access: Access) {
538 if let Err(layout) = self.try_leave_shape_mode(access) {
539 std::alloc::handle_alloc_error(layout);
540 }
541 }
542
543 unsafe fn try_leave_shape_mode(&self, access: Access) -> Result<(), Layout> {
545 if self.shape().is_none() {
546 return Ok(());
547 }
548 let n = self.raw_len();
549 let (ents, idx) = Self::alloc_pair((n + 1).max(4) * W, index_cap(n + 1))?;
550 let (ew, iw) = (WordBuf::slots(ents), WordBuf::slots(idx));
551 for i in 0..n {
552 let e = self.entry(i);
553 ew[i * W].store(e.hash, Ordering::Relaxed);
554 ew[i * W + 1].store(e.key.0, Ordering::Relaxed);
555 ew[i * W + 2].store(e.value.0, Ordering::Relaxed);
556 Self::place(iw, e.hash, i);
557 }
558 release(self.indices.swap(idx, Ordering::AcqRel), access);
559 let old = self.entries.swap(ents, Ordering::AcqRel);
560 self.shape.store(std::ptr::null_mut(), Ordering::Release);
561 release(old, access);
562 Ok(())
563 }
564
565 pub unsafe fn remove_at(&self, idx: usize, access: Access) -> Value {
570 let s = self.begin(access);
571 self.leave_shape_mode(access);
572 let mut v = Value::NIL;
573 if let Some(w) = self.ents().get(idx * W..idx * W + W) {
574 v = Value(w[2].load(Ordering::Relaxed));
575 w[1].store(Value::TOMBSTONE.0, Ordering::Relaxed);
577 w[2].store(Value::NIL.0, Ordering::Relaxed);
578 self.live.fetch_sub(1, Ordering::Relaxed);
579 }
580 self.end(s);
581 v
582 }
583
584 pub unsafe fn clear(&self, access: Access) {
587 let s = self.begin(access);
588 if self.shape().is_some() {
589 self.shape.store(*ROOT_SHAPE as *const Shape as *mut Shape, Ordering::Release);
590 }
591 self.used.store(0, Ordering::Relaxed);
592 self.live.store(0, Ordering::Relaxed);
593 release(self.entries.swap(WordBuf::empty(), Ordering::AcqRel), access);
594 release(self.indices.swap(WordBuf::empty(), Ordering::AcqRel), access);
595 self.end(s);
596 }
597}
598
599impl Drop for DictTable {
600 fn drop(&mut self) {
601 unsafe {
602 WordBuf::free(*self.entries.get_mut() as *mut u8);
603 WordBuf::free(*self.indices.get_mut() as *mut u8);
604 }
605 }
606}
607
608impl Clone for DictTable {
609 fn clone(&self) -> Self {
610 self.try_clone().unwrap_or_else(|layout| std::alloc::handle_alloc_error(layout))
611 }
612}
613
614impl DictTable {
615 pub fn try_clone(&self) -> Result<Self, Layout> {
617 let n = self.raw_len();
618 if let Some(shape) = self.shape() {
619 let t = DictTable::try_record(n)?;
620 for (d, w) in t.ents().iter().zip(&self.ents()[..n.min(self.ents().len())]) {
621 d.store(w.load(Ordering::Relaxed), Ordering::Relaxed);
622 }
623 t.shape.store(shape as *const Shape as *mut Shape, Ordering::Relaxed);
624 t.used.store(n, Ordering::Relaxed);
625 t.live.store(self.len(), Ordering::Relaxed);
626 return Ok(t);
627 }
628 let t = DictTable::try_with_capacity(self.len())?;
629 for e in self.iter() {
630 unsafe { t.try_insert_new(e.hash, e.key, e.value, Access::Local)? };
632 }
633 Ok(t)
634 }
635}