Allow imports & Variant::Function functions in stubs

compiler/ast/src/passes/reconstructor.rs

@@ -428,6 +428,7 @@ pub trait ProgramReconstructor: StatementReconstructor {
 
     fn reconstruct_stub(&mut self, input: Stub) -> Stub {
         Stub {
+            imports: input.imports,
             stub_id: input.stub_id,
             consts: input.consts,
             structs: input.structs,

compiler/ast/src/stub/mod.rs

@@ -17,19 +17,22 @@
 //! A stub contains function templates as well as definitions for mappings, structs, records, and constants.
 
 pub mod function_stub;
+
 pub use function_stub::*;
 
-use crate::{ConstDeclaration, Mapping, ProgramId, ProgramScope, Struct};
+use crate::{ConstDeclaration, Identifier, Mapping, ProgramId, Struct};
 use leo_span::{Span, Symbol};
 use serde::{Deserialize, Serialize};
 use std::fmt;
 
 /// Stores the Leo stub abstract syntax tree.
 #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 pub struct Stub {
+    /// A vector of imported programs.
+    pub imports: Vec<Identifier>,
     /// The stub id
     pub stub_id: ProgramId,
-    /// A vector of const definitions
+    /// A vector of const definitions.
     pub consts: Vec<(Symbol, ConstDeclaration)>,
     /// A vector of struct definitions.
     pub structs: Vec<(Symbol, Struct)>,
@@ -41,23 +44,6 @@ pub struct Stub {
     pub span: Span,
 }
 
-impl From<ProgramScope> for Stub {
-    fn from(program_scope: ProgramScope) -> Self {
-        Self {
-            stub_id: program_scope.program_id,
-            consts: program_scope.consts,
-            structs: program_scope.structs,
-            mappings: program_scope.mappings,
-            functions: program_scope
-                .functions
-                .into_iter()
-                .map(|(symbol, function)| (symbol, FunctionStub::from(function)))
-                .collect(),
-            span: program_scope.span,
-        }
-    }
-}
-
 impl fmt::Display for Stub {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         writeln!(f, "stub {} {{", self.stub_id)?;

compiler/parser/src/parser/file.rs

@@ -120,7 +120,7 @@ impl ParserContext<'_> {
     }
 
     /// Parses a program body `credits.aleo { ... }`
-    fn parse_program_body(&mut self, start: Span) -> Result<ProgramScope> {
+    fn parse_program_body(&mut self, start: Span) -> Result<(ProgramScope, Vec<Identifier>)> {
         // Parse the program name.
         let name = self.expect_identifier()?;
 
@@ -142,9 +142,18 @@ impl ParserContext<'_> {
         let mut functions: Vec<(Symbol, Function)> = Vec::new();
         let mut structs: Vec<(Symbol, Struct)> = Vec::new();
         let mut mappings: Vec<(Symbol, Mapping)> = Vec::new();
+        let mut imports: Vec<Identifier> = Vec::new();
 
         while self.has_next() {
             match &self.token.token {
+                Token::Import => {
+                    self.expect(&Token::Import)?;
+                    let import_name = self.expect_identifier()?;
+                    self.expect(&Token::Dot)?;
+                    self.expect(&Token::Aleo)?;
+                    self.expect(&Token::Semicolon)?;
+                    imports.push(import_name);
+                }
                 Token::Const => {
                     let declaration = self.parse_const_declaration_statement()?;
                     consts.push((Symbol::intern(&declaration.place.to_string()), declaration));
@@ -180,24 +189,39 @@ impl ParserContext<'_> {
         // Parse `}`.
         let end = self.expect(&Token::RightCurly)?;
 
-        Ok(ProgramScope { program_id, consts, functions, structs, mappings, span: start + end })
+        Ok((ProgramScope { program_id, consts, functions, structs, mappings, span: start + end }, imports))
     }
 
     /// Parses a stub `stub credits.aleo { ... }`.
     fn parse_stub(&mut self) -> Result<(Symbol, Stub)> {
         // Parse `stub` keyword.
         let start = self.expect(&Token::Stub)?;
-        let stub = Stub::from(self.parse_program_body(start)?);
-
-        Ok((stub.stub_id.name.name, stub))
+        let (program_body, imports) = self.parse_program_body(start)?;
+
+        Ok((program_body.program_id.name.name, Stub {
+            imports,
+            stub_id: program_body.program_id,
+            consts: program_body.consts,
+            structs: program_body.structs,
+            mappings: program_body.mappings,
+            functions: program_body
+                .functions
+                .into_iter()
+                .map(|(symbol, function)| (symbol, FunctionStub::from(function)))
+                .collect(),
+            span: program_body.span,
+        }))
     }
 
     /// Parses a program scope `program foo.aleo { ... }`.
     fn parse_program_scope(&mut self) -> Result<ProgramScope> {
         // Parse `program` keyword.
         let start = self.expect(&Token::Program)?;
-
-        self.parse_program_body(start)
+        let (program_scope, imports) = self.parse_program_body(start)?;
+        if imports.len() != 0 {
+            self.emit_err(ParserError::cannot_import_inside_program_body(self.token.span));
+        }
+        Ok(program_scope)
     }
 
     /// Returns a [`Vec<Member>`] AST node if the next tokens represent a struct member.

compiler/passes/src/type_checking/check_program.rs

@@ -88,8 +88,8 @@ impl<'a> ProgramVisitor<'a> for TypeChecker<'a> {
         }
 
         // Must be transition functions
-        if input.variant != Variant::Transition {
-            self.emit_err(TypeCheckerError::stub_functions_must_be_transitions(input.span));
+        if input.variant == Variant::Inline {
+            self.emit_err(TypeCheckerError::stub_functions_must_not_be_inlines(input.span));
         }
 
         // Must be empty

errors/src/errors/parser/parser_errors.rs

@@ -308,6 +308,14 @@ create_messages!(
         help: None,
     }
 
+    /// Enforce that cannot use import in program scope
+    @formatted
+    cannot_import_inside_program_body {
+        args: (),
+        msg: format!("Cannot use import inside program body."),
+        help: None,
+    }
+
     @formatted
     array_must_have_at_least_one_element {
         args: (kind: impl Display),

errors/src/errors/type_checker/type_checker_error.rs

@@ -688,14 +688,14 @@ create_messages!(
     @formatted
     stubs_can_only_have_records_and_transitions {
         args: (found: impl Display),
-        msg: format!("Stubs can only have records and transitions -- found {found}"),
+        msg: format!("Stubs can only have records, transitions, functions and imports -- found {found}"),
         help: None,
     }
 
     @formatted
-    stub_functions_must_be_transitions {
+    stub_functions_must_not_be_inlines {
         args: (),
-        msg: format!("Function stubs must be transitions"),
+        msg: format!("Function stubs must be transitions or function variants not inlines"),
         help: None,
     }
 

tests/expectations/parser/program/illegal_import_fail.out

@@ -0,0 +1,5 @@
+---
+namespace: Parse
+expectation: Fail
+outputs:
+  - "Error [EPAR0370036]: Cannot use import inside program body.\n    --> test:13:1\n     |\n  13 | }\n     | ^"