From 250071939bad8affa6d6e8f595dd2614b01ec424 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20=C5=BD=C3=A1dn=C3=ADk?= <kubouch@gmail.com>
Date: Sun, 7 May 2023 14:41:40 +0300
Subject: [PATCH] Reuse parsed modules (#9125)

---
 crates/nu-parser/src/parse_keywords.rs        | 20 +++++++++--
 crates/nu-protocol/src/engine/engine_state.rs | 35 ++++++++++++-------
 tests/modules/mod.rs                          | 10 +++++-
 3 files changed, 50 insertions(+), 15 deletions(-)

diff --git a/crates/nu-parser/src/parse_keywords.rs b/crates/nu-parser/src/parse_keywords.rs
index 4ca3d1edaa..b6cecf229c 100644
--- a/crates/nu-parser/src/parse_keywords.rs
+++ b/crates/nu-parser/src/parse_keywords.rs
@@ -1727,6 +1727,10 @@ fn parse_module_file(
     let file_id = working_set.add_file(path.to_string_lossy().to_string(), &contents);
     let new_span = working_set.get_span_for_file(file_id);
 
+    if let Some(module_id) = working_set.find_module_by_span(new_span) {
+        return Some(module_id);
+    }
+
     // Change the currently parsed directory
     let prev_currently_parsed_cwd = if let Some(parent) = path.parent() {
         let prev = working_set.currently_parsed_cwd.clone();
@@ -1834,13 +1838,25 @@ pub fn parse_module_file_or_dir(
                     path_span,
                     name_override.or(Some(module_name)),
                 ) {
-                    let module = working_set.get_module_mut(module_id);
+                    let mut module = working_set.get_module(module_id).clone();
 
                     for (submodule_name, submodule_id) in submodules {
                         module.add_submodule(submodule_name, submodule_id);
                     }
 
-                    Some(module_id)
+                    let module_name = String::from_utf8_lossy(&module.name).to_string();
+
+                    let module_comments =
+                        if let Some(comments) = working_set.get_module_comments(module_id) {
+                            comments.to_vec()
+                        } else {
+                            vec![]
+                        };
+
+                    let new_module_id =
+                        working_set.add_module(&module_name, module, module_comments);
+
+                    Some(new_module_id)
                 } else {
                     None
                 }
diff --git a/crates/nu-protocol/src/engine/engine_state.rs b/crates/nu-protocol/src/engine/engine_state.rs
index 1c81c1ec2a..64b9c561d7 100644
--- a/crates/nu-protocol/src/engine/engine_state.rs
+++ b/crates/nu-protocol/src/engine/engine_state.rs
@@ -1322,6 +1322,13 @@ impl<'a> StateWorkingSet<'a> {
         module_id
     }
 
+    pub fn get_module_comments(&self, module_id: ModuleId) -> Option<&[Span]> {
+        self.delta
+            .usage
+            .get_module_comments(module_id)
+            .or_else(|| self.permanent_state.get_module_comments(module_id))
+    }
+
     pub fn next_span_start(&self) -> usize {
         let permanent_span_start = self.permanent_state.next_span_start();
 
@@ -1779,18 +1786,6 @@ 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 {
@@ -1997,6 +1992,22 @@ impl<'a> StateWorkingSet<'a> {
 
         None
     }
+
+    pub fn find_module_by_span(&self, span: Span) -> Option<ModuleId> {
+        for (id, module) in self.delta.modules.iter().enumerate() {
+            if Some(span) == module.span {
+                return Some(self.permanent_state.num_modules() + id);
+            }
+        }
+
+        for (module_id, module) in self.permanent_state.modules.iter().enumerate() {
+            if Some(span) == module.span {
+                return Some(module_id);
+            }
+        }
+
+        None
+    }
 }
 
 impl Default for EngineState {
diff --git a/tests/modules/mod.rs b/tests/modules/mod.rs
index fb53bfba93..29bbebd1a5 100644
--- a/tests/modules/mod.rs
+++ b/tests/modules/mod.rs
@@ -1,6 +1,6 @@
 use nu_test_support::fs::Stub::FileWithContentToBeTrimmed;
 use nu_test_support::playground::Playground;
-use nu_test_support::{nu, pipeline};
+use nu_test_support::{nu, nu_repl_code, pipeline};
 use pretty_assertions::assert_eq;
 
 #[test]
@@ -641,6 +641,14 @@ fn module_dir() {
     assert_eq!(actual.out, "spambaz");
 }
 
+#[test]
+fn module_dir_import_twice_no_panic() {
+    let import = "use samples/spam";
+    let inp = &[import, import, "spam"];
+    let actual_repl = nu!(cwd: "tests/modules", nu_repl_code(inp));
+    assert_eq!(actual_repl.out, "spam");
+}
+
 #[test]
 fn not_allowed_submodule_file() {
     let inp = &["use samples/not_allowed"];