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
30pub struct RegistryCache {
32 pub registry_packets: Arc<[EncodedPacket]>,
34 pub tags_packet: Arc<EncodedPacket>,
36}
37
38impl RegistryCache {
39 #[must_use]
41 pub fn new(compression: Option<CompressionInfo>) -> Self {
42 let registry_packets = Self::build_registry_packets(®ISTRY);
43 let tags_by_registry_packet = Self::build_tags_packet(®ISTRY);
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 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 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 CUpdateTags::new(tags_by_registry)
151 }
152}
153
154fn 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#[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}