Skip to main content

pfx_steam/
lockstep.rs

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}