Skip to main content

miden_standards/code_builder/
mod.rs

1use alloc::sync::Arc;
2use alloc::vec::Vec;
3
4use miden_protocol::account::AccountComponentCode;
5use miden_protocol::assembly::{
6    Assembler,
7    DefaultSourceManager,
8    Library,
9    Parse,
10    ParseOptions,
11    Path,
12    SourceManagerSync,
13};
14use miden_protocol::note::NoteScript;
15use miden_protocol::transaction::{TransactionKernel, TransactionScript};
16use miden_protocol::vm::AdviceMap;
17use miden_protocol::{Felt, Word};
18
19use crate::errors::CodeBuilderError;
20use crate::standards_lib::StandardsLib;
21
22// CODE BUILDER
23// ================================================================================================
24
25/// A builder for compiling account components, note scripts, and transaction scripts with optional
26/// library dependencies.
27///
28/// The [`CodeBuilder`] simplifies the process of creating transaction scripts by providing:
29/// - A clean API for adding multiple libraries with static or dynamic linking
30/// - Automatic assembler configuration with all added libraries
31/// - Debug mode support
32/// - Builder pattern support for method chaining
33///
34/// ## Static vs Dynamic Linking
35///
36/// **Static Linking** (`link_static_library()` / `with_statically_linked_library()`):
37/// - Use when you control and know the library code
38/// - The library code is copied into the script code
39/// - Best for most user-written libraries and dependencies
40/// - Results in larger script size but ensures the code is always available
41///
42/// **Dynamic Linking** (`link_dynamic_library()` / `with_dynamically_linked_library()`):
43/// - Use when making Foreign Procedure Invocation (FPI) calls
44/// - The library code is available on-chain and referenced, not copied
45/// - Results in smaller script size but requires the code to be available on-chain
46///
47/// ## Typical Workflow
48///
49/// 1. Create a new CodeBuilder with debug mode preference
50/// 2. Add any required modules using `link_module()` or `with_linked_module()`
51/// 3. Add libraries using `link_static_library()` / `link_dynamic_library()` as appropriate
52/// 4. Compile your script with `compile_note_script()` or `compile_tx_script()`
53///
54/// Note that the compiling methods consume the CodeBuilder, so if you need to compile
55/// multiple scripts with the same configuration, you should clone the builder first.
56///
57/// ## Builder Pattern Example
58///
59/// ```no_run
60/// # use anyhow::Context;
61/// # use miden_standards::code_builder::CodeBuilder;
62/// # use miden_protocol::assembly::Library;
63/// # use miden_protocol::CoreLibrary;
64/// # fn example() -> anyhow::Result<()> {
65/// # let module_code = "pub proc test push.1 add end";
66/// # let script_code = "begin nop end";
67/// # // Create sample libraries for the example
68/// # let my_lib: Library = CoreLibrary::default().into(); // Convert CoreLibrary to Library
69/// # let fpi_lib: Library = CoreLibrary::default().into();
70/// let script = CodeBuilder::default()
71///     .with_linked_module("my::module", module_code).context("failed to link module")?
72///     .with_statically_linked_library(&my_lib).context("failed to link static library")?
73///     .with_dynamically_linked_library(&fpi_lib).context("failed to link dynamic library")?  // For FPI calls
74///     .compile_tx_script(script_code).context("failed to parse tx script")?;
75/// # Ok(())
76/// # }
77/// ```
78///
79/// # Note
80/// The CodeBuilder automatically includes the `miden` and `std` libraries, which
81/// provide access to transaction kernel procedures. Due to being available on-chain
82/// these libraries are linked dynamically and do not add to the size of built script.
83#[derive(Clone)]
84pub struct CodeBuilder {
85    assembler: Assembler,
86    source_manager: Arc<dyn SourceManagerSync>,
87    advice_map: AdviceMap,
88}
89
90impl CodeBuilder {
91    // CONSTRUCTORS
92    // --------------------------------------------------------------------------------------------
93
94    /// Creates a new CodeBuilder.
95    pub fn new() -> Self {
96        Self::with_source_manager(Arc::new(DefaultSourceManager::default()))
97    }
98
99    /// Creates a new CodeBuilder with the specified source manager.
100    ///
101    /// # Arguments
102    /// * `source_manager` - The source manager to use with the internal `Assembler`
103    pub fn with_source_manager(source_manager: Arc<dyn SourceManagerSync>) -> Self {
104        let assembler = TransactionKernel::assembler_with_source_manager(source_manager.clone())
105            .with_dynamic_library(StandardsLib::default())
106            .expect("linking std lib should work");
107        Self {
108            assembler,
109            source_manager,
110            advice_map: AdviceMap::default(),
111        }
112    }
113
114    // LIBRARY MANAGEMENT
115    // --------------------------------------------------------------------------------------------
116
117    /// Parses and links a module to the code builder.
118    ///
119    /// This method compiles the provided module code and adds it directly to the assembler
120    /// for use in script compilation.
121    ///
122    /// # Arguments
123    /// * `module_path` - The path identifier for the module (e.g., "my_lib::my_module")
124    /// * `module_code` - The source code of the module to compile and link
125    ///
126    /// # Errors
127    /// Returns an error if:
128    /// - The module path is invalid
129    /// - The module code cannot be parsed
130    /// - The module cannot be assembled
131    pub fn link_module(
132        &mut self,
133        module_path: impl AsRef<str>,
134        module_code: impl Parse,
135    ) -> Result<(), CodeBuilderError> {
136        let mut parse_options = ParseOptions::for_library();
137        parse_options.path = Some(Path::new(module_path.as_ref()).into());
138
139        let module = module_code.parse_with_options(self.source_manager(), parse_options).map_err(
140            |err| CodeBuilderError::build_error_with_report("failed to parse module code", err),
141        )?;
142
143        self.assembler.compile_and_statically_link(module).map_err(|err| {
144            CodeBuilderError::build_error_with_report("failed to assemble module", err)
145        })?;
146
147        Ok(())
148    }
149
150    /// Statically links the given library.
151    ///
152    /// Static linking means the library code is copied into the script code.
153    /// Use this for most libraries that are not available on-chain.
154    ///
155    /// # Arguments
156    /// * `library` - The compiled library to statically link
157    ///
158    /// # Errors
159    /// Returns an error if:
160    /// - adding the library to the assembler failed
161    pub fn link_static_library(&mut self, library: &Library) -> Result<(), CodeBuilderError> {
162        self.assembler.link_static_library(library).map_err(|err| {
163            CodeBuilderError::build_error_with_report("failed to add static library", err)
164        })
165    }
166
167    /// Dynamically links a library.
168    ///
169    /// This is useful to dynamically link the [`Library`] of a foreign account
170    /// that is invoked using foreign procedure invocation (FPI). Its code is available
171    /// on-chain and so it does not have to be copied into the script code.
172    ///
173    /// For all other use cases not involving FPI, link the library statically.
174    ///
175    /// # Arguments
176    /// * `library` - The compiled library to dynamically link
177    ///
178    /// # Errors
179    /// Returns an error if the library cannot be added to the assembler
180    pub fn link_dynamic_library(&mut self, library: &Library) -> Result<(), CodeBuilderError> {
181        self.assembler.link_dynamic_library(library).map_err(|err| {
182            CodeBuilderError::build_error_with_report("failed to add dynamic library", err)
183        })
184    }
185
186    /// Builder-style method to statically link a library and return the modified builder.
187    ///
188    /// This enables method chaining for convenient builder patterns.
189    ///
190    /// # Arguments
191    /// * `library` - The compiled library to statically link
192    ///
193    /// # Errors
194    /// Returns an error if the library cannot be added to the assembler
195    pub fn with_statically_linked_library(
196        mut self,
197        library: &Library,
198    ) -> Result<Self, CodeBuilderError> {
199        self.link_static_library(library)?;
200        Ok(self)
201    }
202
203    /// Builder-style method to dynamically link a library and return the modified builder.
204    ///
205    /// This enables method chaining for convenient builder patterns.
206    ///
207    /// # Arguments
208    /// * `library` - The compiled library to dynamically link
209    ///
210    /// # Errors
211    /// Returns an error if the library cannot be added to the assembler
212    pub fn with_dynamically_linked_library(
213        mut self,
214        library: impl AsRef<Library>,
215    ) -> Result<Self, CodeBuilderError> {
216        self.link_dynamic_library(library.as_ref())?;
217        Ok(self)
218    }
219
220    /// Builder-style method to link a module and return the modified builder.
221    ///
222    /// This enables method chaining for convenient builder patterns.
223    ///
224    /// # Arguments
225    /// * `module_path` - The path identifier for the module (e.g., "my_lib::my_module")
226    /// * `module_code` - The source code of the module to compile and link
227    ///
228    /// # Errors
229    /// Returns an error if the module cannot be compiled or added to the assembler
230    pub fn with_linked_module(
231        mut self,
232        module_path: impl AsRef<str>,
233        module_code: impl Parse,
234    ) -> Result<Self, CodeBuilderError> {
235        self.link_module(module_path, module_code)?;
236        Ok(self)
237    }
238
239    // ADVICE MAP MANAGEMENT
240    // --------------------------------------------------------------------------------------------
241
242    /// Adds an entry to the advice map that will be included in compiled scripts.
243    ///
244    /// The advice map allows passing non-deterministic inputs to the VM that can be
245    /// accessed using `adv.push_mapval` instruction.
246    ///
247    /// # Arguments
248    /// * `key` - The key for the advice map entry (a Word)
249    /// * `value` - The values to associate with this key
250    pub fn add_advice_map_entry(&mut self, key: Word, value: impl Into<Vec<Felt>>) {
251        self.advice_map.insert(key, value.into());
252    }
253
254    /// Builder-style method to add an advice map entry.
255    ///
256    /// # Arguments
257    /// * `key` - The key for the advice map entry (a Word)
258    /// * `value` - The values to associate with this key
259    pub fn with_advice_map_entry(mut self, key: Word, value: impl Into<Vec<Felt>>) -> Self {
260        self.add_advice_map_entry(key, value);
261        self
262    }
263
264    /// Extends the advice map with entries from another advice map.
265    ///
266    /// # Arguments
267    /// * `advice_map` - The advice map to merge into this builder's advice map
268    pub fn extend_advice_map(&mut self, advice_map: AdviceMap) {
269        self.advice_map.extend(advice_map);
270    }
271
272    /// Builder-style method to extend the advice map.
273    ///
274    /// # Arguments
275    /// * `advice_map` - The advice map to merge into this builder's advice map
276    pub fn with_extended_advice_map(mut self, advice_map: AdviceMap) -> Self {
277        self.extend_advice_map(advice_map);
278        self
279    }
280
281    // PRIVATE HELPERS
282    // --------------------------------------------------------------------------------------------
283
284    /// Applies the advice map to a program if it's non-empty.
285    ///
286    /// This avoids cloning the MAST forest when there are no advice map entries.
287    fn apply_advice_map(
288        advice_map: AdviceMap,
289        program: miden_protocol::vm::Program,
290    ) -> miden_protocol::vm::Program {
291        if advice_map.is_empty() {
292            program
293        } else {
294            program.with_advice_map(advice_map)
295        }
296    }
297
298    /// Applies the advice map to a library if it's non-empty.
299    ///
300    /// This avoids cloning the MAST forest when there are no advice map entries.
301    fn apply_advice_map_to_library(advice_map: AdviceMap, library: Library) -> Library {
302        if advice_map.is_empty() {
303            library
304        } else {
305            library.with_advice_map(advice_map)
306        }
307    }
308
309    // COMPILATION
310    // --------------------------------------------------------------------------------------------
311
312    /// Compiles the provided module path and MASM code into an [`AccountComponentCode`].
313    /// The resulting code can be used to create account components.
314    ///
315    /// # Arguments
316    /// * `component_path` - The path to the account code module (e.g., `my_account::my_module`)
317    /// * `component_code` - The account component source code
318    ///
319    /// # Errors
320    /// Returns an error if:
321    /// - Compiling the account component code fails
322    pub fn compile_component_code(
323        self,
324        component_path: impl AsRef<str>,
325        component_code: impl Parse,
326    ) -> Result<AccountComponentCode, CodeBuilderError> {
327        let CodeBuilder { assembler, source_manager, advice_map } = self;
328
329        let mut parse_options = ParseOptions::for_library();
330        parse_options.path = Some(Path::new(component_path.as_ref()).into());
331
332        let module =
333            component_code
334                .parse_with_options(source_manager, parse_options)
335                .map_err(|err| {
336                    CodeBuilderError::build_error_with_report("failed to parse component code", err)
337                })?;
338
339        let library = assembler.assemble_library([module]).map_err(|err| {
340            CodeBuilderError::build_error_with_report("failed to parse component code", err)
341        })?;
342
343        Ok(AccountComponentCode::from(Self::apply_advice_map_to_library(
344            advice_map, library,
345        )))
346    }
347
348    /// Compiles the provided MASM code into a [`TransactionScript`].
349    ///
350    /// The parsed script will have access to all modules that have been added to this builder.
351    ///
352    /// # Arguments
353    /// * `tx_script` - The transaction script source code
354    ///
355    /// # Errors
356    /// Returns an error if:
357    /// - The transaction script compiling fails
358    pub fn compile_tx_script(
359        self,
360        tx_script: impl Parse,
361    ) -> Result<TransactionScript, CodeBuilderError> {
362        let CodeBuilder { assembler, advice_map, .. } = self;
363
364        let program = assembler.assemble_program(tx_script).map_err(|err| {
365            CodeBuilderError::build_error_with_report("failed to parse transaction script", err)
366        })?;
367
368        Ok(TransactionScript::new(Self::apply_advice_map(advice_map, program)))
369    }
370
371    /// Compiles the provided MASM code into a [`NoteScript`].
372    ///
373    /// The parsed script will have access to all modules that have been added to this builder.
374    ///
375    /// # Arguments
376    /// * `program` - The note script source code
377    ///
378    /// # Errors
379    /// Returns an error if:
380    /// - The note script compiling fails
381    pub fn compile_note_script(self, source: impl Parse) -> Result<NoteScript, CodeBuilderError> {
382        let CodeBuilder { assembler, advice_map, .. } = self;
383
384        let program = assembler.assemble_program(source).map_err(|err| {
385            CodeBuilderError::build_error_with_report("failed to parse note script", err)
386        })?;
387
388        Ok(NoteScript::new(Self::apply_advice_map(advice_map, program)))
389    }
390
391    // ACCESSORS
392    // --------------------------------------------------------------------------------------------
393
394    /// Access the [`Assembler`]'s [`SourceManagerSync`].
395    pub fn source_manager(&self) -> Arc<dyn SourceManagerSync> {
396        self.source_manager.clone()
397    }
398
399    // TESTING CONVENIENCE FUNCTIONS
400    // --------------------------------------------------------------------------------------------
401
402    /// Returns a [`CodeBuilder`] with the transaction kernel as a library.
403    ///
404    /// This assembler is the same as [`TransactionKernel::assembler`] but additionally includes the
405    /// kernel library on the namespace of `$kernel`. The `$kernel` library is added separately
406    /// because even though the library (`api.masm`) and the kernel binary (`main.masm`) include
407    /// this code, it is not otherwise accessible. By adding it separately, we can invoke procedures
408    /// from the kernel library to test them individually.
409    #[cfg(any(feature = "testing", test))]
410    pub fn with_kernel_library(source_manager: Arc<dyn SourceManagerSync>) -> Self {
411        let mut builder = Self::with_source_manager(source_manager);
412        builder
413            .link_dynamic_library(&TransactionKernel::library())
414            .expect("failed to link kernel library");
415        builder
416    }
417
418    /// Returns a [`CodeBuilder`] with the `mock::{account, faucet, util}` libraries.
419    ///
420    /// This assembler includes:
421    /// - [`MockAccountCodeExt::mock_account_library`][account_lib],
422    /// - [`MockAccountCodeExt::mock_faucet_library`][faucet_lib],
423    /// - [`mock_util_library`][util_lib]
424    ///
425    /// [account_lib]: crate::testing::mock_account_code::MockAccountCodeExt::mock_account_library
426    /// [faucet_lib]: crate::testing::mock_account_code::MockAccountCodeExt::mock_faucet_library
427    /// [util_lib]: miden_protocol::testing::mock_util_lib::mock_util_library
428    #[cfg(any(feature = "testing", test))]
429    pub fn with_mock_libraries() -> Self {
430        Self::with_mock_libraries_with_source_manager(Arc::new(DefaultSourceManager::default()))
431    }
432
433    /// Returns the mock account and faucet libraries used in testing.
434    #[cfg(any(feature = "testing", test))]
435    pub fn mock_libraries() -> impl Iterator<Item = Library> {
436        use miden_protocol::account::AccountCode;
437
438        use crate::testing::mock_account_code::MockAccountCodeExt;
439
440        vec![AccountCode::mock_account_library(), AccountCode::mock_faucet_library()].into_iter()
441    }
442
443    #[cfg(any(feature = "testing", test))]
444    pub fn with_mock_libraries_with_source_manager(
445        source_manager: Arc<dyn SourceManagerSync>,
446    ) -> Self {
447        use miden_protocol::testing::mock_util_lib::mock_util_library;
448
449        // Start with the builder linking against the transaction kernel, protocol library and
450        // standards library.
451        let mut builder = Self::with_source_manager(source_manager);
452
453        // Expose kernel procedures under `$kernel` for testing.
454        builder
455            .link_dynamic_library(&TransactionKernel::library())
456            .expect("failed to link kernel library");
457
458        // Add mock account/faucet libs (built in debug mode) and mock util.
459        for library in Self::mock_libraries() {
460            builder
461                .link_dynamic_library(&library)
462                .expect("failed to link mock account libraries");
463        }
464        builder
465            .link_static_library(&mock_util_library())
466            .expect("failed to link mock util library");
467
468        builder
469    }
470}
471
472impl Default for CodeBuilder {
473    fn default() -> Self {
474        Self::new()
475    }
476}
477
478impl From<CodeBuilder> for Assembler {
479    fn from(builder: CodeBuilder) -> Self {
480        builder.assembler
481    }
482}
483
484// TESTS
485// ================================================================================================
486
487#[cfg(test)]
488mod tests {
489    use anyhow::Context;
490    use miden_protocol::assembly::diagnostics::NamedSource;
491
492    use super::*;
493
494    #[test]
495    fn test_code_builder_new() {
496        let _builder = CodeBuilder::default();
497        // Test that the builder can be created successfully
498    }
499
500    #[test]
501    fn test_code_builder_basic_script_compiling() -> anyhow::Result<()> {
502        let builder = CodeBuilder::default();
503        builder
504            .compile_tx_script("begin nop end")
505            .context("failed to parse basic tx script")?;
506        Ok(())
507    }
508
509    #[test]
510    fn test_create_library_and_create_tx_script() -> anyhow::Result<()> {
511        let script_code = "
512            use external_contract::counter_contract
513
514            begin
515                call.counter_contract::increment
516            end
517        ";
518
519        let account_code = "
520            use miden::protocol::active_account
521            use miden::protocol::native_account
522            use miden::core::sys
523
524            pub proc increment
525                push.0
526                exec.active_account::get_item
527                push.1 add
528                push.0
529                exec.native_account::set_item
530                exec.sys::truncate_stack
531            end
532        ";
533
534        let library_path = "external_contract::counter_contract";
535
536        let mut builder_with_lib = CodeBuilder::default();
537        builder_with_lib
538            .link_module(library_path, account_code)
539            .context("failed to link module")?;
540        builder_with_lib
541            .compile_tx_script(script_code)
542            .context("failed to parse tx script")?;
543
544        Ok(())
545    }
546
547    #[test]
548    fn test_parse_library_and_add_to_builder() -> anyhow::Result<()> {
549        let script_code = "
550            use external_contract::counter_contract
551
552            begin
553                call.counter_contract::increment
554            end
555        ";
556
557        let account_code = "
558            use miden::protocol::active_account
559            use miden::protocol::native_account
560            use miden::core::sys
561
562            pub proc increment
563                push.0
564                exec.active_account::get_item
565                push.1 add
566                push.0
567                exec.native_account::set_item
568                exec.sys::truncate_stack
569            end
570        ";
571
572        let library_path = "external_contract::counter_contract";
573
574        // Test single library
575        let mut builder_with_lib = CodeBuilder::default();
576        builder_with_lib
577            .link_module(library_path, account_code)
578            .context("failed to link module")?;
579        builder_with_lib
580            .compile_tx_script(script_code)
581            .context("failed to parse tx script")?;
582
583        // Test multiple libraries
584        let mut builder_with_libs = CodeBuilder::default();
585        builder_with_libs
586            .link_module(library_path, account_code)
587            .context("failed to link first module")?;
588        builder_with_libs
589            .link_module("test::lib", "pub proc test nop end")
590            .context("failed to link second module")?;
591        builder_with_libs
592            .compile_tx_script(script_code)
593            .context("failed to parse tx script with multiple libraries")?;
594
595        Ok(())
596    }
597
598    #[test]
599    fn test_builder_style_chaining() -> anyhow::Result<()> {
600        let script_code = "
601            use external_contract::counter_contract
602
603            begin
604                call.counter_contract::increment
605            end
606        ";
607
608        let account_code = "
609            use miden::protocol::active_account
610            use miden::protocol::native_account
611            use miden::core::sys
612
613            pub proc increment
614                push.0
615                exec.active_account::get_item
616                push.1 add
617                push.0
618                exec.native_account::set_item
619                exec.sys::truncate_stack
620            end
621        ";
622
623        // Test builder-style chaining with modules
624        let builder = CodeBuilder::default()
625            .with_linked_module("external_contract::counter_contract", account_code)
626            .context("failed to link module")?;
627
628        builder.compile_tx_script(script_code).context("failed to parse tx script")?;
629
630        Ok(())
631    }
632
633    #[test]
634    fn test_multiple_chained_modules() -> anyhow::Result<()> {
635        let script_code =
636            "use test::lib1 use test::lib2 begin exec.lib1::test1 exec.lib2::test2 end";
637
638        // Test chaining multiple modules
639        let builder = CodeBuilder::default()
640            .with_linked_module("test::lib1", "pub proc test1 push.1 add end")
641            .context("failed to link first module")?
642            .with_linked_module("test::lib2", "pub proc test2 push.2 add end")
643            .context("failed to link second module")?;
644
645        builder.compile_tx_script(script_code).context("failed to parse tx script")?;
646
647        Ok(())
648    }
649
650    #[test]
651    fn test_static_and_dynamic_linking() -> anyhow::Result<()> {
652        let script_code = "
653            use contracts::static_contract
654
655            begin
656                call.static_contract::increment_1
657            end
658        ";
659
660        let account_code_1 = "
661            pub proc increment_1
662                push.0 drop
663            end
664        ";
665
666        let account_code_2 = "
667            pub proc increment_2
668                push.0 drop
669            end
670        ";
671
672        // Create libraries using the assembler
673        let temp_assembler = TransactionKernel::assembler();
674
675        let static_lib = temp_assembler
676            .clone()
677            .assemble_library([NamedSource::new("contracts::static_contract", account_code_1)])
678            .map_err(|e| anyhow::anyhow!("failed to assemble static library: {}", e))?;
679
680        let dynamic_lib = temp_assembler
681            .assemble_library([NamedSource::new("contracts::dynamic_contract", account_code_2)])
682            .map_err(|e| anyhow::anyhow!("failed to assemble dynamic library: {}", e))?;
683
684        // Test linking both static and dynamic libraries
685        let builder = CodeBuilder::default()
686            .with_statically_linked_library(&static_lib)
687            .context("failed to link static library")?
688            .with_dynamically_linked_library(&dynamic_lib)
689            .context("failed to link dynamic library")?;
690
691        builder
692            .compile_tx_script(script_code)
693            .context("failed to parse tx script with static and dynamic libraries")?;
694
695        Ok(())
696    }
697
698    #[test]
699    fn test_code_builder_with_advice_map_entry() -> anyhow::Result<()> {
700        let key = Word::from([1u32, 2, 3, 4]);
701        let value = vec![Felt::new(42), Felt::new(43)];
702
703        let script = CodeBuilder::default()
704            .with_advice_map_entry(key, value.clone())
705            .compile_tx_script("begin nop end")
706            .context("failed to compile tx script with advice map")?;
707
708        let mast = script.mast();
709        let stored_value = mast.advice_map().get(&key).expect("advice map entry should be present");
710        assert_eq!(stored_value.as_ref(), value.as_slice());
711
712        Ok(())
713    }
714
715    #[test]
716    fn test_code_builder_extend_advice_map() -> anyhow::Result<()> {
717        let key1 = Word::from([1u32, 0, 0, 0]);
718        let key2 = Word::from([2u32, 0, 0, 0]);
719
720        let mut advice_map = AdviceMap::default();
721        advice_map.insert(key1, vec![Felt::new(1)]);
722        advice_map.insert(key2, vec![Felt::new(2)]);
723
724        let script = CodeBuilder::default()
725            .with_extended_advice_map(advice_map)
726            .compile_tx_script("begin nop end")
727            .context("failed to compile tx script")?;
728
729        let mast = script.mast();
730        assert!(mast.advice_map().get(&key1).is_some(), "key1 should be present");
731        assert!(mast.advice_map().get(&key2).is_some(), "key2 should be present");
732
733        Ok(())
734    }
735
736    #[test]
737    fn test_code_builder_advice_map_in_note_script() -> anyhow::Result<()> {
738        let key = Word::from([5u32, 6, 7, 8]);
739        let value = vec![Felt::new(100)];
740
741        let script = CodeBuilder::default()
742            .with_advice_map_entry(key, value.clone())
743            .compile_note_script("begin nop end")
744            .context("failed to compile note script with advice map")?;
745
746        let mast = script.mast();
747        let stored_value = mast
748            .advice_map()
749            .get(&key)
750            .expect("advice map entry should be present in note script");
751        assert_eq!(stored_value.as_ref(), value.as_slice());
752
753        Ok(())
754    }
755
756    #[test]
757    fn test_code_builder_advice_map_in_component_code() -> anyhow::Result<()> {
758        let key = Word::from([11u32, 22, 33, 44]);
759        let value = vec![Felt::new(500)];
760
761        let component_code = CodeBuilder::default()
762            .with_advice_map_entry(key, value.clone())
763            .compile_component_code("test::component", "pub proc test nop end")
764            .context("failed to compile component code with advice map")?;
765
766        let mast = component_code.mast_forest();
767        let stored_value = mast
768            .advice_map()
769            .get(&key)
770            .expect("advice map entry should be present in component code");
771        assert_eq!(stored_value.as_ref(), value.as_slice());
772
773        Ok(())
774    }
775}