From 0d346d58d77e5323a0e605026794d2ca89087f74 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20=C5=BD=C3=A1dn=C3=ADk?= <kubouch@gmail.com>
Date: Wed, 5 Jun 2024 22:48:43 +0300
Subject: [PATCH] Add GetSpan trait

---
 crates/nu-protocol/src/ast/expression.rs      |  2 +-
 crates/nu-protocol/src/engine/engine_state.rs | 16 ++++----
 .../src/engine/state_working_set.rs           | 38 ++++++++++++-------
 crates/nu-protocol/src/span.rs                |  5 +++
 4 files changed, 40 insertions(+), 21 deletions(-)

diff --git a/crates/nu-protocol/src/ast/expression.rs b/crates/nu-protocol/src/ast/expression.rs
index 832bc1f438..3ea0962dff 100644
--- a/crates/nu-protocol/src/ast/expression.rs
+++ b/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;
diff --git a/crates/nu-protocol/src/engine/engine_state.rs b/crates/nu-protocol/src/engine/engine_state.rs
index e81d3360ab..3960f6d9d0 100644
--- a/crates/nu-protocol/src/engine/engine_state.rs
+++ b/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 {
diff --git a/crates/nu-protocol/src/engine/state_working_set.rs b/crates/nu-protocol/src/engine/state_working_set.rs
index 31600b5eb9..93728f1cdf 100644
--- a/crates/nu-protocol/src/engine/state_working_set.rs
+++ b/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")
     }
 }
 
diff --git a/crates/nu-protocol/src/span.rs b/crates/nu-protocol/src/span.rs
index 3d32aa4ddf..0d280eaa9d 100644
--- a/crates/nu-protocol/src/span.rs
+++ b/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> {