KCL: Refactor pipe bodies

PipeExpressions should always have at least two child expressions. The first expression (the head) is a little different from the following subexpressions, because the head can be any value. But the remaining expressions (tail) must be function calls. This refactors the AST node for PipeExpression so that information is encoded in the type system, not just checked at runtime. Instead of a Vec<Value>, it should have a separate head Value and tail Vec<Value>.
2024-04-03 10:38:43 -05:00
11 changed files with 273 additions and 260 deletions
--- a/src/wasm-lib/grackle/src/lib.rs
+++ b/src/wasm-lib/grackle/src/lib.rs
@ -465,7 +465,7 @@ impl Planner {
                }
            }
            SingleValue::PipeExpression(expr) => {
-                let mut bodies = expr.body.into_iter();
+                let mut bodies = (*expr).all_children();

                // Get the first expression (i.e. body) of the pipeline.
                let first = bodies.next().expect("Pipe expression must have > 1 item");
--- a/src/wasm-lib/kcl/src/ast/types.rs
+++ b/src/wasm-lib/kcl/src/ast/types.rs
@ -2532,7 +2532,10 @@ pub struct PipeExpression {
    pub end: usize,
    // TODO: Only the first body expression can be any Value.
    // The rest will be CallExpression, and the AST type should reflect this.
-    pub body: Vec<Value>,
+    /// First child expression.
+    pub head: Value,
+    /// Remaining child expressions.
+    pub tail: Vec<Value>,
    pub non_code_meta: NonCodeMeta,
 }

@ -2545,23 +2548,40 @@ impl From<PipeExpression> for Value {
 }

 impl PipeExpression {
-    pub fn new(body: Vec<Value>) -> Self {
+    pub fn new(mut body: Vec<Value>) -> Self {
+        let head = body.remove(0);
        Self {
            start: 0,
            end: 0,
-            body,
+            head,
+            tail: body,
            non_code_meta: Default::default(),
        }
    }

+    /// Iterate over all child expressions in this pipe expression.
+    pub fn iter(&self) -> impl Iterator<Item = &Value> {
+        std::iter::once(&self.head).chain(self.tail.iter())
+    }
+
+    /// Iterate over all child expressions in this pipe expression.
+    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut Value> {
+        std::iter::once(&mut self.head).chain(self.tail.iter_mut())
+    }
+
+    /// Iterate over all child expressions in this pipe expression.
+    pub fn all_children(self) -> impl Iterator<Item = Value> {
+        std::iter::once(self.head).chain(self.tail)
+    }
+
    pub fn replace_value(&mut self, source_range: SourceRange, new_value: Value) {
-        for value in &mut self.body {
+        for value in self.iter_mut() {
            value.replace_value(source_range, new_value.clone());
        }
    }

    pub fn get_constraint_level(&self) -> ConstraintLevel {
-        if self.body.is_empty() {
+        if self.tail.is_empty() {
            return ConstraintLevel::Ignore {
                source_ranges: vec![self.into()],
            };
@ -2569,7 +2589,7 @@ impl PipeExpression {

        // Iterate over all body expressions.
        let mut constraint_levels = ConstraintLevels::new();
-        for expression in &self.body {
+        for expression in self.iter() {
            constraint_levels.push(expression.get_constraint_level());
        }

@ -2577,8 +2597,7 @@ impl PipeExpression {
    }

    fn recast(&self, options: &FormatOptions, indentation_level: usize) -> String {
-        self.body
-            .iter()
+        self.iter()
            .enumerate()
            .map(|(index, statement)| {
                let indentation = options.get_indentation(indentation_level + 1);
@ -2596,7 +2615,7 @@ impl PipeExpression {
                    }
                }

-                if index != self.body.len() - 1 {
+                if index != self.tail.len() - 1 {
                    s += "\n";
                    s += &indentation;
                    s += PIPE_OPERATOR;
@ -2609,7 +2628,7 @@ impl PipeExpression {

    /// Returns a hover value that includes the given character position.
    pub fn get_hover_value_for_position(&self, pos: usize, code: &str) -> Option<Hover> {
-        for b in &self.body {
+        for b in self.iter() {
            let b_source_range: SourceRange = b.into();
            if b_source_range.contains(pos) {
                return b.get_hover_value_for_position(pos, code);
@ -2625,12 +2644,12 @@ impl PipeExpression {
        pipe_info: &PipeInfo,
        ctx: &ExecutorContext,
    ) -> Result<MemoryItem, KclError> {
-        execute_pipe_body(memory, &self.body, pipe_info, self.into(), ctx).await
+        execute_pipe_body(memory, &self.head, &self.tail, pipe_info, ctx).await
    }

    /// Rename all identifiers that have the old name to the new given name.
    fn rename_identifiers(&mut self, old_name: &str, new_name: &str) {
-        for statement in &mut self.body {
+        for statement in self.iter_mut() {
            statement.rename_identifiers(old_name, new_name);
        }
    }
@ -2639,31 +2658,24 @@ impl PipeExpression {
 #[async_recursion::async_recursion(?Send)]
 async fn execute_pipe_body(
    memory: &mut ProgramMemory,
-    body: &[Value],
+    head: &Value,
+    tail: &[Value],
    pipe_info: &PipeInfo,
-    source_range: SourceRange,
    ctx: &ExecutorContext,
 ) -> Result<MemoryItem, KclError> {
-    let mut body = body.iter();
-    let first = body.next().ok_or_else(|| {
-        KclError::Semantic(KclErrorDetails {
-            message: "Pipe expressions cannot be empty".to_owned(),
-            source_ranges: vec![source_range],
-        })
-    })?;
    // Evaluate the first element in the pipeline.
    // They use the `pipe_info` from some AST node above this, so that if pipe expression is nested in a larger pipe expression,
    // they use the % from the parent. After all, this pipe expression hasn't been executed yet, so it doesn't have any % value
    // of its own.
-    let output = match first {
+    let output = match head {
        Value::BinaryExpression(binary_expression) => binary_expression.get_result(memory, pipe_info, ctx).await?,
        Value::CallExpression(call_expression) => call_expression.execute(memory, pipe_info, ctx).await?,
        Value::Identifier(identifier) => memory.get(&identifier.name, identifier.into())?.clone(),
        _ => {
            // Return an error this should not happen.
            return Err(KclError::Semantic(KclErrorDetails {
-                message: format!("PipeExpression not implemented here: {:?}", first),
-                source_ranges: vec![first.into()],
+                message: format!("PipeExpression not implemented here: {head:?}"),
+                source_ranges: vec![head.into()],
            }));
        }
    };
@ -2673,7 +2685,7 @@ async fn execute_pipe_body(
    let mut new_pipe_info = PipeInfo::new();
    new_pipe_info.previous_results = Some(output);
    // Evaluate remaining elements.
-    for expression in body {
+    for expression in tail {
        let output = match expression {
            Value::BinaryExpression(binary_expression) => {
                binary_expression.get_result(memory, &new_pipe_info, ctx).await?
--- a/src/wasm-lib/kcl/src/parser/parser_impl.rs
+++ b/src/wasm-lib/kcl/src/parser/parser_impl.rs
@ -136,7 +136,7 @@ fn pipe_expression(i: TokenSlice) -> PResult<PipeExpression> {
    if let Some(nc) = noncode {
        non_code_meta.insert(0, nc);
    }
-    let mut values = vec![head];
+    let mut values = Vec::new();
    let value_surrounded_by_comments = (
        repeat(0.., preceded(opt(whitespace), non_code_node)), // Before the value
        preceded(opt(whitespace), fn_call),                    // The value
@ -185,7 +185,8 @@ fn pipe_expression(i: TokenSlice) -> PResult<PipeExpression> {
    Ok(PipeExpression {
        start: values.first().unwrap().start(),
        end: values.last().unwrap().end().max(max_noncode_end),
-        body: values,
+        head,
+        tail: values,
        non_code_meta,
    })
 }
@ -1584,7 +1585,7 @@ const mySk1 = startSketchAt([0, 0])"#;
        let tokens = crate::token::lexer(test_input);
        let mut slice = tokens.as_slice();
        let PipeExpression {
-            body, non_code_meta, ..
+            tail, non_code_meta, ..
        } = pipe_expression.parse_next(&mut slice).unwrap();
        assert_eq!(non_code_meta.non_code_nodes.len(), 1);
        assert_eq!(
@ -1594,7 +1595,7 @@ const mySk1 = startSketchAt([0, 0])"#;
                style: CommentStyle::Line,
            }
        );
-        assert_eq!(body.len(), 4);
+        assert_eq!(tail.len(), 3);
    }

    #[test]
--- a/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__a.snap
+++ b/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__a.snap
@ -25,48 +25,48 @@ expression: actual
          "init": {
            "type": "PipeExpression",
            "type": "PipeExpression",
-            "start": 18,
+            "start": 47,
            "end": 143,
-            "body": [
-              {
-                "type": "CallExpression",
-                "type": "CallExpression",
+            "head": {
+              "type": "CallExpression",
+              "type": "CallExpression",
+              "start": 18,
+              "end": 39,
+              "callee": {
+                "type": "Identifier",
                "start": 18,
-                "end": 39,
-                "callee": {
-                  "type": "Identifier",
-                  "start": 18,
-                  "end": 31,
-                  "name": "startSketchAt"
-                },
-                "arguments": [
-                  {
-                    "type": "ArrayExpression",
-                    "type": "ArrayExpression",
-                    "start": 32,
-                    "end": 38,
-                    "elements": [
-                      {
-                        "type": "Literal",
-                        "type": "Literal",
-                        "start": 33,
-                        "end": 34,
-                        "value": 0,
-                        "raw": "0"
-                      },
-                      {
-                        "type": "Literal",
-                        "type": "Literal",
-                        "start": 36,
-                        "end": 37,
-                        "value": 0,
-                        "raw": "0"
-                      }
-                    ]
-                  }
-                ],
-                "optional": false
+                "end": 31,
+                "name": "startSketchAt"
              },
+              "arguments": [
+                {
+                  "type": "ArrayExpression",
+                  "type": "ArrayExpression",
+                  "start": 32,
+                  "end": 38,
+                  "elements": [
+                    {
+                      "type": "Literal",
+                      "type": "Literal",
+                      "start": 33,
+                      "end": 34,
+                      "value": 0,
+                      "raw": "0"
+                    },
+                    {
+                      "type": "Literal",
+                      "type": "Literal",
+                      "start": 36,
+                      "end": 37,
+                      "value": 0,
+                      "raw": "0"
+                    }
+                  ]
+                }
+              ],
+              "optional": false
+            },
+            "tail": [
              {
                "type": "CallExpression",
                "type": "CallExpression",
--- a/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__af.snap
+++ b/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__af.snap
@ -25,48 +25,48 @@ expression: actual
          "init": {
            "type": "PipeExpression",
            "type": "PipeExpression",
-            "start": 17,
+            "start": 49,
            "end": 167,
-            "body": [
-              {
-                "type": "CallExpression",
-                "type": "CallExpression",
+            "head": {
+              "type": "CallExpression",
+              "type": "CallExpression",
+              "start": 17,
+              "end": 37,
+              "callee": {
+                "type": "Identifier",
                "start": 17,
-                "end": 37,
-                "callee": {
-                  "type": "Identifier",
-                  "start": 17,
-                  "end": 30,
-                  "name": "startSketchAt"
-                },
-                "arguments": [
-                  {
-                    "type": "ArrayExpression",
-                    "type": "ArrayExpression",
-                    "start": 31,
-                    "end": 36,
-                    "elements": [
-                      {
-                        "type": "Literal",
-                        "type": "Literal",
-                        "start": 32,
-                        "end": 33,
-                        "value": 0,
-                        "raw": "0"
-                      },
-                      {
-                        "type": "Literal",
-                        "type": "Literal",
-                        "start": 34,
-                        "end": 35,
-                        "value": 0,
-                        "raw": "0"
-                      }
-                    ]
-                  }
-                ],
-                "optional": false
+                "end": 30,
+                "name": "startSketchAt"
              },
+              "arguments": [
+                {
+                  "type": "ArrayExpression",
+                  "type": "ArrayExpression",
+                  "start": 31,
+                  "end": 36,
+                  "elements": [
+                    {
+                      "type": "Literal",
+                      "type": "Literal",
+                      "start": 32,
+                      "end": 33,
+                      "value": 0,
+                      "raw": "0"
+                    },
+                    {
+                      "type": "Literal",
+                      "type": "Literal",
+                      "start": 34,
+                      "end": 35,
+                      "value": 0,
+                      "raw": "0"
+                    }
+                  ]
+                }
+              ],
+              "optional": false
+            },
+            "tail": [
              {
                "type": "CallExpression",
                "type": "CallExpression",
--- a/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__ag.snap
+++ b/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__ag.snap
@ -25,48 +25,48 @@ expression: actual
          "init": {
            "type": "PipeExpression",
            "type": "PipeExpression",
-            "start": 17,
+            "start": 41,
            "end": 70,
-            "body": [
-              {
-                "type": "CallExpression",
-                "type": "CallExpression",
+            "head": {
+              "type": "CallExpression",
+              "type": "CallExpression",
+              "start": 17,
+              "end": 37,
+              "callee": {
+                "type": "Identifier",
                "start": 17,
-                "end": 37,
-                "callee": {
-                  "type": "Identifier",
-                  "start": 17,
-                  "end": 30,
-                  "name": "startSketchAt"
-                },
-                "arguments": [
-                  {
-                    "type": "ArrayExpression",
-                    "type": "ArrayExpression",
-                    "start": 31,
-                    "end": 36,
-                    "elements": [
-                      {
-                        "type": "Literal",
-                        "type": "Literal",
-                        "start": 32,
-                        "end": 33,
-                        "value": 0,
-                        "raw": "0"
-                      },
-                      {
-                        "type": "Literal",
-                        "type": "Literal",
-                        "start": 34,
-                        "end": 35,
-                        "value": 0,
-                        "raw": "0"
-                      }
-                    ]
-                  }
-                ],
-                "optional": false
+                "end": 30,
+                "name": "startSketchAt"
              },
+              "arguments": [
+                {
+                  "type": "ArrayExpression",
+                  "type": "ArrayExpression",
+                  "start": 31,
+                  "end": 36,
+                  "elements": [
+                    {
+                      "type": "Literal",
+                      "type": "Literal",
+                      "start": 32,
+                      "end": 33,
+                      "value": 0,
+                      "raw": "0"
+                    },
+                    {
+                      "type": "Literal",
+                      "type": "Literal",
+                      "start": 34,
+                      "end": 35,
+                      "value": 0,
+                      "raw": "0"
+                    }
+                  ]
+                }
+              ],
+              "optional": false
+            },
+            "tail": [
              {
                "type": "CallExpression",
                "type": "CallExpression",
--- a/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__ai.snap
+++ b/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__ai.snap
@ -25,32 +25,32 @@ expression: actual
          "init": {
            "type": "PipeExpression",
            "type": "PipeExpression",
-            "start": 14,
+            "start": 22,
            "end": 29,
-            "body": [
-              {
-                "type": "CallExpression",
-                "type": "CallExpression",
+            "head": {
+              "type": "CallExpression",
+              "type": "CallExpression",
+              "start": 14,
+              "end": 18,
+              "callee": {
+                "type": "Identifier",
                "start": 14,
-                "end": 18,
-                "callee": {
-                  "type": "Identifier",
-                  "start": 14,
-                  "end": 15,
-                  "name": "f"
-                },
-                "arguments": [
-                  {
-                    "type": "Literal",
-                    "type": "Literal",
-                    "start": 16,
-                    "end": 17,
-                    "value": 1,
-                    "raw": "1"
-                  }
-                ],
-                "optional": false
+                "end": 15,
+                "name": "f"
              },
+              "arguments": [
+                {
+                  "type": "Literal",
+                  "type": "Literal",
+                  "start": 16,
+                  "end": 17,
+                  "value": 1,
+                  "raw": "1"
+                }
+              ],
+              "optional": false
+            },
+            "tail": [
              {
                "type": "CallExpression",
                "type": "CallExpression",
--- a/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__aj.snap
+++ b/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__aj.snap
@ -25,31 +25,31 @@ expression: actual
          "init": {
            "type": "PipeExpression",
            "type": "PipeExpression",
-            "start": 14,
+            "start": 34,
            "end": 49,
-            "body": [
-              {
-                "type": "CallExpression",
-                "type": "CallExpression",
+            "head": {
+              "type": "CallExpression",
+              "type": "CallExpression",
+              "start": 14,
+              "end": 30,
+              "callee": {
+                "type": "Identifier",
                "start": 14,
-                "end": 30,
-                "callee": {
-                  "type": "Identifier",
-                  "start": 14,
-                  "end": 27,
-                  "name": "startSketchAt"
-                },
-                "arguments": [
-                  {
-                    "type": "Identifier",
-                    "type": "Identifier",
-                    "start": 28,
-                    "end": 29,
-                    "name": "p"
-                  }
-                ],
-                "optional": false
+                "end": 27,
+                "name": "startSketchAt"
              },
+              "arguments": [
+                {
+                  "type": "Identifier",
+                  "type": "Identifier",
+                  "start": 28,
+                  "end": 29,
+                  "name": "p"
+                }
+              ],
+              "optional": false
+            },
+            "tail": [
              {
                "type": "CallExpression",
                "type": "CallExpression",
--- a/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__au.snap
+++ b/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__au.snap
@ -25,32 +25,32 @@ expression: actual
          "init": {
            "type": "PipeExpression",
            "type": "PipeExpression",
-            "start": 17,
+            "start": 44,
            "end": 86,
-            "body": [
-              {
-                "type": "CallExpression",
-                "type": "CallExpression",
+            "head": {
+              "type": "CallExpression",
+              "type": "CallExpression",
+              "start": 17,
+              "end": 36,
+              "callee": {
+                "type": "Identifier",
                "start": 17,
-                "end": 36,
-                "callee": {
-                  "type": "Identifier",
-                  "start": 17,
-                  "end": 30,
-                  "name": "startSketchOn"
-                },
-                "arguments": [
-                  {
-                    "type": "Literal",
-                    "type": "Literal",
-                    "start": 31,
-                    "end": 35,
-                    "value": "XY",
-                    "raw": "'XY'"
-                  }
-                ],
-                "optional": false
+                "end": 30,
+                "name": "startSketchOn"
              },
+              "arguments": [
+                {
+                  "type": "Literal",
+                  "type": "Literal",
+                  "start": 31,
+                  "end": 35,
+                  "value": "XY",
+                  "raw": "'XY'"
+                }
+              ],
+              "optional": false
+            },
+            "tail": [
              {
                "type": "CallExpression",
                "type": "CallExpression",
--- a/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__d.snap
+++ b/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__d.snap
@ -25,32 +25,32 @@ expression: actual
          "init": {
            "type": "PipeExpression",
            "type": "PipeExpression",
-            "start": 14,
+            "start": 23,
            "end": 36,
-            "body": [
-              {
-                "type": "BinaryExpression",
-                "type": "BinaryExpression",
+            "head": {
+              "type": "BinaryExpression",
+              "type": "BinaryExpression",
+              "start": 14,
+              "end": 19,
+              "operator": "+",
+              "left": {
+                "type": "Literal",
+                "type": "Literal",
                "start": 14,
-                "end": 19,
-                "operator": "+",
-                "left": {
-                  "type": "Literal",
-                  "type": "Literal",
-                  "start": 14,
-                  "end": 15,
-                  "value": 5,
-                  "raw": "5"
-                },
-                "right": {
-                  "type": "Literal",
-                  "type": "Literal",
-                  "start": 18,
-                  "end": 19,
-                  "value": 6,
-                  "raw": "6"
-                }
+                "end": 15,
+                "value": 5,
+                "raw": "5"
              },
+              "right": {
+                "type": "Literal",
+                "type": "Literal",
+                "start": 18,
+                "end": 19,
+                "value": 6,
+                "raw": "6"
+              }
+            },
+            "tail": [
              {
                "type": "CallExpression",
                "type": "CallExpression",
--- a/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__z.snap
+++ b/src/wasm-lib/kcl/src/parser/snapshots/kcl_libparserparser_impl__snapshot_tests__z.snap
@ -25,31 +25,31 @@ expression: actual
          "init": {
            "type": "PipeExpression",
            "type": "PipeExpression",
-            "start": 11,
+            "start": 33,
            "end": 53,
-            "body": [
-              {
-                "type": "CallExpression",
-                "type": "CallExpression",
+            "head": {
+              "type": "CallExpression",
+              "type": "CallExpression",
+              "start": 11,
+              "end": 29,
+              "callee": {
+                "type": "Identifier",
                "start": 11,
-                "end": 29,
-                "callee": {
-                  "type": "Identifier",
-                  "start": 11,
-                  "end": 24,
-                  "name": "startSketchAt"
-                },
-                "arguments": [
-                  {
-                    "type": "Identifier",
-                    "type": "Identifier",
-                    "start": 25,
-                    "end": 28,
-                    "name": "pos"
-                  }
-                ],
-                "optional": false
+                "end": 24,
+                "name": "startSketchAt"
              },
+              "arguments": [
+                {
+                  "type": "Identifier",
+                  "type": "Identifier",
+                  "start": 25,
+                  "end": 28,
+                  "name": "pos"
+                }
+              ],
+              "optional": false
+            },
+            "tail": [
              {
                "type": "CallExpression",
                "type": "CallExpression",