Span ID Refactor - Step 1 (#12960)

# Description
First part of SpanID refactoring series. This PR adds a `SpanId` type
and a corresponding `span_id` field to `Expression`. Parser creating
expressions will now add them to an array in `StateWorkingSet`,
generates a corresponding ID and saves the ID to the Expression. The IDs
are not used anywhere yet.

For the rough overall plan, see
https://github.com/nushell/nushell/issues/12963.

# User-Facing Changes
Hopefully none. This is only a refactor of Nushell's internals that
shouldn't have visible side effects.

# Tests + Formatting

# After Submitting

crates/nu-protocol/src/ast/call.rs

@@ -338,9 +338,13 @@ impl Call {
 #[cfg(test)]
 mod test {
     use super::*;
+    use crate::engine::EngineState;
 
     #[test]
     fn argument_span_named() {
+        let engine_state = EngineState::new();
+        let mut working_set = StateWorkingSet::new(&engine_state);
+
         let named = Spanned {
             item: "named".to_string(),
             span: Span::new(2, 3),
@@ -349,7 +353,7 @@ mod test {
             item: "short".to_string(),
             span: Span::new(5, 7),
         };
-        let expr = Expression::garbage(Span::new(11, 13));
+        let expr = Expression::garbage(&mut working_set, Span::new(11, 13));
 
         let arg = Argument::Named((named.clone(), None, None));
 
@@ -370,8 +374,11 @@ mod test {
 
     #[test]
     fn argument_span_positional() {
+        let engine_state = EngineState::new();
+        let mut working_set = StateWorkingSet::new(&engine_state);
+
         let span = Span::new(2, 3);
-        let expr = Expression::garbage(span);
+        let expr = Expression::garbage(&mut working_set, span);
         let arg = Argument::Positional(expr);
 
         assert_eq!(span, arg.span());
@@ -379,8 +386,11 @@ mod test {
 
     #[test]
     fn argument_span_unknown() {
+        let engine_state = EngineState::new();
+        let mut working_set = StateWorkingSet::new(&engine_state);
+
         let span = Span::new(2, 3);
-        let expr = Expression::garbage(span);
+        let expr = Expression::garbage(&mut working_set, span);
         let arg = Argument::Unknown(expr);
 
         assert_eq!(span, arg.span());
@@ -388,9 +398,12 @@ mod test {
 
     #[test]
     fn call_arguments_span() {
+        let engine_state = EngineState::new();
+        let mut working_set = StateWorkingSet::new(&engine_state);
+
         let mut call = Call::new(Span::new(0, 1));
-        call.add_positional(Expression::garbage(Span::new(2, 3)));
-        call.add_positional(Expression::garbage(Span::new(5, 7)));
+        call.add_positional(Expression::garbage(&mut working_set, Span::new(2, 3)));
+        call.add_positional(Expression::garbage(&mut working_set, Span::new(5, 7)));
 
         assert_eq!(Span::new(2, 7), call.arguments_span());
     }

crates/nu-protocol/src/ast/expression.rs

@@ -1,7 +1,7 @@
 use crate::{
     ast::{Argument, Block, Expr, ExternalArgument, ImportPattern, MatchPattern, RecordItem},
-    engine::StateWorkingSet,
-    BlockId, DeclId, Signature, Span, Type, VarId, IN_VARIABLE_ID,
+    engine::{EngineState, StateWorkingSet},
+    BlockId, DeclId, Signature, Span, SpanId, Type, VarId, IN_VARIABLE_ID,
 };
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
@@ -10,15 +10,18 @@ use std::sync::Arc;
 pub struct Expression {
     pub expr: Expr,
     pub span: Span,
+    pub span_id: SpanId,
     pub ty: Type,
     pub custom_completion: Option<DeclId>,
 }
 
 impl Expression {
-    pub fn garbage(span: Span) -> Expression {
+    pub fn garbage(working_set: &mut StateWorkingSet, span: Span) -> Expression {
+        let span_id = working_set.add_span(span);
         Expression {
             expr: Expr::Garbage,
             span,
+            span_id,
             ty: Type::Any,
             custom_completion: None,
         }
@@ -471,4 +474,49 @@ impl Expression {
             Expr::VarDecl(_) => {}
         }
     }
+
+    pub fn new(working_set: &mut StateWorkingSet, expr: Expr, span: Span, ty: Type) -> Expression {
+        let span_id = working_set.add_span(span);
+        Expression {
+            expr,
+            span,
+            span_id,
+            ty,
+            custom_completion: None,
+        }
+    }
+
+    pub fn new_existing(expr: Expr, span: Span, span_id: SpanId, ty: Type) -> Expression {
+        Expression {
+            expr,
+            span,
+            span_id,
+            ty,
+            custom_completion: None,
+        }
+    }
+
+    pub fn new_unknown(expr: Expr, span: Span, ty: Type) -> Expression {
+        Expression {
+            expr,
+            span,
+            span_id: SpanId(0),
+            ty,
+            custom_completion: None,
+        }
+    }
+
+    pub fn with_span_id(self, span_id: SpanId) -> Expression {
+        Expression {
+            expr: self.expr,
+            span: self.span,
+            span_id,
+            ty: self.ty,
+            custom_completion: self.custom_completion,
+        }
+    }
+
+    pub fn span(&self, engine_state: &EngineState) -> Span {
+        engine_state.get_span(self.span_id)
+    }
 }

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

@@ -8,7 +8,7 @@ use crate::{
     },
     eval_const::create_nu_constant,
     BlockId, Category, Config, DeclId, FileId, HistoryConfig, Module, ModuleId, OverlayId,
-    ShellError, Signature, Span, Type, Value, VarId, VirtualPathId,
+    ShellError, Signature, Span, SpanId, Type, Value, VarId, VirtualPathId,
 };
 use fancy_regex::Regex;
 use lru::LruCache;
@@ -81,6 +81,7 @@ pub struct EngineState {
     // especially long, so it helps
     pub(super) blocks: Arc<Vec<Arc<Block>>>,
     pub(super) modules: Arc<Vec<Arc<Module>>>,
+    pub spans: Vec<Span>,
     usage: Usage,
     pub scope: ScopeFrame,
     pub ctrlc: Option<Arc<AtomicBool>>,
@@ -115,6 +116,9 @@ pub const IN_VARIABLE_ID: usize = 1;
 pub const ENV_VARIABLE_ID: usize = 2;
 // NOTE: If you add more to this list, make sure to update the > checks based on the last in the list
 
+// The first span is unknown span
+pub const UNKNOWN_SPAN_ID: SpanId = SpanId(0);
+
 impl EngineState {
     pub fn new() -> Self {
         Self {
@@ -132,6 +136,7 @@ impl EngineState {
             modules: Arc::new(vec![Arc::new(Module::new(
                 DEFAULT_OVERLAY_NAME.as_bytes().to_vec(),
             ))]),
+            spans: vec![Span::unknown()],
             usage: Usage::new(),
             // make sure we have some default overlay:
             scope: ScopeFrame::with_empty_overlay(
@@ -184,6 +189,7 @@ impl EngineState {
         self.files.extend(delta.files);
         self.virtual_paths.extend(delta.virtual_paths);
         self.vars.extend(delta.vars);
+        self.spans.extend(delta.spans);
         self.usage.merge_with(delta.usage);
 
         // Avoid potentially cloning the Arcs if we aren't adding anything
@@ -565,6 +571,9 @@ impl EngineState {
         self.modules.len()
     }
 
+    pub fn num_spans(&self) -> usize {
+        self.spans.len()
+    }
     pub fn print_vars(&self) {
         for var in self.vars.iter().enumerate() {
             println!("var{}: {:?}", var.0, var.1);
@@ -1019,6 +1028,25 @@ impl EngineState {
             )));
         }
     }
+
+    /// Add new span and return its ID
+    pub fn add_span(&mut self, span: Span) -> SpanId {
+        self.spans.push(span);
+        SpanId(self.num_spans() - 1)
+    }
+
+    /// Get existing span
+    pub fn get_span(&self, span_id: SpanId) -> Span {
+        *self
+            .spans
+            .get(span_id.0)
+            .expect("internal error: missing span")
+    }
+
+    /// Find ID of a span (should be avoided if possible)
+    pub fn find_span_id(&self, span: Span) -> Option<SpanId> {
+        self.spans.iter().position(|sp| sp == &span).map(SpanId)
+    }
 }
 
 impl Default for EngineState {

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

@@ -4,7 +4,7 @@ use crate::{
         usage::Usage, CachedFile, Command, EngineState, OverlayFrame, ScopeFrame, Variable,
         VirtualPath,
     },
-    Module,
+    Module, Span,
 };
 use std::sync::Arc;
 
@@ -21,6 +21,7 @@ pub struct StateDelta {
     pub(super) decls: Vec<Box<dyn Command>>, // indexed by DeclId
     pub blocks: Vec<Arc<Block>>,             // indexed by BlockId
     pub(super) modules: Vec<Arc<Module>>,    // indexed by ModuleId
+    pub spans: Vec<Span>,                    // indexed by SpanId
     pub(super) usage: Usage,
     pub scope: Vec<ScopeFrame>,
     #[cfg(feature = "plugin")]
@@ -45,6 +46,7 @@ impl StateDelta {
             decls: vec![],
             blocks: vec![],
             modules: vec![],
+            spans: vec![],
             scope: vec![scope_frame],
             usage: Usage::new(),
             #[cfg(feature = "plugin")]

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

@@ -5,7 +5,7 @@ use crate::{
         StateDelta, Variable, VirtualPath, Visibility,
     },
     BlockId, Category, Config, DeclId, FileId, Module, ModuleId, ParseError, ParseWarning, Span,
-    Type, Value, VarId, VirtualPathId,
+    SpanId, Type, Value, VarId, VirtualPathId,
 };
 use core::panic;
 use std::{
@@ -1013,6 +1013,25 @@ impl<'a> StateWorkingSet<'a> {
                 .expect("internal error: missing virtual path")
         }
     }
+
+    pub fn add_span(&mut self, span: Span) -> SpanId {
+        let num_permanent_spans = self.permanent_state.spans.len();
+        self.delta.spans.push(span);
+        SpanId(num_permanent_spans + self.delta.spans.len() - 1)
+    }
+
+    pub fn get_span(&self, span_id: SpanId) -> Span {
+        let num_permanent_spans = self.permanent_state.num_spans();
+        if span_id.0 < num_permanent_spans {
+            self.permanent_state.get_span(span_id)
+        } else {
+            *self
+                .delta
+                .spans
+                .get(span_id.0 - num_permanent_spans)
+                .expect("internal error: missing span")
+        }
+    }
 }
 
 impl<'a> miette::SourceCode for &StateWorkingSet<'a> {

crates/nu-protocol/src/id.rs

@@ -1,3 +1,5 @@
+use serde::{Deserialize, Serialize};
+
 pub type VarId = usize;
 pub type DeclId = usize;
 pub type BlockId = usize;
@@ -5,3 +7,5 @@ pub type ModuleId = usize;
 pub type OverlayId = usize;
 pub type FileId = usize;
 pub type VirtualPathId = usize;
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
+pub struct SpanId(pub usize); // more robust ID style used in the new parser