Repository: sharkdp/numbat (rust)

Issue:
Bring back (and improve) "readable types"
For example: `let v = 1m/s` would previously print `let v: Velocity = …`. Similarly, it would be cool if we would annotate `v` in `fn f(v)=v+1 knot` with a `Velocity` type

Pull request description:
Readable types
closes #490 

Patch:
diff --git a/numbat/src/bytecode_interpreter.rs b/numbat/src/bytecode_interpreter.rs
index e8afe5e..4e43aa6 100644
--- a/numbat/src/bytecode_interpreter.rs
+++ b/numbat/src/bytecode_interpreter.rs
@@ -310,7 +310,14 @@ impl BytecodeInterpreter {
                 self.compile_expression_with_simplify(expr)?;
                 self.vm.add_op(Op::Return);
             }
-            Statement::DefineVariable(identifier, decorators, expr, _annotation, _type) => {
+            Statement::DefineVariable(
+                identifier,
+                decorators,
+                expr,
+                _annotation,
+                _type,
+                _readable_type,
+            ) => {
                 let current_depth = self.current_depth();
 
                 // For variables, we ignore the prefix info and only use the names
@@ -342,6 +349,8 @@ impl BytecodeInterpreter {
                 parameters,
                 Some(expr),
                 _return_type,
+                _return_type_annotation,
+                _readable_return_type,
             ) => {
                 self.vm.begin_function(name);
 
@@ -372,6 +381,8 @@ impl BytecodeInterpreter {
                 parameters,
                 None,
                 _return_type,
+                _return_type_annotation,
+                _readable_return_type,
             ) => {
                 // Declaring a foreign function does not generate any bytecode. But we register
                 // its name and arity here to be able to distinguish it from normal functions.
@@ -422,7 +433,14 @@ impl BytecodeInterpreter {
                         .insert(name.into(), constant_idx);
                 }
             }
-            Statement::DefineDerivedUnit(unit_name, expr, decorators, annotation, type_) => {
+            Statement::DefineDerivedUnit(
+                unit_name,
+                expr,
+                decorators,
+                annotation,
+                type_,
+                _readable_type,
+            ) => {
                 let aliases = decorator::name_and_aliases(unit_name, decorators)
                     .map(|(name, ap)| (name.clone(), ap))
                     .collect();
diff --git a/numbat/src/lib.rs b/numbat/src/lib.rs
index d8950d4..03babf8 100644
--- a/numbat/src/lib.rs
+++ b/numbat/src/lib.rs
@@ -724,6 +724,8 @@ impl Context {
 
         let typed_statements = result?;
 
+        let interpreter_old = self.interpreter.clone();
+
         let result = self.interpreter.interpret_statements(
             settings,
             &typed_statements,
@@ -743,6 +745,7 @@ impl Context {
             //
             self.prefix_transformer = prefix_transformer_old;
             self.typechecker = typechecker_old;
+            self.interpreter = interpreter_old;
         }
 
         let result = result.map_err(NumbatError::RuntimeError)?;
diff --git a/numbat/src/traversal.rs b/numbat/src/traversal.rs
index 21656d8..573a699 100644
--- a/numbat/src/traversal.rs
+++ b/numbat/src/traversal.rs
@@ -79,11 +79,11 @@ impl ForAllTypeSchemes for Statement {
     fn for_all_type_schemes(&mut self, f: &mut dyn FnMut(&mut TypeScheme)) {
         match self {
             Statement::Expression(expr) => expr.for_all_type_schemes(f),
-            Statement::DefineVariable(_, _, expr, _annotation, type_) => {
+            Statement::DefineVariable(_, _, expr, _annotation, type_, _) => {
                 expr.for_all_type_schemes(f);
                 f(type_);
             }
-            Statement::DefineFunction(_, _, _, _, body, fn_type) => {
+            Statement::DefineFunction(_, _, _, _, body, fn_type, _, _) => {
                 if let Some(body) = body {
                     body.for_all_type_schemes(f);
                 }
@@ -93,7 +93,7 @@ impl ForAllTypeSchemes for Statement {
             Statement::DefineBaseUnit(_, _, _annotation, type_) => {
                 f(type_);
             }
-            Statement::DefineDerivedUnit(_, expr, _, _annotation, type_) => {
+            Statement::DefineDerivedUnit(_, expr, _, _annotation, type_, _) => {
                 expr.for_all_type_schemes(f);
                 f(type_);
             }
@@ -115,15 +115,15 @@ impl ForAllExpressions for Statement {
     fn for_all_expressions(&self, f: &mut dyn FnMut(&Expression)) {
         match self {
             Statement::Expression(expr) => expr.for_all_expressions(f),
-            Statement::DefineVariable(_, _, expr, _, _) => expr.for_all_expressions(f),
-            Statement::DefineFunction(_, _, _, _, body, _) => {
+            Statement::DefineVariable(_, _, expr, _, _, _) => expr.for_all_expressions(f),
+            Statement::DefineFunction(_, _, _, _, body, _, _, _) => {
                 if let Some(body) = body {
                     body.for_all_expressions(f);
                 }
             }
             Statement::DefineDimension(_, _) => {}
             Statement::DefineBaseUnit(_, _, _, _) => {}
-            Statement::DefineDerivedUnit(_, expr, _, _, _) => expr.for_all_expressions(f),
+            Statement::DefineDerivedUnit(_, expr, _, _, _, _) => expr.for_all_expressions(f),
             Statement::ProcedureCall(_, args) => {
                 for arg in args {
                     arg.for_all_expressions(f);
diff --git a/numbat/src/typechecker/mod.rs b/numbat/src/typechecker/mod.rs
index e0e1228..26c5cf5 100644
--- a/numbat/src/typechecker/mod.rs
+++ b/numbat/src/typechecker/mod.rs
@@ -1152,6 +1152,7 @@ impl TypeChecker {
                     expr_checked,
                     type_annotation.clone(),
                     TypeScheme::concrete(type_deduced),
+                    crate::markup::empty(),
                 )
             }
             ast::Statement::DefineBaseUnit(span, unit_name, type_annotation, decorators) => {
@@ -1271,6 +1272,7 @@ impl TypeChecker {
                     decorators.clone(),
                     type_annotation.clone(),
                     TypeScheme::Concrete(type_deduced),
+                    crate::markup::empty(),
                 )
             }
             ast::Statement::DefineFunction {
@@ -1353,7 +1355,12 @@ impl TypeChecker {
                         *parameter_span,
                         false,
                     );
-                    typed_parameters.push((*parameter_span, parameter.clone(), parameter_type));
+                    typed_parameters.push((
+                        *parameter_span,
+                        parameter.clone(),
+                        parameter_type,
+                        type_annotation,
+                    ));
                 }
 
                 let annotated_return_type = return_type_annotation
@@ -1370,11 +1377,11 @@ impl TypeChecker {
 
                 let parameters: Vec<_> = typed_parameters
                     .iter()
-                    .map(|(span, name, _)| (*span, name.clone()))
+                    .map(|(span, name, _, _)| (*span, name.clone()))
                     .collect();
                 let parameter_types = typed_parameters
                     .iter()
-                    .map(|(_, _, type_)| type_.clone())
+                    .map(|(_, _, type_, _)| type_.clone())
                     .collect();
 
                 let fn_type =
@@ -1495,10 +1502,19 @@ impl TypeChecker {
                         .collect(),
                     typed_parameters
                         .iter()
-                        .map(|(span, name, _)| (*span, name.clone()))
+                        .map(|(span, name, _, ref type_annotation)| {
+                            (
+                                *span,
+                                name.clone(),
+                                (*type_annotation).clone(),
+                                crate::markup::empty(),
+                            )
+                        })
                         .collect(),
                     body_checked,
                     fn_type,
+                    return_type_annotation.clone(),
+                    crate::markup::empty(),
                 )
             }
             ast::Statement::DefineDimension(name_span, name, dexprs) => {
@@ -1709,7 +1725,7 @@ impl TypeChecker {
             TypeCheckError::SubstitutionError(elaborated_statement.pretty_print().to_string(), e)
         })?;
 
-        if let typed_ast::Statement::DefineDerivedUnit(_, expr, _, _annotation, type_) =
+        if let typed_ast::Statement::DefineDerivedUnit(_, expr, _, _annotation, type_, _) =
             &elaborated_statement
         {
             if !type_.unsafe_as_concrete().is_closed() {
@@ -1761,6 +1777,9 @@ impl TypeChecker {
         }
 
         elaborated_statement.generalize_types(&dtype_variables);
+
+        elaborated_statement.update_readable_types(&self.registry);
+
         self.env.generalize_types(&dtype_variables);
 
         // Check if there is a typed hole in the statement
diff --git a/numbat/src/typechecker/substitutions.rs b/numbat/src/typechecker/substitutions.rs
index 31d0c28..67533d7 100644
--- a/numbat/src/typechecker/substitutions.rs
+++ b/numbat/src/typechecker/substitutions.rs
@@ -214,11 +214,11 @@ impl ApplySubstitution for Statement {
     fn apply(&mut self, s: &Substitution) -> Result<(), SubstitutionError> {
         match self {
             Statement::Expression(e) => e.apply(s),
-            Statement::DefineVariable(_, _, e, _annotation, type_) => {
+            Statement::DefineVariable(_, _, e, _annotation, type_, _) => {
                 e.apply(s)?;
                 type_.apply(s)
             }
-            Statement::DefineFunction(_, _, _, _, body, fn_type) => {
+            Statement::DefineFunction(_, _, _, _, body, fn_type, _, _) => {
                 if let Some(body) = body {
                     body.apply(s)?;
                 }
@@ -226,7 +226,7 @@ impl ApplySubstitution for Statement {
             }
             Statement::DefineDimension(_, _) => Ok(()),
             Statement::DefineBaseUnit(_, _, _annotation, type_) => type_.apply(s),
-            Statement::DefineDerivedUnit(_, e, _, _annotation, type_) => {
+            Statement::DefineDerivedUnit(_, e, _, _annotation, type_, _) => {
                 e.apply(s)?;
                 type_.apply(s)
             }
diff --git a/numbat/src/typed_ast.rs b/numbat/src/typed_ast.rs
index 2b2b1ef..398198b 100644
--- a/numbat/src/typed_ast.rs
+++ b/numbat/src/typed_ast.rs
@@ -575,6 +575,7 @@ pub enum Statement {
         Expression,
         Option<TypeAnnotation>,
         TypeScheme,
+        Markup,
     ),
     DefineFunction(
         String,
@@ -582,11 +583,15 @@ pub enum Statement {
         Vec<(String, Option<TypeParameterBound>)>, // type parameters
         Vec<(
             // parameters:
-            Span,   // span of the parameter
-            String, // parameter name
+            Span,                   // span of the parameter
+            String,                 // parameter name
+            Option<TypeAnnotation>, // parameter type annotation
+            Markup,                 // readable parameter type
         )>,
-        Option<Expression>, // function body
-        TypeScheme,         // function type
+        Option<Expression>,     // function body
+        TypeScheme,             // function type
+        Option<TypeAnnotation>, // return type annotation
+        Markup,                 // readable return type
     ),
     DefineDimension(String, Vec<TypeExpression>),
     DefineBaseUnit(String, Vec<Decorator>, Option<TypeAnnotation>, TypeScheme),
@@ -596,6 +601,7 @@ pub enum Statement {
         Vec<Decorator>,
         Option<TypeAnnotation>,
         TypeScheme,
+        Markup,
     ),
     ProcedureCall(crate::ast::ProcedureKind, Vec<Expression>),
     DefineStruct(StructInfo),
@@ -614,6 +620,66 @@ impl Statement {
         self.for_all_type_schemes(&mut |type_: &mut TypeScheme| type_.generalize(dtype_variables));
     }
 
+    fn create_readable_type(
+        registry: &DimensionRegistry,
+        type_: &TypeScheme,
+        annotation: &Option<TypeAnnotation>,
+    ) -> Markup {
+        if let Some(annotation) = annotation {
+            annotation.pretty_print()
+        } else {
+            type_.to_readable_type(registry)
+        }
+    }
+
+    pub(crate) fn update_readable_types(&mut self, registry: &DimensionRegistry) {
+        match self {
+            Statement::Expression(_) => {}
+            Statement::DefineVariable(_, _, _, type_annotation, type_, readable_type) => {
+                *readable_type = Self::create_readable_type(registry, type_, type_annotation);
+            }
+            Statement::DefineFunction(
+                _,
+                _,
+                _,
+                parameters,
+                _,
+                fn_type,
+                return_type_annotation,
+                readable_return_type,
+            ) => {
+                let (fn_type, _) = fn_type.instantiate_for_printing();
+
+                let Type::Fn(parameter_types, return_type) = fn_type.inner else {
+                    unreachable!("Expected a function type")
+                };
+
+                *readable_return_type = Self::create_readable_type(
+                    registry,
+                    &TypeScheme::concrete(*return_type),
+                    return_type_annotation,
+                );
+
+                for ((_, _, type_annotation, readable_parameter_type), parameter_type) in
+                    parameters.iter_mut().zip(parameter_types.iter())
+                {
+                    *readable_parameter_type = Self::create_readable_type(
+                        registry,
+                        &TypeScheme::concrete(parameter_type.clone()),
+                        type_annotation,
+                    );
+                }
+            }
+            Statement::DefineDimension(_, _) => {}
+            Statement::DefineBaseUnit(_, _, _, _) => {}
+            Statement::DefineDerivedUnit(_, _, _, type_annotation, type_, readable_type) => {
+                *readable_type = Self::create_readable_type(registry, type_, type_annotation);
+            }
+            Statement::ProcedureCall(_, _) => {}
+            Statement::DefineStruct(_) => {}
+        }
+    }
+
     pub(crate) fn exponents_for(&mut self, tv: &TypeVariable) -> Vec<Exponent> {
         // TODO: things to not need to be mutable in this function
         let mut exponents = vec![];
@@ -773,16 +839,20 @@ fn decorator_markup(decorators: &Vec<Decorator>) -> Markup {
 impl PrettyPrint for Statement {
     fn pretty_print(&self) -> Markup {
         match self {
-            Statement::DefineVariable(identifier, _decs, expr, annotation, type_) => {
+            Statement::DefineVariable(
+                identifier,
+                _decs,
+                expr,
+                _annotation,
+                _type,
+                readable_type,
+            ) => {
                 m::keyword("let")
                     + m::space()
                     + m::identifier(identifier)
                     + m::operator(":")
                     + m::space()
-                    + annotation
-                        .as_ref()
-                        .map(|a| a.pretty_print())
-                        .unwrap_or(type_.pretty_print())
+                    + readable_type.clone()
                     + m::space()
                     + m::operator("=")
                     + m::space()
@@ -795,13 +865,11 @@ impl PrettyPrint for Statement {
                 parameters,
                 body,
                 fn_type,
+                _return_type_annotation,
+                readable_return_type,
             ) => {
                 let (fn_type, type_parameters) = fn_type.instantiate_for_printing();
 
-                let Type::Fn(parameter_types, return_type) = fn_type.inner else {
-                    unreachable!("Expected a function type")
-                };
-
                 let markup_type_parameters = if type_parameters.is_empty() {
                     m::empty()
                 } else {
@@ -822,20 +890,15 @@ impl PrettyPrint for Statement {
                 };
 
                 let markup_parameters = Itertools::intersperse(
-                    parameters.iter().zip(parameter_types.iter()).map(
-                        |((_span, name), parameter_type)| {
-                            m::identifier(name)
-                                + m::operator(":")
-                                + m::space()
-                                + parameter_type.pretty_print()
-                        },
-                    ),
+                    parameters.iter().map(|(_span, name, _, parameter_type)| {
+                        m::identifier(name) + m::operator(":") + m::space() + parameter_type.clone()
+                    }),
                     m::operator(", "),
                 )
                 .sum();
 
                 let markup_return_type =
-                    m::space() + m::operator("->") + m::space() + return_type.pretty_print();
+                    m::space() + m::operator("->") + m::space() + readable_return_type.clone();
 
                 m::keyword("fn")
                     + m::space()
@@ -879,17 +942,21 @@ impl PrettyPrint for Statement {
                         .map(|a| a.pretty_print())
                         .unwrap_or(type_.pretty_print())
             }
-            Statement::DefineDerivedUnit(identifier, expr, decorators, annotation, type_) => {
+            Statement::DefineDerivedUnit(
+                identifier,
+                expr,
+                decorators,
+                _annotation,
+                _type,
+                readable_type,
+            ) => {
                 decorator_markup(decorators)
                     + m::keyword("unit")
                     + m::space()
                     + m::unit(identifier)
                     + m::operator(":")
                     + m::space()
-                    + annotation
-                        .as_ref()
-                        .map(|a| a.pretty_print())
-                        .unwrap_or(type_.pretty_print())
+                    + readable_type.clone()
                     + m::space()
                     + m::operator("=")
                     + m::space()


Test patch:
diff --git a/numbat-cli/tests/integration.rs b/numbat-cli/tests/integration.rs
index bfea11f..f00f4c2 100644
--- a/numbat-cli/tests/integration.rs
+++ b/numbat-cli/tests/integration.rs
@@ -125,18 +125,18 @@ fn without_prelude() {
         .stderr(predicates::str::contains("unknown identifier"));
 }
 
-// #[test]
-// fn pretty_printing() {
-//     numbat()
-//         .arg("--pretty-print=always")
-//         .arg("--expression")
-//         .arg("let v=30km/h")
-//         .assert()
-//         .success()
-//         .stdout(predicates::str::contains(
-//             "let v: Velocity = 30 kilometre / hour",
-//         ));
-// }
+#[test]
+fn pretty_printing() {
+    numbat()
+        .arg("--pretty-print=always")
+        .arg("--expression")
+        .arg("let v=30km/h")
+        .assert()
+        .success()
+        .stdout(predicates::str::contains(
+            "let v: Velocity = 30 kilometre / hour",
+        ));
+}
 
 #[test]
 fn help_text() {
diff --git a/numbat/src/typechecker/tests/mod.rs b/numbat/src/typechecker/tests/mod.rs
index 6faf94c..1b8519f 100644
--- a/numbat/src/typechecker/tests/mod.rs
+++ b/numbat/src/typechecker/tests/mod.rs
@@ -68,7 +68,7 @@ fn assert_successful_typecheck(input: &str) {
 fn get_inferred_fn_type(input: &str) -> TypeScheme {
     let statement = run_typecheck(input).expect("Input was expected to type-check");
     match statement {
-        Statement::DefineFunction(_, _, _, _, _, fn_type) => fn_type,
+        Statement::DefineFunction(_, _, _, _, _, fn_type, _, _) => fn_type,
         _ => {
             unreachable!();
         }
diff --git a/numbat/tests/interpreter.rs b/numbat/tests/interpreter.rs
index dd292a7..d79195d 100644
--- a/numbat/tests/interpreter.rs
+++ b/numbat/tests/interpreter.rs
@@ -46,6 +46,21 @@ fn expect_output(code: &str, expected_output: impl AsRef<str>) {
     expect_output_with_context(&mut ctx, code, expected_output)
 }
 
+#[track_caller]
+fn expect_pretty_print(code: &str, expected_pretty_print_output: impl AsRef<str>) {
+    let mut ctx = get_test_context();
+
+    let (statements, _result) = ctx.interpret(code, CodeSource::Internal).unwrap();
+
+    assert_eq!(statements.len(), 1);
+
+    let statement = &statements[0];
+    assert_eq!(
+        statement.pretty_print().to_string(),
+        expected_pretty_print_output.as_ref()
+    );
+}
+
 #[track_caller]
 fn expect_failure_with_context(ctx: &mut Context, code: &str, msg_part: &str) {
     if let Err(e) = ctx.interpret(code, CodeSource::Internal) {
@@ -719,3 +734,82 @@ fn test_user_errors() {
         "User error: test",
     );
 }
+
+#[test]
+fn test_recovery_after_runtime_error() {
+    {
+        let mut ctx = get_test_context();
+
+        expect_failure_with_context(&mut ctx, "let x = 1/0", "Division by zero");
+        expect_failure_with_context(&mut ctx, "x", "Unknown identifier 'x'");
+    }
+    {
+        let mut ctx = get_test_context();
+
+        expect_failure_with_context(&mut ctx, "let x = 1/0", "Division by zero");
+        assert!(ctx
+            .interpret("let x = 1", CodeSource::Internal)
+            .unwrap()
+            .1
+            .is_continue());
+        expect_output_with_context(&mut ctx, "x", "1");
+    }
+}
+
+#[test]
+fn test_statement_pretty_printing() {
+    // Let definitions
+    expect_pretty_print("let v = 10 m/s", "let v: Velocity = 10 metre / second");
+    expect_pretty_print(
+        "let v: Length * Frequency = 10 m/s",
+        "let v: Length × Frequency = 10 metre / second",
+    );
+
+    expect_pretty_print("let x = 0 + 1 m", "let x: Length = 0 + 1 metre");
+
+    expect_pretty_print("let x: Length = 0", "let x: Length = 0");
+    expect_pretty_print("let x = 0", "let x: forall A: Dim. A = 0"); // TODO: This is not ideal. 'forall' is not valid Numbat syntax.
+
+    // Derived unit definitions
+    expect_pretty_print(
+        "unit my_length_base_unit: Length",
+        "unit my_length_base_unit: Length",
+    );
+    expect_pretty_print(
+        "unit my_custom_base_unit",
+        "unit my_custom_base_unit: MyCustomBaseUnit",
+    );
+    expect_pretty_print(
+        "unit my_speed_unit = 10 m/s",
+        "unit my_speed_unit: Velocity = 10 metre / second",
+    );
+
+    // Function definitions
+    expect_pretty_print(
+        "fn f(v)=(v+1 knot)/1s",
+        "fn f(v: Velocity) -> Acceleration = (v + 1 knot) / 1 second",
+    );
+
+    expect_pretty_print("fn f(x) = x", "fn f<A>(x: A) -> A = x");
+
+    expect_pretty_print("fn f(x, y) = y", "fn f<A, B>(x: A, y: B) -> B = y");
+    expect_pretty_print("fn f(x, y) = x", "fn f<A, B>(x: B, y: A) -> B = x"); // TODO: This is correct, but it would be nice to associate 'x' to 'A', not 'B'.
+
+    expect_pretty_print("fn f(x) = 2 x", "fn f<A: Dim>(x: A) -> A = 2 x");
+
+    // Partially annotated functions
+    expect_pretty_print(
+        "fn f() -> Length * Frequency = c",
+        "fn f() -> Length × Frequency = c",
+    );
+    expect_pretty_print(
+        "fn f(v: Length * Frequency) = 1",
+        "fn f(v: Length × Frequency) -> Scalar = 1",
+    );
+
+    expect_pretty_print("fn f(x: Length) = 2 x", "fn f(x: Length) -> Length = 2 x");
+    expect_pretty_print("fn f(x) -> Length = 2 x", "fn f(x: Length) -> Length = 2 x");
+
+    // TODO:
+    // expect_pretty_print("fn f<Z>(z: Z) = z", "fn f<Z>(z: Z) -> Z = z");
+}