Fix to use correct units for array range expressions (#6837)

* Fix to use correct units for array range expressions

* Update output

Author: jtran

rust/kcl-lib/src/execution/exec_ast.rs

@@ -16,6 +16,7 @@ use crate::{
         BodyType, EnvironmentRef, ExecState, ExecutorContext, KclValue, Metadata, PlaneType, TagEngineInfo,
         TagIdentifier,
     },
+    fmt,
     modules::{ModuleId, ModulePath, ModuleRepr},
     parsing::ast::types::{
         Annotation, ArrayExpression, ArrayRangeExpression, AscribedExpression, BinaryExpression, BinaryOperator,
@@ -1611,7 +1612,7 @@ impl Node<ArrayRangeExpression> {
     #[async_recursion]
     pub async fn execute(&self, exec_state: &mut ExecState, ctx: &ExecutorContext) -> Result<KclValue, KclError> {
         let metadata = Metadata::from(&self.start_element);
-        let start = ctx
+        let start_val = ctx
             .execute_expr(
                 &self.start_element,
                 exec_state,
@@ -1620,19 +1621,30 @@ impl Node<ArrayRangeExpression> {
                 StatementKind::Expression,
             )
             .await?;
-        let start = start.as_int().ok_or(KclError::Semantic(KclErrorDetails {
+        let (start, start_ty) = start_val.as_int_with_ty().ok_or(KclError::Semantic(KclErrorDetails {
             source_ranges: vec![self.into()],
-            message: format!("Expected int but found {}", start.human_friendly_type()),
+            message: format!("Expected int but found {}", start_val.human_friendly_type()),
         }))?;
         let metadata = Metadata::from(&self.end_element);
-        let end = ctx
+        let end_val = ctx
             .execute_expr(&self.end_element, exec_state, &metadata, &[], StatementKind::Expression)
             .await?;
-        let end = end.as_int().ok_or(KclError::Semantic(KclErrorDetails {
+        let (end, end_ty) = end_val.as_int_with_ty().ok_or(KclError::Semantic(KclErrorDetails {
             source_ranges: vec![self.into()],
-            message: format!("Expected int but found {}", end.human_friendly_type()),
+            message: format!("Expected int but found {}", end_val.human_friendly_type()),
         }))?;
 
+        if start_ty != end_ty {
+            let start = start_val.as_ty_f64().unwrap_or(TyF64 { n: 0.0, ty: start_ty });
+            let start = fmt::human_display_number(start.n, start.ty);
+            let end = end_val.as_ty_f64().unwrap_or(TyF64 { n: 0.0, ty: end_ty });
+            let end = fmt::human_display_number(end.n, end.ty);
+            return Err(KclError::Semantic(KclErrorDetails {
+                source_ranges: vec![self.into()],
+                message: format!("Range start and end must be of the same type, but found {start} and {end}"),
+            }));
+        }
+
         if end < start {
             return Err(KclError::Semantic(KclErrorDetails {
                 source_ranges: vec![self.into()],
@@ -1655,11 +1667,11 @@ impl Node<ArrayRangeExpression> {
                 .into_iter()
                 .map(|num| KclValue::Number {
                     value: num as f64,
-                    ty: NumericType::count(),
+                    ty: start_ty.clone(),
                     meta: meta.clone(),
                 })
                 .collect(),
-            ty: RuntimeType::Primitive(PrimitiveType::Number(NumericType::count())),
+            ty: RuntimeType::Primitive(PrimitiveType::Number(start_ty)),
         })
     }
 }

rust/kcl-lib/src/execution/kcl_value.rs

@@ -404,6 +404,13 @@ impl KclValue {
         }
     }
 
+    pub fn as_int_with_ty(&self) -> Option<(i64, NumericType)> {
+        match self {
+            KclValue::Number { value, ty, .. } => crate::try_f64_to_i64(*value).map(|i| (i, ty.clone())),
+            _ => None,
+        }
+    }
+
     pub fn as_object(&self) -> Option<&KclObjectFields> {
         if let KclValue::Object { value, meta: _ } = &self {
             Some(value)

rust/kcl-lib/src/simulation_tests.rs

@@ -594,6 +594,48 @@ mod array_range_negative_expr {
         super::execute(TEST_NAME, false).await
     }
 }
+mod array_range_with_units {
+    const TEST_NAME: &str = "array_range_with_units";
+
+    /// Test parsing KCL.
+    #[test]
+    fn parse() {
+        super::parse(TEST_NAME)
+    }
+
+    /// Test that parsing and unparsing KCL produces the original KCL input.
+    #[tokio::test(flavor = "multi_thread")]
+    async fn unparse() {
+        super::unparse(TEST_NAME).await
+    }
+
+    /// Test that KCL is executed correctly.
+    #[tokio::test(flavor = "multi_thread")]
+    async fn kcl_test_execute() {
+        super::execute(TEST_NAME, false).await
+    }
+}
+mod array_range_mismatch_units {
+    const TEST_NAME: &str = "array_range_mismatch_units";
+
+    /// Test parsing KCL.
+    #[test]
+    fn parse() {
+        super::parse(TEST_NAME)
+    }
+
+    /// Test that parsing and unparsing KCL produces the original KCL input.
+    #[tokio::test(flavor = "multi_thread")]
+    async fn unparse() {
+        super::unparse(TEST_NAME).await
+    }
+
+    /// Test that KCL is executed correctly.
+    #[tokio::test(flavor = "multi_thread")]
+    async fn kcl_test_execute() {
+        super::execute(TEST_NAME, false).await
+    }
+}
 mod sketch_in_object {
     const TEST_NAME: &str = "sketch_in_object";