Functions in libraries (#29234)

Co-authored-by: Mohammad Fawaz <mohammadfawaz@gmail.com>

Author: mohammadfawaz

crates/ast/src/library.rs

@@ -16,20 +16,22 @@
 
 use leo_span::Symbol;
 
-use crate::{Composite, ConstDeclaration, Indent};
+use crate::{Composite, ConstDeclaration, Function, Indent};
 use serde::{Deserialize, Serialize};
 use std::fmt;
 
 /// Stores the Leo library abstract syntax tree.
 ///
-/// Libraries may contain `const` declarations and `struct` definitions.
+/// Libraries may contain `const` declarations, `struct` definitions, and `fn` functions.
 #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 pub struct Library {
     pub name: Symbol,
     /// The constants defined in this library.
     pub consts: Vec<(Symbol, ConstDeclaration)>,
     /// The struct definitions in this library.
     pub structs: Vec<(Symbol, Composite)>,
+    /// The function definitions in this library.
+    pub functions: Vec<(Symbol, Function)>,
 }
 
 impl fmt::Display for Library {
@@ -40,6 +42,10 @@ impl fmt::Display for Library {
             writeln!(f, "{}", Indent(struct_def))?;
         }
 
+        for (_, func) in self.functions.iter() {
+            writeln!(f, "{}", Indent(func))?;
+        }
+
         for (_, const_decl) in self.consts.iter() {
             writeln!(f, "{};", Indent(const_decl))?;
         }

crates/ast/src/passes/reconstructor.rs

@@ -607,6 +607,7 @@ pub trait ProgramReconstructor: AstReconstructor {
                 })
                 .collect(),
             structs: input.structs.into_iter().map(|(i, s)| (i, self.reconstruct_composite(s))).collect(),
+            functions: input.functions.into_iter().map(|(i, f)| (i, self.reconstruct_function(f))).collect(),
         }
     }
 

crates/ast/src/passes/visitor.rs

