Skip to main content

leo_compiler/
compiler.rs

1// Copyright (C) 2019-2026 Provable Inc.
2// This file is part of the Leo library.
3
4// The Leo library is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// The Leo library is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
16
17//! The compiler for Leo programs.
18//!
19//! The [`Compiler`] type compiles Leo programs into R1CS circuits.
20
21use crate::{AstSnapshots, CompilerOptions};
22
23pub use leo_ast::{Ast, DiGraph, Program};
24use leo_ast::{Library, NetworkName, NodeBuilder, Stub};
25use leo_errors::{CompilerError, Handler, Result};
26use leo_passes::*;
27use leo_span::{
28    Span,
29    Symbol,
30    file_source::{DiskFileSource, FileSource},
31    source_map::FileName,
32    with_session_globals,
33};
34
35use std::{
36    fs,
37    path::{Path, PathBuf},
38    rc::Rc,
39};
40
41use indexmap::IndexMap;
42
43/// A single compiled program with its bytecode and ABI.
44pub struct CompiledProgram {
45    /// The program name (without `.aleo` suffix).
46    pub name: String,
47    /// The generated Aleo bytecode.
48    pub bytecode: String,
49    /// The ABI describing the program's public interface.
50    pub abi: leo_abi::Program,
51}
52
53/// The result of compiling a Leo program.
54pub struct Compiled {
55    /// The primary program that was compiled.
56    pub primary: CompiledProgram,
57    /// Compiled programs for imports.
58    pub imports: Vec<CompiledProgram>,
59}
60
61/// The primary entry point of the Leo compiler.
62pub struct Compiler {
63    /// The path to where the compiler outputs all generated files.
64    output_directory: PathBuf,
65    /// The program name,
66    pub program_name: Option<String>,
67    /// Options configuring compilation.
68    compiler_options: CompilerOptions,
69    /// State.
70    state: CompilerState,
71    /// The stubs for imported programs.
72    import_stubs: IndexMap<Symbol, Stub>,
73    /// How many statements were in the AST before DCE?
74    pub statements_before_dce: u32,
75    /// How many statements were in the AST after DCE?
76    pub statements_after_dce: u32,
77}
78
79impl Compiler {
80    pub fn network(&self) -> NetworkName {
81        self.state.network
82    }
83
84    /// Parses the given source into a program AST and stores it in the compiler state.
85    ///
86    /// The source file and any provided module sources are first registered in the
87    /// session source map so spans can be resolved correctly. The parser then
88    /// constructs the program AST from the main source and its modules.
89    ///
90    /// After parsing, this verifies that the program scope name matches the expected
91    /// program name (from `program.json` or the test filename). The resulting AST is
92    /// stored in `self.state.ast`, and optionally written to disk if configured.
93    pub fn parse_program(&mut self, source: &str, filename: FileName, modules: &[(&str, FileName)]) -> Result<()> {
94        // Register the source in the source map.
95        let source_file = with_session_globals(|s| s.source_map.new_source(source, filename.clone()));
96
97        // Register the sources of all the modules in the source map.
98        let modules = modules
99            .iter()
100            .map(|(source, filename)| with_session_globals(|s| s.source_map.new_source(source, filename.clone())))
101            .collect::<Vec<_>>();
102
103        // Use the parser to construct the abstract syntax tree (ast).
104        let program = leo_parser::parse_program(
105            self.state.handler.clone(),
106            &self.state.node_builder,
107            &source_file,
108            &modules,
109            self.state.network,
110        )?;
111
112        // Check that the name of its program scope matches the expected name.
113        // Note that parsing enforces that there is exactly one program scope in a file.
114        let program_scope = program.program_scopes.values().next().unwrap();
115        if let Some(program_name) = &self.program_name {
116            if program_name != &program_scope.program_id.as_symbol().to_string() {
117                return Err(CompilerError::program_name_should_match_file_name(
118                    program_scope.program_id.as_symbol(),
119                    // If this is a test, use the filename as the expected name.
120                    if self.state.is_test {
121                        format!(
122                            "`{}` (the test file name)",
123                            filename.to_string().split("/").last().expect("Could not get file name")
124                        )
125                    } else {
126                        format!("`{program_name}` (specified in `program.json`)")
127                    },
128                    program_scope.program_id.span(),
129                )
130                .into());
131            }
132        } else {
133            self.program_name = Some(program_scope.program_id.as_symbol().to_string());
134        }
135
136        self.state.ast = Ast::Program(program);
137
138        if self.compiler_options.initial_ast {
139            self.write_ast_to_json("initial.json")?;
140            self.write_ast("initial.ast")?;
141        }
142
143        Ok(())
144    }
145
146    /// Simple wrapper around `parse_program` that also returns a program AST.
147    pub fn parse_and_return_program(
148        &mut self,
149        source: &str,
150        filename: FileName,
151        modules: &[(&str, FileName)],
152    ) -> Result<Program> {
153        // Parse the program.
154        self.parse_program(source, filename, modules)?;
155
156        match &self.state.ast {
157            Ast::Program(program) => Ok(program.clone()),
158            Ast::Library(_) => unreachable!("expected Program AST"),
159        }
160    }
161
162    /// Simple wrapper around `parse_library` that also returns a library AST.
163    pub fn parse_and_return_library(
164        &mut self,
165        library_name: &str,
166        source: &str,
167        filename: FileName,
168        modules: &[(&str, FileName)],
169    ) -> Result<Library> {
170        self.parse_library(Symbol::intern(library_name), source, filename, modules)?;
171
172        match &self.state.ast {
173            Ast::Program(_) => unreachable!("expected Library AST"),
174            Ast::Library(library) => Ok(library.clone()),
175        }
176    }
177
178    /// Parses a library source (and its submodules) into a library AST.
179    ///
180    /// All source strings are registered in the session source map so span information
181    /// can be resolved correctly. The resulting AST is stored in `self.state.ast`.
182    pub fn parse_library(
183        &mut self,
184        library_name: Symbol,
185        source: &str,
186        filename: FileName,
187        modules: &[(&str, FileName)],
188    ) -> Result<()> {
189        let source_file = with_session_globals(|s| s.source_map.new_source(source, filename.clone()));
190
191        // Register each module source in the source map.
192        let module_files = modules
193            .iter()
194            .map(|(src, name)| with_session_globals(|s| s.source_map.new_source(src, name.clone())))
195            .collect::<Vec<_>>();
196
197        self.state.ast = Ast::Library(leo_parser::parse_library(
198            self.state.handler.clone(),
199            &self.state.node_builder,
200            library_name,
201            &source_file,
202            &module_files,
203            self.state.network,
204        )?);
205
206        Ok(())
207    }
208
209    /// Returns a new Leo compiler.
210    #[allow(clippy::too_many_arguments)]
211    pub fn new(
212        expected_program_name: Option<String>,
213        is_test: bool,
214        handler: Handler,
215        node_builder: Rc<NodeBuilder>,
216        output_directory: PathBuf,
217        compiler_options: Option<CompilerOptions>,
218        import_stubs: IndexMap<Symbol, Stub>,
219        network: NetworkName,
220    ) -> Self {
221        Self {
222            state: CompilerState {
223                handler,
224                node_builder: Rc::clone(&node_builder),
225                is_test,
226                network,
227                ..Default::default()
228            },
229            output_directory,
230            program_name: expected_program_name,
231            compiler_options: compiler_options.unwrap_or_default(),
232            import_stubs,
233            statements_before_dce: 0,
234            statements_after_dce: 0,
235        }
236    }
237
238    pub fn do_pass<P: Pass>(&mut self, input: P::Input) -> Result<P::Output> {
239        let output = P::do_pass(input, &mut self.state)?;
240
241        let write = match &self.compiler_options.ast_snapshots {
242            AstSnapshots::All => true,
243            AstSnapshots::Some(passes) => passes.contains(P::NAME),
244        };
245
246        if write {
247            self.write_ast_to_json(&format!("{}.json", P::NAME))?;
248            self.write_ast(&format!("{}.ast", P::NAME))?;
249        }
250
251        Ok(output)
252    }
253
254    /// Runs all frontend passes: NameValidation through StaticAnalyzing.
255    pub fn frontend_passes(&mut self) -> Result<()> {
256        // Bail out if the parser already found errors.  The error-recovering parser may have
257        // produced ErrExpression nodes in the AST, which would cause panics in later passes.
258        self.state.handler.last_err()?;
259
260        self.do_pass::<NameValidation>(())?;
261        self.do_pass::<GlobalVarsCollection>(())?;
262        self.do_pass::<PathResolution>(())?;
263        self.do_pass::<GlobalItemsCollection>(())?;
264        self.do_pass::<CheckInterfaces>(())?;
265        self.do_pass::<TypeChecking>(TypeCheckingInput::new(self.state.network))?;
266        self.do_pass::<Disambiguate>(())?;
267        self.do_pass::<ProcessingAsync>(TypeCheckingInput::new(self.state.network))?;
268        self.do_pass::<StaticAnalyzing>(())
269    }
270
271    /// Runs the compiler stages.
272    ///
273    /// Returns the generated ABIs (primary and imports), which are captured
274    /// immediately after monomorphisation to ensure all types are resolved,
275    /// but not yet lowered.
276    pub fn intermediate_passes(&mut self) -> Result<(leo_abi::Program, IndexMap<String, leo_abi::Program>)> {
277        let type_checking_config = TypeCheckingInput::new(self.state.network);
278
279        self.frontend_passes()?;
280
281        self.do_pass::<ConstPropUnrollAndMorphing>(type_checking_config.clone())?;
282
283        // Generate ABIs after monomorphization to capture concrete types.
284        // Const generic structs are resolved to their monomorphized versions.
285        let abis = self.generate_abi();
286
287        self.do_pass::<StorageLowering>(type_checking_config.clone())?;
288
289        self.do_pass::<OptionLowering>(type_checking_config)?;
290
291        self.do_pass::<SsaForming>(SsaFormingInput { rename_defs: true })?;
292
293        self.do_pass::<Destructuring>(())?;
294
295        self.do_pass::<SsaForming>(SsaFormingInput { rename_defs: false })?;
296
297        self.do_pass::<WriteTransforming>(())?;
298
299        self.do_pass::<SsaForming>(SsaFormingInput { rename_defs: false })?;
300
301        self.do_pass::<Flattening>(())?;
302
303        self.do_pass::<FunctionInlining>(())?;
304
305        // Flattening may produce ternary expressions not in SSA form.
306        self.do_pass::<SsaForming>(SsaFormingInput { rename_defs: false })?;
307
308        self.do_pass::<SsaConstPropagation>(())?;
309
310        self.do_pass::<SsaForming>(SsaFormingInput { rename_defs: false })?;
311
312        self.do_pass::<CommonSubexpressionEliminating>(())?;
313
314        let output = self.do_pass::<DeadCodeEliminating>(())?;
315        self.statements_before_dce = output.statements_before;
316        self.statements_after_dce = output.statements_after;
317
318        Ok(abis)
319    }
320
321    /// Generates ABIs for the primary program and all imports.
322    ///
323    /// Returns `(primary_abi, import_abis)` where `import_abis` maps program
324    /// names to their ABIs.
325    ///
326    /// This method only expects program ASTs. Library ASTs cause this method to panic.
327    fn generate_abi(&self) -> (leo_abi::Program, IndexMap<String, leo_abi::Program>) {
328        let program = match &self.state.ast {
329            Ast::Program(program) => program,
330            Ast::Library(_) => panic!("expected Program AST"),
331        };
332
333        // Generate primary ABI (pruning happens inside generate).
334        let primary_abi = leo_abi::generate(program);
335
336        // Generate import ABIs from stubs, ignoring libraries.
337        let import_abis: IndexMap<String, leo_abi::Program> = program
338            .stubs
339            .iter()
340            .filter(|(_, stub)| !matches!(stub, Stub::FromLibrary { .. }))
341            .map(|(name, stub)| {
342                let abi = match stub {
343                    Stub::FromLeo { program, .. } => leo_abi::generate(program),
344                    Stub::FromAleo { program, .. } => leo_abi::aleo::generate(program),
345                    Stub::FromLibrary { .. } => unreachable!("filtered out"),
346                };
347                (name.to_string(), abi)
348            })
349            .collect();
350
351        (primary_abi, import_abis)
352    }
353
354    /// Compiles a program from a given source string and a list of module sources.
355    ///
356    /// # Arguments
357    ///
358    /// * `source` - The main source code as a string slice.
359    /// * `filename` - The name of the main source file.
360    /// * `modules` - A vector of tuples where each tuple contains:
361    ///     - A module source as a string slice.
362    ///     - Its associated `FileName`.
363    ///
364    /// # Returns
365    ///
366    /// * `Ok(CompiledPrograms)` containing the generated bytecode and ABI if compilation succeeds.
367    /// * `Err(CompilerError)` if any stage of the pipeline fails.
368    pub fn compile(&mut self, source: &str, filename: FileName, modules: &Vec<(&str, FileName)>) -> Result<Compiled> {
369        // Parse the program.
370        self.parse_program(source, filename, modules)?;
371        // Merge the stubs into the AST.
372        self.add_import_stubs()?;
373        // Run the intermediate compiler stages, which also generates ABIs.
374        let (primary_abi, import_abis) = self.intermediate_passes()?;
375        // Run code generation.
376        let bytecodes = CodeGenerating::do_pass((), &mut self.state)?;
377
378        // Build the primary compiled program.
379        let primary = CompiledProgram {
380            name: self.program_name.clone().unwrap(),
381            bytecode: bytecodes.primary_bytecode,
382            abi: primary_abi,
383        };
384
385        // Build compiled programs for imports, looking up ABIs by name.
386        let imports: Vec<CompiledProgram> = bytecodes
387            .import_bytecodes
388            .into_iter()
389            .map(|bc| {
390                let abi = import_abis.get(&bc.program_name).expect("ABI should exist for all imports").clone();
391                CompiledProgram { name: bc.program_name, bytecode: bc.bytecode, abi }
392            })
393            .collect();
394
395        Ok(Compiled { primary, imports })
396    }
397
398    /// Reads the main source file and all module files in the same directory tree.
399    ///
400    /// This helper walks all `.leo` files under `source_directory` (excluding the main file itself),
401    /// reads their contents, and returns:
402    /// - The main file’s source as a `String`.
403    /// - A vector of module tuples `(String, FileName)` suitable for compilation or parsing.
404    ///
405    /// # Arguments
406    ///
407    /// * `entry_file_path` - The main source file.
408    /// * `source_directory` - The directory root for discovering `.leo` module files.
409    ///
410    /// # Errors
411    ///
412    /// Returns `Err(CompilerError)` if reading any file fails.
413    fn read_sources_and_modules(
414        file_source: &impl FileSource,
415        entry_file_path: impl AsRef<Path>,
416        source_directory: impl AsRef<Path>,
417    ) -> Result<(String, Vec<(String, FileName)>)> {
418        let entry_file_path = entry_file_path.as_ref();
419        let source_directory = source_directory.as_ref();
420
421        // Read the contents of the main source file.
422        let source = file_source
423            .read_file(entry_file_path)
424            .map_err(|e| CompilerError::file_read_error(entry_file_path.display().to_string(), e))?;
425
426        let files = file_source
427            .list_leo_files(source_directory, entry_file_path)
428            .map_err(|e| CompilerError::file_read_error(source_directory.display().to_string(), e))?;
429
430        let mut modules = Vec::with_capacity(files.len());
431        for path in files {
432            let module_source = file_source
433                .read_file(&path)
434                .map_err(|e| CompilerError::file_read_error(path.display().to_string(), e))?;
435            modules.push((module_source, FileName::Real(path)));
436        }
437
438        Ok((source, modules))
439    }
440
441    /// Compiles a program from a source file and its associated module files in the same directory tree.
442    pub fn compile_from_directory(
443        &mut self,
444        entry_file_path: impl AsRef<Path>,
445        source_directory: impl AsRef<Path>,
446    ) -> Result<Compiled> {
447        self.compile_from_directory_with_file_source(entry_file_path, source_directory, &DiskFileSource)
448    }
449
450    /// Compiles a program from a source file using the given file source.
451    pub fn compile_from_directory_with_file_source(
452        &mut self,
453        entry_file_path: impl AsRef<Path>,
454        source_directory: impl AsRef<Path>,
455        file_source: &impl FileSource,
456    ) -> Result<Compiled> {
457        let (source, modules_owned) = Self::read_sources_and_modules(file_source, &entry_file_path, &source_directory)?;
458
459        // Convert owned module sources into temporary (&str, FileName) tuples.
460        let module_refs: Vec<(&str, FileName)> =
461            modules_owned.iter().map(|(src, fname)| (src.as_str(), fname.clone())).collect();
462
463        // Compile the main source along with all collected modules.
464        self.compile(&source, FileName::Real(entry_file_path.as_ref().into()), &module_refs)
465    }
466
467    /// Parses a program from a source file and its associated module files in the same directory tree.
468    pub fn parse_program_from_directory(
469        &mut self,
470        entry_file_path: impl AsRef<Path>,
471        source_directory: impl AsRef<Path>,
472    ) -> Result<Program> {
473        self.parse_program_from_directory_with_file_source(entry_file_path, source_directory, &DiskFileSource)
474    }
475
476    /// Parses a program from a source file using the given file source.
477    pub fn parse_program_from_directory_with_file_source(
478        &mut self,
479        entry_file_path: impl AsRef<Path>,
480        source_directory: impl AsRef<Path>,
481        file_source: &impl FileSource,
482    ) -> Result<Program> {
483        let (source, modules_owned) = Self::read_sources_and_modules(file_source, &entry_file_path, &source_directory)?;
484
485        // Convert owned module sources into temporary (&str, FileName) tuples.
486        let module_refs: Vec<(&str, FileName)> =
487            modules_owned.iter().map(|(src, fname)| (src.as_str(), fname.clone())).collect();
488
489        // Parse the main source along with all collected modules.
490        self.parse_program(&source, FileName::Real(entry_file_path.as_ref().into()), &module_refs)?;
491
492        match &self.state.ast {
493            Ast::Program(program) => Ok(program.clone()),
494            Ast::Library(_) => unreachable!("expected Program AST"),
495        }
496    }
497
498    /// Parses a program from a source file and its associated module files in the same directory tree.
499    pub fn parse_library_from_directory(
500        &mut self,
501        library_name: Symbol,
502        entry_file_path: impl AsRef<Path>,
503        source_directory: impl AsRef<Path>,
504    ) -> Result<Library> {
505        self.parse_library_from_directory_with_file_source(
506            library_name,
507            entry_file_path,
508            source_directory,
509            &DiskFileSource,
510        )
511    }
512
513    /// Parses a library from a source file.
514    pub fn parse_library_from_directory_with_file_source(
515        &mut self,
516        library_name: Symbol,
517        entry_file_path: impl AsRef<Path>,
518        source_directory: impl AsRef<Path>,
519        file_source: &impl FileSource,
520    ) -> Result<Library> {
521        let (source, modules_owned) = Self::read_sources_and_modules(file_source, &entry_file_path, &source_directory)?;
522
523        let module_refs: Vec<(&str, FileName)> =
524            modules_owned.iter().map(|(src, fname)| (src.as_str(), fname.clone())).collect();
525
526        self.parse_library(library_name, &source, FileName::Real(entry_file_path.as_ref().into()), &module_refs)?;
527
528        match &self.state.ast {
529            Ast::Library(library) => Ok(library.clone()),
530            Ast::Program(_) => unreachable!("expected Library AST"),
531        }
532    }
533
534    /// Writes the AST to a JSON file.
535    fn write_ast_to_json(&self, file_suffix: &str) -> Result<()> {
536        match &self.state.ast {
537            Ast::Program(program) => {
538                // Remove `Span`s if they are not enabled.
539                if self.compiler_options.ast_spans_enabled {
540                    program.to_json_file(
541                        self.output_directory.clone(),
542                        &format!("{}.{file_suffix}", self.program_name.as_ref().unwrap()),
543                    )?;
544                } else {
545                    program.to_json_file_without_keys(
546                        self.output_directory.clone(),
547                        &format!("{}.{file_suffix}", self.program_name.as_ref().unwrap()),
548                        &["_span", "span"],
549                    )?;
550                }
551            }
552            Ast::Library(_) => {
553                // no-op for libraries
554            }
555        }
556        Ok(())
557    }
558
559    /// Writes the AST to a file (Leo syntax, not JSON).
560    fn write_ast(&self, file_suffix: &str) -> Result<()> {
561        let filename = format!("{}.{file_suffix}", self.program_name.as_ref().unwrap());
562        let full_filename = self.output_directory.join(&filename);
563
564        let contents = match &self.state.ast {
565            Ast::Program(program) => program.to_string(),
566            Ast::Library(_) => String::new(), // empty for libraries
567        };
568
569        fs::write(&full_filename, contents).map_err(|e| CompilerError::failed_ast_file(full_filename.display(), e))?;
570
571        Ok(())
572    }
573
574    /// Resolves and registers all import stubs for the current program.
575    ///
576    /// This method performs a graph traversal over the program’s import relationships to:
577    /// 1. Establish parent–child relationships between stubs based on imports.
578    /// 2. Collect all reachable stubs in traversal order.
579    /// 3. Store the explored stubs back into `self.state.ast.ast.stubs`.
580    ///
581    /// The traversal starts from the imports of the main program and recursively follows
582    /// their transitive dependencies. Any missing stub during traversal results in an error.
583    ///
584    /// # Returns
585    ///
586    /// * `Ok(())` if all imports are successfully resolved and stubs are collected.
587    /// * `Err(CompilerError)` if any imported program cannot be found.
588    pub fn add_import_stubs(&mut self) -> Result<()> {
589        use indexmap::IndexSet;
590
591        // Track which programs we've already processed.
592        let mut explored = IndexSet::<Symbol>::new();
593
594        // Compute initial imports: explicit program imports + library dependencies
595        let initial_imports: IndexMap<Symbol, Span> = match &self.state.ast {
596            Ast::Program(program) => {
597                let mut map: IndexMap<Symbol, Span> =
598                    program.imports.iter().map(|(name, id)| (*name, id.span())).collect();
599                // Add any libraries that have this program as a parent
600                for (stub_name, stub) in &self.import_stubs {
601                    if matches!(stub, Stub::FromLibrary { .. })
602                        && stub.parents().contains(&Symbol::intern(self.program_name.as_ref().unwrap()))
603                    {
604                        map.insert(
605                            *stub_name,
606                            Span::default(), // library dependencies are implicit
607                        );
608                    }
609                }
610                map
611            }
612            Ast::Library(_) => IndexMap::new(),
613        };
614
615        // Initialize the exploration queue with the root’s direct imports.
616        let mut to_explore: Vec<(Symbol, Span)> = initial_imports.iter().map(|(sym, span)| (*sym, *span)).collect();
617
618        // If this is a named program, set the main program as the parent of its direct imports.
619        if let Some(main_program_name) = self.program_name.clone() {
620            let main_symbol = Symbol::intern(&main_program_name);
621            for import in initial_imports.keys() {
622                if let Some(child_stub) = self.import_stubs.get_mut(import) {
623                    child_stub.add_parent(main_symbol);
624                }
625            }
626        }
627
628        // Traverse the dependency graph breadth-first, populating parents
629        while let Some((import_symbol, span)) = to_explore.pop() {
630            // Mark this import as explored.
631            explored.insert(import_symbol);
632
633            // Look up the corresponding stub.
634            let Some(stub) = self.import_stubs.get(&import_symbol) else {
635                return Err(CompilerError::imported_program_not_found(
636                    self.program_name.as_ref().unwrap(),
637                    import_symbol,
638                    span,
639                )
640                .into());
641            };
642
643            // Combine imports: explicit stub.explicit_imports() + libraries that list this stub as parent
644            let mut combined_imports: IndexMap<Symbol, Span> = stub.explicit_imports().collect();
645            for (lib_name, lib_stub) in &self.import_stubs {
646                if matches!(lib_stub, Stub::FromLibrary { .. }) && lib_stub.parents().contains(&import_symbol) {
647                    combined_imports.insert(
648                        *lib_name,
649                        Span::default(), // library dependencies are implicit
650                    );
651                }
652            }
653
654            for (child_symbol, child_span) in combined_imports {
655                // Record parent relationship
656                if let Some(child_stub) = self.import_stubs.get_mut(&child_symbol) {
657                    child_stub.add_parent(import_symbol);
658                }
659
660                // Schedule child for exploration if not yet visited.
661                if explored.insert(child_symbol) {
662                    to_explore.push((child_symbol, child_span));
663                }
664            }
665        }
666
667        // Only Programs store stubs on the AST (at least for now).
668        if let Ast::Program(program) = &mut self.state.ast {
669            program.stubs = self
670                .import_stubs
671                .iter()
672                .filter(|(symbol, _)| explored.contains(*symbol))
673                .map(|(symbol, stub)| (*symbol, stub.clone()))
674                .collect();
675        }
676
677        Ok(())
678    }
679}
680
681#[cfg(test)]
682mod tests {
683    use super::Compiler;
684
685    use leo_ast::{NetworkName, NodeBuilder};
686    use leo_errors::Handler;
687    use leo_span::{Symbol, create_session_if_not_set_then, file_source::InMemoryFileSource};
688
689    use std::{path::PathBuf, rc::Rc};
690
691    use indexmap::IndexMap;
692
693    #[test]
694    fn parse_library_from_directory_in_memory() {
695        create_session_if_not_set_then(|_| {
696            let mut source = InMemoryFileSource::new();
697            source.set(
698                PathBuf::from("/mylib/src/lib.leo"),
699                concat!("const SCALE: u32 = 10u32;\n", "const OFFSET: u32 = SCALE + 1u32;\n",).into(),
700            );
701
702            let handler = Handler::default();
703            let node_builder = Rc::new(NodeBuilder::default());
704            let mut compiler = Compiler::new(
705                None,
706                false,
707                handler,
708                node_builder,
709                PathBuf::from("/unused"),
710                None,
711                IndexMap::new(),
712                NetworkName::TestnetV0,
713            );
714
715            let library = compiler
716                .parse_library_from_directory_with_file_source(
717                    Symbol::intern("mylib"),
718                    "/mylib/src/lib.leo",
719                    "/mylib/src",
720                    &source,
721                )
722                .unwrap_or_else(|err| panic!("parsing library from in-memory file source failed: {err}"));
723
724            assert_eq!(library.name, Symbol::intern("mylib"));
725            assert_eq!(library.consts.len(), 2, "expected 2 consts, got {}", library.consts.len());
726            assert!(
727                library.consts.iter().any(|(name, _)| *name == Symbol::intern("SCALE")),
728                "expected const `SCALE` in library"
729            );
730            assert!(
731                library.consts.iter().any(|(name, _)| *name == Symbol::intern("OFFSET")),
732                "expected const `OFFSET` in library"
733            );
734        });
735    }
736
737    #[test]
738    fn parse_program_from_directory_in_memory_with_module() {
739        create_session_if_not_set_then(|_| {
740            let mut source = InMemoryFileSource::new();
741            source.set(
742                PathBuf::from("/project/src/main.leo"),
743                concat!(
744                    "program test.aleo {\n",
745                    "  fn main() -> u32 {\n",
746                    "    return utils::helper();\n",
747                    "  }\n",
748                    "}\n",
749                )
750                .into(),
751            );
752            source.set(PathBuf::from("/project/src/utils.leo"), "fn helper() -> u32 {\n  return 42u32;\n}\n".into());
753
754            let handler = Handler::default();
755            let node_builder = Rc::new(NodeBuilder::default());
756            let mut compiler = Compiler::new(
757                Some("test.aleo".into()),
758                false,
759                handler,
760                node_builder,
761                PathBuf::from("/unused"),
762                None,
763                IndexMap::new(),
764                NetworkName::TestnetV0,
765            );
766
767            let ast = compiler
768                .parse_program_from_directory_with_file_source("/project/src/main.leo", "/project/src", &source)
769                .unwrap_or_else(|err| panic!("parsing from in-memory file source failed: {err}"));
770            let utils_key = vec![Symbol::intern("utils")];
771
772            assert!(
773                ast.modules.contains_key(&utils_key),
774                "module `utils` should be loaded from the in-memory file source; found keys: {:?}",
775                ast.modules.keys().collect::<Vec<_>>()
776            );
777        });
778    }
779}