Merge pull request #329 from schungx/master

Make Engine::compile_into_self_contained recursive.

RELEASES.md

@@ -10,11 +10,13 @@ Breaking changes
 * The error variant `EvalAltResult::ErrorInFunctionCall` has a new parameter holding the _source_ of the function.
 * `ParseErrorType::WrongFnDefinition` is renamed `FnWrongDefinition`.
 * Redefining an existing function within the same script now throws a new `ParseErrorType::FnDuplicatedDefinition`. This is to prevent accidental overwriting an earlier function definition.
+* `AST::set_source` is now split into `AST::set_source` and `AST::clear_source`.
 
 New features
 ------------
 
 * `Engine::compile_to_self_contained` compiles a script into an `AST` and _eagerly_ resolves all `import` statements with string literal paths. The resolved modules are directly embedded into the `AST`. When the `AST` is later evaluated, `import` statements directly yield the pre-resolved modules without going through the resolution process once again.
+* `AST::walk`, `Stmt::walk` and `Expr::walk` internal API's to recursively walk an `AST`.
 
 Enhancements
 ------------

src/ast.rs

@@ -218,24 +218,32 @@ impl AST {
             resolver: None,
         }
     }
-    /// Get the source.
+    /// Get the source, if any.
     #[inline(always)]
     pub fn source(&self) -> Option<&str> {
         self.source.as_ref().map(|s| s.as_str())
     }
-    /// Clone the source.
+    /// Clone the source, if any.
     #[inline(always)]
     pub(crate) fn clone_source(&self) -> Option<ImmutableString> {
         self.source.clone()
     }
     /// Set the source.
     #[inline(always)]