@@ -345,6 +345,7 @@ pub trait ProgramVisitor: AstVisitor {
     fn visit_library(&mut self, input: &Library) {
         input.consts.iter().for_each(|(_, c)| self.visit_const(c));
         input.structs.iter().for_each(|(_, s)| self.visit_composite(s));
+        input.functions.iter().for_each(|(_, f)| self.visit_function(f));
     }
 
     fn visit_stub(&mut self, input: &Stub) {

crates/parser/src/rowan.rs

@@ -1802,6 +1802,7 @@ impl<'a> ConversionContext<'a> {
         item: &SyntaxNode,
         consts: &mut Vec<(Symbol, leo_ast::ConstDeclaration)>,
         structs: &mut Vec<(Symbol, leo_ast::Composite)>,
+        functions: &mut Vec<(Symbol, leo_ast::Function)>,
     ) -> Result<()> {
         if is_library_item(item.kind()) {
             match item.kind() {
@@ -1813,14 +1814,18 @@ impl<'a> ConversionContext<'a> {
                     let composite = self.to_composite(item)?;
                     structs.push((composite.identifier.name, composite));
                 }
-                // Future library items (functions, etc.) will be handled here.
+                FUNCTION_DEF => {
+                    // `is_in_program_block = false` so the variant is always `Fn` (not EntryPoint).
+                    let func = self.to_function(item, false)?;
+                    functions.push((func.identifier.name, func));
+                }
                 _ => {}
             }
         } else if is_program_item(item.kind()) {
             // A recognized program-only item appeared in a library file.
             let span = self.to_span(item);
             self.handler.emit_err(ParserError::custom(
-                "Only `const` declarations and `struct` definitions are allowed in a library.",
+                "Only `const` declarations, `struct` definitions, and `fn` functions are allowed in a library.",
                 span,
             ));
         }
@@ -1978,12 +1983,13 @@ impl<'a> ConversionContext<'a> {
     fn to_library(&self, name: Symbol, node: &SyntaxNode) -> Result<leo_ast::Library> {
         let mut consts = Vec::new();
         let mut structs = Vec::new();
+        let mut functions = Vec::new();
 
         for child in children(node) {
-            self.collect_library_item(&child, &mut consts, &mut structs)?;
+            self.collect_library_item(&child, &mut consts, &mut structs, &mut functions)?;
         }
 
-        Ok(leo_ast::Library { name, consts, structs })
+        Ok(leo_ast::Library { name, consts, structs, functions })
     }
 
     /// Extract a ProgramId from an IMPORT node. Guarantees `network` is always present.
@@ -3055,7 +3061,7 @@ fn compute_module_key(name: &FileName, root_dir: Option<&std::path::Path>) -> Op
 /// Currently `const` declarations and `struct` definitions are allowed; this list will grow
 /// as library support expands to include functions and other items.
 fn is_library_item(kind: SyntaxKind) -> bool {
-    matches!(kind, GLOBAL_CONST | STRUCT_DEF)
+    matches!(kind, GLOBAL_CONST | STRUCT_DEF | FUNCTION_DEF)
 }
 
 /// Returns `true` for syntax node kinds that are valid inside a program (`main.leo`).

crates/passes/src/common_subexpression_elimination/program.rs

@@ -16,9 +16,16 @@
 
 use super::CommonSubexpressionEliminatingVisitor;
 
-use leo_ast::{AstReconstructor, Constructor, Function, Module, ProgramReconstructor};
+use leo_ast::{AstReconstructor, Constructor, Function, Library, Module, ProgramReconstructor};
 
 impl ProgramReconstructor for CommonSubexpressionEliminatingVisitor<'_> {
+    fn reconstruct_library(&mut self, input: Library) -> Library {
+        // Library function bodies are no longer relevant after function inlining and may not
+        // satisfy CSE's invariants (e.g. expression type-table entries for raw function bodies).
+        // Pass functions through unchanged.
+        input
+    }
+
     fn reconstruct_program_scope(&mut self, mut input: leo_ast::ProgramScope) -> leo_ast::ProgramScope {
         input.functions = input.functions.into_iter().map(|(i, f)| (i, self.reconstruct_function(f))).collect();
         input.constructor = input.constructor.map(|c| self.reconstruct_constructor(c));

crates/passes/src/const_propagation/program.rs

@@ -99,6 +99,14 @@ impl ProgramReconstructor for ConstPropagationVisitor<'_> {
             *c = declaration;
         }
 
+        // Skip generic functions (those with const parameters): they cannot be fully evaluated
+        // until they are monomorphized into the consuming program's scope.
+        input.functions = input
+            .functions
+            .into_iter()
+            .map(|(i, f)| if f.const_parameters.is_empty() { (i, self.reconstruct_function(f)) } else { (i, f) })
+            .collect();
+
         input
     }
 

crates/passes/src/flattening/program.rs

@@ -21,13 +21,27 @@ use leo_ast::{
     Constructor,
     Expression,
     Function,
+    Library,
     ProgramReconstructor,
     ProgramScope,
     ReturnStatement,
     Statement,
 };
 
 impl ProgramReconstructor for FlatteningVisitor<'_> {
+    fn reconstruct_library(&mut self, input: Library) -> Library {
+        let prev_program = self.program;
+        self.program = input.name;
+        let library = Library {
+            name: input.name,
+            consts: input.consts,
+            structs: input.structs,
+            functions: input.functions.into_iter().map(|(i, f)| (i, self.reconstruct_function(f))).collect(),
+        };
+        self.program = prev_program;
+        library
+    }
+
     /// Flattens a program scope.
     fn reconstruct_program_scope(&mut self, input: ProgramScope) -> ProgramScope {
         self.program = input.program_id.as_symbol();

crates/passes/src/function_inlining/transform.rs

@@ -78,7 +78,11 @@ impl ProgramReconstructor for TransformVisitor<'_> {
             }
         }
 
-        // This is a sanity check to ensure that functions in the program scope have been processed.
+        // Drop any library functions that were not reachable from this program (dead library code).
+        // They are not needed and would fail the sanity check below.
+        self.function_map.retain(|loc, _| !self.state.symbol_table.is_library(loc.program));
+
+        // This is a sanity check to ensure that all non-library functions in the program scope have been processed.
         assert!(self.function_map.is_empty(), "All functions in the program should have been processed.");
 
         // Reconstruct the constructor.
@@ -87,11 +91,15 @@ impl ProgramReconstructor for TransformVisitor<'_> {
 
         // Note that this intentionally clears `self.reconstructed_functions` for the next program scope.
         let functions = core::mem::take(&mut self.reconstructed_functions)
-            .iter()
+            .into_iter()
             .filter_map(|(loc, f)| {
-                // Only consider functions defined at program scope. The rest are not relevant since they should all
-                // have been inlined by now.
-                loc.path.split_last().filter(|(_, rest)| rest.is_empty()).map(|(last, _)| (*last, f.clone()))
+                // Only emit functions that belong to the current program scope and have a single-segment
+                // path (no module prefix). Library functions and module-nested functions have all been
+                // inlined at their call sites and must not appear as standalone functions in output.
+                if loc.program != self.program {
+                    return None;
+                }
+                loc.path.split_last().filter(|(_, rest)| rest.is_empty()).map(|(last, _)| (*last, f))
             })
             .collect();
 
@@ -150,8 +158,8 @@ impl ProgramReconstructor for TransformVisitor<'_> {
     }
 
     fn reconstruct_program(&mut self, input: Program) -> Program {
-        // Populate `self.function_map` using the functions in the program scopes and the modules
-        // Use full Location (program + path) as keys to avoid collisions between programs
+        // Populate `self.function_map` using the functions in the program scopes and the modules.
+        // Use full Location (program + path) as keys to avoid collisions between programs.
         input
             .modules
             .iter()
@@ -169,6 +177,16 @@ impl ProgramReconstructor for TransformVisitor<'_> {
                 self.function_map.insert(location, f);
             });
 
