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
31pub struct RegistryCache {
33 pub registry_packets: Arc<[EncodedPacket]>,
35 pub tags_packet: Arc<EncodedPacket>,
37}
38
39impl RegistryCache {
40 #[must_use]
42 pub fn new(compression: Option<CompressionInfo>) -> Self {
43 let registry_packets = Self::build_registry_packets(®ISTRY);
44 let tags_by_registry_packet = Self::build_tags_packet(®ISTRY);
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 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 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 CUpdateTags::new(tags_by_registry)
155 }
156}
157
158fn 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#[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}