KCL: Array index can be any expression (#7755)

Resolves #6651.

Before, `arr[f(i)]` and `arr[i + 1]` were syntax errors. Now they work as you would expect.

## Implementation

Before, the properties in a MemberExpression were an enum with either Literal or Identifier, named (imaginatively) LiteralIdentifier. Now that enum is gone. The property of a MemberExpression is now just any Expr (an enum of all KCL expressions).

Author: adamchalmers

e2e/playwright/snapshot-tests.spec.ts-snapshots/Draft-segments-should-look-right-3-Google-Chrome-linux.png

@@ -17,14 +17,10 @@ use crate::{
     },
     fmt,
     modules::{ModuleId, ModulePath, ModuleRepr},
-    parsing::{
-        ast::types::{
-            Annotation, ArrayExpression, ArrayRangeExpression, AscribedExpression, BinaryExpression, BinaryOperator,
-            BinaryPart, BodyItem, Expr, IfExpression, ImportPath, ImportSelector, ItemVisibility, LiteralIdentifier,
-            LiteralValue, MemberExpression, Name, Node, NodeRef, ObjectExpression, PipeExpression, Program,
-            TagDeclarator, Type, UnaryExpression, UnaryOperator,
-        },
-        token::NumericSuffix,
+    parsing::ast::types::{
+        Annotation, ArrayExpression, ArrayRangeExpression, AscribedExpression, BinaryExpression, BinaryOperator,
+        BinaryPart, BodyItem, Expr, IfExpression, ImportPath, ImportSelector, ItemVisibility, MemberExpression, Name,
+        Node, NodeRef, ObjectExpression, PipeExpression, Program, TagDeclarator, Type, UnaryExpression, UnaryOperator,
     },
     source_range::SourceRange,
     std::args::TyF64,
@@ -984,10 +980,20 @@ impl Node<Name> {
 
 impl Node<MemberExpression> {
     async fn get_result(&self, exec_state: &mut ExecState, ctx: &ExecutorContext) -> Result<KclValue, KclError> {
-        let property = Property::try_from(self.computed, self.property.clone(), exec_state, self.into())?;
         let meta = Metadata {
             source_range: SourceRange::from(self),
         };
+        let property = Property::try_from(
+            self.computed,
+            self.property.clone(),
+            exec_state,
+            self.into(),
+            ctx,
+            &meta,
+            &[],
+            StatementKind::Expression,
+        )
+        .await?;
         let object = ctx
             .execute_expr(&self.object, exec_state, &meta, &[], StatementKind::Expression)
             .await?;
@@ -1726,87 +1732,62 @@ enum Property {
 }
 
 impl Property {
-    fn try_from(
+    #[allow(clippy::too_many_arguments)]
+    async fn try_from<'a>(
         computed: bool,
-        value: LiteralIdentifier,
-        exec_state: &ExecState,
+        value: Expr,
+        exec_state: &mut ExecState,
         sr: SourceRange,
+        ctx: &ExecutorContext,
+        metadata: &Metadata,
+        annotations: &[Node<Annotation>],
+        statement_kind: StatementKind<'a>,
     ) -> Result<Self, KclError> {
         let property_sr = vec![sr];
-        let property_src: SourceRange = value.clone().into();
-        match value {
-            LiteralIdentifier::Identifier(identifier) => {
-                let name = &identifier.name;
-                if !computed {
-                    // This is dot syntax. Treat the property as a literal.
-                    Ok(Property::String(name.to_string()))
-                } else {
-                    // This is bracket syntax. Actually evaluate memory to
-                    // compute the property.
-                    let prop = exec_state.stack().get(name, property_src)?;
-                    jvalue_to_prop(prop, property_sr, name)
-                }
-            }
-            LiteralIdentifier::Literal(literal) => {
-                let value = literal.value.clone();
-                match value {
-                    n @ LiteralValue::Number { value, suffix } => {
-                        if !matches!(suffix, NumericSuffix::None | NumericSuffix::Count) {
-                            return Err(KclError::new_semantic(KclErrorDetails::new(
-                                format!("{n} is not a valid index, indices must be non-dimensional numbers"),
-                                property_sr,
-                            )));
-                        }
-                        if let Some(x) = crate::try_f64_to_usize(value) {
-                            Ok(Property::UInt(x))
-                        } else {
-                            Err(KclError::new_semantic(KclErrorDetails::new(
-                                format!("{n} is not a valid index, indices must be whole numbers >= 0"),
-                                property_sr,
-                            )))
-                        }
-                    }
-                    _ => Err(KclError::new_semantic(KclErrorDetails::new(
-                        "Only numbers (>= 0) can be indexes".to_owned(),
-                        vec![sr],
-                    ))),
-                }
-            }
+        if !computed {
+            let Expr::Name(identifier) = value else {
+                // Should actually be impossible because the parser would reject it.
+                return Err(KclError::new_semantic(KclErrorDetails::new(
+                    "Object expressions like `obj.property` must use simple identifier names, not complex expressions"
+                        .to_owned(),
+                    property_sr,
+                )));
+            };
+            return Ok(Property::String(identifier.to_string()));
         }
-    }
-}
 
-fn jvalue_to_prop(value: &KclValue, property_sr: Vec<SourceRange>, name: &str) -> Result<Property, KclError> {
-    let make_err =
-        |message: String| Err::<Property, _>(KclError::new_semantic(KclErrorDetails::new(message, property_sr)));
-    match value {
-        n @ KclValue::Number { value: num, ty, .. } => {
-            if !matches!(
-                ty,
-                NumericType::Known(crate::exec::UnitType::Count) | NumericType::Default { .. } | NumericType::Any
-            ) {
-                return make_err(format!(
-                    "arrays can only be indexed by non-dimensioned numbers, found {}",
-                    n.human_friendly_type()
-                ));
-            }
-            let num = *num;
-            if num < 0.0 {
-                return make_err(format!("'{num}' is negative, so you can't index an array with it"));
-            }
-            let nearest_int = crate::try_f64_to_usize(num);
-            if let Some(nearest_int) = nearest_int {
-                Ok(Property::UInt(nearest_int))
-            } else {
-                make_err(format!(
-                    "'{num}' is not an integer, so you can't index an array with it"
-                ))
+        let prop_value = ctx
+            .execute_expr(&value, exec_state, metadata, annotations, statement_kind)
+            .await?;
+        match prop_value {
+            KclValue::Number { value, ty, meta: _ } => {
+                if !matches!(
+                    ty,
+                    NumericType::Unknown
+                        | NumericType::Default { .. }
+                        | NumericType::Known(crate::exec::UnitType::Count)
+                ) {
+                    return Err(KclError::new_semantic(KclErrorDetails::new(
+                        format!(
+                            "{value} is not a valid index, indices must be non-dimensional numbers. If you're sure this is correct, you can add `: number(Count)` to tell KCL this number is an index"
+                        ),
+                        property_sr,
+                    )));
+                }
+                if let Some(x) = crate::try_f64_to_usize(value) {
+                    Ok(Property::UInt(x))
+                } else {
+                    Err(KclError::new_semantic(KclErrorDetails::new(
+                        format!("{value} is not a valid index, indices must be whole numbers >= 0"),
+                        property_sr,
+                    )))
+                }
             }
+            _ => Err(KclError::new_semantic(KclErrorDetails::new(
+                "Only numbers (>= 0) can be indexes".to_owned(),
+                vec![sr],
+            ))),
         }
-        KclValue::String { value: x, meta: _ } => Ok(Property::String(x.to_owned())),
-        _ => make_err(format!(
-            "{name} is not a valid property/index, you can only use a string to get the property of an object, or an int (>= 0) to get an item in an array"
-        )),
     }
 }
 

public/kcl-samples/screenshots/focusrite-scarlett-mounting-bracket.png

@@ -4,9 +4,9 @@ use crate::parsing::ast::types::{
     Annotation, ArrayExpression, ArrayRangeExpression, AscribedExpression, BinaryExpression, BinaryPart, BodyItem,
     CallExpressionKw, DefaultParamVal, ElseIf, Expr, ExpressionStatement, FunctionExpression, FunctionType, Identifier,
     IfExpression, ImportItem, ImportSelector, ImportStatement, ItemVisibility, KclNone, LabelledExpression, Literal,
-    LiteralIdentifier, LiteralValue, MemberExpression, Name, ObjectExpression, ObjectProperty, Parameter,
-    PipeExpression, PipeSubstitution, PrimitiveType, Program, ReturnStatement, TagDeclarator, Type, TypeDeclaration,
-    UnaryExpression, VariableDeclaration, VariableDeclarator, VariableKind,
+    LiteralValue, MemberExpression, Name, ObjectExpression, ObjectProperty, Parameter, PipeExpression,
+    PipeSubstitution, PrimitiveType, Program, ReturnStatement, TagDeclarator, Type, TypeDeclaration, UnaryExpression,
+    VariableDeclaration, VariableDeclarator, VariableKind,
 };
 
 /// Position-independent digest of the AST node.
@@ -170,14 +170,6 @@ impl BinaryPart {
     }
 }
 
-impl LiteralIdentifier {
-    pub fn compute_digest(&mut self) -> Digest {
-        match self {
-            LiteralIdentifier::Identifier(id) => id.compute_digest(),
-            LiteralIdentifier::Literal(lit) => lit.compute_digest(),
-        }
-    }
-}
 impl Type {
     pub fn compute_digest(&mut self) -> Digest {
         let mut hasher = Sha256::new();

public/kcl-samples/screenshots/food-service-spatula.png

@@ -3,7 +3,7 @@ pub mod types;
 
 use crate::{
     ModuleId,
-    parsing::ast::types::{BinaryPart, BodyItem, Expr, LiteralIdentifier},
+    parsing::ast::types::{BinaryPart, BodyItem, Expr},
 };
 
 impl BodyItem {
@@ -59,12 +59,3 @@ impl BinaryPart {
         }
     }
 }
-
-impl LiteralIdentifier {
-    pub fn module_id(&self) -> ModuleId {
-        match self {
-            LiteralIdentifier::Identifier(identifier) => identifier.module_id,
-            LiteralIdentifier::Literal(literal) => literal.module_id,
-        }
-    }
-}

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

@@ -2758,53 +2758,13 @@ impl ObjectProperty {
     }
 }
 
-#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, ts_rs::TS, JsonSchema)]
-#[ts(export)]
-#[serde(tag = "type")]
-pub enum LiteralIdentifier {
-    Identifier(BoxNode<Identifier>),
-    Literal(BoxNode<Literal>),
-}
-
-impl LiteralIdentifier {
-    pub fn start(&self) -> usize {
-        match self {
-            LiteralIdentifier::Identifier(identifier) => identifier.start,
-            LiteralIdentifier::Literal(literal) => literal.start,
-        }
-    }
-
-    pub fn end(&self) -> usize {
-        match self {
-            LiteralIdentifier::Identifier(identifier) => identifier.end,
-            LiteralIdentifier::Literal(literal) => literal.end,
-        }
-    }
-
-    pub(crate) fn contains_range(&self, range: &SourceRange) -> bool {
-        let sr = SourceRange::from(self);
-        sr.contains_range(range)
-    }
-}
-
-impl From<LiteralIdentifier> for SourceRange {
-    fn from(id: LiteralIdentifier) -> Self {
-        Self::new(id.start(), id.end(), id.module_id())
-    }
-}
-
-impl From<&LiteralIdentifier> for SourceRange {
-    fn from(id: &LiteralIdentifier) -> Self {
-        Self::new(id.start(), id.end(), id.module_id())
-    }
-}
-
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, ts_rs::TS, JsonSchema)]
 #[ts(export)]
 #[serde(tag = "type")]
 pub struct MemberExpression {
     pub object: Expr,
-    pub property: LiteralIdentifier,
+    pub property: Expr,
+    /// True if `obj[prop]`, false if obj.prop
     pub computed: bool,
 
     #[serde(default, skip_serializing_if = "Option::is_none")]
@@ -2826,22 +2786,10 @@ impl MemberExpression {
     /// Rename all identifiers that have the old name to the new given name.
     fn rename_identifiers(&mut self, old_name: &str, new_name: &str) {
         self.object.rename_identifiers(old_name, new_name);
-
-        match &mut self.property {
-            LiteralIdentifier::Identifier(identifier) => identifier.rename(old_name, new_name),
-            LiteralIdentifier::Literal(_) => {}
-        }
+        self.property.rename_identifiers(old_name, new_name);
     }
 }
 
-#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, ts_rs::TS, JsonSchema)]
-#[ts(export)]
-pub struct ObjectKeyInfo {
-    pub key: LiteralIdentifier,
-    pub index: usize,
-    pub computed: bool,
-}
-
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, ts_rs::TS, JsonSchema)]
 #[ts(export)]
 #[serde(tag = "type")]

rust/kcl-lib/src/parsing/ast/digest.rs

@@ -266,7 +266,7 @@ impl NodePath {
                 }
                 if node.property.contains_range(&range) {
                     path.push(Step::MemberExpressionProperty);
-                    return Some(path);
+                    return NodePath::from_expr(&node.property, range, path);
                 }
             }
             Expr::UnaryExpression(node) => {

rust/kcl-lib/src/parsing/ast/mod.rs

@@ -28,10 +28,10 @@ use crate::{
             Annotation, ArrayExpression, ArrayRangeExpression, BinaryExpression, BinaryOperator, BinaryPart, BodyItem,
             BoxNode, CallExpressionKw, CommentStyle, DefaultParamVal, ElseIf, Expr, ExpressionStatement,
             FunctionExpression, FunctionType, Identifier, IfExpression, ImportItem, ImportSelector, ImportStatement,
-            ItemVisibility, LabeledArg, Literal, LiteralIdentifier, LiteralValue, MemberExpression, Name, Node,
-            NodeList, NonCodeMeta, NonCodeNode, NonCodeValue, ObjectExpression, ObjectProperty, Parameter,
-            PipeExpression, PipeSubstitution, PrimitiveType, Program, ReturnStatement, Shebang, TagDeclarator, Type,
-            TypeDeclaration, UnaryExpression, UnaryOperator, VariableDeclaration, VariableDeclarator, VariableKind,
+            ItemVisibility, LabeledArg, Literal, LiteralValue, MemberExpression, Name, Node, NodeList, NonCodeMeta,
+            NonCodeNode, NonCodeValue, ObjectExpression, ObjectProperty, Parameter, PipeExpression, PipeSubstitution,
+            PrimitiveType, Program, ReturnStatement, Shebang, TagDeclarator, Type, TypeDeclaration, UnaryExpression,
+            UnaryOperator, VariableDeclaration, VariableDeclarator, VariableKind,
         },
         math::BinaryExpressionToken,
         token::{Token, TokenSlice, TokenType},
@@ -1299,36 +1299,34 @@ fn function_decl(i: &mut TokenSlice) -> ModalResult<Node<FunctionExpression>> {
 }
 
 /// E.g. `person.name`
-fn member_expression_dot(i: &mut TokenSlice) -> ModalResult<(LiteralIdentifier, usize, bool)> {
+fn member_expression_dot(i: &mut TokenSlice) -> ModalResult<(Expr, usize, bool)> {
     period.parse_next(i)?;
     let property = nameable_identifier
         .map(Box::new)
-        .map(LiteralIdentifier::Identifier)
+        .map(|p| {
+            let ni: Node<Identifier> = *p;
+            let nn: Node<Name> = ni.into();
+            Expr::Name(Box::new(nn))
+        })
         .parse_next(i)?;
     let end = property.end();
-    Ok((property, end, false))
+    let computed = false;
+    Ok((property, end, computed))
 }
 
-/// E.g. `people[0]` or `people[i]` or `people['adam']`
-fn member_expression_subscript(i: &mut TokenSlice) -> ModalResult<(LiteralIdentifier, usize, bool)> {
+fn member_expression_subscript_complex_expr(i: &mut TokenSlice) -> ModalResult<(Expr, usize, bool)> {
     let _ = open_bracket.parse_next(i)?;
-    // TODO: This should be an expression, not just a literal or identifier.
-    let property = alt((
-        literal.map(LiteralIdentifier::Literal),
-        nameable_identifier.map(Box::new).map(LiteralIdentifier::Identifier),
-    ))
-    .parse_next(i)?;
-
+    let property = expression.parse_next(i)?;
     let end = close_bracket.parse_next(i)?.end;
-    let computed = matches!(property, LiteralIdentifier::Identifier(_));
+    let computed = true;
     Ok((property, end, computed))
 }
 
 /// Get a property of an object, or an index of an array, or a member of a collection.
 /// Can be arbitrarily nested, e.g. `people[i]['adam'].age`.
 fn build_member_expression(object: Expr, i: &mut TokenSlice) -> ModalResult<Node<MemberExpression>> {
     // Now a sequence of members.
-    let member = alt((member_expression_dot, member_expression_subscript)).context(expected("a member/property, e.g. size.x and size['height'] and size[0] are all different ways to access a member/property of 'size'"));
+    let member = alt((member_expression_dot, member_expression_subscript_complex_expr)).context(expected("a member/property, e.g. size.x and size['height'] and size[0] are all different ways to access a member/property of 'size'"));
     let mut members: Vec<_> = repeat(1.., member)
         .context(expected("a sequence of at least one members/properties"))
         .parse_next(i)?;
@@ -3407,6 +3405,19 @@ mod tests {
         binary_expression.parse(tokens).unwrap();
     }
 
+    #[test]
+    fn expression_in_array_index() {
+        let tokens = crate::parsing::token::lex("arr[x + 1]", ModuleId::default()).unwrap();
+        let tokens = tokens.as_slice();
+        let Expr::MemberExpression(expr) = expression.parse(tokens).unwrap() else {
+            panic!();
+        };
+        let Expr::BinaryExpression(be) = expr.inner.property else {
+            panic!();
+        };
+        assert_eq!(be.inner.operator, BinaryOperator::Add);
+    }
+
     #[test]
     fn parse_binary_operator_on_object() {
         let tokens = crate::parsing::token::lex("{ a = 1 } + 2", ModuleId::default()).unwrap();