Introduce DType::is_scalar function

sharkdp · Nov 23, 2023 · 341ca6d · 341ca6d
1 parent 1651bbc
commit 341ca6d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/numbat/src/typechecker.rs b/numbat/src/typechecker.rs
@@ -457,7 +457,7 @@ impl TypeChecker {
 
                 match *op {
                     ast::UnaryOperator::Factorial => {
-                        if dtype != DType::unity() {
+                        if !dtype.is_scalar() {
                             return Err(TypeCheckError::NonScalarFactorialArgument(
                                 expr.full_span(),
                                 dtype,
@@ -539,15 +539,15 @@ impl TypeChecker {
                     }
                     typed_ast::BinaryOperator::Power => {
                         let exponent_type = dtype(&rhs_checked)?;
-                        if exponent_type != DType::unity() {
+                        if !exponent_type.is_scalar() {
                             return Err(TypeCheckError::NonScalarExponent(
                                 rhs.full_span(),
                                 exponent_type,
                             ));
                         }
 
                         let base_type = dtype(&lhs_checked)?;
-                        if base_type == DType::unity() {
+                        if base_type.is_scalar() {
                             // Skip evaluating the exponent if the lhs is a scalar. This allows
                             // for arbitrary (decimal) exponents, if the base is a scalar.
 

diff --git a/numbat/src/typed_ast.rs b/numbat/src/typed_ast.rs
@@ -15,8 +15,11 @@ use crate::{
 pub type DType = BaseRepresentation;
 
 impl DType {
+    pub fn is_scalar(&self) -> bool {
+        self == &DType::unity()
+    }
     pub fn to_readable_type(&self, registry: &DimensionRegistry) -> m::Markup {
-        if self == &DType::unity() {
+        if self.is_scalar() {
             return m::type_identifier("Scalar");
         }