Add GetSpan trait

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

@@ -1,7 +1,7 @@
 use crate::{
     ast::{Argument, Block, Expr, ExternalArgument, ImportPattern, MatchPattern, RecordItem},
     engine::{EngineState, StateWorkingSet},
-    BlockId, DeclId, Signature, Span, SpanId, Type, VarId, IN_VARIABLE_ID,
+    BlockId, DeclId, GetSpan, Signature, Span, SpanId, Type, VarId, IN_VARIABLE_ID,
 };
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;

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

@@ -7,7 +7,7 @@ use crate::{
         Variable, Visibility, DEFAULT_OVERLAY_NAME,
     },
     eval_const::create_nu_constant,
-    BlockId, Category, Config, DeclId, FileId, HistoryConfig, Module, ModuleId, OverlayId,
+    BlockId, Category, Config, DeclId, FileId, GetSpan, HistoryConfig, Module, ModuleId, OverlayId,
     ShellError, Signature, Span, SpanId, Type, Value, VarId, VirtualPathId,
 };
 use fancy_regex::Regex;
@@ -1035,18 +1035,20 @@ impl EngineState {
         SpanId(self.num_spans() - 1)
     }
 
+    /// 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<'a> GetSpan for &'a EngineState {
     /// Get existing span
-    pub fn get_span(&self, span_id: SpanId) -> Span {
+    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_working_set.rs

@@ -4,8 +4,8 @@ use crate::{
         usage::build_usage, CachedFile, Command, CommandType, EngineState, OverlayFrame,
         StateDelta, Variable, VirtualPath, Visibility,
     },
-    BlockId, Category, Config, DeclId, FileId, Module, ModuleId, ParseError, ParseWarning, Span,
-    SpanId, Type, Value, VarId, VirtualPathId,
+    BlockId, Category, Config, DeclId, FileId, GetSpan, Module, ModuleId, ParseError, ParseWarning,
+    Span, SpanId, Type, Value, VarId, VirtualPathId,
 };
 use core::panic;
 use std::{
@@ -1019,18 +1019,30 @@ impl<'a> StateWorkingSet<'a> {
         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> GetSpan for StateWorkingSet<'a> {
+    fn get_span(&self, span_id: SpanId) -> Span {
+        get_span(self, span_id)
+    }
+}
+
+impl<'a> GetSpan for &'a StateWorkingSet<'a> {
+    fn get_span(&self, span_id: SpanId) -> Span {
+        get_span(self, span_id)
+    }
+}
+
+fn get_span(working_set: &StateWorkingSet, span_id: SpanId) -> Span {
+    let num_permanent_spans = working_set.permanent_state.num_spans();
+    if span_id.0 < num_permanent_spans {
+        working_set.permanent_state.get_span(span_id)
+    } else {
+        *working_set
+            .delta
+            .spans
+            .get(span_id.0 - num_permanent_spans)
+            .expect("internal error: missing span")
     }
 }
 

crates/nu-protocol/src/span.rs

@@ -1,7 +1,12 @@
+use crate::SpanId;
 use miette::SourceSpan;
 use serde::{Deserialize, Serialize};
 use std::ops::Deref;
 
+pub trait GetSpan {
+    fn get_span(&self, span_id: SpanId) -> Span;
+}
+
 /// A spanned area of interest, generic over what kind of thing is of interest
 #[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub struct Spanned<T> {