libtw2_demo/
writer.rs

1use arrayvec::ArrayVec;
2use binrw::BinWrite;
3use libtw2_common::digest::Sha256;
4use libtw2_common::num::Cast;
5use libtw2_huffman::instances::TEEWORLDS as HUFFMAN;
6use libtw2_packer::with_packer;
7use std::io;
8use thiserror::Error;
9
10use crate::format::CappedString;
11use crate::format::ChunkHeader;
12use crate::format::DataKind;
13use crate::format::DemoKind;
14use crate::format::Header;
15use crate::format::MapSha256;
16use crate::format::RawChunk;
17use crate::format::TickMarker;
18use crate::format::TimelineMarkers;
19use crate::format::Version;
20use crate::format::MAX_SNAPSHOT_SIZE;
21
22#[derive(Error, Debug)]
23#[error(transparent)]
24pub struct WriteError(#[from] binrw::Error);
25
26impl WriteError {
27    pub fn io_error(self) -> Result<io::Error, WriteError> {
28        match self.0 {
29            binrw::Error::Io(io) => Ok(io),
30            err => Err(WriteError(err)),
31        }
32    }
33}
34
35pub struct Writer {
36    file: Box<dyn SeekableWrite>,
37    header: Header,
38    prev_tick: Option<i32>,
39    huffman: ArrayVec<[u8; MAX_SNAPSHOT_SIZE]>,
40    buffer2: ArrayVec<[u8; MAX_SNAPSHOT_SIZE]>,
41}
42
43const WRITER_VERSION: Version = Version::V5;
44const WRITER_VERSION_DDNET: Version = Version::V6Ddnet;
45
46pub(crate) trait SeekableWrite: io::Write + io::Seek {}
47impl<T: io::Write + io::Seek> SeekableWrite for T {}
48
49impl Writer {
50    pub fn new<W: io::Write + io::Seek + 'static>(
51        file: W,
52        net_version: &[u8],
53        map_name: &[u8],
54        map_sha256: Option<Sha256>,
55        map_crc: u32,
56        kind: DemoKind,
57        length: i32,
58        timestamp: &[u8],
59        map: &[u8],
60    ) -> Result<Writer, WriteError> {
61        let mut writer = Writer {
62            file: Box::new(file),
63            header: Header {
64                net_version: CappedString::from_raw(net_version),
65                map_name: CappedString::from_raw(map_name),
66                map_size: map.len().assert_i32(),
67                map_crc: map_crc,
68                kind: kind,
69                length,
70                timestamp: CappedString::from_raw(timestamp),
71            },
72            prev_tick: None,
73            huffman: ArrayVec::new(),
74            buffer2: ArrayVec::new(),
75        };
76        writer.write_header(map_sha256.is_some())?;
77        TimelineMarkers {
78            amount: 0,
79            markers: [0; 64],
80        }
81        .write(&mut writer.file)?;
82        if let Some(sha256) = map_sha256 {
83            MapSha256::new(sha256).write_le(&mut writer.file)?;
84        }
85        map.write(&mut writer.file)?;
86        Ok(writer)
87    }
88    fn write_header(&mut self, ddnet: bool) -> Result<(), WriteError> {
89        let version = if ddnet {
90            WRITER_VERSION_DDNET
91        } else {
92            WRITER_VERSION
93        };
94        version.write(&mut self.file)?;
95        self.header.write(&mut self.file)?;
96        Ok(())
97    }
98    pub fn write_chunk(&mut self, chunk: RawChunk) -> Result<(), WriteError> {
99        match chunk {
100            RawChunk::Tick { tick, keyframe } => self.write_tick(keyframe, tick),
101            RawChunk::Snapshot(snapshot) => self.write_snapshot(snapshot),
102            RawChunk::SnapshotDelta(delta) => self.write_snapshot_delta(delta),
103            RawChunk::Message(msg) => self.write_message(msg),
104            RawChunk::Unknown => panic!(),
105        }
106    }
107    pub fn write_tick(&mut self, keyframe: bool, tick: i32) -> Result<(), WriteError> {
108        let tm = TickMarker::new(tick, self.prev_tick, keyframe, WRITER_VERSION);
109        ChunkHeader::Tick {
110            marker: tm,
111            keyframe: keyframe,
112        }
113        .write(&mut self.file, WRITER_VERSION)?;
114        self.prev_tick = Some(tick);
115        Ok(())
116    }
117    fn write_chunk_impl(&mut self, kind: DataKind, data: Option<&[u8]>) -> Result<(), WriteError> {
118        let data = data.unwrap_or(&self.buffer2);
119        self.huffman.clear();
120        HUFFMAN
121            .compress(data, &mut self.huffman)
122            .expect("too long compression");
123        ChunkHeader::Data {
124            kind,
125            size: self.huffman.len().assert_u16(),
126        }
127        .write(&mut self.file, WRITER_VERSION)?;
128        self.file
129            .write_all(&self.huffman)
130            .map_err(binrw::Error::Io)?;
131        Ok(())
132    }
133    pub fn write_snapshot(&mut self, snapshot: &[u8]) -> Result<(), WriteError> {
134        self.write_chunk_impl(DataKind::Snapshot, Some(snapshot))
135    }
136    pub fn write_snapshot_delta(&mut self, delta: &[u8]) -> Result<(), WriteError> {
137        self.write_chunk_impl(DataKind::SnapshotDelta, Some(delta))
138    }
139    pub fn write_message(&mut self, msg: &[u8]) -> Result<(), WriteError> {
140        self.buffer2.clear();
141        with_packer(
142            &mut self.buffer2,
143            |mut p| -> Result<(), buffer::CapacityError> {
144                for b in msg.chunks(4) {
145                    // Get or return 0.
146                    fn g(bytes: &[u8], idx: usize) -> u8 {
147                        bytes.get(idx).cloned().unwrap_or(0)
148                    }
149                    p.write_int(i32::from_le_bytes([g(b, 0), g(b, 1), g(b, 2), g(b, 3)]))?;
150                }
151                Ok(())
152            },
153        )
154        .expect("overlong message");
155        self.write_chunk_impl(DataKind::Message, None)
156    }
157    // TODO: Add a `finalize` function that writes the demo length into the
158    // original header.
159}
libtw2_demo/writer.rs

libtw2_demo/
writer.rs
