Add source spans for assertion error messages

numbat/src/bytecode_interpreter.rs

@@ -487,9 +487,18 @@ impl BytecodeInterpreter {
 
                 let name = &ffi::procedures().get(kind).unwrap().name;
 
-                let idx = self.vm.get_ffi_callable_idx(name).unwrap();
-                self.vm
-                    .add_op2(Op::FFICallProcedure, idx, args.len() as u16); // TODO: check overflow
+                let callable_idx = self.vm.get_ffi_callable_idx(name).unwrap();
+
+                let arg_spans = args.iter().map(|a| a.full_span()).collect();
+                let spans_idx = self.vm.add_procedure_arg_span(arg_spans);
+
+                self.vm.add_op3(
+                    Op::FFICallProcedure,
+                    callable_idx,
+                    args.len() as u16,
+                    spans_idx,
+                );
+                // TODO: check overflow
             }
             Statement::DefineStruct(struct_info) => {
                 self.vm.add_struct_info(struct_info);

numbat/src/diagnostic.rs

@@ -497,8 +497,43 @@ impl ErrorDiagnostic for TypeCheckError {
 
 impl ErrorDiagnostic for RuntimeError {
     fn diagnostics(&self) -> Vec<Diagnostic> {
-        vec![Diagnostic::error()
-            .with_message("runtime error")
-            .with_notes(vec![format!("{self:#}")])]
+        let inner = format!("{self:#}");
+
+        match self {
+            RuntimeError::AssertFailed(span) => vec![Diagnostic::error()
+                .with_message("assertion failed")
+                .with_labels(vec![span
+                    .diagnostic_label(LabelStyle::Primary)
+                    .with_message("assertion failed")])],
+            RuntimeError::AssertEq2Failed(span_lhs, lhs, span_rhs, rhs) => {
+                vec![Diagnostic::error()
+                    .with_message("Assertion failed")
+                    .with_labels(vec![
+                        span_lhs
+                            .diagnostic_label(LabelStyle::Secondary)
+                            .with_message(format!("{lhs}")),
+                        span_rhs
+                            .diagnostic_label(LabelStyle::Primary)
+                            .with_message(format!("{rhs}")),
+                    ])
+                    .with_notes(vec![inner])]
+            }
+            RuntimeError::AssertEq3Failed(span_lhs, lhs, span_rhs, rhs, _) => {
+                vec![Diagnostic::error()
+                    .with_message("Assertion failed")
+                    .with_labels(vec![
+                        span_lhs
+                            .diagnostic_label(LabelStyle::Secondary)
+                            .with_message(format!("{lhs}")),
+                        span_rhs
+                            .diagnostic_label(LabelStyle::Primary)
+                            .with_message(format!("{rhs}")),
+                    ])
+                    .with_notes(vec![format!("{self:#}")])]
+            }
+            _ => vec![Diagnostic::error()
+                .with_message("runtime error")
+                .with_notes(vec![inner])],
+        }
     }
 }

numbat/src/ffi/mod.rs

@@ -11,6 +11,7 @@ mod strings;
 use std::collections::VecDeque;
 
 use crate::interpreter::RuntimeError;
+use crate::span::Span;
 use crate::value::Value;
 use crate::vm::ExecutionContext;
 
@@ -26,7 +27,7 @@ type BoxedFunction = Box<dyn Fn(Args) -> Result<Value> + Send + Sync>;
 
 pub(crate) enum Callable {
     Function(BoxedFunction),
-    Procedure(fn(&mut ExecutionContext, Args) -> ControlFlow),
+    Procedure(fn(&mut ExecutionContext, Args, Vec<Span>) -> ControlFlow),
 }
 
 pub(crate) struct ForeignFunction {

numbat/src/ffi/procedures.rs

@@ -4,7 +4,7 @@ use std::sync::OnceLock;
 
 use super::macros::*;
 use crate::{
-    ast::ProcedureKind, ffi::ControlFlow, pretty_print::PrettyPrint, value::Value,
+    ast::ProcedureKind, ffi::ControlFlow, pretty_print::PrettyPrint, span::Span, value::Value,
     vm::ExecutionContext, RuntimeError,
 };
 
@@ -46,7 +46,7 @@ pub(crate) fn procedures() -> &'static HashMap<ProcedureKind, ForeignFunction> {
     })
 }
 
