feat: function scope alias (#426)

* feat: function scope alias

* fix types

core/expression/src/compiler/compiler.rs

@@ -38,11 +38,28 @@ impl Compiler {
 struct CompilerInner<'arena, 'bytecode_ref> {
     root: &'arena Node<'arena>,
     bytecode: &'bytecode_ref mut Vec<Opcode>,
+    closure_aliases: Vec<Option<&'arena str>>,
 }
 
 impl<'arena, 'bytecode_ref> CompilerInner<'arena, 'bytecode_ref> {
     pub fn new(bytecode: &'bytecode_ref mut Vec<Opcode>, root: &'arena Node<'arena>) -> Self {
-        Self { root, bytecode }
+        Self {
+            root,
+            bytecode,
+            closure_aliases: Vec::new(),
+        }
+    }
+
+    fn lookup_alias(&self, name: &str) -> Option<u32> {
+        self.closure_aliases
+            .iter()
+            .enumerate()
+            .rev()
+            .find_map(|(index, &alias)| {
+                alias.and_then(|a| {
+                    (a == name).then_some((self.closure_aliases.len() - 1 - index) as u32)
+                })
+            })
     }
 
     pub fn compile(&mut self) -> CompilerResult<()> {
@@ -116,10 +133,12 @@ impl<'arena, 'bytecode_ref> CompilerInner<'arena, 'bytecode_ref> {
     fn compile_member_fast(&mut self, node: &'arena Node<'arena>) -> Option<Vec<FetchFastTarget>> {
         match node {
             Node::Root => Some(vec![FetchFastTarget::Root]),
-            Node::Identifier(v) => Some(vec![
-                FetchFastTarget::Begin,
-                FetchFastTarget::String(Arc::from(*v)),
-            ]),
+            Node::Identifier(v) => self.lookup_alias(v).is_none().then(|| {
+                vec![
+                    FetchFastTarget::Begin,
+                    FetchFastTarget::String(Arc::from(*v)),
+                ]
+            }),
             Node::Member { node, property } => {
                 let mut path = self.compile_member_fast(node)?;
                 match property {
@@ -149,7 +168,7 @@ impl<'arena, 'bytecode_ref> CompilerInner<'arena, 'bytecode_ref> {
             Node::Bool(v) => Ok(self.emit(Opcode::PushBool(*v))),
             Node::Number(v) => Ok(self.emit(Opcode::PushNumber(*v))),
             Node::String(v) => Ok(self.emit(Opcode::PushString(Arc::from(*v)))),
-            Node::Pointer => Ok(self.emit(Opcode::Pointer)),
+            Node::Pointer => Ok(self.emit(Opcode::Pointer(0))),
             Node::Root => Ok(self.emit(Opcode::FetchRootEnv)),
             Node::Array(v) => {
                 v.iter()
@@ -189,8 +208,16 @@ impl<'arena, 'bytecode_ref> CompilerInner<'arena, 'bytecode_ref> {
                     with_return: output.is_some(),
                 }))
             }
-            Node::Identifier(v) => Ok(self.emit(Opcode::FetchEnv(Arc::from(*v)))),
-            Node::Closure(v) => self.compile_node(v),
+            Node::Identifier(v) => Ok(self.emit(
+                self.lookup_alias(v)
+                    .map_or_else(|| Opcode::FetchEnv(Arc::from(*v)), Opcode::Pointer),
+            )),
+            Node::Closure { body, alias } => {
+                self.closure_aliases.push(*alias);
+                let result = self.compile_node(body);
+                self.closure_aliases.pop();
+                result
+            }
             Node::Parenthesized(v) => self.compile_node(v),
             Node::Member {
                 node: n,
@@ -496,7 +523,7 @@ impl<'arena, 'bytecode_ref> CompilerInner<'arena, 'bytecode_ref> {
                             c.compile_argument(kind, arguments, 1)?;
                             c.emit_cond(|c| {
                                 c.emit(Opcode::IncrementCount);
-                                c.emit(Opcode::Pointer);
+                                c.emit(Opcode::Pointer(0));
                             });
                             Ok(())
                         })?;

core/expression/src/compiler/opcode.rs

@@ -51,7 +51,8 @@ pub enum Opcode {
     IncrementCount,
     GetCount,
     GetLen,
-    Pointer,
+    /// The u32 is the depth from innermost scope (0 = innermost, 1 = parent, etc.)
+    Pointer(u32),
     Begin,
     End,
     CallFunction {

core/expression/src/intellisense/mod.rs

@@ -52,6 +52,7 @@ impl<'arena> IntelliSense<'arena> {
                 pointer_data: data.shallow_clone(),
                 root_data: data.shallow_clone(),
                 current_data: data.shallow_clone(),
+                ..Default::default()
             },
         );
 
@@ -98,6 +99,7 @@ impl<'arena> IntelliSense<'arena> {
                 pointer_data: data.shallow_clone(),
                 root_data: data.shallow_clone(),
                 current_data: data.shallow_clone(),
+                ..Default::default()
             },
         );
 

core/expression/src/intellisense/scope.rs

@@ -1,8 +1,17 @@
 use crate::variable::VariableType;
+use ahash::HashMap;
+use std::rc::Rc;
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Default)]
 pub struct IntelliSenseScope {
     pub root_data: VariableType,
     pub current_data: VariableType,
     pub pointer_data: VariableType,
+    pub aliases: HashMap<Rc<str>, VariableType>,
+}
+
+impl IntelliSenseScope {
+    pub fn get_alias(&self, name: &str) -> Option<&VariableType> {
+        self.aliases.get(name)
+    }
 }

core/expression/src/intellisense/types/provider.rs

@@ -165,7 +165,10 @@ impl TypesProvider {
                 }
             }
 
-            Node::Identifier(i) => TypeInfo::from(scope.root_data.get(i)),
+            Node::Identifier(i) => scope
+                .get_alias(i)
+                .map(|t| TypeInfo::from(t.clone()))
+                .unwrap_or_else(|| TypeInfo::from(scope.root_data.get(i))),
             Node::Member { node, property } => {
                 let node_type = self.determine(node, scope.clone());
                 let property_type = self.determine(property, scope.clone());
@@ -389,15 +392,27 @@ impl TypesProvider {
 
                 if let FunctionKind::Closure(_) = kind {
                     let ptr_type = type_list[0].iterator().unwrap_or_default();
-                    let new_type = self.determine(
-                        arguments[1],
-                        IntelliSenseScope {
-                            pointer_data: ptr_type.deref().clone(),
-                            current_data: scope.current_data,
-                            root_data: scope.root_data,
-                        },
-                    );
+                    let ptr_type_inner = ptr_type.deref().clone();
 
+                    let alias = match arguments[1] {
+                        Node::Closure { alias, .. } => *alias,
+                        _ => None,
+                    };
+
+                    let mut closure_scope = IntelliSenseScope {
+                        pointer_data: ptr_type_inner.clone(),
+                        current_data: scope.current_data.clone(),
+                        root_data: scope.root_data.clone(),
+                        aliases: scope.aliases.clone(),
+                    };
+
+                    if let Some(alias_name) = alias {
+                        closure_scope
+                            .aliases
+                            .insert(Rc::from(alias_name), ptr_type_inner);
+                    }
+
+                    let new_type = self.determine(arguments[1], closure_scope);
                     type_list[1] = new_type.kind;
                 }
 
@@ -493,7 +508,7 @@ impl TypesProvider {
                     error: typecheck.general,
                 }
             }
-            Node::Closure(c) => self.determine(c, scope.clone()),
+            Node::Closure { body, .. } => self.determine(body, scope.clone()),
             Node::Parenthesized(c) => self.determine(c, scope.clone()),
             Node::Error { node, error } => match node {
                 None => TypeInfo {

core/expression/src/parser/ast.rs

@@ -21,7 +21,10 @@ pub enum Node<'a> {
         output: Option<&'a Node<'a>>,
     },
     Identifier(&'a str),
-    Closure(&'a Node<'a>),
+    Closure {
+        body: &'a Node<'a>,
+        alias: Option<&'a str>,
+    },
     Parenthesized(&'a Node<'a>),
     Root,
     Member {
@@ -106,7 +109,7 @@ impl<'a> Node<'a> {
                     output.walk(func.clone());
                 }
             }
-            Node::Closure(closure) => closure.walk(func.clone()),
+            Node::Closure { body, .. } => body.walk(func.clone()),
             Node::Parenthesized(c) => c.walk(func.clone()),
             Node::Member { node, property } => {
                 node.walk(func.clone());

core/expression/src/parser/parser.rs

@@ -638,7 +638,11 @@ impl<'arena, 'token_ref, Flavor> Parser<'arena, 'token_ref, Flavor> {
     }
 
     /// Closure
-    pub(crate) fn closure<F>(&self, expression_parser: &F) -> &'arena Node<'arena>
+    pub(crate) fn closure<F>(
+        &self,
+        expression_parser: &F,
+        alias: Option<&'arena str>,
+    ) -> &'arena Node<'arena>
     where
         F: Fn(ParserContext) -> &'arena Node<'arena>,
     {
@@ -648,7 +652,7 @@ impl<'arena, 'token_ref, Flavor> Parser<'arena, 'token_ref, Flavor> {
         let node = expression_parser(ParserContext::Closure);
         self.depth.set(self.depth.get() - 1);
 
-        self.node(Node::Closure(node), |_| NodeMetadata {
+        self.node(Node::Closure { body: node, alias }, |_| NodeMetadata {
             span: (start, self.prev_token_end()),
         })
     }
@@ -734,11 +738,42 @@ impl<'arena, 'token_ref, Flavor> Parser<'arena, 'token_ref, Flavor> {
                 let mut arguments = BumpVec::new_in(&self.bump);
 
                 arguments.push(expression_parser(ParserContext::Global));
+
+                let alias: Option<&'arena str> =
+                    if self
+                        .current()
+                        .is_some_and(|t| t.kind == TokenKind::Literal && t.value == "as")
+                    {
+                        self.next();
+
+                        let alias_token = self.current();
+                        match alias_token {
+                            Some(t) if t.kind == TokenKind::Literal => {
+                                let alias_str = self.bump.alloc_str(t.value);
+                                self.next();
+                                Some(alias_str)
+                            }
+                            _ => {
+                                arguments.push(self.error(AstNodeError::Custom {
+                                    message: afmt!(self, "Expected identifier after 'as'"),
+                                    span:
+                                        alias_token.map(|t| t.span).unwrap_or((
+                                            self.prev_token_end(),
+                                            self.prev_token_end(),
+                                        )),
+                                }));
+                                None
+                            }
+                        }
+                    } else {
+                        None
+                    };
+
                 if let Some(error) = self.expect(TokenKind::Operator(Operator::Comma)) {
                     arguments.push(error);
                 };
 
-                arguments.push(self.closure(&expression_parser));
+                arguments.push(self.closure(&expression_parser, alias));
                 if let Some(error) = self.expect(TokenKind::Bracket(Bracket::RightParenthesis)) {
                     arguments.push(error);
                 }

core/expression/src/parser/unary.rs

@@ -299,7 +299,7 @@ impl From<&Node<'_>> for UnaryNodeBehaviour {
             Node::Pointer => AsBoolean,
             Node::Array(_) => CompareWithReference(In),
             Node::Identifier(_) => CompareWithReference(Equal),
-            Node::Closure(_) => AsBoolean,
+            Node::Closure { .. } => AsBoolean,
             Node::Member { .. } => CompareWithReference(Equal),
             Node::Slice { .. } => CompareWithReference(In),
             Node::Interval { .. } => CompareWithReference(In),

core/expression/src/vm/vm.rs

@@ -845,8 +845,17 @@ impl<'arena, 'parent_ref, 'bytecode_ref> VMInner<'parent_ref, 'bytecode_ref> {
 
                     self.push(Number(scope.len.into()));
                 }
-                Opcode::Pointer => {
-                    let scope = self.scopes.last().ok_or_else(|| OpcodeErr {
+                Opcode::Pointer(depth) => {
+                    let scope_index = self
+                        .scopes
+                        .len()
+                        .checked_sub(1 + *depth as usize)
+                        .ok_or_else(|| OpcodeErr {
+                            opcode: "Pointer".into(),
+                            message: format!("Scope depth {} out of bounds", depth),
+                        })?;
+
+                    let scope = self.scopes.get(scope_index).ok_or_else(|| OpcodeErr {
                         opcode: "Pointer".into(),
                         message: "Empty scope".into(),
                     })?;

core/expression/tests/data/standard.csv

@@ -613,4 +613,27 @@ round(1.234e2);;123
 "user.name = 'Eve'; user.name";{};'Eve'
 "config.debug = true; config.env = 'dev'; config";{};{"debug": true, "env": "dev"}
 "config.debug = true; config.env = 'dev'; $root";{};{"config": {"debug": true, "env": "dev"}}
-"a = 5; b = 10; a + b";{};15
+"a = 5; b = 10; a + b";{};15
+
+# Closure 'as' alias syntax
+map(products as p, p.price);{"products": [{"name": "A", "price": 10}, {"name": "B", "price": 20}]};[10, 20]
+filter(products as p, p.price > 15);{"products": [{"name": "A", "price": 10}, {"name": "B", "price": 20}]};[{"name": "B", "price": 20}]
+some(items as i, i.active);{"items": [{"active": false}, {"active": true}]};true
+all(items as i, i.value > 0);{"items": [{"value": 1}, {"value": 2}]};true
+none(items as i, i.value < 0);{"items": [{"value": 1}, {"value": 2}]};true
+count(items as i, i.value > 1);{"items": [{"value": 1}, {"value": 2}, {"value": 3}]};2
+one(items as i, i.id == 2);{"items": [{"id": 1}, {"id": 2}, {"id": 3}]};true
+flatMap(groups as g, g.items);{"groups": [{"items": [1, 2]}, {"items": [3, 4]}]};[1, 2, 3, 4]
+
+# Nested closures with 'as' aliases
+map(orders as o, map(o.items as i, i.qty));{"orders": [{"items": [{"qty": 1}, {"qty": 2}]}, {"items": [{"qty": 3}]}]};[[1, 2], [3]]
+map(orders as o, map(o.items as i, o.id + i.qty));{"orders": [{"id": 10, "items": [{"qty": 1}, {"qty": 2}]}]};[[11, 12]]
+filter(orders as o, some(o.items as i, i.qty > 2));{"orders": [{"id": 1, "items": [{"qty": 1}]}, {"id": 2, "items": [{"qty": 3}]}]};[{"id": 2, "items": [{"qty": 3}]}]
+
+# Mixed # and 'as' alias
+map(orders as o, map(o.items, #.qty));{"orders": [{"items": [{"qty": 5}, {"qty": 6}]}]};[[5, 6]]
+map(items, #.value * 2);{"items": [{"value": 1}, {"value": 2}]};[2, 4]
+
+# Backwards compatibility - # still works
+map([1, 2, 3], # * 2);;[2, 4, 6]
+filter([1, 2, 3, 4], # > 2);;[3, 4]