Allow creating modules from directories (#9066)

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

@@ -1202,6 +1202,15 @@ impl<'a> StateWorkingSet<'a> {
         }
     }
 
+    pub fn use_modules(&mut self, modules: Vec<(Vec<u8>, ModuleId)>) {
+        let overlay_frame = self.last_overlay_mut();
+
+        for (name, module_id) in modules {
+            overlay_frame.insert_module(name, module_id);
+            // overlay_frame.visibility.use_module_id(&module_id);  // TODO: Add hiding modules
+        }
+    }
+
     pub fn add_predecl(&mut self, decl: Box<dyn Command>) -> Option<DeclId> {
         let name = decl.name().as_bytes().to_vec();
 
@@ -1770,6 +1779,18 @@ impl<'a> StateWorkingSet<'a> {
         }
     }
 
+    pub fn get_module_mut(&mut self, module_id: ModuleId) -> &mut Module {
+        let num_permanent_modules = self.permanent_state.num_modules();
+        if module_id < num_permanent_modules {
+            panic!("Attempt to mutate a module that is in the permanent (immutable) state")
+        } else {
+            self.delta
+                .modules
+                .get_mut(module_id - num_permanent_modules)
+                .expect("internal error: missing module")
+        }
+    }
+
     pub fn get_block_mut(&mut self, block_id: BlockId) -> &mut Block {
         let num_permanent_blocks = self.permanent_state.num_blocks();
         if block_id < num_permanent_blocks {
@@ -1848,7 +1869,7 @@ impl<'a> StateWorkingSet<'a> {
             let name = self.last_overlay_name().to_vec();
             let origin = overlay_frame.origin;
             let prefixed = overlay_frame.prefixed;
-            self.add_overlay(name, origin, vec![], prefixed);
+            self.add_overlay(name, origin, vec![], vec![], prefixed);
         }
 
         self.delta
@@ -1886,6 +1907,7 @@ impl<'a> StateWorkingSet<'a> {
         name: Vec<u8>,
         origin: ModuleId,
         decls: Vec<(Vec<u8>, DeclId)>,
+        modules: Vec<(Vec<u8>, ModuleId)>,
         prefixed: bool,
     ) {
         let last_scope_frame = self.delta.last_scope_frame_mut();
@@ -1913,6 +1935,7 @@ impl<'a> StateWorkingSet<'a> {
         self.move_predecls_to_overlay();
 
         self.use_decls(decls);
+        self.use_modules(modules);
     }
 
     pub fn remove_overlay(&mut self, name: &[u8], keep_custom: bool) {

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

@@ -206,6 +206,10 @@ impl OverlayFrame {
         self.decls.insert((name, input), decl_id)
     }
 
+    pub fn insert_module(&mut self, name: Vec<u8>, module_id: ModuleId) -> Option<ModuleId> {
+        self.modules.insert(name, module_id)
+    }
+
     pub fn get_decl(&self, name: &[u8], input: &Type) -> Option<DeclId> {
         if let Some(decl) = self.decls.get(&(name, input) as &dyn DeclKey) {
             Some(*decl)

crates/nu-protocol/src/exportable.rs

@@ -1,5 +1,6 @@
-use crate::DeclId;
+use crate::{DeclId, ModuleId};
 
 pub enum Exportable {
     Decl { name: Vec<u8>, id: DeclId },
+    Module { name: Vec<u8>, id: ModuleId },
 }

crates/nu-protocol/src/module.rs

@@ -1,12 +1,26 @@
-use crate::{BlockId, DeclId, Span};
+use crate::{
+    ast::ImportPatternMember, engine::StateWorkingSet, BlockId, DeclId, ModuleId, ParseError, Span,
+};
 
 use indexmap::IndexMap;
 
+pub struct ResolvedImportPattern {
+    pub decls: Vec<(Vec<u8>, DeclId)>,
+    pub modules: Vec<(Vec<u8>, ModuleId)>,
+}
+
+impl ResolvedImportPattern {
+    pub fn new(decls: Vec<(Vec<u8>, DeclId)>, modules: Vec<(Vec<u8>, ModuleId)>) -> Self {
+        ResolvedImportPattern { decls, modules }
+    }
+}
+
 /// Collection of definitions that can be exported from a module
 #[derive(Debug, Clone)]
 pub struct Module {
     pub name: Vec<u8>,
     pub decls: IndexMap<Vec<u8>, DeclId>,
+    pub submodules: IndexMap<Vec<u8>, ModuleId>,
     pub env_block: Option<BlockId>, // `export-env { ... }` block
     pub main: Option<DeclId>,       // `export def main`
     pub span: Option<Span>,
@@ -17,6 +31,7 @@ impl Module {
         Module {
             name,
             decls: IndexMap::new(),
+            submodules: IndexMap::new(),
             env_block: None,
             main: None,
             span: None,
@@ -27,32 +42,29 @@ impl Module {
         Module {
             name,
             decls: IndexMap::new(),
+            submodules: IndexMap::new(),
             env_block: None,
             main: None,
             span: Some(span),
         }
     }
 
+    pub fn name(&self) -> Vec<u8> {
+        self.name.clone()
+    }
+
     pub fn add_decl(&mut self, name: Vec<u8>, decl_id: DeclId) -> Option<DeclId> {
         self.decls.insert(name, decl_id)
     }
 
+    pub fn add_submodule(&mut self, name: Vec<u8>, module_id: ModuleId) -> Option<ModuleId> {
+        self.submodules.insert(name, module_id)
+    }
+
     pub fn add_env_block(&mut self, block_id: BlockId) {
         self.env_block = Some(block_id);
     }
 
-    pub fn extend(&mut self, other: &Module) {
-        self.decls.extend(other.decls.clone());
-    }
-
-    pub fn is_empty(&self) -> bool {
-        self.decls.is_empty()
-    }
-
-    pub fn get_decl_id(&self, name: &[u8]) -> Option<DeclId> {
-        self.decls.get(name).copied()
-    }
-
     pub fn has_decl(&self, name: &[u8]) -> bool {
         if name == self.name && self.main.is_some() {
             return true;
@@ -61,6 +73,120 @@ impl Module {
         self.decls.contains_key(name)
     }
 
+    pub fn resolve_import_pattern(
+        &self,
+        working_set: &StateWorkingSet,
+        self_id: ModuleId,
+        members: &[ImportPatternMember],
+        name_override: Option<&[u8]>, // name under the module was stored (doesn't have to be the
+                                      // same as self.name)
+    ) -> (ResolvedImportPattern, Vec<ParseError>) {
+        let final_name = name_override.unwrap_or(&self.name).to_vec();
+
+        let (head, rest) = if let Some((head, rest)) = members.split_first() {
+            (head, rest)
+        } else {
+            // Import pattern was just name without any members
+            let mut results = vec![];
+            let mut errors = vec![];
+
+            for (_, id) in &self.submodules {
+                let submodule = working_set.get_module(*id);
+                let (sub_results, sub_errors) =
+                    submodule.resolve_import_pattern(working_set, *id, &[], None);
+                errors.extend(sub_errors);
+
+                for (sub_name, sub_decl_id) in sub_results.decls {
+                    let mut new_name = final_name.clone();
+                    new_name.push(b' ');
+                    new_name.extend(sub_name);
+
+                    results.push((new_name, sub_decl_id));
+                }
+            }
+
+            results.extend(self.decls_with_head(&final_name));
+
+            return (
+                ResolvedImportPattern::new(results, vec![(final_name, self_id)]),
+                errors,
+            );
+        };
+
+        match head {
+            ImportPatternMember::Name { name, span } => {
+                if name == b"main" {
+                    if let Some(main_decl_id) = self.main {
+                        (
+                            ResolvedImportPattern::new(vec![(final_name, main_decl_id)], vec![]),
+                            vec![],
+                        )
+                    } else {
+                        (
+                            ResolvedImportPattern::new(vec![], vec![]),
+                            vec![ParseError::ExportNotFound(*span)],
+                        )
+                    }
+                } else if let Some(decl_id) = self.decls.get(name) {
+                    (
+                        ResolvedImportPattern::new(vec![(name.clone(), *decl_id)], vec![]),
+                        vec![],
+                    )
+                } else if let Some(submodule_id) = self.submodules.get(name) {
+                    let submodule = working_set.get_module(*submodule_id);
+                    submodule.resolve_import_pattern(working_set, *submodule_id, rest, None)
+                } else {
+                    (
+                        ResolvedImportPattern::new(vec![], vec![]),
+                        vec![ParseError::ExportNotFound(*span)],
+                    )
+                }
+            }
+            ImportPatternMember::Glob { .. } => {
+                let mut decls = vec![];
+                let mut submodules = vec![];
+                let mut errors = vec![];
+
+                for (_, id) in &self.submodules {
+                    let submodule = working_set.get_module(*id);
+                    let (sub_results, sub_errors) =
+                        submodule.resolve_import_pattern(working_set, *id, &[], None);
+                    decls.extend(sub_results.decls);
+                    submodules.extend(sub_results.modules);
+                    errors.extend(sub_errors);
+                }
+
+                decls.extend(self.decls());
+                submodules.extend(self.submodules());
+
+                (ResolvedImportPattern::new(decls, submodules), errors)
+            }
+            ImportPatternMember::List { names } => {
+                let mut decls = vec![];
+                let mut submodules = vec![];
+                let mut errors = vec![];
+
+                for (name, span) in names {
+                    if name == b"main" {
+                        if let Some(main_decl_id) = self.main {
+                            decls.push((final_name.clone(), main_decl_id));
+                        } else {
+                            errors.push(ParseError::ExportNotFound(*span));
+                        }
+                    } else if let Some(decl_id) = self.decls.get(name) {
+                        decls.push((name.clone(), *decl_id));
+                    } else if let Some(submodule_id) = self.submodules.get(name) {
+                        submodules.push((name.clone(), *submodule_id));
+                    } else {
+                        errors.push(ParseError::ExportNotFound(*span));
+                    }
+                }
+
+                (ResolvedImportPattern::new(decls, submodules), errors)
+            }
+        }
+    }
+
     pub fn decl_name_with_head(&self, name: &[u8], head: &[u8]) -> Option<Vec<u8>> {
         if self.has_decl(name) {
             let mut new_name = head.to_vec();
@@ -124,6 +250,13 @@ impl Module {
         result
     }
 
+    pub fn submodules(&self) -> Vec<(Vec<u8>, ModuleId)> {
+        self.submodules
+            .iter()
+            .map(|(name, id)| (name.clone(), *id))
+            .collect()
+    }
+
     pub fn decl_names(&self) -> Vec<Vec<u8>> {
         let mut result: Vec<Vec<u8>> = self.decls.keys().cloned().collect();
 

crates/nu-protocol/src/parse_error.rs

@@ -200,14 +200,36 @@ pub enum ParseError {
     #[error("Can't export {0} named same as the module.")]
     #[diagnostic(
         code(nu::parser::named_as_module),
-        help("Module {1} can't export {0} named the same as the module. Either change the module name, or export `main` custom command.")
+        help("Module {1} can't export {0} named the same as the module. Either change the module name, or export `{2}` {0}.")
     )]
     NamedAsModule(
+        String,
         String,
         String,
         #[label = "can't export from module {1}"] Span,
     ),
 
+    #[error("Module already contains 'main' command.")]
+    #[diagnostic(
+        code(nu::parser::module_double_main),
+        help("Tried to add 'main' command to module '{0}' but it has already been added.")
+    )]
+    ModuleDoubleMain(
+        String,
+        #[label = "module '{0}' already contains 'main'"] Span,
+    ),
+
+    #[error("Invalid module file name")]
+    #[diagnostic(
+        code(nu::parser::invalid_module_file_name),
+        help("File {0} resolves to module name {1} which is the same as the parent module. Either rename the file or, save it as 'mod.nu' to define the parent module.")
+    )]
+    InvalidModuleFileName(
+        String,
+        String,
+        #[label = "submodule can't have the same name as the parent module"] Span,
+    ),
+
     #[error("Can't export alias defined as 'main'.")]
     #[diagnostic(
         code(nu::parser::export_main_alias_not_allowed),
@@ -371,7 +393,7 @@ pub enum ParseError {
 
     #[error("Wrong import pattern structure.")]
     #[diagnostic(code(nu::parser::missing_import_pattern))]
-    WrongImportPattern(#[label = "invalid import pattern structure"] Span),
+    WrongImportPattern(String, #[label = "{0}"] Span),
 
     #[error("Export not found.")]
     #[diagnostic(code(nu::parser::export_not_found))]
@@ -452,7 +474,9 @@ impl ParseError {
             ParseError::AliasNotValid(s) => *s,
             ParseError::CommandDefNotValid(s) => *s,
             ParseError::ModuleNotFound(s) => *s,
-            ParseError::NamedAsModule(_, _, s) => *s,
+            ParseError::NamedAsModule(_, _, _, s) => *s,
+            ParseError::ModuleDoubleMain(_, s) => *s,
+            ParseError::InvalidModuleFileName(_, _, s) => *s,
             ParseError::ExportMainAliasNotAllowed(s) => *s,
             ParseError::CyclicalModuleImport(_, s) => *s,
             ParseError::ModuleOrOverlayNotFound(s) => *s,
@@ -486,7 +510,7 @@ impl ParseError {
             ParseError::MissingColumns(_, s) => *s,
             ParseError::AssignmentMismatch(_, _, s) => *s,
             ParseError::MissingImportPattern(s) => *s,
-            ParseError::WrongImportPattern(s) => *s,
+            ParseError::WrongImportPattern(_, s) => *s,
             ParseError::ExportNotFound(s) => *s,
             ParseError::SourcedFileNotFound(_, s) => *s,
             ParseError::RegisteredFileNotFound(_, s) => *s,