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