1use std::collections::BTreeMap;
2
3use crate::hash::fnv1a;
4use crate::{Delivery, Error, SteamApi, User};
5
6pub const LOCKSTEP_WINDOW: u64 = 256;
7const INPUT: u8 = 1;
8const SUM: u8 = 2;
9const HEADER: usize = 9;
10const BATCH: usize = 64;
11
12#[derive(Clone, Debug, PartialEq, Eq)]
13pub struct Turn {
14 pub step: u64,
15 pub inputs: [Vec<u8>; 2],
16 pub check: bool,
17}
18
19#[derive(Clone, Copy, Debug, PartialEq, Eq)]
20pub struct Desync {
21 pub step: u64,
22 pub mine: u64,
23 pub theirs: u64,
24}
25
26#[derive(Clone, Debug)]
27pub struct Lockstep {
28 me: User,
29 peer: User,
30 channel: u32,
31 every: u64,
32 next: u64,
33 sent: u64,
34 local: BTreeMap<u64, Vec<u8>>,
35 remote: BTreeMap<u64, Vec<u8>>,
36 mine: BTreeMap<u64, u64>,
37 theirs: BTreeMap<u64, u64>,
38 desync: Option<Desync>,
39}
40
41pub fn checksum(bytes: &[u8]) -> u64 {
42 pfx_core::sim::mix64(fnv1a(bytes))
43}
44
45fn packet(tag: u8, step: u64, payload: &[u8]) -> Vec<u8> {
46 let mut bytes = Vec::with_capacity(HEADER + payload.len());
47 bytes.push(tag);
48 bytes.extend_from_slice(&step.to_le_bytes());
49 bytes.extend_from_slice(payload);
50 bytes
51}
52
53impl Lockstep {
54 pub fn new(me: User, peer: User, channel: u32, every: u64) -> Result<Self, Error> {
55 if me == peer || every == 0 {
56 return Err(Error::Rejected);
57 }
58 Ok(Self {
59 me,
60 peer,
61 channel,
62 every,
63 next: 0,
64 sent: 0,
65 local: BTreeMap::new(),
66 remote: BTreeMap::new(),
67 mine: BTreeMap::new(),
68 theirs: BTreeMap::new(),
69 desync: None,
70 })
71 }
72
73 pub fn me(&self) -> User {
74 self.me
75 }
76
77 pub fn peer(&self) -> User {
78 self.peer
79 }
80
81 pub fn side(&self) -> usize {
82 usize::from(self.me > self.peer)
83 }
84
85 pub fn step(&self) -> u64 {
86 self.next
87 }
88
89 pub fn ahead(&self) -> u64 {
90 self.sent - self.next
91 }
92
93 pub fn desync(&self) -> Option<Desync> {
94 self.desync
95 }
96
97 pub fn is_check(&self, step: u64) -> bool {
98 step % self.every == self.every - 1
99 }
100
101 pub fn send_input(&mut self, steam: &mut impl SteamApi, input: &[u8]) -> Result<u64, Error> {
102 if self.ahead() >= LOCKSTEP_WINDOW {
103 return Err(Error::Rejected);
104 }
105 let step = self.sent;
106 steam.send_message(
107 self.peer,
108 self.channel,
109 &packet(INPUT, step, input),
110 Delivery::Reliable,
111 )?;
112 self.local.insert(step, input.to_vec());
113 self.sent += 1;
114 Ok(step)
115 }
116
117 pub fn send_checksum(
118 &mut self,
119 steam: &mut impl SteamApi,
120 step: u64,
121 sum: u64,
122 ) -> Result<(), Error> {
123 if !self.is_check(step) || step >= self.next {
124 return Err(Error::Rejected);
125 }
126 steam.send_message(
127 self.peer,
128 self.channel,
129 &packet(SUM, step, &sum.to_le_bytes()),
130 Delivery::Reliable,
131 )?;
132 self.mine.insert(step, sum);
133 self.compare(step);
134 Ok(())
135 }
136
137 pub fn receive(&mut self, steam: &mut impl SteamApi) -> Result<usize, Error> {
138 let mut taken = 0;
139 loop {
140 let batch = steam.receive_messages(self.channel, BATCH)?;
141 if batch.is_empty() {
142 return Ok(taken);
143 }
144 for message in batch {
145 if message.peer == self.peer && self.accept(&message.bytes) {
146 taken += 1;
147 }
148 }
149 }
150 }
151
152 pub fn accept(&mut self, bytes: &[u8]) -> bool {
153 if bytes.len() < HEADER {
154 return false;
155 }
156 let mut word = [0u8; 8];
157 word.copy_from_slice(&bytes[1..HEADER]);
158 let step = u64::from_le_bytes(word);
159 let payload = &bytes[HEADER..];
160 if step >= self.next.saturating_add(2 * LOCKSTEP_WINDOW) {
161 return false;
162 }
163 match bytes[0] {
164 INPUT if step >= self.next && !self.remote.contains_key(&step) => {
165 self.remote.insert(step, payload.to_vec());
166 true
167 }
168 SUM if payload.len() == 8
169 && self.is_check(step)
170 && !self.theirs.contains_key(&step) =>
171 {
172 word.copy_from_slice(payload);
173 self.theirs.insert(step, u64::from_le_bytes(word));
174 self.compare(step);
175 true
176 }
177 _ => false,
178 }
179 }
180
181 pub fn turn(&mut self) -> Option<Turn> {
182 if self.desync.is_some() {
183 return None;
184 }
185 let step = self.next;
186 if !self.local.contains_key(&step) || !self.remote.contains_key(&step) {
187 return None;
188 }
189 let local = self.local.remove(&step).unwrap_or_default();
190 let remote = self.remote.remove(&step).unwrap_or_default();
191 self.next += 1;
192 let floor = self.next.saturating_sub(2 * LOCKSTEP_WINDOW);
193 self.mine = self.mine.split_off(&floor);
194 self.theirs = self.theirs.split_off(&floor);
195 let inputs = if self.side() == 0 {
196 [local, remote]
197 } else {
198 [remote, local]
199 };
200 Some(Turn {
201 step,
202 inputs,
203 check: self.is_check(step),
204 })
205 }
206
207 fn compare(&mut self, step: u64) {
208 let (Some(mine), Some(theirs)) = (self.mine.get(&step), self.theirs.get(&step)) else {
209 return;
210 };
211 if mine != theirs && self.desync.is_none() {
212 self.desync = Some(Desync {
213 step,
214 mine: *mine,
215 theirs: *theirs,
216 });
217 }
218 self.mine.remove(&step);
219 self.theirs.remove(&step);
220 }
221}