feat: parse function literals

viddrobnic · May 29, 2024 · 9c5115a · 9c5115a
1 parent 9a406fa
commit 9c5115a
Show file tree

Hide file tree

Showing 3 changed files with 158 additions and 2 deletions.
diff --git a/parser/src/error.rs b/parser/src/error.rs
@@ -16,6 +16,7 @@ pub enum ErrorKind {
     InvalidTokenKind { expected: TokenKind, got: TokenKind },
     InvalidAssignee,
     InvalidRange,
+    InvalidFunctionParameter,
 }
 
 #[derive(Debug, Error, PartialEq)]
@@ -49,6 +50,9 @@ impl Display for ErrorKind {
             }
             ErrorKind::InvalidAssignee => write!(f, "assignee must be identifier, index or array"),
             ErrorKind::InvalidRange => write!(f, "invalid range defined"),
+            ErrorKind::InvalidFunctionParameter => {
+                write!(f, "function parameter must be an identifier")
+            }
         }
     }
 }
diff --git a/parser/src/parser/mod.rs b/parser/src/parser/mod.rs
@@ -201,7 +201,7 @@ impl Parser<'_> {
             TokenKind::Break => (ast::NodeValue::Break, range.end),
             TokenKind::Continue => (ast::NodeValue::Continue, range.end),
             TokenKind::Return => todo!("parse return statement"),
-            TokenKind::Fn => todo!("parse function literal"),
+            TokenKind::Fn => self.parse_fn_literal()?,
             TokenKind::Use => {
                 let token = self.next_token()?;
                 let TokenKind::String(val) = token.kind else {
@@ -351,10 +351,16 @@ impl Parser<'_> {
     fn parse_assign(&mut self, left: ast::Node) -> Result<(ast::NodeValue, Position)> {
         let token = self.next_token()?;
 
-        let right = self.parse_node(token, Precedence::Lowest)?;
+        let mut right = self.parse_node(token, Precedence::Lowest)?;
         validate_node_kind(&right, NodeKind::Expression)?;
         validate_assignee(&left)?;
 
+        if let (ast::NodeValue::Identifier(ident), ast::NodeValue::FunctionLiteral { name, .. }) =
+            (&left.value, &mut right.value)
+        {
+            *name = Some(ident.clone());
+        }
+
         let end = right.range.end;
         Ok((
             ast::NodeValue::Assign {
@@ -536,6 +542,37 @@ impl Parser<'_> {
         ))
     }
 
+    fn parse_fn_literal(&mut self) -> Result<(ast::NodeValue, Position)> {
+        // Read `(`
+        let token = self.next_token()?;
+        validate_token_kind(&token, TokenKind::LBracket)?;
+
+        // Read arguments
+        let (args, _) = self.parse_multiple(
+            TokenKind::RBracket,
+            TokenKind::Comma,
+            |_, token| match token.kind {
+                TokenKind::Ident(ident) => Ok(ident),
+                _ => Err(Error {
+                    kind: ErrorKind::InvalidFunctionParameter,
+                    range: token.range,
+                }),
+            },
+        )?;
+
+        let body_token = self.next_token()?;
+        let (body, end) = self.parse_block(body_token)?;
+
+        Ok((
+            ast::NodeValue::FunctionLiteral {
+                name: None,
+                parameters: args,
+                body,
+            },
+            end,
+        ))
+    }
+
     // Helper function that reads block { ... }.
     // It returns vector of nodes and end position, which is the end
     // position of `}`

diff --git a/parser/src/parser/test.rs b/parser/src/parser/test.rs
@@ -1037,6 +1037,101 @@ fn for_loop() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn fn_literal() -> Result<()> {
+    let tests = [
+        (
+            "fn(){}",
+            ast::Node {
+                value: ast::NodeValue::FunctionLiteral {
+                    name: None,
+                    parameters: vec![],
+                    body: vec![],
+                },
+                range: Range {
+                    start: Position::new(0, 0),
+                    end: Position::new(0, 6),
+                },
+            },
+        ),
+        (
+            "fn(a){}",
+            ast::Node {
+                value: ast::NodeValue::FunctionLiteral {
+                    name: None,
+                    parameters: vec!["a".to_string()],
+                    body: vec![],
+                },
+                range: Range {
+                    start: Position::new(0, 0),
+                    end: Position::new(0, 7),
+                },
+            },
+        ),
+        (
+            "fn(a, b){\n}",
+            ast::Node {
+                value: ast::NodeValue::FunctionLiteral {
+                    name: None,
+                    parameters: vec!["a".to_string(), "b".to_string()],
+                    body: vec![],
+                },
+                range: Range {
+                    start: Position::new(0, 0),
+                    end: Position::new(1, 1),
+                },
+            },
+        ),
+    ];
+
+    for (input, expected) in tests {
+        let program = parse(input)?;
+
+        assert_eq!(program.statements.len(), 1);
+        assert_eq!(program.statements[0], expected);
+    }
+
+    Ok(())
+}
+
+#[test]
+fn fn_literal_named() -> Result<()> {
+    let program = parse("foo = fn(){}")?;
+
+    assert_eq!(program.statements.len(), 1);
+    assert_eq!(
+        program.statements[0],
+        ast::Node {
+            value: ast::NodeValue::Assign {
+                ident: Box::new(ast::Node {
+                    value: ast::NodeValue::Identifier("foo".to_string()),
+                    range: Range {
+                        start: Position::new(0, 0),
+                        end: Position::new(0, 3),
+                    }
+                }),
+                value: Box::new(ast::Node {
+                    value: ast::NodeValue::FunctionLiteral {
+                        name: Some("foo".to_string()),
+                        parameters: vec![],
+                        body: vec![]
+                    },
+                    range: Range {
+                        start: Position::new(0, 6),
+                        end: Position::new(0, 12),
+                    }
+                })
+            },
+            range: Range {
+                start: Position::new(0, 0),
+                end: Position::new(0, 12),
+            }
+        }
+    );
+
+    Ok(())
+}
+
 #[test]
 fn errors() {
     let tests = [
@@ -1060,6 +1155,26 @@ fn errors() {
                 },
             },
         ),
+        (
+            "for (true) {}",
+            Error {
+                kind: ErrorKind::InvalidRange,
+                range: Range {
+                    start: Position::new(0, 4),
+                    end: Position::new(0, 10),
+                },
+            },
+        ),
+        (
+            "fn(1 + 1){}",
+            Error {
+                kind: ErrorKind::InvalidFunctionParameter,
+                range: Range {
+                    start: Position::new(0, 3),
+                    end: Position::new(0, 4),
+                },
+            },
+        ),
     ];
 
     for (input, expected) in tests {