Skip to main content

naia_shared/world/component/
component_kinds.rs

1use std::{any::TypeId, collections::HashMap};
2
3use naia_serde::{BitReader, BitWrite, Serde, SerdeErr};
4
5use crate::{
6    ComponentUpdate, LocalEntityAndGlobalEntityConverter,
7    Replicate, ReplicateBuilder,
8};
9use crate::world::component::replicate::SplitUpdateResult;
10
11type NetId = u16;
12
13/// Wire encoding for `ComponentKind` NetIds is a fixed-width raw bit
14/// field whose width is `ceil(log2(N))`, where N is the number of kinds
15/// registered in the protocol. Both ends share the same registration
16/// order, so both compute the same width and the encoding stays in sync.
17///
18/// This is strictly optimal: every tag is exactly the bits it needs, no
19/// proceed-bit overhead, no varint loop. For a protocol with N kinds:
20///
21/// |     N      | Bits per tag |
22/// |------------|-------------:|
23/// |    0..1    |            0 |
24/// |    2       |            1 |
25/// |    3..4    |            2 |
26/// |    5..8    |            3 |
27/// |    9..16   |            4 |
28/// |   17..32   |            5 |
29/// |   33..64   |            6 |
30/// |  65..128   |            7 |
31/// | 129..256   |            8 |
32///
33/// Width is precomputed at registration time and cached on
34/// `ComponentKinds`, so ser/de pays only an inline u8 read plus N
35/// `write_bit` calls — no struct construction, no HashMap lookup for the
36/// width. Pinned by `benches/tests/component_kind_wire.rs`.
37fn bit_width_for_kind_count(count: NetId) -> u8 {
38    // count <= 1 → 0 bits (degenerate; nothing to disambiguate).
39    // count   N → ceil(log2(N)) bits.
40    if count < 2 {
41        0
42    } else {
43        (count as u32).next_power_of_two().trailing_zeros() as u8
44    }
45}
46
47/// ComponentKind - should be one unique value for each type of Component
48#[derive(Eq, Hash, Copy, Clone, PartialEq, Debug)]
49pub struct ComponentKind {
50    type_id: TypeId,
51}
52
53impl From<TypeId> for ComponentKind {
54    fn from(type_id: TypeId) -> Self {
55        Self { type_id }
56    }
57}
58impl From<ComponentKind> for TypeId {
59    fn from(val: ComponentKind) -> Self {
60        val.type_id
61    }
62}
63
64impl ComponentKind {
65    /// Returns the `ComponentKind` corresponding to the type `C`.
66    pub fn of<C: Replicate>() -> Self {
67        Self {
68            type_id: TypeId::of::<C>(),
69        }
70    }
71
72    /// Serializes this kind's compact net-ID into `writer` using the bit-width in `component_kinds`.
73    pub fn ser(&self, component_kinds: &ComponentKinds, writer: &mut dyn BitWrite) {
74        let net_id = component_kinds.kind_to_net_id(self);
75        let bits = component_kinds.kind_bit_width;
76        for i in 0..bits {
77            writer.write_bit((net_id >> i) & 1 != 0);
78        }
79    }
80
81    /// Deserializes a `ComponentKind` from `reader` using the bit-width in `component_kinds`.
82    pub fn de(component_kinds: &ComponentKinds, reader: &mut BitReader) -> Result<Self, SerdeErr> {
83        let bits = component_kinds.kind_bit_width;
84        let mut net_id: NetId = 0;
85        for i in 0..bits {
86            if bool::de(reader)? {
87                net_id |= 1 << i;
88            }
89        }
90        Ok(component_kinds.net_id_to_kind(&net_id))
91    }
92}
93
94/// A map to hold all component types
95pub struct ComponentKinds {
96    current_net_id: NetId,
97    /// Number of bits needed to encode any registered NetId — recomputed
98    /// on every `add_component` so it always reflects the current count.
99    /// Read directly by `ComponentKind::ser`/`de` on the hot path.
100    kind_bit_width: u8,
101    kind_map: HashMap<ComponentKind, (NetId, Box<dyn ReplicateBuilder>, String)>,
102    net_id_map: HashMap<NetId, ComponentKind>,
103}
104
105impl Clone for ComponentKinds {
106    fn clone(&self) -> Self {
107        let current_net_id = self.current_net_id;
108        let kind_bit_width = self.kind_bit_width;
109        let net_id_map = self.net_id_map.clone();
110
111        let mut kind_map = HashMap::new();
112        for (key, value) in self.kind_map.iter() {
113            kind_map.insert(*key, (value.0, value.1.box_clone(), value.2.clone()));
114        }
115
116        Self {
117            current_net_id,
118            kind_bit_width,
119            kind_map,
120            net_id_map,
121        }
122    }
123}
124
125impl Default for ComponentKinds {
126    fn default() -> Self {
127        Self::new()
128    }
129}
130
131impl ComponentKinds {
132    /// Creates an empty `ComponentKinds` registry.
133    pub fn new() -> Self {
134        Self {
135            current_net_id: 0,
136            kind_bit_width: 0,
137            kind_map: HashMap::new(),
138            net_id_map: HashMap::new(),
139        }
140    }
141
142    /// Registers replicated component type `C`, assigning it the next sequential net-ID.
143    pub fn add_component<C: Replicate>(&mut self) {
144        let component_kind = ComponentKind::of::<C>();
145
146        let net_id = self.current_net_id;
147        // Pre-2026-05-05 there was a `net_id < 64` cap here because the
148        // per-user `DirtyQueue` stored dirty bits in a single `u64`
149        // per entity. The queue is now flat-strided over multiple
150        // `u64`s sized to `ceil(kind_count / 64)`, so there's no
151        // longer a 64-kind ceiling. The wire-format kind tag is a
152        // `u16` NetId (cap 65,535) — that's the real ceiling and well
153        // beyond any realistic protocol size.
154        self.kind_map.insert(
155            component_kind,
156            (net_id, C::create_builder(), C::protocol_name().to_string()),
157        );
158        self.net_id_map.insert(net_id, component_kind);
159        self.current_net_id += 1;
160        self.kind_bit_width = bit_width_for_kind_count(self.current_net_id);
161    }
162
163    /// Number of component kinds currently registered. Used at
164    /// `UserDiffHandler` construction to size the per-user `DirtyQueue`'s
165    /// stride (= `ceil(kind_count / 64)` AtomicU64 words per entity).
166    pub fn kind_count(&self) -> u16 {
167        self.current_net_id
168    }
169
170    /// Reads a component kind tag then deserializes and returns the component from `reader`.
171    pub fn read(
172        &self,
173        reader: &mut BitReader,
174        converter: &dyn LocalEntityAndGlobalEntityConverter,
175    ) -> Result<Box<dyn Replicate>, SerdeErr> {
176        let component_kind: ComponentKind = ComponentKind::de(self, reader)?;
177        self
178            .kind_to_builder(&component_kind)
179            .read(reader, converter)
180    }
181
182    /// Reads a component kind tag then deserializes an initial-create update payload from `reader`.
183    pub fn read_create_update(&self, reader: &mut BitReader) -> Result<ComponentUpdate, SerdeErr> {
184        let component_kind: ComponentKind = ComponentKind::de(self, reader)?;
185        self
186            .kind_to_builder(&component_kind)
187            .read_create_update(reader)
188    }
189
190    /// Splits a full-component update into a waiting portion and a ready-to-apply portion.
191    pub fn split_update(
192        &self,
193        converter: &dyn LocalEntityAndGlobalEntityConverter,
194        component_kind: &ComponentKind,
195        update: ComponentUpdate,
196    ) -> SplitUpdateResult {
197        self
198            .kind_to_builder(component_kind)
199            .split_update(converter, update)
200    }
201
202    /// Returns the protocol name for `component_kind`. Panics if not registered.
203    pub fn kind_to_name(&self, component_kind: &ComponentKind) -> String {
204        self
205            .kind_map
206            .get(component_kind)
207            .expect(
208                "Must properly initialize Component with Protocol via `add_component()` function!",
209            )
210            .2
211            .clone()
212    }
213
214    fn net_id_to_kind(&self, net_id: &NetId) -> ComponentKind {
215        *self.net_id_map.get(net_id).expect(
216            "Must properly initialize Component with Protocol via `add_component()` function!",
217        )
218    }
219
220    fn kind_to_net_id(&self, component_kind: &ComponentKind) -> NetId {
221        self
222            .kind_map
223            .get(component_kind)
224            .expect(
225                "Must properly initialize Component with Protocol via `add_component()` function!",
226            )
227            .0
228    }
229
230    /// Public accessor for a kind's NetId (== bit position in the
231    /// `DirtyQueue` u64 mask, max 64). Returns `None` for unregistered
232    /// kinds.
233    pub fn net_id_of(&self, component_kind: &ComponentKind) -> Option<u16> {
234        self.kind_map.get(component_kind).map(|(net_id, _, _)| *net_id)
235    }
236
237    fn kind_to_builder(&self, component_kind: &ComponentKind) -> &dyn ReplicateBuilder {
238        self
239            .kind_map
240            .get(component_kind)
241            .expect(
242                "Must properly initialize Component with Protocol via `add_component()` function!",
243            )
244            .1
245            .as_ref()
246    }
247
248    /// Returns `true` if the given kind was registered as an immutable component.
249    pub fn kind_is_immutable(&self, component_kind: &ComponentKind) -> bool {
250        self.kind_map
251            .get(component_kind)
252            .map(|(_, builder, _)| builder.is_immutable())
253            .unwrap_or(false)
254    }
255
256    /// Returns a sorted list of all registered component protocol names.
257    pub fn all_names(&self) -> Vec<String> {
258        let mut output = Vec::new();
259        for (_, _, name) in self.kind_map.values() {
260            output.push(name.clone());
261        }
262        output.sort();
263        output
264    }
265}