Allow overlays to import prefixed definitions (#6301)

* WIP

* Fix overlay prefix not preserving correctly

* Work around failing REPL tests

* Remove wrong code when removing with --keep-custom

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

@@ -101,7 +101,11 @@ impl EngineState {
             blocks: vec![],
             modules: vec![Module::new()],
             // make sure we have some default overlay:
-            scope: ScopeFrame::with_empty_overlay(DEFAULT_OVERLAY_NAME.as_bytes().to_vec(), 0),
+            scope: ScopeFrame::with_empty_overlay(
+                DEFAULT_OVERLAY_NAME.as_bytes().to_vec(),
+                0,
+                false,
+            ),
             ctrlc: None,
             env_vars: EnvVars::from([(DEFAULT_OVERLAY_NAME.to_string(), HashMap::new())]),
             previous_env_vars: HashMap::new(),
@@ -832,9 +836,11 @@ pub struct StateDelta {
 
 impl StateDelta {
     pub fn new(engine_state: &EngineState) -> Self {
+        let last_overlay = engine_state.last_overlay(&[]);
         let scope_frame = ScopeFrame::with_empty_overlay(
             engine_state.last_overlay_name(&[]).to_owned(),
-            engine_state.last_overlay(&[]).origin,
+            last_overlay.origin,
+            last_overlay.prefixed,
         );
 
         StateDelta {
@@ -1726,16 +1732,16 @@ impl<'a> StateWorkingSet<'a> {
         self.permanent_state.has_overlay(name)
     }
 
-    pub fn find_overlay_origin(&self, name: &[u8]) -> Option<ModuleId> {
+    pub fn find_overlay(&self, name: &[u8]) -> Option<&OverlayFrame> {
         for scope_frame in self.delta.scope.iter().rev() {
             if let Some(overlay_id) = scope_frame.find_overlay(name) {
-                return Some(scope_frame.get_overlay(overlay_id).origin);
+                return Some(scope_frame.get_overlay(overlay_id));
             }
         }
 
         self.permanent_state
             .find_overlay(name)
-            .map(|id| self.permanent_state.get_overlay(id).origin)
+            .map(|id| self.permanent_state.get_overlay(id))
     }
 
     pub fn last_overlay_name(&self) -> &Vec<u8> {
@@ -1775,9 +1781,11 @@ impl<'a> StateWorkingSet<'a> {
     pub fn last_overlay_mut(&mut self) -> &mut OverlayFrame {
         if self.delta.last_overlay_mut().is_none() {
             // If there is no overlay, automatically activate the last one
+            let overlay_frame = self.last_overlay();
             let name = self.last_overlay_name().to_vec();
-            let origin = self.last_overlay().origin;
-            self.add_overlay(name, origin, vec![], vec![]);
+            let origin = overlay_frame.origin;
+            let prefixed = overlay_frame.prefixed;
+            self.add_overlay(name, origin, vec![], vec![], prefixed);
         }
 
         self.delta
@@ -1841,6 +1849,7 @@ impl<'a> StateWorkingSet<'a> {
         origin: ModuleId,
         decls: Vec<(Vec<u8>, DeclId)>,
         aliases: Vec<(Vec<u8>, AliasId)>,
+        prefixed: bool,
     ) {
         let last_scope_frame = self.delta.last_scope_frame_mut();
 
@@ -1855,7 +1864,7 @@ impl<'a> StateWorkingSet<'a> {
         } else {
             last_scope_frame
                 .overlays
-                .push((name, OverlayFrame::from_origin(origin)));
+                .push((name, OverlayFrame::from_origin(origin, prefixed)));
             last_scope_frame.overlays.len() - 1
         };
 
@@ -1873,7 +1882,7 @@ impl<'a> StateWorkingSet<'a> {
     pub fn remove_overlay(&mut self, name: &[u8], keep_custom: bool) {
         let last_scope_frame = self.delta.last_scope_frame_mut();
 
-        let removed_overlay_origin = if let Some(overlay_id) = last_scope_frame.find_overlay(name) {
+        let maybe_module_id = if let Some(overlay_id) = last_scope_frame.find_overlay(name) {
             last_scope_frame
                 .active_overlays
                 .retain(|id| id != &overlay_id);
@@ -1885,7 +1894,7 @@ impl<'a> StateWorkingSet<'a> {
                 .map(|id| self.permanent_state.get_overlay(id).origin)
         };
 
-        if let Some(module_id) = removed_overlay_origin {
+        if let Some(module_id) = maybe_module_id {
             last_scope_frame.removed_overlays.push(name.to_owned());
 
             if keep_custom {

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

@@ -100,9 +100,9 @@ impl ScopeFrame {
         }
     }
 
-    pub fn with_empty_overlay(name: Vec<u8>, origin: ModuleId) -> Self {
+    pub fn with_empty_overlay(name: Vec<u8>, origin: ModuleId, prefixed: bool) -> Self {
         Self {
-            overlays: vec![(name, OverlayFrame::from_origin(origin))],
+            overlays: vec![(name, OverlayFrame::from_origin(origin, prefixed))],
             active_overlays: vec![0],
             removed_overlays: vec![],
             predecls: HashMap::new(),
@@ -205,10 +205,11 @@ pub struct OverlayFrame {
     pub modules: HashMap<Vec<u8>, ModuleId>,
     pub visibility: Visibility,
     pub origin: ModuleId, // The original module the overlay was created from
+    pub prefixed: bool,   // Whether the overlay has definitions prefixed with its name
 }
 
 impl OverlayFrame {
-    pub fn from_origin(origin: ModuleId) -> Self {
+    pub fn from_origin(origin: ModuleId, prefixed: bool) -> Self {
         Self {
             vars: HashMap::new(),
             predecls: HashMap::new(),
@@ -217,6 +218,7 @@ impl OverlayFrame {
             modules: HashMap::new(),
             visibility: Visibility::new(),
             origin,
+            prefixed,
         }
     }