+        // Populate `function_map` with library functions from FromLibrary stubs.
+        // Library functions are inlined at call sites just like module functions.
+        input.stubs.iter().for_each(|(_, stub)| {
+            if let Stub::FromLibrary { library, .. } = stub {
+                library.functions.iter().for_each(|(name, f)| {
+                    self.function_map.entry(Location::new(library.name, vec![*name])).or_insert_with(|| f.clone());
+                });
+            }
+        });
+
         // Reconstruct program scopes. Inline functions defined in modules will be traversed
         // using the call graph and reconstructed in the right order.
         let program_scopes =
@@ -204,15 +222,17 @@ impl AstReconstructor for TransformVisitor<'_> {
 
     /* Expressions */
     fn reconstruct_call(&mut self, input: CallExpression, _additional: &()) -> (Expression, Self::AdditionalOutput) {
-        // Type checking guarantees that only functions local to the program scope can be inlined.
-        if input.function.expect_global_location().program != self.program {
+        let function_location = input.function.expect_global_location();
+
+        // Pass through calls to external programs (non-library). Library functions and module
+        // functions are always inlined; only true cross-program calls are left as-is.
+        if function_location.program != self.program && !self.state.symbol_table.is_library(function_location.program) {
             return (input.into(), Default::default());
         }
 
         // Lookup the reconstructed callee function.
         // Since this pass processes functions in post-order, the callee function is guaranteed to exist in `self.reconstructed_functions`
         // Use full Location (program + path) to ensure correct lookup across multiple programs.
-        let function_location = input.function.expect_global_location();
         let (_, callee) = self
             .reconstructed_functions
             .iter()
@@ -241,6 +261,10 @@ impl AstReconstructor for TransformVisitor<'_> {
             function_location.program == self.program && function_location.path.len() > 1,
             "this is a module function",
         ) || match callee.variant {
+            // Always inline library functions (they cannot exist as standalone Aleo functions).
+            _ if self.state.symbol_table.is_library(function_location.program) => {
+                mandatory_cond(true, "this is a library function")
+            }
             Variant::FinalFn => mandatory_cond(true, "this is a final fn"),
             Variant::Fn => {
                 mandatory_cond(

crates/passes/src/global_items_collection.rs

@@ -207,6 +207,7 @@ impl ProgramVisitor for GlobalItemsCollectionVisitor<'_> {
 
         input.structs.iter().for_each(|(_, s)| self.visit_composite(s));
         input.consts.iter().for_each(|(_, c)| self.visit_const(c));
+        input.functions.iter().for_each(|(_, f)| self.visit_function(f));
     }
 
     fn visit_aleo_program(&mut self, input: &AleoProgram) {

crates/passes/src/monomorphization/ast.rs

@@ -26,7 +26,7 @@ use leo_ast::{
     Expression,
     Identifier,
     Node as _,
-    ProgramReconstructor,
+    ProgramReconstructor as _,
     Type,
 };
 
@@ -121,8 +121,12 @@ impl AstReconstructor for MonomorphizationVisitor<'_> {
         input_call: CallExpression,
         _additional: &(),
     ) -> (Expression, Self::AdditionalOutput) {
-        // Skip calls to functions from other programs.
-        if input_call.function.expect_global_location().program != self.program {
+        let callee_loc = input_call.function.expect_global_location();
+        let callee_program = callee_loc.program;
+
+        // Skip calls to external programs that are not libraries (cross-program calls).
+        // Library functions and current-program functions (including modules) proceed below.
+        if callee_program != self.program && !self.state.symbol_table.is_library(callee_program) {
             return (input_call.into(), Default::default());
         }