-fn print(ctx: &mut ExecutionContext, mut args: Args) -> ControlFlow {
+fn print(ctx: &mut ExecutionContext, mut args: Args, _: Vec<Span>) -> ControlFlow {
     assert!(args.len() <= 1);
 
     if args.is_empty() {
@@ -61,24 +61,32 @@ fn print(ctx: &mut ExecutionContext, mut args: Args) -> ControlFlow {
     ControlFlow::Continue(())
 }
 
-fn assert(_: &mut ExecutionContext, mut args: Args) -> ControlFlow {
+fn assert(_: &mut ExecutionContext, mut args: Args, arg_spans: Vec<Span>) -> ControlFlow {
     assert!(args.len() == 1);
 
     if arg!(args).unsafe_as_bool() {
         ControlFlow::Continue(())
     } else {
-        ControlFlow::Break(RuntimeError::AssertFailed)
+        ControlFlow::Break(RuntimeError::AssertFailed(arg_spans[0]))
     }
 }
 
-fn assert_eq(_: &mut ExecutionContext, mut args: Args) -> ControlFlow {
+fn assert_eq(_: &mut ExecutionContext, mut args: Args, arg_spans: Vec<Span>) -> ControlFlow {
     assert!(args.len() == 2 || args.len() == 3);
 
+    let span_lhs = arg_spans[0];
+    let span_rhs = arg_spans[1];
+
     if args.len() == 2 {
         let lhs = arg!(args);
         let rhs = arg!(args);
 
-        let error = ControlFlow::Break(RuntimeError::AssertEq2Failed(lhs.clone(), rhs.clone()));
+        let error = ControlFlow::Break(RuntimeError::AssertEq2Failed(
+            span_lhs,
+            lhs.clone(),
+            span_rhs,
+            rhs.clone(),
+        ));
 
         if lhs.is_quantity() {
             let lhs = lhs.unsafe_as_quantity();
@@ -112,7 +120,9 @@ fn assert_eq(_: &mut ExecutionContext, mut args: Args) -> ControlFlow {
                         ControlFlow::Continue(())
                     } else {
                         ControlFlow::Break(RuntimeError::AssertEq3Failed(
+                            span_lhs,
                             lhs.clone(),
+                            span_rhs,
                             rhs.clone(),
                             eps.clone(),
                         ))

numbat/src/interpreter.rs

@@ -3,6 +3,7 @@ use crate::{
     markup::Markup,
     pretty_print::PrettyPrint,
     quantity::{Quantity, QuantityError},
+    span::Span,
     typed_ast::Statement,
     unit_registry::{UnitRegistry, UnitRegistryError},
 };
@@ -26,11 +27,11 @@ pub enum RuntimeError {
     #[error("{0}")]
     QuantityError(QuantityError),
     #[error("Assertion failed")]
-    AssertFailed,
-    #[error("Assertion failed because the following two values are not the same:\n  {0}\n  {1}")]
-    AssertEq2Failed(Value, Value),
-    #[error("Assertion failed because the following two quantities differ by more than {2}:\n  {0}\n  {1}")]
-    AssertEq3Failed(Quantity, Quantity, Quantity),
+    AssertFailed(Span),
+    #[error("Assertion failed because the following two values are not the same:\n  {1}\n  {3}")]
+    AssertEq2Failed(Span, Value, Span, Value),
+    #[error("Assertion failed because the following two quantities differ by more than {4}:\n  {1}\n  {3}")]
+    AssertEq3Failed(Span, Quantity, Span, Quantity, Quantity),
     #[error("Could not load exchange rates from European Central Bank.")]
     CouldNotLoadExchangeRates,
     #[error("User error: {0}")]

numbat/src/vm.rs

@@ -4,6 +4,7 @@ use std::{cmp::Ordering, fmt::Display};
 
 use indexmap::IndexMap;
 
+use crate::span::Span;
 use crate::typed_ast::StructInfo;
 use crate::value::NumbatList;
 use crate::{
@@ -94,6 +95,7 @@ pub enum Op {
     /// Same as above, but call a foreign/native function
     FFICallFunction,
     /// Same as above, but call a procedure which does not return anything (does not push a value onto the stack)
+    /// It has a third argument which is an index to retrieve the source-span of the arguments
     FFICallProcedure,
 
     /// Call a callable object
@@ -123,11 +125,8 @@ pub enum Op {
 impl Op {
     fn num_operands(self) -> usize {
         match self {
-            Op::SetUnitConstant
-            | Op::Call
-            | Op::FFICallFunction
-            | Op::FFICallProcedure
-            | Op::BuildStructInstance => 2,
+            Op::FFICallProcedure => 3,
+            Op::SetUnitConstant | Op::Call | Op::FFICallFunction | Op::BuildStructInstance => 2,
             Op::LoadConstant
             | Op::ApplyPrefix
             | Op::GetLocal
@@ -308,6 +307,10 @@ pub struct Vm {
     /// List of registered native/foreign functions
     ffi_callables: Vec<&'static ForeignFunction>,
 
+    /// Spans for arguments of procedure calls. This is used for
+    /// assertion error messages, for example.
+    procedure_arg_spans: Vec<Vec<Span>>,
+
     /// The call stack
     frames: Vec<CallFrame>,
 
@@ -332,6 +335,7 @@ impl Vm {
             unit_information: vec![],
             last_result: None,
             ffi_callables: ffi::procedures().iter().map(|(_, ff)| ff).collect(),
+            procedure_arg_spans: vec![],
             frames: vec![CallFrame::root()],
             stack: vec![],
             debug: false,
@@ -371,6 +375,14 @@ impl Vm {
         Self::push_u16(current_chunk, arg2);
     }
 
+    pub(crate) fn add_op3(&mut self, op: Op, arg1: u16, arg2: u16, arg3: u16) {
+        let current_chunk = self.current_chunk_mut();
+        current_chunk.push(op as u8);
+        Self::push_u16(current_chunk, arg1);
+        Self::push_u16(current_chunk, arg2);
+        Self::push_u16(current_chunk, arg3);
+    }
+
     pub fn current_offset(&self) -> u16 {
         self.bytecode[self.current_chunk_index].1.len() as u16
     }
@@ -466,6 +478,12 @@ impl Vm {
         Some(position as u16)
     }
 
+    pub(crate) fn add_procedure_arg_span(&mut self, spans: Vec<Span>) -> u16 {
+        self.procedure_arg_spans.push(spans);
+        assert!(self.procedure_arg_spans.len() <= u16::MAX as usize);
+        (self.procedure_arg_spans.len() - 1) as u16
+    }
+
     pub fn disassemble(&self) {
         if !self.debug {
             return;
@@ -832,7 +850,10 @@ impl Vm {
                             self.push(result?);
                         }
                         Callable::Procedure(procedure) => {
-                            let result = (procedure)(ctx, args);
+                            let span_idx = self.read_u16() as usize;
+                            let spans = &self.procedure_arg_spans[span_idx];
+
+                            let result = (procedure)(ctx, args, spans.clone());
 
                             match result {
                                 std::ops::ControlFlow::Continue(()) => {}