Skip to main content

steel_core/server/
registry_cache.rs

1use simdnbt::ToNbtTag;
2use std::sync::Arc;
3
4use steel_protocol::packet_traits::{ClientPacket, EncodedPacket};
5use steel_protocol::packets::common::TagCollection;
6use steel_protocol::{
7    packets::{
8        common::CUpdateTags,
9        config::{CRegistryData, RegistryEntry},
10    },
11    utils::ConnectionProtocol,
12};
13
14use steel_registry::{
15    BANNER_PATTERN_REGISTRY, BIOMES_REGISTRY, BLOCK_STATE_PROVIDER_REGISTRY,
16    BLOCK_TRANSFORMER_REGISTRY, BLOCKS_REGISTRY, CAT_SOUND_VARIANT_REGISTRY, CAT_VARIANT_REGISTRY,
17    CHAT_TYPE_REGISTRY, CHICKEN_SOUND_VARIANT_REGISTRY, CHICKEN_VARIANT_REGISTRY,
18    COW_SOUND_VARIANT_REGISTRY, COW_VARIANT_REGISTRY, DAMAGE_TYPE_REGISTRY,
19    DECORATED_POT_PATTERN_REGISTRY, DIALOG_REGISTRY, DIMENSION_TYPE_REGISTRY, ENCHANTMENT_REGISTRY,
20    ENTITY_TYPE_REGISTRY, FLUID_REGISTRY, FROG_VARIANT_REGISTRY, INSTRUMENT_REGISTRY,
21    ITEMS_REGISTRY, JUKEBOX_SONG_REGISTRY, PAINTING_VARIANT_REGISTRY, PIG_SOUND_VARIANT_REGISTRY,
22    PIG_VARIANT_REGISTRY, REGISTRY, Registry, RegistryEntry as _, TIMELINE_REGISTRY,
23    TRIM_MATERIAL_REGISTRY, TRIM_PATTERN_REGISTRY, TaggedRegistryExt, WOLF_SOUND_VARIANT_REGISTRY,
24    WOLF_VARIANT_REGISTRY, WORLD_CLOCK_REGISTRY, ZOMBIE_NAUTILUS_VARIANT_REGISTRY,
25};
26use steel_utils::Identifier;
27use steel_utils::codec::VarInt;
28
29use steel_protocol::packet_traits::CompressionInfo;
30
31/// Caches compressed registry packets to avoid re-compressing them for every player.
32pub struct RegistryCache {
33    /// The cached registry data packets.
34    pub registry_packets: Arc<[EncodedPacket]>,
35    /// The cached tags packet.
36    pub tags_packet: Arc<EncodedPacket>,
37}
38
39impl RegistryCache {
40    /// Creates a new `RegistryCache` from the given registry.
41    #[must_use]
42    pub fn new(compression: Option<CompressionInfo>) -> Self {
43        let registry_packets = Self::build_registry_packets(&REGISTRY);
44        let tags_by_registry_packet = Self::build_tags_packet(&REGISTRY);
45
46        let (registry_packets, tags_packet) =
47            build_compressed_packets(registry_packets, tags_by_registry_packet, compression);
48
49        Self {
50            registry_packets,
51            tags_packet: Arc::new(tags_packet),
52        }
53    }
54
55    fn build_registry_packets(registry: &Registry) -> Vec<CRegistryData> {
56        let mut packets = Vec::with_capacity(9);
57
58        macro_rules! add_registry {
59            ($reg_key:expr, $field:ident) => {
60                packets.push(CRegistryData::new(
61                    $reg_key,
62                    registry
63                        .$field
64                        .iter()
65                        .map(|(_, entry)| {
66                            RegistryEntry::new(entry.key.clone(), Some(entry.to_nbt_tag()))
67                        })
68                        .collect(),
69                ));
70            };
71        }
72
73        //TODO: For non vanilla entries we need to encode the data into nbt
74
75        add_registry!(BIOMES_REGISTRY, biomes);
76        add_registry!(CHAT_TYPE_REGISTRY, chat_types);
77        add_registry!(TRIM_PATTERN_REGISTRY, trim_patterns);
78        add_registry!(TRIM_MATERIAL_REGISTRY, trim_materials);
79        add_registry!(DECORATED_POT_PATTERN_REGISTRY, decorated_pot_patterns);
80        add_registry!(WOLF_VARIANT_REGISTRY, wolf_variants);
81        add_registry!(WOLF_SOUND_VARIANT_REGISTRY, wolf_sound_variants);
82        add_registry!(PIG_VARIANT_REGISTRY, pig_variants);
83        add_registry!(PIG_SOUND_VARIANT_REGISTRY, pig_sound_variants);
84        add_registry!(FROG_VARIANT_REGISTRY, frog_variants);
85        add_registry!(CAT_VARIANT_REGISTRY, cat_variants);
86        add_registry!(CAT_SOUND_VARIANT_REGISTRY, cat_sound_variants);
87        add_registry!(COW_VARIANT_REGISTRY, cow_variants);
88        add_registry!(COW_SOUND_VARIANT_REGISTRY, cow_sound_variants);
89        add_registry!(CHICKEN_VARIANT_REGISTRY, chicken_variants);
90        add_registry!(CHICKEN_SOUND_VARIANT_REGISTRY, chicken_sound_variants);
91        add_registry!(PAINTING_VARIANT_REGISTRY, painting_variants);
92        add_registry!(DIMENSION_TYPE_REGISTRY, dimension_types);
93        add_registry!(DAMAGE_TYPE_REGISTRY, damage_types);
94        add_registry!(BANNER_PATTERN_REGISTRY, banner_patterns);
95        add_registry!(ZOMBIE_NAUTILUS_VARIANT_REGISTRY, zombie_nautilus_variants);
96
97        add_registry!(ENCHANTMENT_REGISTRY, enchantments);
98
99        add_registry!(JUKEBOX_SONG_REGISTRY, jukebox_songs);
100        add_registry!(INSTRUMENT_REGISTRY, instruments);
101        add_registry!(TIMELINE_REGISTRY, timelines);
102        add_registry!(DIALOG_REGISTRY, dialogs);
103
104        add_registry!(WORLD_CLOCK_REGISTRY, world_clocks);
105        add_registry!(BLOCK_STATE_PROVIDER_REGISTRY, block_state_providers);
106        add_registry!(BLOCK_TRANSFORMER_REGISTRY, block_transformers);
107
108        packets
109    }
110
111    fn build_tags_packet(registry: &Registry) -> CUpdateTags {
112        let mut tags_by_registry: TagCollection = Vec::with_capacity(10);
113        macro_rules! add_tags {
114            ($reg_key:expr, $field:ident) => {
115                let mut tags: Vec<(Identifier, Vec<VarInt>)> =
116                    Vec::with_capacity(registry.$field.tag_keys().count());
117                for tag_key in registry.$field.tag_keys() {
118                    let mut ids = Vec::with_capacity(registry.$field.iter_tag(tag_key).count());
119                    for entry in registry.$field.iter_tag(tag_key) {
120                        ids.push(VarInt::from(entry.id()));
121                    }
122                    tags.push((tag_key.clone(), ids));
123                }
124                tags_by_registry.push(($reg_key, tags));
125            };
126        }
127
128        add_tags!(BLOCKS_REGISTRY, blocks);
129        add_tags!(ITEMS_REGISTRY, items);
130        add_tags!(TIMELINE_REGISTRY, timelines);
131        add_tags!(DIALOG_REGISTRY, dialogs);
132        add_tags!(DAMAGE_TYPE_REGISTRY, damage_types);
133        add_tags!(BANNER_PATTERN_REGISTRY, banner_patterns);
134        add_tags!(ENTITY_TYPE_REGISTRY, entity_types);
135        add_tags!(INSTRUMENT_REGISTRY, instruments);
136        add_tags!(PAINTING_VARIANT_REGISTRY, painting_variants);
137        add_tags!(ENCHANTMENT_REGISTRY, enchantments);
138
139        // Build fluid tags
140        let mut fluid_tags: Vec<(Identifier, Vec<VarInt>)> =
141            Vec::with_capacity(registry.fluids.tag_keys().count());
142        for tag_key in registry.fluids.tag_keys() {
143            let mut fluid_ids = Vec::with_capacity(registry.fluids.iter_tag(tag_key).count());
144
145            for fluid in registry.fluids.iter_tag(tag_key) {
146                fluid_ids.push(VarInt::from(fluid.id() as i32));
147            }
148
149            fluid_tags.push((tag_key.clone(), fluid_ids));
150        }
151        tags_by_registry.push((FLUID_REGISTRY, fluid_tags));
152
153        // Build and return a CUpdateTagsPacket based on the registry data
154        CUpdateTags::new(tags_by_registry)
155    }
156}
157
158/// Compresses a packet.
159fn compress_packet<P: ClientPacket>(
160    packet: P,
161    compression: Option<CompressionInfo>,
162) -> Option<EncodedPacket> {
163    let id = packet.get_id(ConnectionProtocol::Config);
164
165    EncodedPacket::from_bare(packet, compression, ConnectionProtocol::Config)
166        .map_err(|_| {
167            log::error!("Failed to encode packet: {id:?}");
168        })
169        .ok()
170}
171
172/// # Panics
173/// This function will panic if the compression fails.
174#[must_use]
175pub fn build_compressed_packets(
176    registry_packets: Vec<CRegistryData>,
177    tags_packet: CUpdateTags,
178    compression: Option<CompressionInfo>,
179) -> (Arc<[EncodedPacket]>, EncodedPacket) {
180    let mut compressed_packets = Vec::with_capacity(registry_packets.len());
181
182    for packet in registry_packets {
183        compressed_packets
184            .push(compress_packet(packet, compression).expect("Failed to compress packet"));
185    }
186
187    let compressed_tags_packet =
188        compress_packet(tags_packet, compression).expect("Failed to compress tags packet");
189
190    (compressed_packets.into(), compressed_tags_packet)
191}