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