feat(tvix/eval): add data structures for tracking spans in chunks r/4709

This adds a new vector to the chunk data structure which tracks spans
into a codemap. The compiler will emit this information to the chunk
when adding instructions.

The internal representation of the spans is slightly optimised to
avoid storing duplicate spans, as there are cases where many
instructions might be derived from the same span.

Change-Id: I336f8c912e7eb50ea02ed71e6164f651ca3ca790
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6376
Tested-by: BuildkiteCI
Reviewed-by: sterni <sternenseemann@systemli.org>

1 files changed, 50 insertions, 0 deletions

diff --git a/tvix/eval/src/chunk.rs b/tvix/eval/src/chunk.rs
index 7526add169..9075db11b3 100644
--- a/tvix/eval/src/chunk.rs
+++ b/tvix/eval/src/chunk.rs
@@ -1,10 +1,31 @@
 use crate::opcode::{CodeIdx, ConstantIdx, OpCode};
 use crate::value::Value;
 
+/// Represents a source location from which one or more operations
+/// were compiled.
+///
+/// The span itself is an index into a [codemap::Codemap], and the
+/// structure tracks the number of operations that were yielded from
+/// the same span.
+///
+/// At error reporting time, it becomes possible to either just fetch
+/// the textual representation of that span from the codemap, or to
+/// even re-parse the AST using rnix to create more semantically
+/// interesting errors.
+#[derive(Clone, Debug)]
+struct SourceSpan {
+    /// Span into the [codemap::Codemap].
+    span: codemap::Span,
+
+    /// Number of instructions derived from this span.
+    count: usize,
+}
+
 #[derive(Clone, Debug, Default)]
 pub struct Chunk {
     pub code: Vec<OpCode>,
     pub constants: Vec<Value>,
+    spans: Vec<SourceSpan>,
 }
 
 impl Chunk {
@@ -23,4 +44,33 @@ impl Chunk {
     pub fn constant(&self, idx: ConstantIdx) -> &Value {
         &self.constants[idx.0]
     }
+
+    // Span tracking implementation
+
+    fn push_span(&mut self, span: codemap::Span) {
+        match self.spans.last_mut() {
+            // We do not need to insert the same span again, as this
+            // instruction was compiled from the same span as the last
+            // one.
+            Some(last) if last.span == span => last.count += 1,
+
+            // In all other cases, this is a new source span.
+            _ => self.spans.push(SourceSpan { span, count: 1 }),
+        }
+    }
+
+    /// Retrieve the [codemap::Span] from which the instruction at
+    /// `offset` was compiled.
+    pub fn get_span(&self, offset: CodeIdx) -> codemap::Span {
+        let mut pos = 0;
+
+        for span in &self.spans {
+            pos += span.count;
+            if pos > offset.0 {
+                return span.span;
+            }
+        }
+
+        panic!("compiler error: chunk missing span for offset {}", offset.0);
+    }
 }