sov_universal_wallet/schema/
mod.rs

1use std::any::TypeId;
2use std::collections::HashMap;
3use std::fmt::Debug;
4
5use once_cell::sync::OnceCell;
6
7use sha2::{Digest, Sha256};
8pub mod container;
9mod primitive;
10pub mod safe_string;
11pub mod transaction_templates;
12use borsh::{BorshDeserialize, BorshSerialize};
13pub use container::Container;
14use nmt_rs::simple_merkle::db::MemDb;
15use nmt_rs::simple_merkle::tree::MerkleTree;
16use nmt_rs::TmSha2Hasher;
17pub use primitive::Primitive;
18#[cfg(feature = "serde")]
19use serde::{Deserialize, Serialize};
20mod schema_impls;
21
22use thiserror::Error;
23use transaction_templates::TransactionTemplateSet;
24
25use crate::display::{Context as DisplayContext, DisplayVisitor, FormatError};
26#[cfg(feature = "serde")]
27use crate::json_to_borsh::{Context as EncodeContext, EncodeError, EncodeVisitor};
28use crate::ty::byte_display::ByteParseError;
29use crate::ty::{ContainerSerdeMetadata, LinkingScheme, Ty};
30#[cfg(feature = "eip712")]
31use crate::visitors::eip712::{Context as Eip712Context, Eip712Error, Eip712Visitor};
32#[cfg(feature = "eip712")]
33use alloy_dyn_abi::{Eip712Types, Error as AlloyEip712Error, PropertyDef, TypedData};
34
35#[derive(Debug, Error)]
36pub enum SchemaError {
37    #[error(transparent)]
38    FormatError(#[from] FormatError),
39    #[error(transparent)]
40    BorshError(#[from] borsh::io::Error),
41    #[cfg(feature = "serde")]
42    #[error(transparent)]
43    EncodeError(#[from] EncodeError),
44    #[cfg(feature = "serde")]
45    #[error(transparent)]
46    JsonError(#[from] serde_json::Error),
47    #[cfg(feature = "eip712")]
48    #[error(transparent)]
49    Eip712Error(#[from] Eip712Error),
50    #[cfg(feature = "eip712")]
51    #[error(transparent)]
52    AlloyEip712Error(#[from] AlloyEip712Error),
53    #[error(transparent)]
54    Bech32Error(#[from] ByteParseError),
55    #[error("Rollup type {0:?} was missing from schema")]
56    MissingRollupRoot(RollupRoots),
57    #[error("Template {0} not found in schema")]
58    UnknownTemplate(String),
59    #[error("Index {0} not found in schema")]
60    InvalidIndex(usize),
61    #[error("Metadata hash must be provided but was not initialized. The schema was not properly finalized, or the serialized schema was invalid.")]
62    MetadataHashNotInitialized,
63}
64
65#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
66#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
67pub struct IndexLinking;
68
69impl LinkingScheme for IndexLinking {
70    type TypeLink = Link;
71}
72
73// TODO: Some type safety for fully-constructed schemas.
74// It should be possible to use the type system to ensure at compile-time that
75// a) constructed Schemas do not have any Link::Placeholder; and
76// b) it is not possible to call construction methods (the ones that edit the links) on a finished
77// Schema.
78// This could be done with, for example, an intermediate SchemaUnderConstruction type using a
79// ConstrutionIndexLinking, which implements into::<Schema>().
80//
81// Right now this is mostly achieved using member visibility (nobody outside can call the private
82// construction methods) and sanity checking (on a derived schema, if under_construction is empty,
83// there won't be any placeholders); but a separate type would provide a stronger guarantee.
84#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
85#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
86pub enum Link {
87    ByIndex(usize),
88    Immediate(Primitive),
89    Placeholder,
90    /// Placeholder indexed by its place in the parent datastructure
91    IndexedPlaceholder(usize),
92}
93
94#[derive(Debug, Clone, PartialEq, Eq)]
95enum MaybePartialLink {
96    Partial(Link),
97    Complete(Link),
98}
99
100impl MaybePartialLink {
101    fn into_inner(self) -> Link {
102        match self {
103            MaybePartialLink::Partial(link) => link,
104            MaybePartialLink::Complete(link) => link,
105        }
106    }
107}
108
109/// This newtype is mainly necessary to allow the schema to derive Debug ergonomically
110/// Stores both the tree and its root since MerkleTree::root() requires &mut self
111#[derive(Default)]
112#[allow(clippy::type_complexity)] // This is only used internally
113struct ConstructedMerkleTree(OnceCell<(MerkleTree<MemDb<[u8; 32]>, TmSha2Hasher>, [u8; 32])>);
114
115impl Debug for ConstructedMerkleTree {
116    fn fmt(&self, _f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
117        Ok(())
118    }
119}
120
121/// This extra metadata is used in contexts where serde features are enabled; thus, we do not
122/// serialize it in serde formats, as it should be recomputed before using the data committed
123/// using it. (Otherwise a frontend could supply malicious metadata and a mismatching hash to
124/// make the hash pass chain ID checks.)
125/// When serde is disabled (i.e. usecases using borsh serialization), this metadata is unused, so
126/// the committment is the only relevant information and a mismatch is not possible.
127/// TL;DR:
128/// - In borsh: serializes/deserializes the actual hash value, while corresponding metada is empty
129/// - In serde: skips serialization and recalculates on first use, ensuring hash matches the
130///   deserialized metadata
131#[derive(Default)]
132struct MetadataHash(OnceCell<[u8; 32]>);
133
134impl Debug for MetadataHash {
135    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
136        self.0.get().fmt(f)
137    }
138}
139
140impl BorshSerialize for MetadataHash {
141    fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
142        // The hash must be calculated before serialization (via finalize())
143        // It's an error to serialize a schema that hasn't been finalized
144        let hash = self.0.get().copied()
145            .ok_or_else(|| std::io::Error::new(
146                std::io::ErrorKind::InvalidData,
147                "Cannot serialize Schema: metadata_hash not initialized. Call finalize() before serializing"
148            ))?;
149        BorshSerialize::serialize(&hash, writer)
150    }
151}
152
153impl BorshDeserialize for MetadataHash {
154    fn deserialize_reader<R: std::io::Read>(reader: &mut R) -> std::io::Result<Self> {
155        let hash: [u8; 32] = BorshDeserialize::deserialize_reader(reader)?;
156        let metadata_hash = MetadataHash::default();
157        // Set the hash - this should always succeed for a new MetadataHash
158        metadata_hash.0.set(hash).map_err(|_| {
159            std::io::Error::new(
160                std::io::ErrorKind::InvalidData,
161                "Failed to set metadata_hash in OnceCell during deserialization",
162            )
163        })?;
164        Ok(metadata_hash)
165    }
166}
167
168#[derive(Debug, Clone, PartialEq, Eq, Hash)]
169pub struct ItemId(pub TypeId);
170
171impl ItemId {
172    pub fn of<T: 'static + UniversalWallet>() -> Self {
173        T::id_override().unwrap_or(ItemId(TypeId::of::<T>()))
174    }
175}
176
177/// Not enforced in the types, but the expected convention that should be followed when generating
178/// the schema.
179#[derive(Debug, Copy, Clone)]
180pub enum RollupRoots {
181    Transaction = 0,
182    UnsignedTransaction = 1,
183    RuntimeCall = 2,
184    Address = 3,
185}
186
187/// The standard metadata format for every chain. Includes a numeric chain_id and a human-readable
188/// chain name.
189#[derive(Debug, Default, Clone, BorshSerialize, BorshDeserialize)]
190#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
191pub struct ChainData {
192    pub chain_id: u64,
193    pub chain_name: String,
194}
195
196/// A schema, representing set of types (i.e. rust code) as a data structure.
197/// The schema allows any included type's borsh serialization to be displayed as a human readable string,
198/// and the type's JSON serialisation to be re-serialised to borsh without depending on the
199/// original Rust type.
200/// It is also serialisable and therefore, once generated for a rollup, can be imported and used with
201/// non-Rust languages, enabling toolkits in any language to implement the same functionality as above.
202///
203/// A schema can be instantiated for any type that implements either `UniversalWallet` or
204/// `OverrideSchema`. In turn, `UniversalWallet` is intended to be automatically derived using the
205/// `UniversalWallet` macro.
206#[derive(Default, Debug, BorshSerialize, BorshDeserialize)]
207#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
208pub struct Schema {
209    /// The types described by this schema. This is an array of type descriptions, where complex
210    /// types refer to their sub-types by index within the array.
211    /// Any of the types here can be used for schema operations (such as borsh-to-human or
212    /// json-to-borsh reserialisations).
213    types: Vec<Ty<IndexLinking>>,
214
215    /// A mapping from the complex root types that parametrized the schema generation invocation (in
216    /// order, skipping primitives) to the actual indices they ended up at in the type array above.
217    root_type_indices: Vec<usize>,
218
219    /// Global metadata for the chain.
220    chain_data: ChainData,
221
222    /// Extra metadata hash. "Extra" metadata is defined as metadata irrelevant when no additional
223    /// features are enabled.
224    #[cfg_attr(feature = "serde", serde(skip))]
225    extra_metadata_hash: MetadataHash,
226
227    /// The chain hash: the top-level hash committing to the entire schema, including all types and
228    /// all metadata.
229    /// This should be recalculated independently whenever the schema is used, thus is not included
230    /// in serializations. This field caches the results of the calculation for subsequent uses
231    /// after the schema has been deserialized and constructed.
232    #[cfg_attr(feature = "serde", serde(skip))]
233    #[borsh(skip)]
234    chain_hash: OnceCell<[u8; 32]>,
235
236    /// A list of templatable objects that can be constructed from standard input, per root type (in
237    /// order corresponding to root_type_indices). Mapped by template name.
238    /// Should be skipped in binary serialisation for hardware wallet apps.
239    #[borsh(skip)]
240    templates: Vec<TransactionTemplateSet>,
241
242    /// A set of metadata items for each field in the `types` vec, used for serde-compatible
243    /// deserialisation (i.e. `json_to_borsh()`).
244    /// It is separated from the main vec of `Ty` structs to allow non-serde implementations of the
245    /// schema, meaning ones only concerned with `borsh`-serialized interpretation (i.e.
246    /// `display()` functionality), such as hardware wallets, to avoid deserializing this. Not only
247    /// does this save resources but it also allows the format to be modified to implement
248    /// additional serde compatibility features without causing breaking changes for non-serde
249    /// implementations.
250    #[borsh(skip)]
251    serde_metadata: Vec<ContainerSerdeMetadata>,
252
253    /// Cached (lazily-constructed) merkelization of the entire schema.
254    #[cfg_attr(feature = "serde", serde(skip))]
255    #[borsh(skip)]
256    merkle_tree: ConstructedMerkleTree,
257
258    /// A map from the type ID of an item to its index in the types array. Note that primitives and "virtual" structs/tuples
259    /// (i.e. the contents of an enum variant) are not included in this map.
260    /// Only used during schema construction.
261    #[cfg_attr(feature = "serde", serde(skip))]
262    #[borsh(skip)]
263    known_types: HashMap<ItemId, usize>,
264
265    /// Keeps track of all the types which are partially constructed. By the end of schema generation, this
266    /// must be empty.
267    #[cfg_attr(feature = "serde", serde(skip))]
268    #[borsh(skip)]
269    under_construction: HashMap<ItemId, usize>,
270}
271
272impl Schema {
273    /// Instantiate a schema for a single type.
274    /// This root type will be at index 0
275    pub fn of_single_type<T: UniversalWallet>() -> Result<Self, SchemaError> {
276        // TODO: this could easily be implemented with a macro for N types for any N >= 1, if ever needed
277        let mut schema = Self::default();
278        T::make_root_of(&mut schema);
279        schema.finalize()?;
280        Ok(schema)
281    }
282
283    /// Instantiate a schema for a standard set of rollup types: its complete transaction, its
284    /// unsigned transaction, and its call message type.
285    /// The types will be accessible using the indices stored in root_type_indices (in the above
286    /// order); they can also be queried using the `RollupRoots` enum through the `_rollup`-tagged
287    /// functions on the schema
288    pub fn of_rollup_types_with_chain_data<
289        Transaction: UniversalWallet,
290        UnsignedTransaction: UniversalWallet,
291        RuntimeCall: UniversalWallet,
292        Address: UniversalWallet,
293    >(
294        chain_data: ChainData,
295    ) -> Result<Self, SchemaError> {
296        let mut schema = Schema {
297            chain_data,
298            ..Self::default()
299        };
300        Transaction::make_root_of(&mut schema);
301        UnsignedTransaction::make_root_of(&mut schema);
302        RuntimeCall::make_root_of(&mut schema);
303        Address::make_root_of(&mut schema);
304
305        schema.finalize()?;
306        Ok(schema)
307    }
308
309    /// Get the chain data for the schema.
310    pub fn chain_data(&self) -> &ChainData {
311        &self.chain_data
312    }
313
314    #[cfg(not(feature = "serde"))]
315    pub fn metadata_hash(&self) -> Result<[u8; 32], SchemaError> {
316        // In borsh-only context, the hash must have been deserialized
317        // If it's not present, that's a critical error
318        self.extra_metadata_hash
319            .0
320            .get()
321            .copied()
322            .ok_or(SchemaError::MetadataHashNotInitialized)
323    }
324
325    #[cfg(feature = "serde")]
326    pub fn metadata_hash(&self) -> Result<[u8; 32], SchemaError> {
327        // In serde context, calculate on first use if not present
328        self.extra_metadata_hash
329            .0
330            .get_or_try_init(|| self.calculate_metadata_hash())
331            .copied()
332    }
333
334    #[cfg(feature = "serde")]
335    fn calculate_metadata_hash(&self) -> Result<[u8; 32], SchemaError> {
336        let mut hasher = Sha256::new();
337        hasher.update(&borsh::to_vec(&self.templates)?);
338        hasher.update(&borsh::to_vec(&self.serde_metadata)?);
339        Ok(hasher.finalize().into())
340    }
341
342    #[cfg(feature = "serde")]
343    pub fn from_json(input: &str) -> Result<Self, serde_json::Error> {
344        serde_json::from_str(input)
345    }
346
347    pub fn rollup_expected_index(&self, rollup_type: RollupRoots) -> Result<usize, SchemaError> {
348        self.root_type_indices
349            .get(rollup_type as usize)
350            .copied()
351            .ok_or(SchemaError::MissingRollupRoot(rollup_type))
352    }
353
354    /// Use the schema to display the given type using the provided borsh encoded input
355    pub fn display(&self, type_index: usize, input: &[u8]) -> Result<String, SchemaError> {
356        let mut output = String::new();
357        let input = &mut &input[..];
358        let mut visitor = DisplayVisitor::new(input, &mut output);
359        self.types
360            .get(type_index)
361            .ok_or(SchemaError::InvalidIndex(type_index))?
362            .visit(self, &mut visitor, DisplayContext::default())?;
363
364        if !visitor.has_displayed_whole_input() {
365            return Err(FormatError::UnusedInput.into());
366        }
367        Ok(output)
368    }
369
370    /// EIP712-compatible JSON encoding of an object, which can be used directly as a parameter
371    /// for `eth_signTypedData_v4` RPCs.
372    #[cfg(feature = "eip712")]
373    pub fn eip712_json(&self, type_index: usize, input: &[u8]) -> Result<String, SchemaError> {
374        let Some(typed_data) = self.eip712_get_typed_data_inner(type_index, input)? else {
375            return Ok(String::default());
376        };
377        Ok(serde_json::to_string(&typed_data)?)
378    }
379
380    /// The EIP712 signing hash of an object, corresponding to its `eip712_json` encoding.
381    /// This hash should be signed directly with a secp256k1 key to obtain the EIP712 signature.
382    #[cfg(feature = "eip712")]
383    pub fn eip712_signing_hash(
384        &self,
385        type_index: usize,
386        input: &[u8],
387    ) -> Result<[u8; 32], SchemaError> {
388        let Some(typed_data) = self.eip712_get_typed_data_inner(type_index, input)? else {
389            return Ok(Default::default());
390        };
391        Ok(typed_data.eip712_signing_hash()?.into())
392    }
393
394    /// The unashed EIP712 signing digest of an object, corresponding to its `eip712_json` encoding.
395    /// This value needs to be keccak256 hashed then secp256k1 signed to obtain the EIP712 signature.
396    #[cfg(feature = "eip712")]
397    pub fn eip712_signing_digest(
398        &self,
399        type_index: usize,
400        input: &[u8],
401    ) -> Result<[u8; 66], SchemaError> {
402        let Some(typed_data) = self.eip712_get_typed_data_inner(type_index, input)? else {
403            return Ok([0; 66]);
404        };
405        // References:
406        // - https://eips.ethereum.org/EIPS/eip-712#specification-of-the-eth_signtypeddata-json-rpc
407        // - https://github.com/alloy-rs/core/blob/main/crates/dyn-abi/src/eip712/typed_data.rs#L212
408        let mut buf = [0u8; 66];
409        buf[0] = 0x19;
410        buf[1] = 0x01;
411        buf[2..34].copy_from_slice(typed_data.domain.separator().as_slice());
412        buf[34..].copy_from_slice(typed_data.hash_struct()?.as_slice());
413        Ok(buf)
414    }
415
416    #[cfg(feature = "eip712")]
417    fn eip712_get_typed_data_inner(
418        &self,
419        type_index: usize,
420        input: &[u8],
421    ) -> Result<Option<TypedData>, SchemaError> {
422        let mut out_types = Eip712Types::default();
423        let input = &mut &input[..];
424        let mut visitor = Eip712Visitor::new(input, &mut out_types);
425        let root_type = self
426            .types()
427            .get(type_index)
428            .ok_or(SchemaError::InvalidIndex(type_index))?;
429        let Some(visitor_return) = root_type.visit(self, &mut visitor, Eip712Context::default())?
430        else {
431            return Ok(None);
432        };
433        if !visitor.has_displayed_whole_input() {
434            return Err(FormatError::UnusedInput.into());
435        }
436
437        // We manually add the EIP712Domain type outside of the visitor
438        // unwrap: hardcoded types are known to be valid and should never fail to construct
439        out_types.insert(
440            "EIP712Domain".to_string(),
441            vec![
442                PropertyDef::new("string", "name").unwrap(),
443                PropertyDef::new("uint256", "chainId").unwrap(),
444                PropertyDef::new("bytes32", "salt").unwrap(),
445            ],
446        );
447
448        Ok(Some(TypedData {
449            domain: alloy_dyn_abi::Eip712Domain {
450                name: Some(self.chain_data.chain_name.clone().into()),
451                version: None,
452                chain_id: Some(alloy_primitives::U256::from(self.chain_data.chain_id)),
453                // Our chain hash is 32 bytes. We could truncate it to fit in the 20-byte ethereum
454                // Address, but by putting it in the salt we retain the full entropy and security.
455                verifying_contract: None,
456                salt: Some(self.chain_hash()?.into()),
457            },
458            resolver: out_types.into(),
459            primary_type: visitor_return.unique_type_name,
460            message: visitor_return.json_value,
461        }))
462    }
463
464    /// Use the schema to convert a serde-compatible JSON string of the given type into its borsh
465    /// encoding
466    #[cfg(feature = "serde")]
467    pub fn json_to_borsh(&self, type_index: usize, input: &str) -> Result<Vec<u8>, SchemaError> {
468        let mut output = Vec::new();
469
470        let mut visitor = EncodeVisitor::new(&mut output)?;
471
472        self.types
473            .get(type_index)
474            .ok_or(SchemaError::InvalidIndex(type_index))?
475            .visit(self, &mut visitor, EncodeContext::new(input, type_index)?)?;
476
477        Ok(output)
478    }
479
480    /// Use a stub JSON to create a full type using the named template.
481    #[cfg(feature = "serde")]
482    pub fn fill_template_from_json(
483        &self,
484        root_index: usize,
485        template_name: &str,
486        input: &str,
487    ) -> Result<Vec<u8>, SchemaError> {
488        fn serde_to_schema_err(e: serde_json::Error) -> SchemaError {
489            SchemaError::EncodeError(EncodeError::Json(e.to_string()))
490        }
491
492        let template = self
493            .templates
494            .get(root_index)
495            .ok_or(SchemaError::InvalidIndex(root_index))?
496            .0
497            .get(template_name)
498            .ok_or(SchemaError::UnknownTemplate(template_name.to_string()))?;
499
500        // Parse the JSON as a map/object of inputs
501        let mut input_map: serde_json::Map<String, serde_json::Value> =
502            serde_json::from_str::<serde_json::Map<String, serde_json::Value>>(input)
503                .map_err(serde_to_schema_err)?;
504
505        let mut output = template.preencoded_bytes().to_owned();
506
507        // For every input in the template, starting from the end (to preserve the `offset` values
508        // of previous inputs)...
509        for (name, input) in template.inputs().iter().rev() {
510            // ...get its type from the schema,
511            let ty = match input.type_link() {
512                Link::ByIndex(i) => self.types.get(*i).expect("Template {name} contained an invalid link: {i}. This is a major bug with template generation."),
513                Link::Immediate(ty) => &ty.clone().into(),
514                Link::Placeholder | Link::IndexedPlaceholder(_) => panic!("Template {name} contained placeholder link. This is a major bug with template generation.")
515            };
516            // find the corresponding JSON value,
517            let json_value = input_map.remove(name).ok_or(EncodeError::MissingType {
518                name: name.to_owned(),
519            })?;
520            // and use our json_to_borsh functionality to get the bytes for the input.
521            let mut buf = Vec::new();
522            let mut visitor = EncodeVisitor::new(&mut buf)?;
523            ty.visit(
524                self,
525                &mut visitor,
526                EncodeContext::from_val(json_value, input.type_link()),
527            )?;
528
529            // Finally, splice the obtained bytes at the specified offset into the template.
530            output.splice(input.offset()..input.offset(), buf);
531        }
532
533        if !input_map.is_empty() {
534            // Unwrap: we know input_map isn't empty, so it must have at least one entry
535            return Err(SchemaError::EncodeError(EncodeError::UnusedInput {
536                value: input_map.iter().next().unwrap().0.to_owned(),
537            }));
538        }
539
540        Ok(output)
541    }
542
543    /// Lists all templates available for the given root type.
544    #[cfg(feature = "serde")]
545    pub fn templates(&self, index: usize) -> Result<Vec<String>, SchemaError> {
546        Ok(self
547            .templates
548            .get(index)
549            .ok_or(SchemaError::InvalidIndex(index))?
550            .0
551            .keys()
552            .cloned()
553            .collect())
554    }
555
556    /// Returns the chain ID calculated using the merkle root of all the schema types, combined
557    /// with any chain-specific metadata.
558    /// This allows the chain ID to be used for verification of the schema (and thus verification
559    /// that a transaction claiming to correspond to a given schema will have the effect it claims).
560    pub fn chain_hash(&self) -> Result<[u8; 32], SchemaError> {
561        self.chain_hash
562            .get_or_try_init(|| {
563                // First, merkleize the schema
564                let merkle_root = self.merkle_root()?;
565
566                // Then, hash the auxilliary internal data - root indices and chain data
567                let mut hasher = Sha256::new();
568                hasher.update(&borsh::to_vec(&self.root_type_indices)?);
569                hasher.update(&borsh::to_vec(&self.chain_data)?);
570                let internal_data_hash: [u8; 32] = hasher.finalize().into();
571
572                // Get the metadata hash
573                let metadata_hash = self.metadata_hash()?;
574
575                // Finally, combine the three hashes in order to get the final chain hash
576                let mut hasher = Sha256::new();
577                hasher.update(merkle_root);
578                hasher.update(internal_data_hash);
579                hasher.update(metadata_hash);
580
581                let chain_hash: [u8; 32] = hasher.finalize().into();
582                Ok(chain_hash)
583            })
584            .copied()
585    }
586
587    fn merkle_root(&self) -> Result<[u8; 32], SchemaError> {
588        let (_, root) = self.merkle_tree.0.get_or_try_init(|| {
589            let mut tree = MerkleTree::new();
590            for ty in &self.types {
591                tree.push_raw_leaf(&borsh::to_vec(ty)?)
592            }
593            let root = tree.root();
594            Ok::<_, SchemaError>((tree, root))
595        })?;
596        Ok(*root)
597    }
598
599    fn finalize(&self) -> Result<(), SchemaError> {
600        // Ensure both hashes are calculated and cached
601        self.metadata_hash()?;
602        self.chain_hash()?;
603        Ok(())
604    }
605
606    pub fn types(&self) -> &[Ty<IndexLinking>] {
607        &self.types
608    }
609
610    pub fn serde_metadata(&self) -> &[ContainerSerdeMetadata] {
611        &self.serde_metadata
612    }
613
614    pub fn root_types(&self) -> &[usize] {
615        &self.root_type_indices
616    }
617
618    fn find_item_by_id(&self, item_id: &ItemId) -> Option<usize> {
619        self.known_types.get(item_id).copied()
620    }
621
622    /// Link a child type to its parent, panicking if the parent type is not in the schema or if the parent type has no more placeholders.
623    fn link_child_to_parent(&mut self, parent: ItemId, child: Link) {
624        let idx = self.known_types.get(&parent).unwrap_or_else(|| panic!("Tried to link a child to a parent ({parent:?}) that the schema doesn't have. This is a bug in a hand-written schema."));
625
626        let remaining_children = *self.under_construction.get(&parent).unwrap_or_else(|| panic!("Tried to link too many children to parent ({parent:?}). This is a bug in a hand-written schema."));
627        if remaining_children == 1 {
628            self.under_construction.remove(&parent);
629        } else {
630            self.under_construction
631                .insert(parent, remaining_children - 1);
632        }
633        self.types[*idx].fill_next_placholder(child);
634    }
635
636    /// Get a link to the given type, adding it to the top-level schema if necessary.
637    /// Unlike all other methods in this crate, the linked type returned by this method is allowed to be only partially generated.
638    ///
639    /// It is the responsibility of the caller to complete the returned link.
640    fn get_partial_link_to(
641        &mut self,
642        item: Item<IndexLinking>,
643        item_id: ItemId,
644    ) -> MaybePartialLink {
645        match item {
646            Item::Container(c) => {
647                if let Some(location) = self.find_item_by_id(&item_id) {
648                    MaybePartialLink::Complete(Link::ByIndex(location))
649                } else {
650                    let num_children = c.num_children();
651                    let serde_metadata = c.serde();
652                    let location = self.types.len();
653                    self.known_types.insert(item_id.clone(), location);
654                    self.types.push(c.into());
655                    self.serde_metadata.push(serde_metadata);
656                    if num_children != 0 {
657                        self.under_construction.insert(item_id, num_children);
658                        MaybePartialLink::Partial(Link::ByIndex(location))
659                    } else {
660                        MaybePartialLink::Complete(Link::ByIndex(location))
661                    }
662                }
663            }
664            Item::Atom(primitive) => MaybePartialLink::Complete(Link::Immediate(primitive)),
665        }
666    }
667
668    /// After generating a root type, register it with the schema for ease of reference. Sets
669    /// the canonical "root links", so has to be carefully called in the right order (normally,
670    /// immediately after root type construction, with the link to the newly created type).
671    /// No-op for primitive links.
672    /// Panics on placeholder links.
673    fn push_root_link(&mut self, link: Link) {
674        match link {
675            Link::ByIndex(i) => self.root_type_indices.push(i),
676            Link::Immediate(..) => {},
677            Link::Placeholder | Link::IndexedPlaceholder(_) => panic!("Attempted to register a placeholder link as a schema root - are you passing the right link?"),
678        }
679    }
680}
681
682pub enum Item<L: LinkingScheme> {
683    Container(Container<L>),
684    Atom(Primitive),
685}
686
687/// Generate the schema for a type.
688/// For complex types, this should typically be derived with a macro,
689/// rather than implemented by hand.
690/// This is also automatically implemented for all types implementing `OverrideSchema`.
691pub trait UniversalWallet: Sized + 'static {
692    /// Ensure that each type contained in the outer type (i.e. the type of each struct/tuple field) is added to the schema,
693    /// and return a `Link` connecting the child to the parent.
694    ///
695    /// Ideally, this function would return something like `Box<dyn UniversalWallet>`.
696    /// Unfortunately, we need to return *types*, not instances (because we don't want to
697    /// add a `Default` bound on all types that implement UniversalWallet) which Rustc doesn't like.
698    /// So, we have a slightly messier signature where the type is expected to register each of its child
699    /// types with the schema directly rather than returning them to the caller for future registration.
700    fn get_child_links(schema: &mut Schema) -> Vec<Link>;
701
702    /// Generate the "scaffolding" of the item. If the item is a primtive, this is just the corresponding primtive.
703    /// If the type is composed of other types, this is the container with all links set to [`Link::Placeholder`].
704    fn scaffold() -> Item<IndexLinking>;
705
706    /// Writes the type to the schema if it is not already present and returns a link to it.
707    ///
708    /// Any child types will have their schemas generated as well, but the placement of those types is left
709    /// to the discretion of the implementation - they may or may not appear at the top level of the schema.
710    fn write_schema(schema: &mut Schema) -> Link {
711        let item = Self::scaffold();
712        let item_id = ItemId::of::<Self>();
713        match item {
714            Item::Atom(_primitive) => {
715                // When recursively building the schema, primitives get filled in directly as
716                // Link::Immediate and do not get `write_schema` called for them. Thus this can
717                // only happen from a user call.
718                // Forbidding this makes managing metadata significantly easier.
719                panic!("Creating a schema for primitive root types is not supported. If this is necessary, wrap the primitive in a newtype struct. If you did not specify a primitive root type, this may be a bug in schema generation.");
720            }
721            Item::Container(container) => {
722                let link = schema.get_partial_link_to(Item::Container(container), item_id.clone());
723                if let MaybePartialLink::Complete(link) = link {
724                    return link;
725                }
726
727                for child in Self::get_child_links(schema) {
728                    schema.link_child_to_parent(item_id.clone(), child);
729                }
730                link.into_inner()
731            }
732        }
733    }
734
735    /// Writes the type and all its children to the schema, if not already present, and sets the
736    /// type as a root type. Generates any templates defined on that type.
737    fn make_root_of(schema: &mut Schema) {
738        let link = Self::write_schema(schema);
739        assert!(
740            schema.under_construction.is_empty(),
741            "Schema generation left some types partially constructed. This is a bug in the schema. {schema:?}"
742        );
743        schema.push_root_link(link);
744        let templates = Self::get_child_templates(schema);
745        schema.templates.push(templates);
746    }
747
748    /// Empty by default
749    /// When derived by the macro, builds a template set from annotations on the fields + the field
750    /// types' own get_child_templates()
751    fn get_child_templates(_schema: &mut Schema) -> TransactionTemplateSet {
752        Default::default()
753    }
754
755    /// Gets a link to the type, writing the type to the schema if necessary.
756    fn make_linkable(schema: &mut Schema) -> Link {
757        match Self::scaffold() {
758            Item::Container(_) => Self::write_schema(schema),
759            Item::Atom(atom) => Link::Immediate(atom),
760        }
761    }
762
763    /// Override the type ID of the item. This should typically not be written by hand. Instead,
764    /// use the [`OverrideSchema`] trait.
765    fn id_override() -> Option<ItemId> {
766        None
767    }
768}
769
770/// Establish that this type should use the `Output` type to generate its schema.
771/// This is appropriate for cases where different types represent the same kind of data structure.
772/// For instance, HashMap and BTreeMap both represent a `Container::Map` in the data model of the
773/// schema; their internal implementation differences don't affect the shape of their schemas.
774///
775/// Note that, for types to be considered equivalent in the schema, their borsh and JSON
776/// serialisations must both also be equivalent.
777pub trait OverrideSchema {
778    type Output: UniversalWallet;
779}
780
781impl<T: OverrideSchema + 'static> UniversalWallet for T {
782    fn scaffold() -> Item<IndexLinking> {
783        <Self as OverrideSchema>::Output::scaffold()
784    }
785    fn get_child_links(schema: &mut Schema) -> Vec<Link> {
786        <Self as OverrideSchema>::Output::get_child_links(schema)
787    }
788    fn id_override() -> Option<ItemId> {
789        <Self as OverrideSchema>::Output::id_override()
790    }
791    fn make_linkable(schema: &mut Schema) -> Link {
792        <Self as OverrideSchema>::Output::make_linkable(schema)
793    }
794    fn write_schema(schema: &mut Schema) -> Link {
795        <Self as OverrideSchema>::Output::write_schema(schema)
796    }
797}
sov_universal_wallet/schema/mod.rs

sov_universal_wallet/schema/
mod.rs
