Allow adding module blocks to engine state

crates/nu-command/src/lib.rs

@@ -15,8 +15,8 @@ mod let_env;
 mod lines;
 mod list_git_branches;
 mod ls;
-mod run_external;
 mod module;
+mod run_external;
 mod table;
 mod where_;
 
@@ -37,6 +37,6 @@ pub use let_env::LetEnv;
 pub use lines::Lines;
 pub use list_git_branches::ListGitBranches;
 pub use ls::Ls;
-pub use run_external::External;
 pub use module::Module;
+pub use run_external::External;
 pub use table::Table;

crates/nu-parser/src/parser.rs

@@ -2658,9 +2658,13 @@ pub fn parse_module(
 
     // parse_def() equivalent
     if bytes == b"module" && spans.len() >= 3 {
-        let (name_expr, err) = parse_string(working_set, spans[1]);
+        let (module_name_expr, err) = parse_string(working_set, spans[1]);
         error = error.or(err);
 
+        let module_name = module_name_expr
+            .as_string()
+            .expect("internal error: module name is not a string");
+
         // parse_block_expression() equivalent
         let block_span = spans[2];
         let block_bytes = working_set.get_span_contents(block_span);
@@ -2715,6 +2719,8 @@ pub fn parse_module(
             }
         }
 
+        let mut exports: Vec<Vec<u8>> = vec![];
+
         let block: Block = output
             .block
             .iter()
@@ -2726,11 +2732,22 @@ pub fn parse_module(
 
                     let (stmt, err) = match name {
                         // TODO: Here we can add other stuff that's alowed for modules
-                        b"def" => parse_def(working_set, &pipeline.commands[0].parts),
+                        b"def" => {
+                            let (stmt, err) = parse_def(working_set, &pipeline.commands[0].parts);
+
+                            if err.is_none() {
+                                let def_name =
+                                    working_set.get_span_contents(pipeline.commands[0].parts[1]);
+                                // TODO: Later, we want to put this behind 'export'
+                                exports.push(def_name.into());
+                            }
+
+                            (stmt, err)
+                        }
                         _ => (
                             garbage_statement(&pipeline.commands[0].parts),
                             Some(ParseError::Expected("def".into(), block_span)),
-                        )
+                        ),
                     };
 
                     if error.is_none() {
@@ -2745,9 +2762,11 @@ pub fn parse_module(
             })
             .into();
 
+        let block = block.with_exports(exports);
+
         working_set.exit_scope();
 
-        let block_id = working_set.add_block(block);
+        let block_id = working_set.add_module(&module_name, block);
 
         let block_expr = Expression {
             expr: Expr::Block(block_id),
@@ -2763,7 +2782,7 @@ pub fn parse_module(
         let call = Box::new(Call {
             head: spans[0],
             decl_id: module_decl_id,
-            positional: vec![name_expr, block_expr],
+            positional: vec![module_name_expr, block_expr],
             named: vec![],
         });
 

crates/nu-protocol/src/ast/block.rs

@@ -8,6 +8,7 @@ use super::Statement;
 pub struct Block {
     pub signature: Box<Signature>,
     pub stmts: Vec<Statement>,
+    pub exports: Vec<Vec<u8>>, // Assuming just defs for now
 }
 
 impl Block {
@@ -45,6 +46,15 @@ impl Block {
         Self {
             signature: Box::new(Signature::new("")),
             stmts: vec![],
+            exports: vec![],
+        }
+    }
+
+    pub fn with_exports(self, exports: Vec<Vec<u8>>) -> Self {
+        Self {
+            signature: self.signature,
+            stmts: self.stmts,
+            exports,
         }
     }
 }
@@ -57,6 +67,7 @@ where
         Self {
             signature: Box::new(Signature::new("")),
             stmts: stmts.collect(),
+            exports: vec![],
         }
     }
 }

crates/nu-protocol/src/engine/engine_state.rs

@@ -17,6 +17,7 @@ pub struct ScopeFrame {
     vars: HashMap<Vec<u8>, VarId>,
     decls: HashMap<Vec<u8>, DeclId>,
     aliases: HashMap<Vec<u8>, Vec<Span>>,
+    modules: HashMap<Vec<u8>, BlockId>,
 }
 
 impl ScopeFrame {
@@ -25,6 +26,7 @@ impl ScopeFrame {
             vars: HashMap::new(),
             decls: HashMap::new(),
             aliases: HashMap::new(),
+            modules: HashMap::new(),
         }
     }
 
@@ -290,6 +292,23 @@ impl<'a> StateWorkingSet<'a> {
         self.num_blocks() - 1
     }
 
+    pub fn add_module(&mut self, name: &str, block: Block) -> BlockId {
+        let name = name.as_bytes().to_vec();
+
+        self.delta.blocks.push(block);
+        let block_id = self.num_blocks() - 1;
+
+        let scope_frame = self
+            .delta
+            .scope
+            .last_mut()
+            .expect("internal error: missing required scope frame");
+
+        scope_frame.modules.insert(name, block_id);
+
+        block_id
+    }
+
     pub fn next_span_start(&self) -> usize {
         self.permanent_state.next_span_start() + self.delta.file_contents.len()
     }