From 9a482ce2844008e4e284e9662fb8dd345793d99d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20=C5=BD=C3=A1dn=C3=ADk?= Date: Wed, 25 May 2022 00:22:17 +0300 Subject: [PATCH] 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 --- .../src/core_commands/overlay/add.rs | 93 ++++------- .../src/core_commands/overlay/remove.rs | 47 +++++- crates/nu-parser/src/parse_keywords.rs | 13 +- crates/nu-protocol/src/engine/engine_state.rs | 90 ++++++++-- crates/nu-protocol/src/engine/overlay.rs | 47 +----- crates/nu-protocol/src/engine/stack.rs | 60 ++++++- crates/nu-protocol/src/shell_error.rs | 2 +- tests/overlays/mod.rs | 156 ++++++++++++++++++ 8 files changed, 363 insertions(+), 145 deletions(-) diff --git a/crates/nu-command/src/core_commands/overlay/add.rs b/crates/nu-command/src/core_commands/overlay/add.rs index d17f39c985..3cbde45267 100644 --- a/crates/nu-command/src/core_commands/overlay/add.rs +++ b/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 { let name_arg: Spanned = 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)) diff --git a/crates/nu-command/src/core_commands/overlay/remove.rs b/crates/nu-command/src/core_commands/overlay/remove.rs index 2cd216e18c..a64a4010b3 100644 --- a/crates/nu-command/src/core_commands/overlay/remove.rs +++ b/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 { - // let module_name: Spanned = call.req(engine_state, stack, 0)?; - - let module_name: Spanned = if let Some(name) = call.opt(engine_state, stack, 0)? { + let overlay_name: Spanned = 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)) } diff --git a/crates/nu-parser/src/parse_keywords.rs b/crates/nu-parser/src/parse_keywords.rs index 8f17a0475b..9ada7f47c5 100644 --- a/crates/nu-parser/src/parse_keywords.rs +++ b/crates/nu-parser/src/parse_keywords.rs @@ -2063,6 +2063,8 @@ pub fn parse_overlay_remove( ) }; + let keep_custom = call.has_flag("keep-custom"); + let pipeline = Pipeline::from_vec(vec![Expression { expr: Expr::Call(call), span: span(spans), @@ -2097,16 +2099,7 @@ pub fn parse_overlay_remove( ); } - // let original_module = if call.has_flag("discard") { - // None - // } else if let Some(module_id) = working_set.find_module(overlay_name.as_bytes()) { - // // TODO: Remove clone - // Some(working_set.get_module(module_id).clone()) - // } else { - // Some(Module::new()) - // }; - - working_set.remove_overlay(overlay_name.as_bytes()); + working_set.remove_overlay(overlay_name.as_bytes(), keep_custom); (pipeline, None) } diff --git a/crates/nu-protocol/src/engine/engine_state.rs b/crates/nu-protocol/src/engine/engine_state.rs index d4aa735d6f..98bc7709d0 100644 --- a/crates/nu-protocol/src/engine/engine_state.rs +++ b/crates/nu-protocol/src/engine/engine_state.rs @@ -1692,6 +1692,56 @@ impl<'a> StateWorkingSet<'a> { .expect("internal error: missing added overlay") } + /// Collect all decls that belong to an overlay + pub fn decls_of_overlay(&self, name: &[u8]) -> HashMap, DeclId> { + let mut result = HashMap::new(); + + if let Some(overlay_id) = self.permanent_state.find_overlay(name) { + let overlay_frame = self.permanent_state.get_overlay(overlay_id); + + for (decl_name, decl_id) in &overlay_frame.decls { + result.insert(decl_name.to_owned(), *decl_id); + } + } + + for scope_frame in self.delta.scope.iter() { + if let Some(overlay_id) = scope_frame.find_overlay(name) { + let overlay_frame = scope_frame.get_overlay(overlay_id); + + for (decl_name, decl_id) in &overlay_frame.decls { + result.insert(decl_name.to_owned(), *decl_id); + } + } + } + + result + } + + /// Collect all aliases that belong to an overlay + pub fn aliases_of_overlay(&self, name: &[u8]) -> HashMap, DeclId> { + let mut result = HashMap::new(); + + if let Some(overlay_id) = self.permanent_state.find_overlay(name) { + let overlay_frame = self.permanent_state.get_overlay(overlay_id); + + for (alias_name, alias_id) in &overlay_frame.aliases { + result.insert(alias_name.to_owned(), *alias_id); + } + } + + for scope_frame in self.delta.scope.iter() { + if let Some(overlay_id) = scope_frame.find_overlay(name) { + let overlay_frame = scope_frame.get_overlay(overlay_id); + + for (alias_name, alias_id) in &overlay_frame.aliases { + result.insert(alias_name.to_owned(), *alias_id); + } + } + } + + result + } + pub fn add_overlay( &mut self, name: Vec, @@ -1712,7 +1762,7 @@ impl<'a> StateWorkingSet<'a> { } else { last_scope_frame .overlays - .push((name, OverlayFrame::from(origin))); + .push((name, OverlayFrame::from_origin(origin))); last_scope_frame.overlays.len() - 1 }; @@ -1727,33 +1777,43 @@ impl<'a> StateWorkingSet<'a> { self.use_aliases(aliases); } - pub fn remove_overlay(&mut self, name: &[u8]) { + pub fn remove_overlay(&mut self, name: &[u8], keep_custom: bool) { let last_scope_frame = self.delta.last_scope_frame_mut(); - let removed_overlay = if let Some(overlay_id) = last_scope_frame.find_overlay(name) { + let removed_overlay_origin = if let Some(overlay_id) = last_scope_frame.find_overlay(name) { last_scope_frame .active_overlays .retain(|id| id != &overlay_id); - Some(last_scope_frame.get_overlay(overlay_id).clone()) + Some(last_scope_frame.get_overlay(overlay_id).origin) } else { self.permanent_state .find_overlay(name) - .map(|id| self.permanent_state.get_overlay(id).clone()) + .map(|id| self.permanent_state.get_overlay(id).origin) }; - if removed_overlay.is_some() { + if let Some(module_id) = removed_overlay_origin { last_scope_frame.removed_overlays.push(name.to_owned()); + + if keep_custom { + let origin_module = self.get_module(module_id); + + let decls = self + .decls_of_overlay(name) + .into_iter() + .filter(|(n, _)| !origin_module.has_decl(n)) + .collect(); + + let aliases = self + .aliases_of_overlay(name) + .into_iter() + .filter(|(n, _)| !origin_module.has_alias(n)) + .collect(); + + self.use_decls(decls); + self.use_aliases(aliases); + } } - - // if let Some(module) = original_module { - // let last_overlay_name = self.last_overlay_name().to_owned(); - - // if let Some(overlay) = removed_overlay { - // let (diff_decls, diff_aliases) = overlay.diff(&module); - // self.add_overlay(last_overlay_name, diff_decls, diff_aliases); - // } - // } } pub fn render(self) -> StateDelta { diff --git a/crates/nu-protocol/src/engine/overlay.rs b/crates/nu-protocol/src/engine/overlay.rs index 3c3e0655e9..7d04be0f4c 100644 --- a/crates/nu-protocol/src/engine/overlay.rs +++ b/crates/nu-protocol/src/engine/overlay.rs @@ -101,7 +101,7 @@ impl ScopeFrame { pub fn with_empty_overlay(name: Vec, origin: ModuleId) -> Self { Self { - overlays: vec![(name, OverlayFrame::from(origin))], + overlays: vec![(name, OverlayFrame::from_origin(origin))], active_overlays: vec![0], removed_overlays: vec![], predecls: HashMap::new(), @@ -195,8 +195,6 @@ impl ScopeFrame { } } -// type OverlayDiff = (Vec<(Vec, DeclId)>, Vec<(Vec, AliasId)>); - #[derive(Debug, Clone)] pub struct OverlayFrame { pub vars: HashMap, VarId>, @@ -209,7 +207,7 @@ pub struct OverlayFrame { } impl OverlayFrame { - pub fn from(origin: ModuleId) -> Self { + pub fn from_origin(origin: ModuleId) -> Self { Self { vars: HashMap::new(), predecls: HashMap::new(), @@ -220,45 +218,4 @@ impl OverlayFrame { origin, } } - - // Find out which definitions are custom compared to the origin module - // pub fn diff(&self, engine_state: &EngineState) -> OverlayDiff { - // let module = engine_state.get_module(self.origin); - - // let decls = self - // .decls - // .iter() - // .filter(|(name, decl_id)| { - // if self.visibility.is_decl_id_visible(decl_id) { - // if let Some(original_id) = module.get_decl_id(name) { - // &original_id != *decl_id - // } else { - // true - // } - // } else { - // false - // } - // }) - // .map(|(name, decl_id)| (name.to_owned(), *decl_id)) - // .collect(); - - // let aliases = self - // .aliases - // .iter() - // .filter(|(name, alias_id)| { - // if self.visibility.is_alias_id_visible(alias_id) { - // if let Some(original_id) = module.get_alias_id(name) { - // &original_id != *alias_id - // } else { - // true - // } - // } else { - // false - // } - // }) - // .map(|(name, alias_id)| (name.to_owned(), *alias_id)) - // .collect(); - - // (decls, aliases) - // } } diff --git a/crates/nu-protocol/src/engine/stack.rs b/crates/nu-protocol/src/engine/stack.rs index 8b7a35c4b7..935669fcb9 100644 --- a/crates/nu-protocol/src/engine/stack.rs +++ b/crates/nu-protocol/src/engine/stack.rs @@ -178,6 +178,39 @@ impl Stack { result } + /// Flatten the env var scope frames into one frame, only from one overlay + pub fn get_overlay_env_vars( + &self, + engine_state: &EngineState, + overlay_name: &str, + ) -> HashMap { + let mut result = HashMap::new(); + + // for active_overlay in self.active_overlays.iter() { + if let Some(active_overlay) = self.active_overlays.iter().find(|n| n == &overlay_name) { + if let Some(env_vars) = engine_state.env_vars.get(active_overlay) { + result.extend( + env_vars + .iter() + .filter(|(k, _)| { + if let Some(env_hidden) = self.env_hidden.get(active_overlay) { + !env_hidden.contains(*k) + } else { + // nothing has been hidden in this overlay + true + } + }) + .map(|(k, v)| (k.clone(), v.clone())) + .collect::>(), + ); + } + } + + result.extend(self.get_stack_overlay_env_vars(overlay_name)); + + result + } + /// Get flattened environment variables only from the stack pub fn get_stack_env_vars(&self) -> HashMap { let mut result = HashMap::new(); @@ -193,6 +226,21 @@ impl Stack { result } + /// Get flattened environment variables only from the stack and one overlay + pub fn get_stack_overlay_env_vars(&self, overlay_name: &str) -> HashMap { + let mut result = HashMap::new(); + + for scope in &self.env_vars { + if let Some(active_overlay) = self.active_overlays.iter().find(|n| n == &overlay_name) { + if let Some(env_vars) = scope.get(active_overlay) { + result.extend(env_vars.clone()); + } + } + } + + result + } + /// Same as get_env_vars, but returns only the names as a HashSet pub fn get_env_var_names(&self, engine_state: &EngineState) -> HashSet { let mut result = HashSet::new(); @@ -326,6 +374,10 @@ impl Stack { engine_state.env_vars.contains_key(name) } + pub fn is_overlay_active(&self, name: &String) -> bool { + self.active_overlays.contains(name) + } + pub fn add_overlay(&mut self, name: String) { self.env_hidden.remove(&name); @@ -333,14 +385,8 @@ impl Stack { self.active_overlays.push(name); } - pub fn remove_overlay(&mut self, name: &String, span: &Span) -> Result<(), ShellError> { - if !self.active_overlays.contains(name) { - return Err(ShellError::OverlayNotFoundAtRuntime(name.into(), *span)); - } - + pub fn remove_overlay(&mut self, name: &String) { self.active_overlays.retain(|o| o != name); - - Ok(()) } } diff --git a/crates/nu-protocol/src/shell_error.rs b/crates/nu-protocol/src/shell_error.rs index 052b87dea1..5d01efc2fd 100644 --- a/crates/nu-protocol/src/shell_error.rs +++ b/crates/nu-protocol/src/shell_error.rs @@ -234,7 +234,7 @@ pub enum ShellError { /// /// Check the module name. Did you typo it? Did you forget to declare it? Is the casing right? #[error("Module or overlay'{0}' not found")] - #[diagnostic(code(nu::shell::module_not_found), url(docsrs))] + #[diagnostic(code(nu::shell::module_or_overlay_not_found), url(docsrs))] ModuleOrOverlayNotFoundAtRuntime(String, #[label = "not a module or overlay"] Span), /// A referenced overlay was not found at runtime. diff --git a/tests/overlays/mod.rs b/tests/overlays/mod.rs index fdc5966f4b..d8772738ad 100644 --- a/tests/overlays/mod.rs +++ b/tests/overlays/mod.rs @@ -307,6 +307,162 @@ fn remove_overlay_discard_env() { assert!(actual_repl.err.contains("DidYouMean")); } +#[test] +fn remove_overlay_keep_decl() { + let inp = &[ + r#"overlay add samples/spam.nu"#, + r#"def bagr [] { "bagr" }"#, + r#"overlay remove --keep-custom spam"#, + r#"bagr"#, + ]; + + let actual = nu!(cwd: "tests/overlays", pipeline(&inp.join("; "))); + let actual_repl = nu_repl("tests/overlays", inp); + + assert!(actual.out.contains("bagr")); + assert!(actual_repl.out.contains("bagr")); +} + +#[test] +fn remove_overlay_keep_alias() { + let inp = &[ + r#"overlay add samples/spam.nu"#, + r#"alias bagr = "bagr""#, + r#"overlay remove --keep-custom spam"#, + r#"bagr"#, + ]; + + let actual = nu!(cwd: "tests/overlays", pipeline(&inp.join("; "))); + let actual_repl = nu_repl("tests/overlays", inp); + + assert!(actual.out.contains("bagr")); + assert!(actual_repl.out.contains("bagr")); +} + +#[test] +fn remove_overlay_keep_env() { + let inp = &[ + r#"overlay add samples/spam.nu"#, + r#"let-env BAGR = "bagr""#, + r#"overlay remove --keep-custom spam"#, + r#"$env.BAGR"#, + ]; + + let actual = nu!(cwd: "tests/overlays", pipeline(&inp.join("; "))); + let actual_repl = nu_repl("tests/overlays", inp); + + assert!(actual.out.contains("bagr")); + assert!(actual_repl.out.contains("bagr")); +} + +#[test] +fn remove_overlay_keep_discard_overwritten_decl() { + let inp = &[ + r#"overlay add samples/spam.nu"#, + r#"def foo [] { 'bar' }"#, + r#"overlay remove --keep-custom spam"#, + r#"foo"#, + ]; + + let actual = nu!(cwd: "tests/overlays", pipeline(&inp.join("; "))); + let actual_repl = nu_repl("tests/overlays", inp); + + assert!(!actual.err.is_empty()); + #[cfg(windows)] + assert!(actual_repl.out != "bagr"); + #[cfg(not(windows))] + assert!(!actual_repl.err.is_empty()); +} + +#[test] +fn remove_overlay_keep_discard_overwritten_alias() { + let inp = &[ + r#"overlay add samples/spam.nu"#, + r#"alias bar = 'baz'"#, + r#"overlay remove --keep-custom spam"#, + r#"bar"#, + ]; + + let actual = nu!(cwd: "tests/overlays", pipeline(&inp.join("; "))); + let actual_repl = nu_repl("tests/overlays", inp); + + assert!(!actual.err.is_empty()); + #[cfg(windows)] + assert!(actual_repl.out != "bagr"); + #[cfg(not(windows))] + assert!(!actual_repl.err.is_empty()); +} + +#[test] +fn remove_overlay_keep_discard_overwritten_env() { + let inp = &[ + r#"overlay add samples/spam.nu"#, + r#"let-env BAZ = "bagr""#, + r#"overlay remove --keep-custom spam"#, + r#"$env.BAZ"#, + ]; + + let actual = nu!(cwd: "tests/overlays", pipeline(&inp.join("; "))); + let actual_repl = nu_repl("tests/overlays", inp); + + assert!(actual.err.contains("did you mean")); + assert!(actual_repl.err.contains("DidYouMean")); +} + +#[test] +fn remove_overlay_keep_decl_in_latest_overlay() { + let inp = &[ + r#"overlay add samples/spam.nu"#, + r#"def bagr [] { 'bagr' }"#, + r#"module eggs { }"#, + r#"overlay add eggs"#, + r#"overlay remove --keep-custom spam"#, + r#"bagr"#, + ]; + + let actual = nu!(cwd: "tests/overlays", pipeline(&inp.join("; "))); + let actual_repl = nu_repl("tests/overlays", inp); + + assert!(actual.out.contains("bagr")); + assert!(actual_repl.out.contains("bagr")); +} + +#[test] +fn remove_overlay_keep_alias_in_latest_overlay() { + let inp = &[ + r#"overlay add samples/spam.nu"#, + r#"alias bagr = 'bagr'"#, + r#"module eggs { }"#, + r#"overlay add eggs"#, + r#"overlay remove --keep-custom spam"#, + r#"bagr"#, + ]; + + let actual = nu!(cwd: "tests/overlays", pipeline(&inp.join("; "))); + let actual_repl = nu_repl("tests/overlays", inp); + + assert!(actual.out.contains("bagr")); + assert!(actual_repl.out.contains("bagr")); +} + +#[test] +fn remove_overlay_keep_env_in_latest_overlay() { + let inp = &[ + r#"overlay add samples/spam.nu"#, + r#"let-env BAGR = "bagr""#, + r#"module eggs { }"#, + r#"overlay add eggs"#, + r#"overlay remove --keep-custom spam"#, + r#"$env.BAGR"#, + ]; + + let actual = nu!(cwd: "tests/overlays", pipeline(&inp.join("; "))); + let actual_repl = nu_repl("tests/overlays", inp); + + assert!(actual.out.contains("bagr")); + assert!(actual_repl.out.contains("bagr")); +} + #[test] fn preserve_overrides() { let inp = &[