-    pub fn set_source<S: Into<ImmutableString>>(&mut self, source: Option<S>) {
+    pub fn set_source(&mut self, source: impl Into<ImmutableString>) -> &mut Self {
-        self.source = source.map(|s| s.into());
+        self.source = Some(source.into());
 
         if let Some(module) = Shared::get_mut(&mut self.functions) {
             module.set_id(self.source.clone());
         }
+
+        self
+    }
+    /// Clear the source.
+    #[inline(always)]
+    pub fn clear_source(&mut self) -> &mut Self {
+        self.source = None;
+        self
     }
     /// Get the statements.
     #[cfg(not(feature = "internals"))]
@@ -674,6 +682,44 @@ impl AST {
     pub fn clear_statements(&mut self) {
         self.statements = vec![];
     }
+    /// Recursively walk the [`AST`], including function bodies (if any).
+    #[cfg(not(feature = "internals"))]
+    #[cfg(not(feature = "no_module"))]
+    #[inline(always)]
+    pub(crate) fn walk(&self, on_node: &mut impl FnMut(&[ASTNode])) {
+        self.statements()
+            .iter()
+            .chain({
+                #[cfg(not(feature = "no_function"))]
+                {
+                    self.iter_fn_def().map(|f| &f.body)
+                }
+                #[cfg(feature = "no_function")]
+                {
+                    crate::stdlib::iter::empty()
+                }
+            })
+            .for_each(|stmt| stmt.walk(&mut Default::default(), on_node));
+    }
+    /// _(INTERNALS)_ Recursively walk the [`AST`], including function bodies (if any).
+    /// Exported under the `internals` feature only.
+    #[cfg(feature = "internals")]
+    #[inline(always)]
+    pub fn walk(&self, on_node: &mut impl FnMut(&[ASTNode])) {
+        self.statements()
+            .iter()
+            .chain({
+                #[cfg(not(feature = "no_function"))]
+                {
+                    self.iter_fn_def().map(|f| &f.body)
+                }
+                #[cfg(feature = "no_function")]
+                {
+                    crate::stdlib::iter::empty()
+                }
+            })
+            .for_each(|stmt| stmt.walk(&mut Default::default(), on_node));
+    }
 }
 
 impl<A: AsRef<AST>> Add<A> for &AST {
@@ -741,6 +787,30 @@ pub enum ReturnType {
     Exception,
 }
 
+/// _(INTERNALS)_ An [`AST`] node, consisting of either an [`Expr`] or a [`Stmt`].
+/// Exported under the `internals` feature only.
+///
+/// # WARNING
+///
+/// This type is volatile and may change.
+#[derive(Debug, Clone, Hash)]
+pub enum ASTNode<'a> {
+    Stmt(&'a Stmt),
+    Expr(&'a Expr),
+}
+
+impl<'a> From<&'a Stmt> for ASTNode<'a> {
+    fn from(stmt: &'a Stmt) -> Self {
+        Self::Stmt(stmt)
+    }
+}
+
+impl<'a> From<&'a Expr> for ASTNode<'a> {
+    fn from(expr: &'a Expr) -> Self {
+        Self::Expr(expr)
+    }
+}
+
 /// _(INTERNALS)_ A statement.
 /// Exported under the `internals` feature only.
 ///
@@ -941,50 +1011,50 @@ impl Stmt {
     }
     /// Recursively walk this statement.
     #[inline(always)]
-    pub fn walk(&self, process_stmt: &mut impl FnMut(&Stmt), process_expr: &mut impl FnMut(&Expr)) {
+    pub fn walk<'a>(&'a self, path: &mut Vec<ASTNode<'a>>, on_node: &mut impl FnMut(&[ASTNode])) {
-        process_stmt(self);
+        path.push(self.into());
+        on_node(path);
 
         match self {
-            Self::Let(_, Some(e), _, _) | Self::Const(_, Some(e), _, _) => {
+            Self::Let(_, Some(e), _, _) | Self::Const(_, Some(e), _, _) => e.walk(path, on_node),
-                e.walk(process_stmt, process_expr)
-            }
             Self::If(e, x, _) => {
-                e.walk(process_stmt, process_expr);
+                e.walk(path, on_node);
-                x.0.walk(process_stmt, process_expr);
+                x.0.walk(path, on_node);
                 if let Some(ref s) = x.1 {
-                    s.walk(process_stmt, process_expr);
+                    s.walk(path, on_node);
                 }
             }
             Self::Switch(e, x, _) => {
-                e.walk(process_stmt, process_expr);
+                e.walk(path, on_node);
-                x.0.values()
+                x.0.values().for_each(|s| s.walk(path, on_node));
-                    .for_each(|s| s.walk(process_stmt, process_expr));
                 if let Some(ref s) = x.1 {
-                    s.walk(process_stmt, process_expr);
+                    s.walk(path, on_node);
                 }
             }
             Self::While(e, s, _) | Self::Do(s, e, _, _) => {
-                e.walk(process_stmt, process_expr);
+                e.walk(path, on_node);
-                s.walk(process_stmt, process_expr);
+                s.walk(path, on_node);
             }
             Self::For(e, x, _) => {
-                e.walk(process_stmt, process_expr);
+                e.walk(path, on_node);
-                x.1.walk(process_stmt, process_expr);
+                x.1.walk(path, on_node);
             }
             Self::Assignment(x, _) => {
-                x.0.walk(process_stmt, process_expr);
+                x.0.walk(path, on_node);
-                x.2.walk(process_stmt, process_expr);
+                x.2.walk(path, on_node);
             }
-            Self::Block(x, _) => x.iter().for_each(|s| s.walk(process_stmt, process_expr)),
+            Self::Block(x, _) => x.iter().for_each(|s| s.walk(path, on_node)),
             Self::TryCatch(x, _, _) => {
-                x.0.walk(process_stmt, process_expr);
+                x.0.walk(path, on_node);
-                x.2.walk(process_stmt, process_expr);
+                x.2.walk(path, on_node);
             }
-            Self::Expr(e) | Self::Return(_, Some(e), _) => e.walk(process_stmt, process_expr),
+            Self::Expr(e) | Self::Return(_, Some(e), _) => e.walk(path, on_node),
             #[cfg(not(feature = "no_module"))]
-            Self::Import(e, _, _) => e.walk(process_stmt, process_expr),
+            Self::Import(e, _, _) => e.walk(path, on_node),
             _ => (),
         }
+
+        path.pop().unwrap();
     }
 }
 
@@ -1408,25 +1478,23 @@ impl Expr {
     }
     /// Recursively walk this expression.
     #[inline(always)]
-    pub fn walk(&self, process_stmt: &mut impl FnMut(&Stmt), process_expr: &mut impl FnMut(&Expr)) {
+    pub fn walk<'a>(&'a self, path: &mut Vec<ASTNode<'a>>, on_node: &mut impl FnMut(&[ASTNode])) {
-        process_expr(self);
+        path.push(self.into());
+        on_node(path);
 
         match self {
-            Self::Stmt(x, _) => x.iter().for_each(|s| s.walk(process_stmt, process_expr)),
+            Self::Stmt(x, _) => x.iter().for_each(|s| s.walk(path, on_node)),
-            Self::Array(x, _) => x.iter().for_each(|e| e.walk(process_stmt, process_expr)),
+            Self::Array(x, _) => x.iter().for_each(|e| e.walk(path, on_node)),
-            Self::Map(x, _) => x
+            Self::Map(x, _) => x.iter().for_each(|(_, e)| e.walk(path, on_node)),
-                .iter()
-                .for_each(|(_, e)| e.walk(process_stmt, process_expr)),
             Self::Index(x, _) | Expr::In(x, _) | Expr::And(x, _) | Expr::Or(x, _) => {
-                x.lhs.walk(process_stmt, process_expr);
+                x.lhs.walk(path, on_node);
-                x.rhs.walk(process_stmt, process_expr);
+                x.rhs.walk(path, on_node);
             }
-            Self::Custom(x, _) => x
+            Self::Custom(x, _) => x.keywords.iter().for_each(|e| e.walk(path, on_node)),
-                .keywords
-                .iter()
-                .for_each(|e| e.walk(process_stmt, process_expr)),
             _ => (),
         }
+
+        path.pop().unwrap();
     }
 }
 

