Repository: sharkdp/numbat (rust)

Issue:
Overwrite constants and functions in the REPL
It's a very frequent use case to go "up arrow" and change something in a let/fn definition.

Pull request description:
Allow overwriting of functions
relates to #149 

Patch:
diff --git a/examples/name_resolution_error/duplicate_variable.nbt b/examples/name_resolution_error/duplicate_variable.nbt
deleted file mode 100644
index 4dc90d4..0000000
--- a/examples/name_resolution_error/duplicate_variable.nbt
+++ /dev/null
@@ -1,2 +0,0 @@
-let foo = 2
-let foo = 3
diff --git a/examples/name_resolution_error/function_name_clash.nbt b/examples/name_resolution_error/function_name_clash.nbt
deleted file mode 100644
index 5379ae8..0000000
--- a/examples/name_resolution_error/function_name_clash.nbt
+++ /dev/null
@@ -1,3 +0,0 @@
-let x = 2
-
-fn x() = 2
diff --git a/examples/name_resolution_error/variable_clashes_with_parameter.nbt b/examples/name_resolution_error/variable_clashes_with_parameter.nbt
deleted file mode 100644
index bffe7a1..0000000
--- a/examples/name_resolution_error/variable_clashes_with_parameter.nbt
+++ /dev/null
@@ -1,3 +0,0 @@
-let foo = 2
-
-fn bar(foo: Scalar) -> Scalar = foo
diff --git a/numbat/src/diagnostic.rs b/numbat/src/diagnostic.rs
index 4d0be65..2dc5b7d 100644
--- a/numbat/src/diagnostic.rs
+++ b/numbat/src/diagnostic.rs
@@ -269,14 +269,6 @@ impl ErrorDiagnostic for TypeCheckError {
                     .diagnostic_label(LabelStyle::Primary)
                     .with_message("Incompatible types in 'assert_eq' call"),
             ]),
-            TypeCheckError::ForeignFunctionNeedsTypeAnnotations(span, _)
-            | TypeCheckError::UnknownForeignFunction(span, _)
-            | TypeCheckError::NonRationalExponent(span)
-            | TypeCheckError::OverflowInConstExpr(span)
-            | TypeCheckError::ExpectedDimensionType(span, _)
-            | TypeCheckError::ExpectedBool(span) => d.with_labels(vec![span
-                .diagnostic_label(LabelStyle::Primary)
-                .with_message(inner_error)]),
             TypeCheckError::IncompatibleTypesInAnnotation(
                 what,
                 what_span,
@@ -295,6 +287,31 @@ impl ErrorDiagnostic for TypeCheckError {
                     .diagnostic_label(LabelStyle::Primary)
                     .with_message(format!("Incompatible types in {what}")),
             ]),
+            TypeCheckError::NameAlreadyUsedBy(_, definition_span, previous_definition_span) => {
+                let mut labels = vec![];
+
+                if let Some(span) = previous_definition_span {
+                    labels.push(
+                        span.diagnostic_label(LabelStyle::Secondary)
+                            .with_message("Previously defined here"),
+                    );
+                }
+
+                labels.push(
+                    definition_span
+                        .diagnostic_label(LabelStyle::Primary)
+                        .with_message(inner_error),
+                );
+                d.with_labels(labels)
+            }
+            TypeCheckError::ForeignFunctionNeedsTypeAnnotations(span, _)
+            | TypeCheckError::UnknownForeignFunction(span, _)
+            | TypeCheckError::NonRationalExponent(span)
+            | TypeCheckError::OverflowInConstExpr(span)
+            | TypeCheckError::ExpectedDimensionType(span, _)
+            | TypeCheckError::ExpectedBool(span) => d.with_labels(vec![span
+                .diagnostic_label(LabelStyle::Primary)
+                .with_message(inner_error)]),
         }
     }
 }
diff --git a/numbat/src/prefix_parser.rs b/numbat/src/prefix_parser.rs
index 507cfdf..dbd9fa3 100644
--- a/numbat/src/prefix_parser.rs
+++ b/numbat/src/prefix_parser.rs
@@ -140,13 +140,20 @@ impl PrefixParser {
         }
     }
 
