Skip to main content

celox_state_layout/
trace.rs

1//! Write activity shared by generated code and waveform observers.
2//!
3//! An object is assigned by its stable home, so aliases share notifications
4//! and partial/dynamic stores also notify observers of the complete object.
5use serde::{Deserialize, Serialize};
6
7pub const TRACE_GROUP_BYTES: usize = 64;
8const GROUPS_PER_SUMMARY: usize = 64;
9
10#[derive(Clone, Debug, Serialize, Deserialize)]
11pub struct TraceLayout {
12    pub flags_offset: usize,
13    pub group_count: usize,
14    pub summary_offset: usize,
15    pub summary_count: usize,
16    /// Disjoint physical homes, including both planes and element padding.
17    homes: Vec<(usize, usize)>,
18}
19
20impl TraceLayout {
21    pub(crate) fn new(offset: usize, stable_size: usize, mut homes: Vec<(usize, usize)>) -> Self {
22        homes.sort_unstable();
23        let mut merged: Vec<(usize, usize)> = Vec::new();
24        for (start, end) in homes {
25            if let Some(last) = merged.last_mut()
26                && last.0 == start
27            {
28                last.1 = last.1.max(end);
29            } else {
30                merged.push((start, end));
31            }
32        }
33        let group_count = stable_size.div_ceil(TRACE_GROUP_BYTES);
34        let summary_count = group_count.div_ceil(GROUPS_PER_SUMMARY);
35        Self {
36            flags_offset: offset,
37            group_count,
38            summary_offset: offset + group_count,
39            summary_count,
40            homes: merged,
41        }
42    }
43
44    pub fn end_offset(&self) -> usize {
45        self.summary_offset + self.summary_count
46    }
47
48    pub fn validate(
49        &self,
50        stable_size: usize,
51        metadata_start: usize,
52        scratch_start: usize,
53    ) -> bool {
54        self.group_count == stable_size.div_ceil(TRACE_GROUP_BYTES)
55            && self.summary_count == self.group_count.div_ceil(GROUPS_PER_SUMMARY)
56            && self.flags_offset >= metadata_start
57            && self.flags_offset.checked_add(self.group_count) == Some(self.summary_offset)
58            && self
59                .summary_offset
60                .checked_add(self.summary_count)
61                .is_some_and(|end| end <= scratch_start)
62            && self
63                .homes
64                .iter()
65                .all(|&(start, end)| start <= end && end <= stable_size)
66            && self.homes.windows(2).all(|homes| homes[0].1 <= homes[1].0)
67    }
68
69    /// Two byte stores suffice: no read/modify/write in generated hot paths.
70    pub fn notification_offsets(&self, stable_home: usize) -> [usize; 2] {
71        let group = stable_home / TRACE_GROUP_BYTES;
72        [
73            self.flags_offset + group,
74            self.summary_offset + group / GROUPS_PER_SUMMARY,
75        ]
76    }
77
78    /// # Safety
79    /// `memory` must exclusively reference the complete writable state image.
80    pub unsafe fn mark_home(&self, memory: *mut u8, home: usize) {
81        for offset in self.notification_offsets(home) {
82            unsafe {
83                *memory.add(offset) = 1;
84            }
85        }
86    }
87
88    /// Mark host writes, including an element view into a larger array.
89    /// # Safety
90    /// `memory` must exclusively reference the complete writable state image.
91    pub unsafe fn mark_range(&self, memory: *mut u8, offset: usize, len: usize) {
92        if len == 0 {
93            return;
94        }
95        let first = self
96            .homes
97            .partition_point(|&(start, _)| start <= offset)
98            .saturating_sub(1);
99        for &(start, end) in &self.homes[first..] {
100            if start >= offset.saturating_add(len) {
101                break;
102            }
103            if end > offset {
104                unsafe {
105                    self.mark_home(memory, start);
106                }
107            }
108        }
109    }
110
111    /// Consume only nonempty groups. Clock trigger clearing never touches this
112    /// region. The caller owns the reusable result allocation.
113    pub fn take(&self, memory: &mut [u8], groups: &mut Vec<usize>) {
114        groups.clear();
115        let (prefix, summaries) = memory.split_at_mut(self.summary_offset);
116        let flags = &mut prefix[self.flags_offset..][..self.group_count];
117        take_nonzero(&mut summaries[..self.summary_count], |summary| {
118            let start = summary * GROUPS_PER_SUMMARY;
119            let flags = &mut flags[start..];
120            let len = flags.len().min(GROUPS_PER_SUMMARY);
121            take_nonzero(&mut flags[..len], |index| groups.push(start + index));
122        });
123    }
124}
125
126/// Clear and visit nonzero bytes in address order. Notifications stay byte
127/// stores, while the consumer skips eight empty bytes with one word test.
128fn take_nonzero(bytes: &mut [u8], mut visit: impl FnMut(usize)) {
129    let (words, tail) = bytes.as_chunks_mut::<8>();
130    for (index, bytes) in words.iter_mut().enumerate() {
131        // A safe unaligned load; little endian keeps bit order in address order
132        // even on big-endian hosts.
133        let word = u64::from_le_bytes(*bytes);
134        if word == 0 {
135            continue;
136        }
137        bytes.fill(0);
138        // Set exactly the high bit of each nonzero byte, including flags other
139        // than 1. Adding 0x7f to the low seven bits cannot carry between bytes.
140        const LOW_BITS: u64 = 0x7f7f_7f7f_7f7f_7f7f;
141        let mut active = (((word & LOW_BITS) + LOW_BITS) | word) & !LOW_BITS;
142        while active != 0 {
143            visit(index * 8 + active.trailing_zeros() as usize / 8);
144            active &= active - 1;
145        }
146    }
147    let base = words.len() * 8;
148    for (index, byte) in tail.iter_mut().enumerate() {
149        if *byte != 0 {
150            *byte = 0;
151            visit(base + index);
152        }
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159
160    #[test]
161    fn nonzero_bytes_preserve_address_order() {
162        for alignment in 0..8 {
163            // Exercise every possible byte value, including adjacent zero and
164            // nonzero bytes, word boundaries, and short tails.
165            for len in 0..=512 {
166                let mut bytes = vec![0xa5; alignment + len + 8];
167                for (index, byte) in bytes[alignment..][..len].iter_mut().enumerate() {
168                    *byte = if index % 2 == 0 { (index / 2) as u8 } else { 0 };
169                }
170                let expected: Vec<_> = bytes[alignment..][..len]
171                    .iter()
172                    .enumerate()
173                    .filter_map(|(index, &byte)| (byte != 0).then_some(index))
174                    .collect();
175                let mut actual = vec![];
176                take_nonzero(&mut bytes[alignment..][..len], |index| actual.push(index));
177                assert_eq!(actual, expected);
178                assert!(bytes[..alignment].iter().all(|&byte| byte == 0xa5));
179                assert!(bytes[alignment..][..len].iter().all(|&byte| byte == 0));
180                assert!(bytes[alignment + len..].iter().all(|&byte| byte == 0xa5));
181            }
182        }
183    }
184
185    #[test]
186    fn activity_boundaries_match_scalar_collection() {
187        for group_count in [
188            0, 1, 7, 8, 9, 63, 64, 65, 127, 128, 447, 448, 449, 511, 512, 513, 575, 576, 577, 1024,
189        ] {
190            for alignment in 0..8 {
191                // Include partial stable groups and unaligned metadata.
192                let stable_size = (group_count * TRACE_GROUP_BYTES).saturating_sub(3);
193                let trace = TraceLayout::new(stable_size + alignment, stable_size, vec![]);
194                assert_eq!(trace.group_count, group_count);
195                for pattern in 0..5 {
196                    let mut memory = vec![0xa5; trace.end_offset() + 8];
197                    memory[trace.flags_offset..trace.end_offset()].fill(0);
198                    for group in 0..group_count {
199                        let flag = match pattern {
200                            0 => 0,
201                            1 => 1,
202                            2 => u8::from(group % 2 == 0),
203                            3 => {
204                                u8::from(group == 0 || group + 1 == group_count || group % 64 == 0)
205                            }
206                            _ => [0, 2, 0, 0x80, 0, 0xff, 1][group % 7],
207                        };
208                        memory[trace.flags_offset + group] = flag;
209                        if flag != 0 {
210                            memory[trace.summary_offset + group / GROUPS_PER_SUMMARY] = flag;
211                        }
212                    }
213                    // Both an empty marked summary and an unmarked summary with
214                    // stale flags must retain the scalar collector's behavior.
215                    if trace.summary_count > 1 && matches!(pattern, 0 | 4) {
216                        memory[trace.summary_offset] = 0xff;
217                        memory[trace.summary_offset + 1] = 0;
218                    }
219                    let mut expected_memory = memory.clone();
220                    let mut expected_groups = vec![];
221                    for summary in 0..trace.summary_count {
222                        if expected_memory[trace.summary_offset + summary] == 0 {
223                            continue;
224                        }
225                        expected_memory[trace.summary_offset + summary] = 0;
226                        for group in 0..group_count {
227                            if group / GROUPS_PER_SUMMARY == summary
228                                && expected_memory[trace.flags_offset + group] != 0
229                            {
230                                expected_memory[trace.flags_offset + group] = 0;
231                                expected_groups.push(group);
232                            }
233                        }
234                    }
235                    let mut groups = vec![usize::MAX];
236                    trace.take(&mut memory, &mut groups);
237                    assert_eq!(groups, expected_groups);
238                    // This also checks untouched state/clock bytes and the byte
239                    // immediately following the last (possibly partial) group.
240                    assert_eq!(memory, expected_memory);
241                    groups.push(usize::MAX);
242                    trace.take(&mut memory, &mut groups);
243                    assert!(groups.is_empty());
244                    assert_eq!(memory, expected_memory);
245                }
246            }
247        }
248    }
249
250    #[test]
251    fn aliases_partial_writes_and_separate_consumers() {
252        let trace = TraceLayout::new(
253            8192,
254            8192,
255            vec![(0, 512), (0, 64), (512, 520), (4096, 4104)],
256        );
257        let mut memory = vec![0; trace.end_offset()];
258        // A partial write near the end of an array notifies its stable home.
259        unsafe {
260            trace.mark_range(memory.as_mut_ptr(), 500, 2);
261            trace.mark_range(memory.as_mut_ptr(), 510, 4);
262            trace.mark_home(memory.as_mut_ptr(), 4096);
263        }
264        let mut groups = vec![];
265        trace.take(&mut memory, &mut groups);
266        assert_eq!(groups, [0, 8, 64]);
267        trace.take(&mut memory, &mut groups);
268        assert!(groups.is_empty());
269    }
270}