1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

use naia_derive::MessageFragment;
use naia_serde::{BitReader, BitWrite, ConstBitLength, Serde, SerdeErr, UnsignedInteger};

const FRAGMENT_ID_BITS: u8 = 10;
const FRAGMENT_ID_LIMIT: u16 = 2 ^ (FRAGMENT_ID_BITS as u16);
const FRAGMENT_INDEX_BITS: u8 = 20;
const FRAGMENT_INDEX_LIMIT: u32 = 2 ^ (FRAGMENT_INDEX_BITS as u32);

// FragmentId
#[derive(Copy, Clone, PartialEq, Eq, Hash)]
pub struct FragmentId {
    inner: u16,
}

impl FragmentId {
    pub(crate) fn zero() -> Self {
        Self { inner: 0 }
    }

    pub(crate) fn increment(&mut self) {
        self.inner += 1;
        if self.inner >= FRAGMENT_ID_LIMIT {
            self.inner = 0;
        }
    }
}

impl Serde for FragmentId {
    fn ser(&self, writer: &mut dyn BitWrite) {
        let integer = UnsignedInteger::<FRAGMENT_ID_BITS>::new(self.inner as u64);
        integer.ser(writer);
    }

    fn de(reader: &mut BitReader) -> Result<Self, SerdeErr> {
        let integer = UnsignedInteger::<FRAGMENT_ID_BITS>::de(reader)?;
        let inner = integer.get() as u16;
        Ok(Self { inner })
    }

    fn bit_length(&self) -> u32 {
        <Self as ConstBitLength>::const_bit_length()
    }
}

impl ConstBitLength for FragmentId {
    fn const_bit_length() -> u32 {
        FRAGMENT_ID_BITS as u32
    }
}

// FragmentIndex
#[derive(Copy, Clone, PartialEq)]
pub struct FragmentIndex {
    inner: u32,
}

impl FragmentIndex {
    pub(crate) fn zero() -> Self {
        Self { inner: 0 }
    }

    pub(crate) fn increment(&mut self) {
        self.inner += 1;
        if self.inner >= FRAGMENT_INDEX_LIMIT {
            panic!("Attempting to fragment large message, but hit fragment limit of {FRAGMENT_INDEX_LIMIT}. This means you're trying to transmit about 500 megabytes, which is a bad idea.")
        }
    }

    pub fn as_usize(&self) -> usize {
        self.inner as usize
    }
}

impl Serde for FragmentIndex {
    fn ser(&self, writer: &mut dyn BitWrite) {
        let integer = UnsignedInteger::<FRAGMENT_INDEX_BITS>::new(self.inner as u64);
        integer.ser(writer);
    }

    fn de(reader: &mut BitReader) -> Result<Self, SerdeErr> {
        let integer = UnsignedInteger::<FRAGMENT_INDEX_BITS>::de(reader)?;
        let inner = integer.get() as u32;
        Ok(Self { inner })
    }

    fn bit_length(&self) -> u32 {
        <Self as ConstBitLength>::const_bit_length()
    }
}

impl ConstBitLength for FragmentIndex {
    fn const_bit_length() -> u32 {
        FRAGMENT_INDEX_BITS as u32
    }
}

#[derive(MessageFragment)]
pub struct FragmentedMessage {
    id: FragmentId,
    index: FragmentIndex,
    total: FragmentIndex,
    bytes: Box<[u8]>,
}

impl FragmentedMessage {
    pub fn new(id: FragmentId, index: FragmentIndex, bytes: Box<[u8]>) -> Self {
        Self {
            id,
            index,
            bytes,
            total: FragmentIndex::zero(),
        }
    }

    pub(crate) fn set_total(&mut self, total: FragmentIndex) {
        self.total = total;
    }

    pub(crate) fn id(&self) -> FragmentId {
        self.id
    }

    pub(crate) fn index(&self) -> FragmentIndex {
        self.index
    }

    pub(crate) fn total(&self) -> FragmentIndex {
        self.total
    }

    pub(crate) fn to_payload(self) -> Box<[u8]> {
        self.bytes
    }
}