src/bin/rhai-repl.rs

@@ -92,7 +92,7 @@ fn main() {
                 .compile(&contents)
                 .map_err(|err| err.into())
                 .and_then(|mut ast| {
-                    ast.set_source(Some(&filename));
+                    ast.set_source(&filename);
                     Module::eval_ast_as_new(Default::default(), &ast, &engine)
                 }) {
                 Err(err) => {

src/engine_api.rs

@@ -903,47 +903,53 @@ impl Engine {
         script: &str,
     ) -> Result<AST, Box<EvalAltResult>> {
         use crate::{
-            ast::{Expr, Stmt},
+            ast::{ASTNode, Expr, Stmt},
             fn_native::shared_take_or_clone,
             module::resolvers::StaticModuleResolver,
-            stdlib::collections::HashMap,
+            stdlib::collections::HashSet,
             ImmutableString,
         };
 
-        let mut ast = self.compile_scripts_with_scope(scope, &[script])?;
+        fn collect_imports(
-        let mut imports = HashMap::<ImmutableString, Position>::new();
+            ast: &AST,
-
+            resolver: &StaticModuleResolver,
-        ast.statements()
+            imports: &mut HashSet<ImmutableString>,
-            .iter()
+        ) {
-            .chain({
+            ast.walk(&mut |path| match path.last().unwrap() {
-                #[cfg(not(feature = "no_function"))]
+                // Collect all `import` statements with a string constant path
+                ASTNode::Stmt(Stmt::Import(Expr::StringConstant(s, _), _, _))
+                    if !resolver.contains_path(s) && !imports.contains(s) =>
                 {
-                    ast.iter_fn_def().map(|f| &f.body)
+                    imports.insert(s.clone());
-                }
-                #[cfg(feature = "no_function")]
-                {
-                    crate::stdlib::iter::empty()
-                }
-            })
-            .for_each(|stmt| {
-                stmt.walk(
-                    &mut |stmt| match stmt {
-                        Stmt::Import(Expr::StringConstant(s, pos), _, _)
-                            if !imports.contains_key(s) =>
-                        {
-                            imports.insert(s.clone(), *pos);
                 }
                 _ => (),
-                    },
-                    &mut |_| {},
-                )
             });
+        }
+
+        let mut resolver = StaticModuleResolver::new();
+        let mut ast = self.compile_scripts_with_scope(scope, &[script])?;
+        let mut imports = HashSet::<ImmutableString>::new();
+
+        collect_imports(&ast, &mut resolver, &mut imports);
 
         if !imports.is_empty() {
-            let mut resolver = StaticModuleResolver::new();
+            while let Some(path) = imports.iter().next() {
-            for (path, pos) in imports {
+                let path = path.clone();
-                let module = self.module_resolver.resolve(self, &path, pos)?;
+
-                let module = shared_take_or_clone(module);
+                if let Some(module_ast) =
+                    self.module_resolver
+                        .resolve_ast(self, &path, Position::NONE)?
+                {
+                    collect_imports(&module_ast, &mut resolver, &mut imports);
+                }
+
+                let module = shared_take_or_clone(self.module_resolver.resolve(
+                    self,
+                    &path,
+                    Position::NONE,
+                )?);
+
+                imports.remove(&path);
                 resolver.insert(path, module);
             }
             ast.set_resolver(resolver);
@@ -1957,7 +1963,7 @@ impl Engine {
     ///                ));
     ///
     /// let mut ast = engine.compile(r#"let x = "hello"; debug(x);"#)?;
-    /// ast.set_source(Some("world"));
+    /// ast.set_source("world");
     /// engine.consume_ast(&ast)?;
     ///
     /// assert_eq!(*result.read().unwrap(), r#"world @ 1:18 > "hello""#);

src/lib.rs

@@ -184,7 +184,8 @@ pub use token::{get_next_token, parse_string_literal, InputStream, Token, Tokeni
 #[cfg(feature = "internals")]
 #[deprecated = "this type is volatile and may change"]
 pub use ast::{
-    BinaryExpr, CustomExpr, Expr, FloatWrapper, FnCallExpr, Ident, ReturnType, ScriptFnDef, Stmt,
+    ASTNode, BinaryExpr, CustomExpr, Expr, FloatWrapper, FnCallExpr, Ident, ReturnType,
+    ScriptFnDef, Stmt,
 };
 
 #[cfg(feature = "internals")]

src/module/resolvers/file.rs

@@ -212,7 +212,7 @@ impl ModuleResolver for FileModuleResolver {
                 _ => Box::new(EvalAltResult::ErrorInModule(path.to_string(), err, pos)),
             })?;
 
-        ast.set_source(Some(path));
+        ast.set_source(path);
 
         // Make a module from the AST
         let m: Shared<Module> = Module::eval_ast_as_new(scope, &ast, engine)
@@ -227,4 +227,28 @@ impl ModuleResolver for FileModuleResolver {
 
         Ok(m)
     }
+
+    fn resolve_ast(
+        &self,
+        engine: &Engine,
+        path: &str,
+        pos: Position,
+    ) -> Result<Option<crate::AST>, Box<EvalAltResult>> {
+        // Construct the script file path
+        let mut file_path = self.base_path.clone();
+        file_path.push(path);
+        file_path.set_extension(&self.extension); // Force extension
+
+        // Load the script file and compile it
+        let mut ast = engine.compile_file(file_path).map_err(|err| match *err {
+            EvalAltResult::ErrorSystem(_, err) if err.is::<IoError>() => {
+                Box::new(EvalAltResult::ErrorModuleNotFound(path.to_string(), pos))
+            }
+            _ => Box::new(EvalAltResult::ErrorInModule(path.to_string(), err, pos)),
+        })?;
+
+        ast.set_source(path);
+
+        Ok(Some(ast))
+    }
 }

src/module/resolvers/mod.rs

@@ -1,6 +1,6 @@
 use crate::fn_native::SendSync;
 use crate::stdlib::boxed::Box;
-use crate::{Engine, EvalAltResult, Module, Position, Shared};
+use crate::{Engine, EvalAltResult, Module, Position, Shared, AST};
 
 mod dummy;
 pub use dummy::DummyModuleResolver;
@@ -28,4 +28,23 @@ pub trait ModuleResolver: SendSync {
         path: &str,
         pos: Position,
     ) -> Result<Shared<Module>, Box<EvalAltResult>>;
+
+    /// Resolve a module into an `AST` based on a path string.
+    ///
+    /// Returns [`None`] (default) if such resolution is not supported
+    /// (e.g. if the module is Rust-based).
+    ///
+    /// ## Low-Level API
+    ///
+    /// Override the default implementation of this method if the module resolver
+    /// serves modules based on compiled Rhai scripts.
+    #[allow(unused_variables)]
+    fn resolve_ast(
+        &self,
+        engine: &Engine,
+        path: &str,
+        pos: Position,
+    ) -> Result<Option<AST>, Box<EvalAltResult>> {
+        Ok(None)
+    }
 }

tests/modules.rs

@@ -247,9 +247,14 @@ fn test_module_resolver() -> Result<(), Box<EvalAltResult>> {
         )?;
     }
 
+    #[cfg(not(feature = "no_function"))]
+    {
         let script = r#"
+            fn foo() {
                 import "hello" as h;
                 h::answer
+            }
+            foo() + { import "hello" as h; h::answer }
         "#;
         let mut scope = Scope::new();
 
@@ -257,9 +262,10 @@ fn test_module_resolver() -> Result<(), Box<EvalAltResult>> {
 
         engine.set_module_resolver(DummyModuleResolver::new());
 
-    assert_eq!(engine.eval_ast::<INT>(&ast)?, 42);
+        assert_eq!(engine.eval_ast::<INT>(&ast)?, 84);
 
         assert!(engine.eval::<INT>(script).is_err());
+    }
 
     Ok(())
 }

tests/print.rs

@@ -25,7 +25,7 @@ fn test_print_debug() -> Result<(), Box<EvalAltResult>> {
     // Evaluate script
     engine.consume("print(40 + 2)")?;
     let mut ast = engine.compile(r#"let x = "hello!"; debug(x)"#)?;
-    ast.set_source(Some("world"));
+    ast.set_source("world");
     engine.consume_ast(&ast)?;
 
     // 'logbook' captures all the 'print' and 'debug' output