Turn exponents into integers

numbat/src/typechecker/mod.rs

@@ -14,6 +14,7 @@ pub mod type_scheme;
 use std::collections::HashMap;
 use std::ops::Deref;
 
+use crate::arithmetic::Exponent;
 use crate::ast::{self, BinaryOperator, ProcedureKind, StringPart, TypeAnnotation, TypeExpression};
 use crate::dimension::DimensionRegistry;
 use crate::name_resolution::Namespace;
@@ -36,7 +37,7 @@ use num_traits::Zero;
 pub use error::{Result, TypeCheckError};
 pub use incompatible_dimensions::IncompatibleDimensionsError;
 use qualified_type::{Bound, QualifiedType};
-use substitutions::ApplySubstitution;
+use substitutions::{ApplySubstitution, Substitution};
 use type_scheme::TypeScheme;
 
 fn dtype(e: &Expression) -> Result<DType> {
@@ -1664,23 +1665,23 @@ impl TypeChecker {
         // within type_, and then multiply the corresponding exponents with the least common
         // multiple of the denominators of the exponents. For example, this will turn
         // T0^(1/3) -> T0^(1/5) -> T0 into T0^5 -> T0^3 -> T0^15.
-        // for tv in &dtype_variables {
-        //     let exponents = elaborated_statement.exponents_for(&tv);
-        //     let lcm = exponents
-        //         .iter()
-        //         .fold(1, |acc, e| num_integer::lcm(acc, *e.denom()));
-
-        //     if lcm != 1 {
-        //         let s = Substitution::single(
-        //             tv.clone(),
-        //             Type::Dimension(
-        //                 DType::from_type_variable(tv.clone()).power(Exponent::from_integer(lcm)),
-        //             ),
-        //         );
-
-        //         elaborated_statement.apply(&s).unwrap();
-        //     }
-        // }
+        for tv in &dtype_variables {
+            let exponents = elaborated_statement.exponents_for(&tv);
+            let lcm = exponents
+                .iter()
+                .fold(1, |acc, e| num_integer::lcm(acc, *e.denom()));
+
+            if lcm != 1 {
+                let s = Substitution::single(
+                    tv.clone(),
+                    Type::Dimension(
+                        DType::from_type_variable(tv.clone()).power(Exponent::from_integer(lcm)),
+                    ),
+                );
+
+                elaborated_statement.apply(&s).unwrap();
+            }
+        }
 
         elaborated_statement.generalize_types(&dtype_variables);
         self.env.generalize_types(&dtype_variables);

numbat/src/typed_ast.rs

@@ -554,6 +554,25 @@ impl Statement {
     pub(crate) fn generalize_types(&mut self, dtype_variables: &[TypeVariable]) {
         self.for_all_type_schemes(&mut |type_: &mut TypeScheme| type_.generalize(dtype_variables));
     }
+
+    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![];
+        self.for_all_type_schemes(&mut |type_: &mut TypeScheme| {
+            match type_.to_concrete_type() {
+                // TODO: do not use to_concrete_type here
+                Type::Dimension(dtype) => {
+                    for (factor, exp) in dtype.factors {
+                        if factor == DTypeFactor::TVar(tv.clone()) {
+                            exponents.push(exp)
+                        }
+                    }
+                }
+                _ => {}
+            }
+        });
+        exponents
+    }
 }
 
 impl Expression {