Repository: reubeno/brush (rust)

Issue:
support exporting/importing functions across shell instances
`bash` supports exporting functions to (and inherited exported functions from) other instances of `bash`. Under the hood, exported functions get serialized to env variables with the prefix `BASH_FUNC_`.

`brush` should be able to:

* Inherit functions exported by `bash`.
* Export functions in a way that will be used by a child invocation of `bash`.
* Export functions in a way that will be used by a child invocation of `brush`.

Pull request description:
feat(funcs): implement function exporting/importing
* Implements `export -f`
* Implements `export -n`
* Implements serialization/deserialization of exported functions via `BASH_FUNC_<name>%%` variables

Resolves #489

Patch:
diff --git a/brush-core/benches/shell.rs b/brush-core/benches/shell.rs
index 83b133c..58fa94a 100644
--- a/brush-core/benches/shell.rs
+++ b/brush-core/benches/shell.rs
@@ -1,7 +1,5 @@
 #[cfg(unix)]
 mod unix {
-    use std::sync::Arc;
-
     use criterion::{black_box, Criterion};
 
     async fn instantiate_shell() -> brush_core::Shell {
@@ -107,7 +105,7 @@ mod unix {
         let mut shell = rt.block_on(instantiate_shell());
         shell.funcs.update(
             String::from("testfunc"),
-            Arc::new(brush_parser::ast::FunctionDefinition {
+            brush_parser::ast::FunctionDefinition {
                 fname: String::from("testfunc"),
                 body: brush_parser::ast::FunctionBody(
                     brush_parser::ast::CompoundCommand::BraceGroup(
@@ -118,7 +116,8 @@ mod unix {
                     None,
                 ),
                 source: String::from("/some/path"),
-            }),
+            }
+            .into(),
         );
         c.bench_function("function_call", |b| {
             b.iter_batched_ref(
diff --git a/brush-core/src/builtins/export.rs b/brush-core/src/builtins/export.rs
index dc521c1..6447b90 100644
--- a/brush-core/src/builtins/export.rs
+++ b/brush-core/src/builtins/export.rs
@@ -5,7 +5,7 @@ use std::io::Write;
 use crate::{
     builtins, commands,
     env::{EnvironmentLookup, EnvironmentScope},
-    variables,
+    error, variables,
 };
 
 /// Add or update exported shell variables.
@@ -40,70 +40,114 @@ impl builtins::DeclarationCommand for ExportCommand {
 impl builtins::Command for ExportCommand {
     async fn execute(
         &self,
-        context: commands::ExecutionContext<'_>,
+        mut context: commands::ExecutionContext<'_>,
     ) -> Result<crate::builtins::ExitCode, crate::error::Error> {
-        if !self.declarations.is_empty() {
-            for decl in &self.declarations {
-                match decl {
-                    commands::CommandArg::String(s) => {
-                        // Try to find the variable already present; if we find it, then mark it
-                        // exported.
-                        if let Some((_, variable)) = context.shell.env.get_mut(s) {
-                            variable.export();
-                        }
-                    }
-                    commands::CommandArg::Assignment(assignment) => {
-                        let name = match &assignment.name {
-                            brush_parser::ast::AssignmentName::VariableName(name) => name,
-                            brush_parser::ast::AssignmentName::ArrayElementName(_, _) => {
-                                writeln!(context.stderr(), "not a valid variable name")?;
-                                return Ok(builtins::ExitCode::InvalidUsage);
-                            }
-                        };
+        if self.declarations.is_empty() {
+            display_all_exported_vars(&context)?;
+            return Ok(builtins::ExitCode::Success);
+        }
 
-                        let value = match &assignment.value {
-                            brush_parser::ast::AssignmentValue::Scalar(s) => {
-                                variables::ShellValueLiteral::Scalar(s.flatten())
-                            }
-                            brush_parser::ast::AssignmentValue::Array(a) => {
-                                variables::ShellValueLiteral::Array(variables::ArrayLiteral(
-                                    a.iter()
-                                        .map(|(k, v)| {
-                                            (k.as_ref().map(|k| k.flatten()), v.flatten())
-                                        })
-                                        .collect(),
-                                ))
-                            }
-                        };
+        let mut exit_code = builtins::ExitCode::Success;
+        for decl in &self.declarations {
+            let current_exit_code = self.process_decl(&mut context, decl)?;
+            if !matches!(current_exit_code, builtins::ExitCode::Success) {
+                exit_code = current_exit_code;
+            }
+        }
 
-                        // Update the variable with the provided value and then mark it exported.
-                        context.shell.env.update_or_add(
-                            name,
-                            value,
-                            |var| {
-                                var.export();
-                                Ok(())
-                            },
-                            EnvironmentLookup::Anywhere,
-                            EnvironmentScope::Global,
-                        )?;
+        Ok(exit_code)
+    }
+}
+
+impl ExportCommand {
+    fn process_decl(
+        &self,
+        context: &mut commands::ExecutionContext<'_>,
+        decl: &commands::CommandArg,
+    ) -> Result<builtins::ExitCode, error::Error> {
+        match decl {
+            commands::CommandArg::String(s) => {
+                // See if this is supposed to be a function name.
+                if self.names_are_functions {
+                    // Try to find the function already present; if we find it, then mark it
+                    // exported.
+                    if let Some(func) = context.shell.funcs.get_mut(s) {
+                        if self.unexport {
+                            func.unexport();
+                        } else {
+                            func.export();
+                        }
+                    } else {
+                        writeln!(context.stderr(), "{s}: not a function")?;
+                        return Ok(builtins::ExitCode::InvalidUsage);
                     }
                 }
-            }
-        } else {
-            // Enumerate variables, sorted by key.
-            for (name, variable) in context.shell.env.iter().sorted_by_key(|v| v.0) {
-                if variable.is_exported() {
-                    let value = variable.value().try_get_cow_str(context.shell);
-                    if let Some(value) = value {
-                        writeln!(context.stdout(), "declare -x {name}=\"{value}\"")?;
+                // Try to find the variable already present; if we find it, then mark it
+                // exported.
+                else if let Some((_, variable)) = context.shell.env.get_mut(s) {
+                    if self.unexport {
+                        variable.unexport();
                     } else {
-                        writeln!(context.stdout(), "declare -x {name}")?;
+                        variable.export();
                     }
                 }
             }
+            commands::CommandArg::Assignment(assignment) => {
+                let name = match &assignment.name {
+                    brush_parser::ast::AssignmentName::VariableName(name) => name,
+                    brush_parser::ast::AssignmentName::ArrayElementName(_, _) => {
+                        writeln!(context.stderr(), "not a valid variable name")?;
+                        return Ok(builtins::ExitCode::InvalidUsage);
+                    }
+                };
+
+                let value = match &assignment.value {
+                    brush_parser::ast::AssignmentValue::Scalar(s) => {
+                        variables::ShellValueLiteral::Scalar(s.flatten())
+                    }
+                    brush_parser::ast::AssignmentValue::Array(a) => {
+                        variables::ShellValueLiteral::Array(variables::ArrayLiteral(
+                            a.iter()
+                                .map(|(k, v)| (k.as_ref().map(|k| k.flatten()), v.flatten()))
+                                .collect(),
+                        ))
+                    }
+                };
+
+                // Update the variable with the provided value and then mark it exported.
+                context.shell.env.update_or_add(
+                    name,
+                    value,
+                    |var| {
+                        if self.unexport {
+                            var.unexport();
+                        } else {
+                            var.export();
+                        }
+                        Ok(())
+                    },
+                    EnvironmentLookup::Anywhere,
+                    EnvironmentScope::Global,
+                )?;
+            }
         }
 
         Ok(builtins::ExitCode::Success)
     }
 }
+
+fn display_all_exported_vars(context: &commands::ExecutionContext<'_>) -> Result<(), error::Error> {
+    // Enumerate variables, sorted by key.
+    for (name, variable) in context.shell.env.iter().sorted_by_key(|v| v.0) {
+        if variable.is_exported() {
+            let value = variable.value().try_get_cow_str(context.shell);
+            if let Some(value) = value {
+                writeln!(context.stdout(), "declare -x {name}=\"{value}\"")?;
+            } else {
+                writeln!(context.stdout(), "declare -x {name}")?;
+            }
+        }
+    }
+
+    Ok(())
+}
diff --git a/brush-core/src/commands.rs b/brush-core/src/commands.rs
index 34ce765..a0584b2 100644
--- a/brush-core/src/commands.rs
+++ b/brush-core/src/commands.rs
@@ -217,6 +217,17 @@ pub(crate) fn compose_std_command<S: AsRef<OsStr>>(
         }
     }
 
+    // Add in exported functions.
+    if !empty_env {
+        for (func_name, registration) in shell.funcs.iter() {
+            if registration.is_exported() {
+                let var_name = std::format!("BASH_FUNC_{func_name}%%");
+                let value = std::format!("() {}", registration.definition.body);
+                cmd.env(var_name, value);
+            }
+        }
+    }
+
     // Redirect stdin, if applicable.
     match open_files.files.remove(&0) {
         Some(OpenFile::Stdin) | None => (),
diff --git a/brush-core/src/error.rs b/brush-core/src/error.rs
index 46cf192..d31381d 100644
--- a/brush-core/src/error.rs
+++ b/brush-core/src/error.rs
@@ -139,6 +139,10 @@ pub enum Error {
     #[error("{0}")]
     FormattingError(#[from] std::fmt::Error),
 
+    /// An error occurred while parsing.
+    #[error("{0}")]
+    ParseError(#[from] brush_parser::ParseError),
+
     /// An error occurred while parsing a word.
     #[error("{0}")]
     WordParseError(#[from] brush_parser::WordParseError),
diff --git a/brush-core/src/functions.rs b/brush-core/src/functions.rs
index 851d0b8..e37b4fe 100644
--- a/brush-core/src/functions.rs
+++ b/brush-core/src/functions.rs
@@ -18,6 +18,16 @@ impl FunctionEnv {
         self.functions.get(name)
     }
 
+    /// Tries to retrieve a mutable reference to the registration for a
+    /// function by name.
+    ///
+    /// # Arguments
+    ///
+    /// * `name` - The name of the function to retrieve.
+    pub fn get_mut(&mut self, name: &str) -> Option<&mut FunctionRegistration> {
+        self.functions.get_mut(name)
+    }
+
     /// Unregisters a function from the environment.
     ///
     /// # Arguments
@@ -32,10 +42,9 @@ impl FunctionEnv {
     /// # Arguments
     ///
     /// * `name` - The name of the function to update.
-    /// * `definition` - The new definition for the function.
-    pub fn update(&mut self, name: String, definition: Arc<brush_parser::ast::FunctionDefinition>) {
-        self.functions
-            .insert(name, FunctionRegistration { definition });
+    /// * `registration` - The new registration for the function.
+    pub fn update(&mut self, name: String, registration: FunctionRegistration) {
+        self.functions.insert(name, registration);
     }
 
     /// Returns an iterator over the functions registered in this environment.
@@ -49,4 +58,32 @@ impl FunctionEnv {
 pub struct FunctionRegistration {
     /// The definition of the function.
     pub(crate) definition: Arc<brush_parser::ast::FunctionDefinition>,
+    /// Whether or not this function definition should be exported to children.
+    exported: bool,
+}
+
+impl From<brush_parser::ast::FunctionDefinition> for FunctionRegistration {
+    fn from(definition: brush_parser::ast::FunctionDefinition) -> Self {
+        FunctionRegistration {
+            definition: Arc::new(definition),
+            exported: false,
+        }
+    }
+}
+
+impl FunctionRegistration {
+    /// Marks the function for export.
+    pub fn export(&mut self) {
+        self.exported = true;
+    }
+
+    /// Unmarks the function for export.
+    pub fn unexport(&mut self) {
+        self.exported = false;
+    }
+
+    /// Returns whether this function is exported.
+    pub fn is_exported(&self) -> bool {
+        self.exported
+    }
 }
diff --git a/brush-core/src/interp.rs b/brush-core/src/interp.rs
index 7785b37..54d0637 100644
--- a/brush-core/src/interp.rs
+++ b/brush-core/src/interp.rs
@@ -7,7 +7,6 @@ use std::os::fd::AsFd;
 #[cfg(unix)]
 use std::os::unix::process::ExitStatusExt;
 use std::path::{Path, PathBuf};
-use std::sync::Arc;
 
 use crate::arithmetic::{self, ExpandAndEvaluate};
 use crate::commands::{self, CommandArg, CommandSpawnResult};
@@ -17,7 +16,9 @@ use crate::shell::Shell;
 use crate::variables::{
     ArrayLiteral, ShellValue, ShellValueLiteral, ShellValueUnsetType, ShellVariable,
 };
-use crate::{error, expansion, extendedtests, jobs, openfiles, processes, sys, timing, traps};
+use crate::{
+    error, expansion, extendedtests, functions, jobs, openfiles, processes, sys, timing, traps,
+};
 
 /// Encapsulates the result of executing a command.
 #[derive(Debug, Default)]
@@ -907,9 +908,10 @@ impl Execute for ast::FunctionDefinition {
         shell: &mut Shell,
         _params: &ExecutionParameters,
     ) -> Result<ExecutionResult, error::Error> {
-        shell
-            .funcs
-            .update(self.fname.clone(), Arc::new(self.clone()));
+        shell.funcs.update(
+            self.fname.clone(),
+            functions::FunctionRegistration::from(self.clone()),
+        );
 
         let result = ExecutionResult::success();
         shell.last_exit_status = result.exit_code;
diff --git a/brush-core/src/shell.rs b/brush-core/src/shell.rs
index 1ff146b..1075025 100644
--- a/brush-core/src/shell.rs
+++ b/brush-core/src/shell.rs
@@ -257,12 +257,44 @@ impl Shell {
         Ok(shell)
     }
 
+    fn initialize_exported_func(
+        &mut self,
+        func_name: &str,
+        body_text: &str,
+    ) -> Result<(), error::Error> {
+        let mut parser = create_parser(body_text, &self.parser_options());
+        let func_body = parser.parse_function_parens_and_body()?;
+
+        let func_def = brush_parser::ast::FunctionDefinition {
+            fname: func_name.to_owned(),
+            body: func_body,
+            source: String::new(),
+        };
+
+        let mut registration = functions::FunctionRegistration::from(func_def);
+        registration.export();
+
+        self.funcs.update(func_name.to_owned(), registration);
+
+        Ok(())
+    }
+
     #[allow(clippy::too_many_lines)]
     #[allow(clippy::unwrap_in_result)]
     fn initialize_vars(&mut self, options: &CreateOptions) -> Result<(), error::Error> {
         // Seed parameters from environment (unless requested not to do so).
         if !options.do_not_inherit_env {
             for (k, v) in std::env::vars() {
+                // See if it's a function exported by an ancestor process.
+                if let Some(func_name) = k.strip_prefix("BASH_FUNC_") {
+                    if let Some(func_name) = func_name.strip_suffix("%%") {
+                        // Intentionally best-effort; don't fail out of the shell if we can't
+                        // parse an incoming function.
+                        let _ = self.initialize_exported_func(func_name, v.as_str());
+                        continue;
+                    }
+                }
+
                 let mut var = ShellVariable::new(ShellValue::String(v));
                 var.export();
                 self.env.set_global(k, var)?;
@@ -850,7 +882,7 @@ impl Shell {
             brush_parser::Parser::new(&mut reader, &self.parser_options(), source_info);
 
         tracing::debug!(target: trace_categories::PARSE, "Parsing sourced file: {}", source_info.source);
-        let parse_result = parser.parse();
+        let parse_result = parser.parse_program();
 
         let mut other_positional_parameters = args.map(|s| s.as_ref().to_owned()).collect();
         let mut other_shell_name = Some(source_info.source.clone());
@@ -1688,15 +1720,22 @@ fn parse_string_impl(
     s: String,
     parser_options: brush_parser::ParserOptions,
 ) -> Result<brush_parser::ast::Program, brush_parser::ParseError> {
-    let mut reader = std::io::BufReader::new(s.as_bytes());
+    let mut parser = create_parser(s.as_str(), &parser_options);
+
+    tracing::debug!(target: trace_categories::PARSE, "Parsing string as program...");
+    parser.parse_program()
+}
+
+fn create_parser<'a>(
+    s: &'a str,
+    parser_options: &brush_parser::ParserOptions,
+) -> brush_parser::Parser<std::io::BufReader<&'a [u8]>> {
+    let reader = std::io::BufReader::new(s.as_bytes());
     let source_info = brush_parser::SourceInfo {
         source: String::from("main"),
     };
-    let mut parser: brush_parser::Parser<&mut std::io::BufReader<&[u8]>> =
-        brush_parser::Parser::new(&mut reader, &parser_options, &source_info);
 
-    tracing::debug!(target: trace_categories::PARSE, "Parsing string as program...");
-    parser.parse()
+    brush_parser::Parser::new(reader, parser_options, &source_info)
 }
 
 fn repeated_char_str(c: char, count: usize) -> String {
diff --git a/brush-parser/src/parser.rs b/brush-parser/src/parser.rs
index edcb4b7..58e3022 100644
--- a/brush-parser/src/parser.rs
+++ b/brush-parser/src/parser.rs
@@ -61,7 +61,7 @@ impl<R: std::io::BufRead> Parser<R> {
     }
 
     /// Parses the input into an abstract syntax tree (AST) of a shell program.
-    pub fn parse(&mut self) -> Result<ast::Program, error::ParseError> {
+    pub fn parse_program(&mut self) -> Result<ast::Program, error::ParseError> {
         //
         // References:
         //   * https://www.gnu.org/software/bash/manual/bash.html#Shell-Syntax
@@ -70,6 +70,25 @@ impl<R: std::io::BufRead> Parser<R> {
         //   * https://pubs.opengroup.org/onlinepubs/9699919799/utilities/V3_chap02.html
         //
 
+        let tokens = self.tokenize()?;
+        parse_tokens(&tokens, &self.options, &self.source_info)
+    }
+
+    /// Parses a function definition body from the input. The body is expected to be
+    /// preceded by "()", but no function name.
+    pub fn parse_function_parens_and_body(
+        &mut self,
+    ) -> Result<ast::FunctionBody, error::ParseError> {
+        let tokens = self.tokenize()?;
+        let parse_result = token_parser::function_parens_and_body(
+            &Tokens { tokens: &tokens },
+            &self.options,
+            &self.source_info,
+        );
+        parse_result_to_error(parse_result, &tokens)
+    }
+
+    fn tokenize(&mut self) -> Result<Vec<Token>, error::ParseError> {
         // First we tokenize the input, according to the policy implied by provided options.
         let mut tokenizer = Tokenizer::new(&mut self.reader, &self.options.tokenizer_options());
 
@@ -100,7 +119,7 @@ impl<R: std::io::BufRead> Parser<R> {
 
         tracing::debug!(target: "tokenize", "  => {} token(s)", tokens.len());
 
-        parse_tokens(&tokens, &self.options, &self.source_info)
+        Ok(tokens)
     }
 }
 
@@ -117,7 +136,16 @@ pub fn parse_tokens(
     source_info: &SourceInfo,
 ) -> Result<ast::Program, error::ParseError> {
     let parse_result = token_parser::program(&Tokens { tokens }, options, source_info);
+    parse_result_to_error(parse_result, tokens)
+}
 
+fn parse_result_to_error<R>(
+    parse_result: Result<R, peg::error::ParseError<usize>>,
+    tokens: &Vec<Token>,
+) -> Result<R, error::ParseError>
+where
+    R: std::fmt::Debug,
+{
     let result = match parse_result {
         Ok(program) => {
             tracing::debug!(target: "parse", "PROG: {:?}", program);
@@ -539,7 +567,7 @@ peg::parser! {
 
         // N.B. Non-sh extensions allows use of the 'function' word to indicate a function definition.
         rule function_definition() -> ast::FunctionDefinition =
-            specific_word("function")? fname:fname() specific_operator("(") specific_operator(")") linebreak() body:function_body() {
+            specific_word("function")? fname:fname() body:function_parens_and_body() {
                 ast::FunctionDefinition { fname: fname.to_owned(), body, source: source_info.source.clone() }
             } /
             specific_word("function") fname:fname() linebreak() body:function_body() {
@@ -547,6 +575,9 @@ peg::parser! {
             } /
             expected!("function definition")
 
+        pub(crate) rule function_parens_and_body() -> ast::FunctionBody =
+            specific_operator("(") specific_operator(")") linebreak() body:function_body() { body }
+
         rule function_body() -> ast::FunctionBody =
             c:compound_command() r:redirect_list()? { ast::FunctionBody(c, r) }
 


Test patch:
diff --git a/brush-shell/tests/cases/builtins/export.yaml b/brush-shell/tests/cases/builtins/export.yaml
index 3212608..68b78f5 100644
--- a/brush-shell/tests/cases/builtins/export.yaml
+++ b/brush-shell/tests/cases/builtins/export.yaml
@@ -18,3 +18,11 @@ cases:
     stdin: |
       export arr=(a 1 2)
       declare -p arr
+
+  - name: "Unexporting variable"
+    stdin: |
+      MY_TEST_VAR="value"
+      export MY_TEST_VAR="value"
+      export -n MY_TEST_VAR
+
+      env | grep MY_TEST_VAR
diff --git a/brush-shell/tests/cases/functions.yaml b/brush-shell/tests/cases/functions.yaml
index 87e2ff7..e3a4f0e 100644
--- a/brush-shell/tests/cases/functions.yaml
+++ b/brush-shell/tests/cases/functions.yaml
@@ -90,3 +90,13 @@ cases:
 
       echo "Calling inner from toplevel"
       inner
+
+  - name: "Exporting functions to child instance"
+    stdin: |
+      mytestfunc() {
+          echo "In mytestfunc"
+      }
+
+      export -f mytestfunc
+
+      $0 -c 'mytestfunc'