1use super::*;
2use sbor::rust::prelude::*;
3
4pub fn generate_full_schema_from_single_type<
5 T: Describe<S::CustomAggregatorTypeKind> + ?Sized,
6 S: CustomSchema,
7>() -> (LocalTypeId, VersionedSchema<S>) {
8 let mut aggregator = TypeAggregator::new();
9 let type_id = aggregator.add_child_type_and_descendents::<T>();
10 (type_id, generate_full_schema(aggregator))
11}
12
13pub fn generate_single_type_schema<
14 T: Describe<S::CustomAggregatorTypeKind> + ?Sized,
15 S: CustomSchema,
16>() -> SingleTypeSchema<S> {
17 let (type_id, schema) = generate_full_schema_from_single_type::<T, S>();
18 SingleTypeSchema::new(schema, type_id)
19}
20
21pub fn generate_full_schema<S: CustomSchema>(
26 aggregator: TypeAggregator<S::CustomAggregatorTypeKind>,
27) -> VersionedSchema<S> {
28 generate_schema_from_types(aggregator.types)
29}
30
31fn generate_schema_from_types<S: CustomSchema>(
32 types: IndexMap<TypeHash, AggregatorTypeData<S>>,
33) -> VersionedSchema<S> {
34 let type_count = types.len();
35 let type_indices = IndexSet::from_iter(types.keys().map(|k| k.clone()));
36
37 let mut type_kinds = Vec::with_capacity(type_count);
38 let mut type_metadata = Vec::with_capacity(type_count);
39 let mut type_validations = Vec::with_capacity(type_count);
40 for (_type_hash, type_data) in types {
41 type_kinds.push(linearize::<S>(type_data.kind, &type_indices));
42 type_metadata.push(type_data.metadata);
43 type_validations.push(type_data.validation);
44 }
45
46 Schema {
47 type_kinds,
48 type_metadata,
49 type_validations,
50 }
51 .into_versioned()
52}
53
54pub fn localize_well_known_type_data<S: CustomSchema>(
55 type_data: AggregatorTypeData<S>,
56) -> LocalTypeData<S> {
57 let TypeData {
58 kind,
59 metadata,
60 validation,
61 } = type_data;
62 TypeData {
63 kind: linearize::<S>(kind, &indexset!()),
64 metadata,
65 validation,
66 }
67}
68
69pub fn localize_well_known<S: CustomSchema>(type_kind: AggregatorTypeKind<S>) -> LocalTypeKind<S> {
70 linearize::<S>(type_kind, &indexset!())
71}
72
73fn linearize<S: CustomSchema>(
74 type_kind: AggregatorTypeKind<S>,
75 type_indices: &IndexSet<TypeHash>,
76) -> LocalTypeKind<S> {
77 match type_kind {
78 TypeKind::Any => TypeKind::Any,
79 TypeKind::Bool => TypeKind::Bool,
80 TypeKind::I8 => TypeKind::I8,
81 TypeKind::I16 => TypeKind::I16,
82 TypeKind::I32 => TypeKind::I32,
83 TypeKind::I64 => TypeKind::I64,
84 TypeKind::I128 => TypeKind::I128,
85 TypeKind::U8 => TypeKind::U8,
86 TypeKind::U16 => TypeKind::U16,
87 TypeKind::U32 => TypeKind::U32,
88 TypeKind::U64 => TypeKind::U64,
89 TypeKind::U128 => TypeKind::U128,
90 TypeKind::String => TypeKind::String,
91 TypeKind::Array { element_type } => TypeKind::Array {
92 element_type: resolve_local_type_id(type_indices, &element_type),
93 },
94 TypeKind::Tuple { field_types } => TypeKind::Tuple {
95 field_types: field_types
96 .into_iter()
97 .map(|t| resolve_local_type_id(type_indices, &t))
98 .collect(),
99 },
100 TypeKind::Enum { variants } => TypeKind::Enum {
101 variants: variants
102 .into_iter()
103 .map(|(variant_index, field_types)| {
104 let new_field_types = field_types
105 .into_iter()
106 .map(|t| resolve_local_type_id(type_indices, &t))
107 .collect();
108 (variant_index, new_field_types)
109 })
110 .collect(),
111 },
112 TypeKind::Map {
113 key_type,
114 value_type,
115 } => TypeKind::Map {
116 key_type: resolve_local_type_id(type_indices, &key_type),
117 value_type: resolve_local_type_id(type_indices, &value_type),
118 },
119 TypeKind::Custom(custom_type_kind) => {
120 TypeKind::Custom(S::linearize_type_kind(custom_type_kind, type_indices))
121 }
122 }
123}
124
125pub fn resolve_local_type_id(
126 type_indices: &IndexSet<TypeHash>,
127 type_id: &RustTypeId,
128) -> LocalTypeId {
129 match type_id {
130 RustTypeId::WellKnown(well_known_type_id) => LocalTypeId::WellKnown(*well_known_type_id),
131 RustTypeId::Novel(type_hash) => {
132 LocalTypeId::SchemaLocalIndex(resolve_index(type_indices, type_hash))
133 }
134 }
135}
136
137fn resolve_index(type_indices: &IndexSet<TypeHash>, type_hash: &TypeHash) -> usize {
138 type_indices.get_index_of(type_hash).unwrap_or_else(|| {
139 panic!(
140 "Fatal error in the type aggregation process - this is likely due to a type impl missing a dependent type in add_all_dependencies. The following type hash wasn't added in add_all_dependencies: {:?}",
141 type_hash
142 )
143 })
144}
145
146pub struct TypeAggregator<C: CustomTypeKind<RustTypeId>> {
147 already_read_dependencies: IndexSet<TypeHash>,
148 named_root_types: IndexMap<String, LocalTypeId>,
149 types: IndexMap<TypeHash, TypeData<C, RustTypeId>>,
150}
151
152impl<C: CustomTypeKind<RustTypeId>> TypeAggregator<C> {
153 pub fn new() -> Self {
154 Self {
155 already_read_dependencies: index_set_new(),
156 named_root_types: IndexMap::default(),
157 types: IndexMap::default(),
158 }
159 }
160
161 pub fn add_root_type<T: Describe<C> + ?Sized>(
167 &mut self,
168 name: impl Into<String>,
169 ) -> LocalTypeId {
170 let local_type_id = self.add_child_type_and_descendents::<T>();
171 self.named_root_types.insert(name.into(), local_type_id);
172 local_type_id
173 }
174
175 pub fn add_child_type_and_descendents<T: Describe<C> + ?Sized>(&mut self) -> LocalTypeId {
177 let schema_type_id = self.add_child_type(T::TYPE_ID, || T::type_data());
178 self.add_schema_descendents::<T>();
179 schema_type_id
180 }
181
182 pub fn add_child_type(
194 &mut self,
195 type_id: RustTypeId,
196 get_type_data: impl FnOnce() -> TypeData<C, RustTypeId>,
197 ) -> LocalTypeId {
198 let complex_type_hash = match type_id {
199 RustTypeId::WellKnown(well_known_type_id) => {
200 return LocalTypeId::WellKnown(well_known_type_id);
201 }
202 RustTypeId::Novel(complex_type_hash) => complex_type_hash,
203 };
204
205 if let Some(index) = self.types.get_index_of(&complex_type_hash) {
206 return LocalTypeId::SchemaLocalIndex(index);
207 }
208
209 let new_index = self.types.len();
210 self.types.insert(complex_type_hash, get_type_data());
211 LocalTypeId::SchemaLocalIndex(new_index)
212 }
213
214 pub fn add_schema_descendents<T: Describe<C> + ?Sized>(&mut self) -> bool {
224 let RustTypeId::Novel(complex_type_hash) = T::TYPE_ID else {
225 return false;
226 };
227
228 if self.already_read_dependencies.contains(&complex_type_hash) {
229 return false;
230 }
231
232 self.already_read_dependencies.insert(complex_type_hash);
233
234 T::add_all_dependencies(self);
235
236 return true;
237 }
238
239 pub fn generate_type_collection_schema<S: CustomSchema<CustomAggregatorTypeKind = C>>(
240 self,
241 ) -> TypeCollectionSchema<S> {
242 TypeCollectionSchema::new(
243 generate_schema_from_types(self.types),
244 self.named_root_types,
245 )
246 }
247}