Overlay keep (#5629)

* Allow env vars to be kept from removed overlay

* Rename --keep to --keep-custom; Add new test

* Rename some symbols

* (WIP) Start working on --keep for defs and aliases

* Fix decls/aliases not melting properly

* Use id instead of the whole cloned overlay

* Rewrite overlay remove for no reason

Doesn't fix the bug but at least looks better.

* Rename variable

* Fix adding overlay env vars

* Add more tests; Fmt + Clippy

crates/nu-command/src/core_commands/overlay/add.rs

@@ -51,26 +51,38 @@ https://www.nushell.sh/book/thinking_in_nushell.html#parsing-and-evaluation-are-
     ) -> Result<PipelineData, ShellError> {
         let name_arg: Spanned<String> = call.req(engine_state, stack, 0)?;
 
-        // TODO: This logic is duplicated in the parser.
-        if stack.has_env_overlay(&name_arg.item, engine_state) {
-            // Activate existing overlay
-            stack.add_overlay(name_arg.item.clone());
+        let maybe_overlay_name = if engine_state
+            .find_overlay(name_arg.item.as_bytes())
+            .is_some()
+        {
+            Some(name_arg.item.clone())
+        } else if let Some(os_str) = Path::new(&name_arg.item).file_stem() {
+            os_str.to_str().map(|name| name.to_string())
+        } else {
+            None
+        };
 
-            if let Some(module_id) = engine_state
-                .find_overlay(name_arg.item.as_bytes())
-                .map(|id| engine_state.get_overlay(id).origin)
-            {
-                if let Some(new_module_id) = engine_state.find_module(name_arg.item.as_bytes(), &[])
+        if let Some(overlay_name) = maybe_overlay_name {
+            if let Some(overlay_id) = engine_state.find_overlay(overlay_name.as_bytes()) {
+                let old_module_id = engine_state.get_overlay(overlay_id).origin;
+
+                stack.add_overlay(overlay_name.clone());
+
+                if let Some(new_module_id) = engine_state.find_module(overlay_name.as_bytes(), &[])
                 {
-                    if module_id != new_module_id {
-                        // The origin module of an overlay changed => update it
+                    if !stack.has_env_overlay(&overlay_name, engine_state)
+                        || (old_module_id != new_module_id)
+                    {
+                        // Add environment variables only if:
+                        // a) adding a new overlay
+                        // b) refreshing an active overlay (the origin module changed)
                         let module = engine_state.get_module(new_module_id);
 
                         for (name, block_id) in module.env_vars() {
                             let name = if let Ok(s) = String::from_utf8(name.clone()) {
                                 s
                             } else {
-                                return Err(ShellError::NonUtf8(name_arg.span));
+                                return Err(ShellError::NonUtf8(call.head));
                             };
 
                             let block = engine_state.get_block(block_id);
@@ -90,57 +102,16 @@ https://www.nushell.sh/book/thinking_in_nushell.html#parsing-and-evaluation-are-
                     }
                 }
             } else {
+                return Err(ShellError::OverlayNotFoundAtRuntime(
+                    name_arg.item,
+                    name_arg.span,
+                ));
             }
         } else {
-            // Create a new overlay from a module
-            let (overlay_name, module) =
-                if let Some(module_id) = engine_state.find_module(name_arg.item.as_bytes(), &[]) {
-                    (name_arg.item, engine_state.get_module(module_id))
-                } else if let Some(os_str) = Path::new(&name_arg.item).file_stem() {
-                    let name = if let Some(s) = os_str.to_str() {
-                        s.to_string()
-                    } else {
-                        return Err(ShellError::NonUtf8(name_arg.span));
-                    };
-
-                    if let Some(module_id) = engine_state.find_module(name.as_bytes(), &[]) {
-                        (name, engine_state.get_module(module_id))
-                    } else {
-                        return Err(ShellError::ModuleOrOverlayNotFoundAtRuntime(
-                            name_arg.item,
-                            name_arg.span,
-                        ));
-                    }
-                } else {
-                    return Err(ShellError::ModuleOrOverlayNotFoundAtRuntime(
-                        name_arg.item,
-                        name_arg.span,
-                    ));
-                };
-
-            stack.add_overlay(overlay_name);
-
-            for (name, block_id) in module.env_vars() {
-                let name = if let Ok(s) = String::from_utf8(name.clone()) {
-                    s
-                } else {
-                    return Err(ShellError::NonUtf8(name_arg.span));
-                };
-
-                let block = engine_state.get_block(block_id);
-
-                let val = eval_block(
-                    engine_state,
-                    stack,
-                    block,
-                    PipelineData::new(call.head),
-                    false,
-                    true,
-                )?
-                .into_value(call.head);
-
-                stack.add_env_var(name, val);
-            }
+            return Err(ShellError::OverlayNotFoundAtRuntime(
+                name_arg.item,
+                name_arg.span,
+            ));
         }
 
         Ok(PipelineData::new(call.head))

crates/nu-command/src/core_commands/overlay/remove.rs

@@ -1,7 +1,9 @@
 use nu_engine::CallExt;
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{Command, EngineState, Stack};
-use nu_protocol::{Category, Example, PipelineData, Signature, Spanned, SyntaxShape};
+use nu_protocol::{
+    Category, Example, PipelineData, ShellError, Signature, Spanned, SyntaxShape, Value,
+};
 
 #[derive(Clone)]
 pub struct OverlayRemove;
@@ -18,6 +20,11 @@ impl Command for OverlayRemove {
     fn signature(&self) -> nu_protocol::Signature {
         Signature::build("overlay remove")
             .optional("name", SyntaxShape::String, "Overlay to remove")
+            .switch(
+                "keep-custom",
+                "Keep newly added symbols within the next activated overlay",
+                Some('k'),
+            )
             .category(Category::Core)
     }
 
@@ -37,9 +44,7 @@ https://www.nushell.sh/book/thinking_in_nushell.html#parsing-and-evaluation-are-
         call: &Call,
         _input: PipelineData,
     ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
-        // let module_name: Spanned<String> = call.req(engine_state, stack, 0)?;
-
-        let module_name: Spanned<String> = if let Some(name) = call.opt(engine_state, stack, 0)? {
+        let overlay_name: Spanned<String> = if let Some(name) = call.opt(engine_state, stack, 0)? {
             name
         } else {
             Spanned {
@@ -48,8 +53,38 @@ https://www.nushell.sh/book/thinking_in_nushell.html#parsing-and-evaluation-are-
             }
         };
 
-        // TODO: Add env merging
-        stack.remove_overlay(&module_name.item, &module_name.span)?;
+        if !stack.is_overlay_active(&overlay_name.item) {
+            return Err(ShellError::OverlayNotFoundAtRuntime(
+                overlay_name.item,
+                overlay_name.span,
+            ));
+        }
+
+        if call.has_flag("keep-custom") {
+            if let Some(overlay_id) = engine_state.find_overlay(overlay_name.item.as_bytes()) {
+                let overlay_frame = engine_state.get_overlay(overlay_id);
+                let origin_module = engine_state.get_module(overlay_frame.origin);
+
+                let env_vars_to_keep: Vec<(String, Value)> = stack
+                    .get_overlay_env_vars(engine_state, &overlay_name.item)
+                    .into_iter()
+                    .filter(|(name, _)| !origin_module.has_env_var(name.as_bytes()))
+                    .collect();
+
+                stack.remove_overlay(&overlay_name.item);
+
+                for (name, val) in env_vars_to_keep {
+                    stack.add_env_var(name, val);
+                }
+            } else {
+                return Err(ShellError::OverlayNotFoundAtRuntime(
+                    overlay_name.item,
+                    overlay_name.span,
+                ));
+            }
+        } else {
+            stack.remove_overlay(&overlay_name.item);
+        }
 
         Ok(PipelineData::new(call.head))
     }