Skip to main content

miden_standards/account/interface/
extension.rs

1use alloc::collections::BTreeSet;
2use alloc::sync::Arc;
3use alloc::vec::Vec;
4
5use miden_processor::MastNodeExt;
6use miden_protocol::Word;
7use miden_protocol::account::{Account, AccountCode, AccountId, AccountProcedureRoot};
8use miden_protocol::assembly::mast::{MastForest, MastNode, MastNodeId};
9use miden_protocol::note::{Note, NoteScript};
10
11use crate::AuthScheme;
12use crate::account::components::{
13    WellKnownComponent,
14    basic_fungible_faucet_library,
15    basic_wallet_library,
16    ecdsa_k256_keccak_acl_library,
17    ecdsa_k256_keccak_library,
18    ecdsa_k256_keccak_multisig_library,
19    falcon_512_rpo_acl_library,
20    falcon_512_rpo_library,
21    falcon_512_rpo_multisig_library,
22    network_fungible_faucet_library,
23    no_auth_library,
24};
25use crate::account::interface::{
26    AccountComponentInterface,
27    AccountInterface,
28    NoteAccountCompatibility,
29};
30use crate::note::WellKnownNote;
31
32// ACCOUNT INTERFACE EXTENSION TRAIT
33// ================================================================================================
34
35/// An extension for [`AccountInterface`] that allows instantiation from higher-level types.
36pub trait AccountInterfaceExt {
37    /// Creates a new [`AccountInterface`] instance from the provided account ID, authentication
38    /// schemes and account code.
39    fn from_code(account_id: AccountId, auth: Vec<AuthScheme>, code: &AccountCode) -> Self;
40
41    /// Creates a new [`AccountInterface`] instance from the provided [`Account`].
42    fn from_account(account: &Account) -> Self;
43
44    /// Returns [NoteAccountCompatibility::Maybe] if the provided note is compatible with the
45    /// current [AccountInterface], and [NoteAccountCompatibility::No] otherwise.
46    fn is_compatible_with(&self, note: &Note) -> NoteAccountCompatibility;
47
48    /// Returns the set of digests of all procedures from all account component interfaces.
49    fn get_procedure_digests(&self) -> BTreeSet<Word>;
50}
51
52impl AccountInterfaceExt for AccountInterface {
53    fn from_code(account_id: AccountId, auth: Vec<AuthScheme>, code: &AccountCode) -> Self {
54        let components = AccountComponentInterface::from_procedures(code.procedures());
55
56        Self::new(account_id, auth, components)
57    }
58
59    fn from_account(account: &Account) -> Self {
60        let components = AccountComponentInterface::from_procedures(account.code().procedures());
61        let mut auth = Vec::new();
62
63        // Find the auth component and extract all auth schemes from it
64        // An account should have only one auth component
65        for component in components.iter() {
66            if component.is_auth_component() {
67                auth = component.get_auth_schemes(account.storage());
68                break;
69            }
70        }
71
72        Self::new(account.id(), auth, components)
73    }
74
75    /// Returns [NoteAccountCompatibility::Maybe] if the provided note is compatible with the
76    /// current [AccountInterface], and [NoteAccountCompatibility::No] otherwise.
77    fn is_compatible_with(&self, note: &Note) -> NoteAccountCompatibility {
78        if let Some(well_known_note) = WellKnownNote::from_note(note) {
79            if well_known_note.is_compatible_with(self) {
80                NoteAccountCompatibility::Maybe
81            } else {
82                NoteAccountCompatibility::No
83            }
84        } else {
85            verify_note_script_compatibility(note.script(), self.get_procedure_digests())
86        }
87    }
88
89    fn get_procedure_digests(&self) -> BTreeSet<Word> {
90        let mut component_proc_digests = BTreeSet::new();
91        for component in self.components.iter() {
92            match component {
93                AccountComponentInterface::BasicWallet => {
94                    component_proc_digests
95                        .extend(basic_wallet_library().mast_forest().procedure_digests());
96                },
97                AccountComponentInterface::BasicFungibleFaucet => {
98                    component_proc_digests
99                        .extend(basic_fungible_faucet_library().mast_forest().procedure_digests());
100                },
101                AccountComponentInterface::NetworkFungibleFaucet => {
102                    component_proc_digests.extend(
103                        network_fungible_faucet_library().mast_forest().procedure_digests(),
104                    );
105                },
106                AccountComponentInterface::AuthEcdsaK256Keccak => {
107                    component_proc_digests
108                        .extend(ecdsa_k256_keccak_library().mast_forest().procedure_digests());
109                },
110                AccountComponentInterface::AuthEcdsaK256KeccakAcl => {
111                    component_proc_digests
112                        .extend(ecdsa_k256_keccak_acl_library().mast_forest().procedure_digests());
113                },
114                AccountComponentInterface::AuthEcdsaK256KeccakMultisig => {
115                    component_proc_digests.extend(
116                        ecdsa_k256_keccak_multisig_library().mast_forest().procedure_digests(),
117                    );
118                },
119                AccountComponentInterface::AuthFalcon512Rpo => {
120                    component_proc_digests
121                        .extend(falcon_512_rpo_library().mast_forest().procedure_digests());
122                },
123                AccountComponentInterface::AuthFalcon512RpoAcl => {
124                    component_proc_digests
125                        .extend(falcon_512_rpo_acl_library().mast_forest().procedure_digests());
126                },
127                AccountComponentInterface::AuthFalcon512RpoMultisig => {
128                    component_proc_digests.extend(
129                        falcon_512_rpo_multisig_library().mast_forest().procedure_digests(),
130                    );
131                },
132                AccountComponentInterface::AuthNoAuth => {
133                    component_proc_digests
134                        .extend(no_auth_library().mast_forest().procedure_digests());
135                },
136                AccountComponentInterface::Custom(custom_procs) => {
137                    component_proc_digests
138                        .extend(custom_procs.iter().map(|info| *info.mast_root()));
139                },
140            }
141        }
142
143        component_proc_digests
144    }
145}
146
147/// An extension for [`AccountComponentInterface`] that allows instantiation from a set of procedure
148/// roots.
149pub trait AccountComponentInterfaceExt {
150    /// Creates a vector of [`AccountComponentInterface`] instances from the provided set of
151    /// procedures.
152    fn from_procedures(procedures: &[AccountProcedureRoot]) -> Vec<AccountComponentInterface>;
153}
154
155impl AccountComponentInterfaceExt for AccountComponentInterface {
156    fn from_procedures(procedures: &[AccountProcedureRoot]) -> Vec<Self> {
157        let mut component_interface_vec = Vec::new();
158
159        let mut procedures = BTreeSet::from_iter(procedures.iter().copied());
160
161        // Well known component interfaces
162        // ----------------------------------------------------------------------------------------
163
164        // Get all available well known components which could be constructed from the
165        // `procedures` map and push them to the `component_interface_vec`
166        WellKnownComponent::extract_well_known_components(
167            &mut procedures,
168            &mut component_interface_vec,
169        );
170
171        // Custom component interfaces
172        // ----------------------------------------------------------------------------------------
173
174        // All remaining procedures are put into the custom bucket.
175        component_interface_vec
176            .push(AccountComponentInterface::Custom(procedures.into_iter().collect()));
177
178        component_interface_vec
179    }
180}
181
182// HELPER FUNCTIONS
183// ------------------------------------------------------------------------------------------------
184
185/// Verifies that the provided note script is compatible with the target account interfaces.
186///
187/// This is achieved by checking that at least one execution branch in the note script is compatible
188/// with the account procedures vector.
189///
190/// This check relies on the fact that account procedures are the only procedures that are `call`ed
191/// from note scripts, while kernel procedures are `sycall`ed.
192fn verify_note_script_compatibility(
193    note_script: &NoteScript,
194    account_procedures: BTreeSet<Word>,
195) -> NoteAccountCompatibility {
196    // collect call branches of the note script
197    let branches = collect_call_branches(note_script);
198
199    // if none of the branches are compatible with the target account, return a `CheckResult::No`
200    if !branches.iter().any(|call_targets| call_targets.is_subset(&account_procedures)) {
201        return NoteAccountCompatibility::No;
202    }
203
204    NoteAccountCompatibility::Maybe
205}
206
207/// Collect call branches by recursively traversing through program execution branches and
208/// accumulating call targets.
209fn collect_call_branches(note_script: &NoteScript) -> Vec<BTreeSet<Word>> {
210    let mut branches = vec![BTreeSet::new()];
211
212    let entry_node = note_script.entrypoint();
213    recursively_collect_call_branches(entry_node, &mut branches, &note_script.mast());
214    branches
215}
216
217/// Generates a list of calls invoked in each execution branch of the provided code block.
218fn recursively_collect_call_branches(
219    mast_node_id: MastNodeId,
220    branches: &mut Vec<BTreeSet<Word>>,
221    note_script_forest: &Arc<MastForest>,
222) {
223    let mast_node = &note_script_forest[mast_node_id];
224
225    match mast_node {
226        MastNode::Block(_) => {},
227        MastNode::Join(join_node) => {
228            recursively_collect_call_branches(join_node.first(), branches, note_script_forest);
229            recursively_collect_call_branches(join_node.second(), branches, note_script_forest);
230        },
231        MastNode::Split(split_node) => {
232            let current_branch = branches.last().expect("at least one execution branch").clone();
233            recursively_collect_call_branches(split_node.on_false(), branches, note_script_forest);
234
235            // If the previous branch had additional calls we need to create a new branch
236            if branches.last().expect("at least one execution branch").len() > current_branch.len()
237            {
238                branches.push(current_branch);
239            }
240
241            recursively_collect_call_branches(split_node.on_true(), branches, note_script_forest);
242        },
243        MastNode::Loop(loop_node) => {
244            recursively_collect_call_branches(loop_node.body(), branches, note_script_forest);
245        },
246        MastNode::Call(call_node) => {
247            if call_node.is_syscall() {
248                return;
249            }
250
251            let callee_digest = note_script_forest[call_node.callee()].digest();
252
253            branches
254                .last_mut()
255                .expect("at least one execution branch")
256                .insert(callee_digest);
257        },
258        MastNode::Dyn(_) => {},
259        MastNode::External(_) => {},
260    }
261}