-    fn ensure_name_is_available(&self, name: &str, conflict_span: Span) -> Result<()> {
+    fn ensure_name_is_available(
+        &self,
+        name: &str,
+        conflict_span: Span,
+        clash_with_other_identifiers: bool,
+    ) -> Result<()> {
         if self.reserved_identifiers.contains(&name) {
             return Err(NameResolutionError::ReservedIdentifier(conflict_span));
         }
 
-        if let Some(original_span) = self.other_identifiers.get(name) {
-            return Err(self.identifier_clash_error(name, conflict_span, *original_span));
+        if clash_with_other_identifiers {
+            if let Some(original_span) = self.other_identifiers.get(name) {
+                return Err(self.identifier_clash_error(name, conflict_span, *original_span));
+            }
         }
 
         match self.parse(name) {
@@ -166,7 +173,7 @@ impl PrefixParser {
         full_name: &str,
         definition_span: Span,
     ) -> Result<()> {
-        self.ensure_name_is_available(unit_name, definition_span)?;
+        self.ensure_name_is_available(unit_name, definition_span, true)?;
 
         for (prefix_long, prefix_short, prefix) in Self::prefixes() {
             if !(prefix.is_metric() && metric || prefix.is_binary() && binary) {
@@ -177,12 +184,14 @@ impl PrefixParser {
                 self.ensure_name_is_available(
                     &format!("{}{}", prefix_long, unit_name),
                     definition_span,
+                    true,
                 )?;
             }
             if accepts_prefix.short {
                 self.ensure_name_is_available(
                     &format!("{}{}", prefix_short, unit_name),
                     definition_span,
+                    true,
                 )?;
             }
         }
@@ -201,15 +210,11 @@ impl PrefixParser {
     }
 
     pub fn add_other_identifier(&mut self, identifier: &str, definition_span: Span) -> Result<()> {
-        self.ensure_name_is_available(identifier, definition_span)?;
-
-        if let Some(original_span) = self.other_identifiers.get(identifier) {
-            Err(self.identifier_clash_error(identifier, definition_span, *original_span))
-        } else {
-            self.other_identifiers
-                .insert(identifier.into(), definition_span);
-            Ok(())
-        }
+        self.ensure_name_is_available(identifier, definition_span, false)?;
+
+        self.other_identifiers
+            .insert(identifier.into(), definition_span);
+        Ok(())
     }
 
     pub fn parse(&self, input: &str) -> PrefixParserResult {
diff --git a/numbat/src/typechecker.rs b/numbat/src/typechecker.rs
index 2d14129..53e9e2d 100644
--- a/numbat/src/typechecker.rs
+++ b/numbat/src/typechecker.rs
@@ -253,6 +253,9 @@ pub enum TypeCheckError {
 
     #[error("Incompatible types in comparison operator")]
     IncompatibleTypesInComparison(Span, Type, Span, Type, Span),
+
+    #[error("This name is already used by {0}")]
+    NameAlreadyUsedBy(&'static str, Span, Option<Span>),
 }
 
 type Result<T> = std::result::Result<T, TypeCheckError>;
@@ -357,38 +360,32 @@ fn evaluate_const_expr(expr: &typed_ast::Expression) -> Result<Exponent> {
     }
 }
 
-#[derive(Clone, PartialEq)]
-enum IdentifierKind {
-    Variable,
-    Other,
+#[derive(Clone)]
+struct FunctionSignature {
+    definition_span: Span,
+    type_parameters: Vec<(Span, String)>,
+    parameter_types: Vec<(Span, Type)>,
+    is_variadic: bool,
+    return_type: Type,
+    is_foreign: bool,
 }
 
 #[derive(Clone, Default)]
 pub struct TypeChecker {
-    identifiers: HashMap<String, (Type, IdentifierKind)>,
-    function_signatures: HashMap<
-        String,
-        (
-            Span,                // span of the function definition
-            Vec<(Span, String)>, // type parameters
-            Vec<(Span, Type)>,   // parameter types
-            bool,                // whether or not the function is variadic
-            Type,                // return type
-        ),
-    >,
+    identifiers: HashMap<String, (Type, Option<Span>)>,
+    function_signatures: HashMap<String, FunctionSignature>,
     registry: DimensionRegistry,
 }
 
 impl TypeChecker {
-    fn identifier_type_and_kind(&self, span: Span, name: &str) -> Result<&(Type, IdentifierKind)> {
-        self.identifiers.get(name).ok_or_else(|| {
-            let suggestion = suggestion::did_you_mean(self.identifiers.keys(), name);
-            TypeCheckError::UnknownIdentifier(span, name.into(), suggestion)
-        })
-    }
-
     fn identifier_type(&self, span: Span, name: &str) -> Result<&Type> {
-        self.identifier_type_and_kind(span, name).map(|(t, _)| t)
+        self.identifiers
+            .get(name)
+            .ok_or_else(|| {
+                let suggestion = suggestion::did_you_mean(self.identifiers.keys(), name);
+                TypeCheckError::UnknownIdentifier(span, name.into(), suggestion)
+            })
+            .map(|(type_, _)| type_)
     }
 
     pub(crate) fn check_expression(&self, ast: &ast::Expression) -> Result<typed_ast::Expression> {
@@ -562,13 +559,14 @@ impl TypeChecker {
                 )
             }
             ast::Expression::FunctionCall(span, full_span, function_name, args) => {
-                let (
-                    callable_definition_span,
+                let FunctionSignature {
+                    definition_span,
                     type_parameters,
                     parameter_types,
                     is_variadic,
                     return_type,
-                ) = self
+                    is_foreign: _,
+                } = self
                     .function_signatures
                     .get(function_name)
                     .ok_or_else(|| TypeCheckError::UnknownCallable(*span, function_name.clone()))?;
@@ -583,7 +581,7 @@ impl TypeChecker {
                     return Err(TypeCheckError::WrongArity {
                         callable_span: *span,
                         callable_name: function_name.clone(),
-                        callable_definition_span: Some(*callable_definition_span),
+                        callable_definition_span: Some(*definition_span),
                         arity: arity_range,
                         num_args: args.len(),
                     });
@@ -725,7 +723,7 @@ impl TypeChecker {
 
                     return Err(TypeCheckError::CanNotInferTypeParameters(
                         *span,
-                        *callable_definition_span,
+                        *definition_span,
                         function_name.clone(),
                         remaining.join(", "),
                     ));
@@ -797,10 +795,8 @@ impl TypeChecker {
             ast::Statement::Expression(expr) => {
                 let checked_expr = self.check_expression(expr)?;
                 for &identifier in LAST_RESULT_IDENTIFIERS {
-                    self.identifiers.insert(
-                        identifier.into(),
-                        (checked_expr.get_type(), IdentifierKind::Variable),
-                    );
+                    self.identifiers
+                        .insert(identifier.into(), (checked_expr.get_type(), None));
                 }
                 typed_ast::Statement::Expression(checked_expr)
             }
@@ -810,6 +806,24 @@ impl TypeChecker {
                 expr,
                 type_annotation,
             } => {
+                // Make sure that identifier does not clash with a function name. We do not
+                // check for clashes with unit names, as this is handled by the prefix parser.
+                if let Some(ref signature) = self.function_signatures.get(identifier) {
+                    return Err(TypeCheckError::NameAlreadyUsedBy(
+                        "a function",
+                        *identifier_span,
+                        Some(signature.definition_span),
+                    ));
+                }
+
+                if let Some(entry) = self.identifiers.get(identifier) {
+                    return Err(TypeCheckError::NameAlreadyUsedBy(
+                        "a constant",
+                        *identifier_span,
+                        entry.1,
+                    ));
+                }
+
                 let expr_checked = self.check_expression(expr)?;
                 let type_deduced = expr_checked.get_type();
 
@@ -856,7 +870,7 @@ impl TypeChecker {
 
                 self.identifiers.insert(
                     identifier.clone(),
-                    (type_deduced.clone(), IdentifierKind::Variable),
+                    (type_deduced.clone(), Some(*identifier_span)),
                 );
 
                 typed_ast::Statement::DefineVariable(
@@ -869,7 +883,7 @@ impl TypeChecker {
                     type_deduced,
                 )
             }
-            ast::Statement::DefineBaseUnit(_span, unit_name, type_annotation, decorators) => {
+            ast::Statement::DefineBaseUnit(span, unit_name, type_annotation, decorators) => {
                 let type_specified = if let Some(dexpr) = type_annotation {
                     self.registry
                         .get_base_representation(dexpr)
@@ -886,10 +900,7 @@ impl TypeChecker {
                 for (name, _) in decorator::name_and_aliases(unit_name, decorators) {
                     self.identifiers.insert(
                         name.clone(),
-                        (
-                            Type::Dimension(type_specified.clone()),
-                            IdentifierKind::Other,
-                        ),
+                        (Type::Dimension(type_specified.clone()), Some(*span)),
                     );
                 }
                 typed_ast::Statement::DefineBaseUnit(
@@ -944,7 +955,10 @@ impl TypeChecker {
                 for (name, _) in decorator::name_and_aliases(identifier, decorators) {
                     self.identifiers.insert(
                         name.clone(),
-                        (Type::Dimension(type_deduced.clone()), IdentifierKind::Other),
+                        (
+                            Type::Dimension(type_deduced.clone()),
+                            Some(*identifier_span),
+                        ),
                     );
                 }
                 typed_ast::Statement::DefineDerivedUnit(
@@ -966,6 +980,26 @@ impl TypeChecker {
                 return_type_annotation_span,
                 return_type_annotation,
             } => {
+                // Make sure that function name does not clash with an identifier. We do not
+                // check for clashes with unit names, as this is handled by the prefix parser.
+                if let Some((_, span)) = self.identifiers.get(function_name) {
+                    return Err(TypeCheckError::NameAlreadyUsedBy(
+                        "a constant",
+                        *function_name_span,
+                        span.clone(),
+                    ));
+                }
+
+                if let Some(ref signature) = self.function_signatures.get(function_name) {
+                    if signature.is_foreign {
+                        return Err(TypeCheckError::NameAlreadyUsedBy(
+                            "a foreign function",
+                            *function_name_span,
+                            Some(signature.definition_span),
+                        ));
+                    }
+                }
+
                 let mut typechecker_fn = self.clone();
                 let is_ffi_function = body.is_none();
                 let mut type_parameters = type_parameters.clone();
@@ -1020,7 +1054,7 @@ impl TypeChecker {
 
                     typechecker_fn.identifiers.insert(
                         parameter.clone(),
-                        (parameter_type.clone(), IdentifierKind::Variable),
+                        (parameter_type.clone(), Some(*parameter_span)),
                     );
                     typed_parameters.push((
                         *parameter_span,
@@ -1045,29 +1079,35 @@ impl TypeChecker {
                     .map(|annotation| typechecker_fn.type_from_annotation(annotation))
                     .transpose()?;
 
-                let add_function_signature = |tc: &mut TypeChecker, return_type: Type| {
-                    let parameter_types = typed_parameters
-                        .iter()
-                        .map(|(span, _, _, _, t)| (*span, t.clone()))
-                        .collect();
-                    tc.function_signatures.insert(
-                        function_name.clone(),
-                        (
-                            *function_name_span,
-                            type_parameters.clone(),
-                            parameter_types,
-                            is_variadic,
-                            return_type,
-                        ),
-                    );
-                };
+                let add_function_signature =
+                    |tc: &mut TypeChecker, return_type: Type, is_foreign: bool| {
+                        let parameter_types = typed_parameters
+                            .iter()
+                            .map(|(span, _, _, _, t)| (*span, t.clone()))
+                            .collect();
+                        tc.function_signatures.insert(
+                            function_name.clone(),
+                            FunctionSignature {
+                                definition_span: *function_name_span,
+                                type_parameters: type_parameters.clone(),
+                                parameter_types,
+                                is_variadic,
+                                return_type,
+                                is_foreign,
+                            },
+                        );
+                    };
 
                 if let Some(ref return_type_specified) = return_type_specified {
                     // This is needed for recursive functions. If the return type
                     // has been specified, we can already provide a function
                     // signature before we check the body of the function. This
                     // way, the 'typechecker_fn' can resolve the recursive call.
-                    add_function_signature(&mut typechecker_fn, return_type_specified.clone());
+                    add_function_signature(
+                        &mut typechecker_fn,
+                        return_type_specified.clone(),
+                        body.is_none(),
+                    );
                 }
 
                 let body_checked = body
@@ -1140,7 +1180,7 @@ impl TypeChecker {
                     })?
                 };
 
-                add_function_signature(self, return_type.clone());
+                add_function_signature(self, return_type.clone(), body.is_none());
 
                 typed_ast::Statement::DefineFunction(
                     function_name.clone(),
diff --git a/numbat/src/vm.rs b/numbat/src/vm.rs
index 0e16846..328e358 100644
--- a/numbat/src/vm.rs
+++ b/numbat/src/vm.rs
@@ -365,7 +365,15 @@ impl Vm {
     }
 
     pub(crate) fn get_function_idx(&self, name: &str) -> u16 {
-        let position = self.bytecode.iter().position(|(n, _)| n == name).unwrap();
+        // We search backwards to allow for functions
+        // to be overwritten.
+        let rev_position = self
+            .bytecode
+            .iter()
+            .rev()
+            .position(|(n, _)| n == name)
+            .unwrap();
+        let position = self.bytecode.len() - 1 - rev_position;
         assert!(position <= u16::MAX as usize);
         position as u16
     }


Test patch:
diff --git a/numbat/tests/interpreter.rs b/numbat/tests/interpreter.rs
index 541756a..8f664c7 100644
--- a/numbat/tests/interpreter.rs
+++ b/numbat/tests/interpreter.rs
@@ -305,26 +305,42 @@ fn test_prefixes() {
 }
 
 #[test]
-fn test_error_messages() {
+fn test_parse_errors() {
     expect_failure(
         "3kg+",
         "Expected one of: number, identifier, parenthesized expression",
     );
-    expect_failure("let kg=2", "Identifier is already in use: 'kg'");
-    expect_failure("let pi=2", "Identifier is already in use: 'pi'");
-    expect_failure("let sin=2", "Identifier is already in use: 'sin'");
     expect_failure("let print=2", "Expected identifier after 'let' keyword");
-    expect_failure("fn kg(x: Scalar) = 1", "Identifier is already in use: 'kg'");
-    expect_failure("fn pi(x: Scalar) = 1", "Identifier is already in use: 'pi'");
-    expect_failure(
-        "fn sin(x: Scalar) = 1",
-        "Identifier is already in use: 'sin'",
-    );
     expect_failure(
         "fn print(x: Scalar) = 1",
         "Expected identifier after 'fn' keyword",
     );
+}
+
+#[test]
+fn test_name_clash_errors() {
+    expect_failure("let kg=2", "Identifier is already in use: 'kg'");
+    expect_failure("fn kg(x: Scalar) = 1", "Identifier is already in use: 'kg'");
+    expect_failure("fn _()=0", "Reserved identifier");
+}
+
+#[test]
+fn test_type_check_errors() {
     expect_failure("foo", "Unknown identifier 'foo'");
+
+    expect_failure("let sin=2", "This name is already used by a function");
+    expect_failure("fn pi() = 1", "This name is already used by a constant");
+    expect_failure(
+        "fn sin(x)=0",
+        "This name is already used by a foreign function",
+    );
+
+    // TODO: this restriction should be lifted in the future:
+    expect_failure("let pi = 1", "This name is already used by a constant");
+}
+
+#[test]
+fn test_runtime_errors() {
     expect_failure("1/0", "Division by zero");
 }
 
@@ -368,3 +384,28 @@ fn test_string_interpolation() {
     expect_output("\"pi = {pi}!\"", "pi = 3.14159!");
     expect_output("\"1 + 2 = {1 + 2}\"", "1 + 2 = 3");
 }
+
+#[test]
+fn test_overwrite_regular_function() {
+    expect_output(
+        "
+        fn f(x)=0
+        fn f(x)=1
+        f(2)",
+        "1",
+    );
+}
+
+#[test]
+fn test_overwrite_inner_function() {
+    expect_output(
+        "
+        fn inner() = 0
+        fn outer() = inner()
+
+        fn inner(x) = 1
+        outer()
+        ",
+        "0",
+    );
+}