sbor/schema/
type_aggregator.rs

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
21/// You may wish to use the newer `aggregator.generate_type_collection_schema()`
22/// which, in tandom with `add_named_root_type_and_descendents`
23/// also captures named root types to give more structure to enable schema
24/// comparisons over time.
25pub 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    /// Adds the type (and its dependencies) to the `TypeAggregator`.
162    /// Also tracks it as a named root type, which can be used e.g. in schema comparisons.
163    ///
164    /// This is only intended for use when adding root types to schemas,
165    /// and should not be called from inside `Describe` implementations.
166    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    /// Adds the dependent type (and its dependencies) to the `TypeAggregator`.
176    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    /// Adds the type's `TypeData` to the `TypeAggregator`.
183    ///
184    /// If the type is well known or already in the aggregator, this returns early with the existing index.
185    ///
186    /// Typically you should use [`add_child_type_and_descendents`], unless you're replacing/mutating
187    /// the child types somehow. In which case, you'll likely wish to call [`add_child_type`] and
188    /// [`add_schema_descendents`] separately.
189    ///
190    /// [`add_child_type`]: #method.add_child_type
191    /// [`add_schema_descendents`]: #method.add_schema_descendents
192    /// [`add_child_type_and_descendents`]: #method.add_child_type_and_descendents
193    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    /// Adds the type's descendent types to the `TypeAggregator`, if they've not already been added.
215    ///
216    /// Typically you should use [`add_child_type_and_descendents`], unless you're replacing/mutating
217    /// the child types somehow. In which case, you'll likely wish to call [`add_child_type`] and
218    /// [`add_schema_descendents`] separately.
219    ///
220    /// [`add_child_type`]: #method.add_child_type
221    /// [`add_schema_descendents`]: #method.add_schema_descendents
222    /// [`add_child_type_and_descendents`]: #method.add_child_type_and_descendents
223    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}