1use std::fmt;
2
3use super::STREAM_VERSION;
4
5pub const JOURNAL_MAGIC: [u8; 4] = *b"PFXJ";
6pub const JOURNAL_FORMAT: u16 = 1;
7pub const MAX_FIELDS: usize = 63;
8
9const HEADER_BYTES: usize = 56;
10const CHECKPOINT_BYTES: usize = 24;
11const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
12const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
13
14#[derive(Clone, Copy, Debug, PartialEq, Eq)]
15pub struct Digest(u64);
16
17impl Digest {
18 pub const fn new() -> Self {
19 Self(FNV_OFFSET)
20 }
21
22 pub fn bytes(&mut self, bytes: &[u8]) -> &mut Self {
23 for &byte in bytes {
24 self.0 = (self.0 ^ byte as u64).wrapping_mul(FNV_PRIME);
25 }
26 self
27 }
28
29 pub fn u64(&mut self, value: u64) -> &mut Self {
30 self.bytes(&value.to_le_bytes())
31 }
32
33 pub fn i64(&mut self, value: i64) -> &mut Self {
34 self.bytes(&value.to_le_bytes())
35 }
36
37 #[must_use]
38 pub const fn value(&self) -> u64 {
39 self.0
40 }
41}
42
43impl Default for Digest {
44 fn default() -> Self {
45 Self::new()
46 }
47}
48
49#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
50pub enum Field {
51 Flag,
52 Unsigned(u8),
53 Signed(u8),
54}
55
56impl Field {
57 pub const fn min(self) -> i64 {
58 match self {
59 Self::Flag | Self::Unsigned(_) => 0,
60 Self::Signed(bits) => i64::MIN >> (64 - bits as u32),
61 }
62 }
63
64 pub const fn max(self) -> i64 {
65 match self {
66 Self::Flag => 1,
67 Self::Unsigned(bits) => (u64::MAX >> (64 - bits as u32)) as i64,
68 Self::Signed(bits) => i64::MAX >> (64 - bits as u32),
69 }
70 }
71
72 const fn is_valid(self) -> bool {
73 match self {
74 Self::Flag => true,
75 Self::Unsigned(bits) => bits >= 1 && bits <= 63,
76 Self::Signed(bits) => bits >= 1 && bits <= 64,
77 }
78 }
79
80 const fn code(self) -> u8 {
81 match self {
82 Self::Flag => 0x00,
83 Self::Unsigned(bits) => 0x40 | (bits - 1),
84 Self::Signed(bits) => 0x80 | (bits - 1),
85 }
86 }
87
88 const fn from_code(code: u8) -> Option<Self> {
89 let bits = (code & 0x3f) + 1;
90 let field = match code >> 6 {
91 0 if code == 0 => Self::Flag,
92 1 => Self::Unsigned(bits),
93 2 => Self::Signed(bits),
94 _ => return None,
95 };
96 if field.is_valid() { Some(field) } else { None }
97 }
98}
99
100#[derive(Clone, Debug, PartialEq, Eq)]
101pub struct Schema {
102 fields: Box<[Field]>,
103}
104
105impl Schema {
106 pub fn new(fields: &[Field]) -> Result<Self, JournalError> {
107 if fields.is_empty() || fields.len() > MAX_FIELDS {
108 return Err(JournalError::Schema);
109 }
110 if !fields.iter().all(|field| field.is_valid()) {
111 return Err(JournalError::Schema);
112 }
113 Ok(Self {
114 fields: fields.into(),
115 })
116 }
117
118 pub fn fields(&self) -> &[Field] {
119 &self.fields
120 }
121
122 #[must_use]
123 pub fn hash(&self) -> u64 {
124 let mut digest = Digest::new();
125 digest.bytes(&(self.fields.len() as u16).to_le_bytes());
126 for field in &self.fields {
127 digest.bytes(&[field.code()]);
128 }
129 digest.value()
130 }
131}
132
133#[derive(Clone, Copy, Debug, PartialEq, Eq)]
134pub struct Header {
135 pub seed: u64,
136 pub sim_version: u32,
137 pub stream_version: u32,
138 pub schema_hash: u64,
139 pub ticks: u64,
140 pub checkpoint_every: u32,
141}
142
143#[derive(Clone, Copy, Debug, PartialEq, Eq)]
144pub struct Checkpoint {
145 pub tick: u64,
146 pub digest: u64,
147 pub offset: u64,
148}
149
150#[derive(Clone, Copy, Debug, PartialEq, Eq)]
151pub enum JournalError {
152 Schema,
153 InputLength {
154 found: usize,
155 expected: usize,
156 },
157 Input {
158 tick: u64,
159 field: usize,
160 },
161 Truncated,
162 Magic,
163 Format {
164 found: u16,
165 },
166 Checksum,
167 Corrupt {
168 at: u64,
169 },
170 SimVersion {
171 found: u32,
172 expected: u32,
173 },
174 StreamVersion {
175 found: u32,
176 },
177 SchemaMismatch {
178 found: u64,
179 expected: u64,
180 },
181 Checkpoint {
182 tick: u64,
183 found: u64,
184 expected: u64,
185 },
186 Beyond {
187 tick: u64,
188 ticks: u64,
189 },
190}
191
192impl fmt::Display for JournalError {
193 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
194 match self {
195 Self::Schema => write!(
196 f,
197 "an input schema takes 1 to {MAX_FIELDS} fields, unsigned 1 to 63 bits, signed 1 to 64"
198 ),
199 Self::InputLength { found, expected } => {
200 write!(
201 f,
202 "input has {found} values, the schema declares {expected}"
203 )
204 }
205 Self::Input { tick, field } => {
206 write!(
207 f,
208 "tick {tick}: field {field} is outside its declared width"
209 )
210 }
211 Self::Truncated => write!(f, "journal is truncated"),
212 Self::Magic => write!(f, "not a pfx journal"),
213 Self::Format { found } => write!(
214 f,
215 "journal format {found}, this build reads {JOURNAL_FORMAT}"
216 ),
217 Self::Checksum => write!(f, "journal checksum does not match its bytes"),
218 Self::Corrupt { at } => write!(f, "journal is corrupt at byte {at}"),
219 Self::SimVersion { found, expected } => write!(
220 f,
221 "journal was recorded by sim version {found}, this game runs {expected}"
222 ),
223 Self::StreamVersion { found } => write!(
224 f,
225 "journal was recorded with rng stream version {found}, this build runs {STREAM_VERSION}"
226 ),
227 Self::SchemaMismatch { found, expected } => write!(
228 f,
229 "journal input schema {found:016x} differs from the game's {expected:016x}"
230 ),
231 Self::Checkpoint {
232 tick,
233 found,
234 expected,
235 } => write!(
236 f,
237 "replay diverged by tick {tick}: digest {found:016x}, journal has {expected:016x}"
238 ),
239 Self::Beyond { tick, ticks } => {
240 write!(f, "tick {tick} is past the journal's {ticks} ticks")
241 }
242 }
243 }
244}
245
246impl std::error::Error for JournalError {}
247
248#[derive(Clone, Debug)]
249pub struct JournalWriter {
250 schema: Schema,
251 seed: u64,
252 sim_version: u32,
253 every: u32,
254 ticks: u64,
255 run: u64,
256 base: Vec<i64>,
257 stream: Vec<u8>,
258 checkpoints: Vec<Checkpoint>,
259}
260
261impl JournalWriter {
262 pub fn new(schema: &Schema, seed: u64, sim_version: u32, checkpoint_every: u32) -> Self {
263 Self {
264 schema: schema.clone(),
265 seed,
266 sim_version,
267 every: checkpoint_every,
268 ticks: 0,
269 run: 0,
270 base: vec![0; schema.fields.len()],
271 stream: Vec::new(),
272 checkpoints: Vec::new(),
273 }
274 }
275
276 pub fn ticks(&self) -> u64 {
277 self.ticks
278 }
279
280 pub fn push(
281 &mut self,
282 input: &[i64],
283 digest: impl FnOnce() -> u64,
284 ) -> Result<(), JournalError> {
285 let fields = &self.schema.fields;
286 if input.len() != fields.len() {
287 return Err(JournalError::InputLength {
288 found: input.len(),
289 expected: fields.len(),
290 });
291 }
292 let mut mask = 0u64;
293 for (i, ((field, &value), &base)) in fields.iter().zip(input).zip(&self.base).enumerate() {
294 if value < field.min() || value > field.max() {
295 return Err(JournalError::Input {
296 tick: self.ticks,
297 field: i,
298 });
299 }
300 mask |= ((value != base) as u64) << i;
301 }
302 if mask == 0 {
303 self.run += 1;
304 } else {
305 flush_run(&mut self.stream, &mut self.run);
306 put_varint(&mut self.stream, (mask << 1) | 1);
307 let mut bits = mask;
308 while bits != 0 {
309 let i = bits.trailing_zeros() as usize;
310 bits &= bits - 1;
311 if fields[i] != Field::Flag {
312 put_varint(
313 &mut self.stream,
314 zigzag(input[i].wrapping_sub(self.base[i])),
315 );
316 }
317 }
318 self.base.copy_from_slice(input);
319 }
320 self.ticks += 1;
321 if self.every != 0 && self.ticks.is_multiple_of(self.every as u64) {
322 flush_run(&mut self.stream, &mut self.run);
323 self.checkpoints.push(Checkpoint {
324 tick: self.ticks,
325 digest: digest(),
326 offset: self.stream.len() as u64,
327 });
328 self.base.fill(0);
329 }
330 Ok(())
331 }
332
333 #[must_use]
334 pub fn finish(self) -> Vec<u8> {
335 self.bytes()
336 }
337
338 #[must_use]
339 pub fn bytes(&self) -> Vec<u8> {
340 let mut out = Vec::new();
341 self.bytes_into(&mut out);
342 out
343 }
344
345 pub fn bytes_into(&self, out: &mut Vec<u8>) {
346 let mut unsealed = self.snapshot_into(std::mem::take(out));
347 unsealed.seal_in_place();
348 *out = unsealed.0;
349 }
350
351 #[must_use]
352 pub fn snapshot(&self) -> Unsealed {
353 self.snapshot_into(Vec::new())
354 }
355
356 #[must_use]
357 pub fn snapshot_into(&self, mut out: Vec<u8>) -> Unsealed {
358 let mut tail = Vec::new();
359 let mut run = self.run;
360 flush_run(&mut tail, &mut run);
361 let codes = self.schema.fields.len();
362 let stream_len = self.stream.len() + tail.len();
363 out.clear();
364 out.reserve(
365 HEADER_BYTES + codes + stream_len + self.checkpoints.len() * CHECKPOINT_BYTES + 8,
366 );
367 out.extend_from_slice(&JOURNAL_MAGIC);
368 out.extend_from_slice(&JOURNAL_FORMAT.to_le_bytes());
369 out.extend_from_slice(&(codes as u16).to_le_bytes());
370 out.extend_from_slice(&self.seed.to_le_bytes());
371 out.extend_from_slice(&self.sim_version.to_le_bytes());
372 out.extend_from_slice(&STREAM_VERSION.to_le_bytes());
373 out.extend_from_slice(&self.schema.hash().to_le_bytes());
374 out.extend_from_slice(&self.ticks.to_le_bytes());
375 out.extend_from_slice(&self.every.to_le_bytes());
376 out.extend_from_slice(&(self.checkpoints.len() as u32).to_le_bytes());
377 out.extend_from_slice(&(stream_len as u64).to_le_bytes());
378 out.extend(self.schema.fields.iter().map(|field| field.code()));
379 out.extend_from_slice(&self.stream);
380 out.extend_from_slice(&tail);
381 for checkpoint in &self.checkpoints {
382 out.extend_from_slice(&checkpoint.tick.to_le_bytes());
383 out.extend_from_slice(&checkpoint.digest.to_le_bytes());
384 out.extend_from_slice(&checkpoint.offset.to_le_bytes());
385 }
386 Unsealed(out)
387 }
388}
389
390#[derive(Clone, Debug, Default, PartialEq, Eq)]
391pub struct Unsealed(Vec<u8>);
392
393impl Unsealed {
394 #[must_use]
395 pub fn seal(mut self) -> Vec<u8> {
396 self.seal_in_place();
397 self.0
398 }
399
400 fn seal_in_place(&mut self) {
401 let sum = Digest::new().bytes(&self.0).value();
402 self.0.extend_from_slice(&sum.to_le_bytes());
403 }
404}
405
406#[derive(Clone, Debug)]
407pub struct Journal<'a> {
408 header: Header,
409 schema: Schema,
410 stream: &'a [u8],
411 stream_at: u64,
412 checkpoints: Vec<Checkpoint>,
413}
414
415impl<'a> Journal<'a> {
416 pub fn open(bytes: &'a [u8]) -> Result<Self, JournalError> {
417 if bytes.len() < JOURNAL_MAGIC.len() {
418 return Err(JournalError::Truncated);
419 }
420 if bytes[..4] != JOURNAL_MAGIC {
421 return Err(JournalError::Magic);
422 }
423 if bytes.len() < HEADER_BYTES {
424 return Err(JournalError::Truncated);
425 }
426 let u16_at = |at: usize| u16::from_le_bytes([bytes[at], bytes[at + 1]]);
427 let u32_at = |at: usize| u32::from_le_bytes(bytes[at..at + 4].try_into().unwrap());
428 let u64_at = |at: usize| u64::from_le_bytes(bytes[at..at + 8].try_into().unwrap());
429 let format = u16_at(4);
430 if format != JOURNAL_FORMAT {
431 return Err(JournalError::Format { found: format });
432 }
433 let count = u16_at(6) as usize;
434 let stream_len = u64_at(48);
435 let checkpoint_count = u32_at(44) as u64;
436 let need = (HEADER_BYTES as u64 + count as u64 + 8)
437 .checked_add(stream_len)
438 .and_then(|n| n.checked_add(checkpoint_count * CHECKPOINT_BYTES as u64))
439 .ok_or(JournalError::Corrupt { at: 48 })?;
440 if (bytes.len() as u64) < need {
441 return Err(JournalError::Truncated);
442 }
443 if bytes.len() as u64 > need {
444 return Err(JournalError::Corrupt { at: need });
445 }
446 let end = bytes.len() - 8;
447 if Digest::new().bytes(&bytes[..end]).value() != u64_at(end) {
448 return Err(JournalError::Checksum);
449 }
450 if count == 0 || count > MAX_FIELDS {
451 return Err(JournalError::Corrupt { at: 6 });
452 }
453 let fields = bytes[HEADER_BYTES..HEADER_BYTES + count]
454 .iter()
455 .enumerate()
456 .map(|(i, &code)| {
457 Field::from_code(code).ok_or(JournalError::Corrupt {
458 at: (HEADER_BYTES + i) as u64,
459 })
460 })
461 .collect::<Result<_, _>>()?;
462 let schema = Schema { fields };
463 let header = Header {
464 seed: u64_at(8),
465 sim_version: u32_at(16),
466 stream_version: u32_at(20),
467 schema_hash: u64_at(24),
468 ticks: u64_at(32),
469 checkpoint_every: u32_at(40),
470 };
471 if header.schema_hash != schema.hash() {
472 return Err(JournalError::Corrupt { at: 24 });
473 }
474 let stream_at = HEADER_BYTES + count;
475 let table_at = stream_at + stream_len as usize;
476 let every = header.checkpoint_every as u64;
477 let expected_count = header.ticks.checked_div(every).unwrap_or(0);
478 if checkpoint_count != expected_count {
479 return Err(JournalError::Corrupt { at: 44 });
480 }
481 let mut checkpoints = Vec::with_capacity(checkpoint_count as usize);
482 let mut last = 0;
483 for k in 0..checkpoint_count as usize {
484 let at = table_at + k * CHECKPOINT_BYTES;
485 let checkpoint = Checkpoint {
486 tick: u64_at(at),
487 digest: u64_at(at + 8),
488 offset: u64_at(at + 16),
489 };
490 if checkpoint.tick != (k as u64 + 1) * every
491 || checkpoint.offset < last
492 || checkpoint.offset > stream_len
493 || (checkpoint.tick == header.ticks && checkpoint.offset != stream_len)
494 {
495 return Err(JournalError::Corrupt { at: at as u64 });
496 }
497 last = checkpoint.offset;
498 checkpoints.push(checkpoint);
499 }
500 Ok(Self {
501 header,
502 schema,
503 stream: &bytes[stream_at..table_at],
504 stream_at: stream_at as u64,
505 checkpoints,
506 })
507 }
508
509 pub fn header(&self) -> &Header {
510 &self.header
511 }
512
513 pub fn schema(&self) -> &Schema {
514 &self.schema
515 }
516
517 pub fn checkpoints(&self) -> &[Checkpoint] {
518 &self.checkpoints
519 }
520
521 pub fn check(&self, schema: &Schema, sim_version: u32) -> Result<(), JournalError> {
522 if self.header.stream_version != STREAM_VERSION {
523 return Err(JournalError::StreamVersion {
524 found: self.header.stream_version,
525 });
526 }
527 if self.header.sim_version != sim_version {
528 return Err(JournalError::SimVersion {
529 found: self.header.sim_version,
530 expected: sim_version,
531 });
532 }
533 let expected = schema.hash();
534 if self.header.schema_hash != expected {
535 return Err(JournalError::SchemaMismatch {
536 found: self.header.schema_hash,
537 expected,
538 });
539 }
540 Ok(())
541 }
542
543 pub fn reader(&self) -> Reader<'_> {
544 self.reader_at(0, 0)
545 }
546
547 pub fn seek(&self, tick: u64) -> Result<Reader<'_>, JournalError> {
548 if tick > self.header.ticks {
549 return Err(JournalError::Beyond {
550 tick,
551 ticks: self.header.ticks,
552 });
553 }
554 let every = self.header.checkpoint_every as u64;
555 let k = tick
556 .checked_div(every)
557 .map_or(0, |k| k.min(self.checkpoints.len() as u64) as usize);
558 let mut reader = match k {
559 0 => self.reader_at(0, 0),
560 _ => {
561 let checkpoint = self.checkpoints[k - 1];
562 self.reader_at(checkpoint.tick, checkpoint.offset as usize)
563 }
564 };
565 while reader.tick < tick {
566 reader.read()?;
567 }
568 Ok(reader)
569 }
570
571 pub fn verify(
572 &self,
573 schema: &Schema,
574 sim_version: u32,
575 mut step: impl FnMut(&[i64]) -> u64,
576 ) -> Result<Option<u64>, JournalError> {
577 self.verify_sparse(schema, sim_version, |input, _| step(input))
578 }
579
580 pub fn verify_sparse(
581 &self,
582 schema: &Schema,
583 sim_version: u32,
584 mut step: impl FnMut(&[i64], bool) -> u64,
585 ) -> Result<Option<u64>, JournalError> {
586 self.check(schema, sim_version)?;
587 let mut reader = self.reader();
588 let mut last = None;
589 let mut next = self.checkpoints.iter().peekable();
590 let total = self.header.ticks;
591 let mut tick = 0;
592 while let Some(input) = reader.read()? {
593 tick += 1;
594 let checkpoint = next.next_if(|checkpoint| checkpoint.tick == tick);
595 let wanted = checkpoint.is_some() || tick == total;
596 let digest = step(input, wanted);
597 if tick == total {
598 last = Some(digest);
599 }
600 if let Some(checkpoint) = checkpoint
601 && checkpoint.digest != digest
602 {
603 return Err(JournalError::Checkpoint {
604 tick: checkpoint.tick,
605 found: digest,
606 expected: checkpoint.digest,
607 });
608 }
609 }
610 Ok(last)
611 }
612
613 fn reader_at(&self, tick: u64, at: usize) -> Reader<'_> {
614 Reader {
615 fields: &self.schema.fields,
616 stream: self.stream,
617 stream_at: self.stream_at,
618 checkpoints: &self.checkpoints,
619 every: self.header.checkpoint_every as u64,
620 total: self.header.ticks,
621 tick,
622 at,
623 run: 0,
624 values: vec![0; self.schema.fields.len()],
625 }
626 }
627}
628
629#[derive(Clone, Debug)]
630pub struct Reader<'r> {
631 fields: &'r [Field],
632 stream: &'r [u8],
633 stream_at: u64,
634 checkpoints: &'r [Checkpoint],
635 every: u64,
636 total: u64,
637 tick: u64,
638 at: usize,
639 run: u64,
640 values: Vec<i64>,
641}
642
643impl Reader<'_> {
644 pub fn tick(&self) -> u64 {
645 self.tick
646 }
647
648 pub fn read(&mut self) -> Result<Option<&[i64]>, JournalError> {
649 if self.every != 0
650 && self.tick != 0
651 && self.tick.is_multiple_of(self.every)
652 && self.run == 0
653 {
654 if self.checkpoints[(self.tick / self.every - 1) as usize].offset as usize != self.at {
655 return Err(self.corrupt(self.at));
656 }
657 self.values.fill(0);
658 }
659 if self.tick == self.total {
660 if self.run != 0 || self.at != self.stream.len() {
661 return Err(self.corrupt(self.at));
662 }
663 return Ok(None);
664 }
665 if self.run > 0 {
666 self.run -= 1;
667 self.tick += 1;
668 return Ok(Some(&self.values));
669 }
670 let start = self.at;
671 let head = self.varint()?;
672 if head & 1 == 0 {
673 let limit = match self.every {
674 0 => self.total,
675 every => (self.tick / every + 1)
676 .saturating_mul(every)
677 .min(self.total),
678 };
679 if (head >> 1) >= limit - self.tick {
680 return Err(self.corrupt(start));
681 }
682 self.run = head >> 1;
683 self.tick += 1;
684 return Ok(Some(&self.values));
685 }
686 let mask = head >> 1;
687 if mask == 0 || mask >> self.fields.len() != 0 {
688 return Err(self.corrupt(start));
689 }
690 let mut bits = mask;
691 while bits != 0 {
692 let i = bits.trailing_zeros() as usize;
693 bits &= bits - 1;
694 let field = self.fields[i];
695 if field == Field::Flag {
696 self.values[i] ^= 1;
697 continue;
698 }
699 let delta_at = self.at;
700 let delta = unzigzag(self.varint()?);
701 let value = self.values[i].wrapping_add(delta);
702 if delta == 0 || value < field.min() || value > field.max() {
703 return Err(self.corrupt(delta_at));
704 }
705 self.values[i] = value;
706 }
707 self.tick += 1;
708 Ok(Some(&self.values))
709 }
710
711 fn varint(&mut self) -> Result<u64, JournalError> {
712 let start = self.at;
713 let mut value = 0u64;
714 let mut shift = 0;
715 loop {
716 let Some(&byte) = self.stream.get(self.at) else {
717 return Err(self.corrupt(start));
718 };
719 self.at += 1;
720 if shift == 63 && byte > 1 {
721 return Err(self.corrupt(start));
722 }
723 value |= ((byte & 0x7f) as u64) << shift;
724 if byte & 0x80 == 0 {
725 if byte == 0 && shift > 0 {
726 return Err(self.corrupt(start));
727 }
728 return Ok(value);
729 }
730 shift += 7;
731 }
732 }
733
734 fn corrupt(&self, at: usize) -> JournalError {
735 JournalError::Corrupt {
736 at: self.stream_at + at as u64,
737 }
738 }
739}
740
741fn flush_run(stream: &mut Vec<u8>, run: &mut u64) {
742 if *run > 0 {
743 put_varint(stream, (*run - 1) << 1);
744 *run = 0;
745 }
746}
747
748fn put_varint(out: &mut Vec<u8>, mut value: u64) {
749 while value >= 0x80 {
750 out.push(value as u8 | 0x80);
751 value >>= 7;
752 }
753 out.push(value as u8);
754}
755
756const fn zigzag(value: i64) -> u64 {
757 ((value << 1) ^ (value >> 63)) as u64
758}
759
760const fn unzigzag(value: u64) -> i64 {
761 (value >> 1) as i64 ^ -((value & 1) as i64)
762}