Add short lambdas

GRAMMAR.txt

@@ -106,7 +106,7 @@ null                 : 'null'
 
 list                 : '[' (arguments ','?)? ']'
 
-function_expression  : 'fn' '(' (parameters ','?)? ')' block
+function_expression  : 'fn' '(' (parameters ','?)? ')' (block | '=>' expression)
 
 call                 : '(' (arguments ','?)? ')'
 

README.md

@@ -299,6 +299,18 @@ println(apply(2, fn(x) {
 }))
 ```
 
+For a single expression, use `=>` to return its value:
+
+```rust
+square = fn(x) => x ^ 2
+println(square(5))
+println(apply(2, fn(x) => x * 3))
+```
+
+The body includes the full expression after `=>`, so `fn(x) => x + 1`
+returns `x + 1`. Parenthesize a short lambda to call it immediately:
+`(fn(x) => x ^ 2)(5)`.
+
 #### Null
 
 Represents the absence of a value.

src/highlighter.rs

@@ -207,7 +207,7 @@ impl<'src> Highlighter<'src> {
   }
 
   fn scan_operator(&self, start: usize) -> Option<usize> {
-    for operator in [">=", "<=", "==", "!=", "&&", "||"] {
+    for operator in [">=", "<=", "==", "!=", "&&", "||", "=>"] {
       if self.content[start..].starts_with(operator) {
         return Some(start + operator.len());
       }
@@ -332,6 +332,21 @@ mod tests {
     case("1.5e2");
   }
 
+  #[test]
+  fn short_lambda() {
+    assert_eq!(
+      Highlighter::new("fn(foo) => foo").collect_highlight_spans(),
+      [
+        HighlightSpan::new(0, 2, HighlightKind::Keyword),
+        HighlightSpan::new(2, 3, HighlightKind::Operator),
+        HighlightSpan::new(3, 6, HighlightKind::Identifier),
+        HighlightSpan::new(6, 7, HighlightKind::Operator),
+        HighlightSpan::new(8, 10, HighlightKind::Operator),
+        HighlightSpan::new(11, 14, HighlightKind::Identifier),
+      ]
+    );
+  }
+
   #[test]
   fn string_contents_are_not_highlighted_as_tokens() {
     let highlighter = Highlighter::new("\"if\" + 'else'");

src/parser.rs

@@ -294,7 +294,16 @@ where
         comma_separated_parser(identifier_parser())
           .delimited_by(padded_parser(just('(')), padded_parser(just(')'))),
       )
-      .then(statement_block.clone())
+      .then(
+        statement_block.clone().or(
+          padded_parser(just("=>"))
+            .ignore_then(expression.clone())
+            .map(|expression: Spanned<Expression>| {
+              let span = expression.1;
+              vec![(Statement::Return(Some(expression)), span)]
+            }),
+        ),
+      )
       .map(|(params, body)| Expression::Function(params, body))
       .map_with(|ast, error| (ast, error.span()));
 
@@ -807,6 +816,66 @@ mod tests {
     case("1e0", "1");
   }
 
+  #[test]
+  fn short_lambda() {
+    #[track_caller]
+    fn case(program: &str, ast: &str) {
+      Test::new().program(program).ast(ast).run();
+    }
+
+    case(
+      "fn(foo) => foo ^ 2 + 1",
+      "statements(expression(function([foo], block(return(binary_op(+, binary_op(^, identifier(foo), number(2)), number(1)))))))",
+    );
+    case(
+      "[fn() => 1, fn(foo,) => foo]",
+      "statements(expression(list(function([], block(return(number(1)))), function([foo], block(return(identifier(foo)))))))",
+    );
+    case(
+      "fn(foo) => fn(bar) => foo + bar",
+      "statements(expression(function([foo], block(return(function([bar], block(return(binary_op(+, identifier(foo), identifier(bar))))))))))",
+    );
+    case(
+      "(fn(foo) => foo)(1)",
+      "statements(expression(function_call(function([foo], block(return(identifier(foo)))), number(1))))",
+    );
+    case(
+      "foo = fn() => 1; bar = 2",
+      "statements(assignment(identifier(foo), function([], block(return(number(1))))), assignment(identifier(bar), number(2)))",
+    );
+  }
+
+  #[test]
+  fn short_lambda_body_is_required() {
+    #[track_caller]
+    fn case(program: &str) {
+      assert!(parse(program).is_err(), "{program}");
+    }
+
+    case("fn() =>");
+    case("fn() => ;");
+    case("fn() => return 1");
+    case("fn() => {}");
+  }
+
+  #[test]
+  fn short_lambda_spans() {
+    let program = parse("fn(foo) => foo + 1").unwrap();
+    let Program::Statements(statements) = program.0;
+    let Statement::Expression((Expression::Function(_, body), span)) =
+      &statements[0].0
+    else {
+      panic!("expected function expression");
+    };
+    let (Statement::Return(Some(expression)), body_span) = &body[0] else {
+      panic!("expected return statement");
+    };
+
+    assert_eq!(*span, SimpleSpan::from(0..18));
+    assert_eq!(*body_span, SimpleSpan::from(11..18));
+    assert_eq!(expression.1, *body_span);
+  }
+
   #[test]
   fn unclosed_string() {
     Test::new()

tests/integration.rs

@@ -2864,6 +2864,110 @@ fn secant() -> Result {
     .run()
 }
 
+#[test]
+fn short_lambda_as_argument() -> Result {
+  Test::new()?
+    .program(indoc! {
+      "
+      fn foo(bar, baz) {
+        baz(bar)
+      }
+
+      println(foo(2, fn(bar) => bar * 3))
+      "
+    })
+    .expected_stdout(Exact("6\n"))
+    .run()
+}
+
+#[test]
+fn short_lambda_assignment() -> Result {
+  Test::new()?
+    .program(indoc! {
+      "
+      foo = fn(bar) => bar ^ 2
+      println(foo(5))
+      "
+    })
+    .expected_stdout(Exact("25\n"))
+    .run()
+}
+
+#[test]
+fn short_lambda_body_calls_function() -> Result {
+  Test::new()?
+    .program(indoc! {
+      "
+      foo = fn() => println('bar')
+      foo()
+      "
+    })
+    .expected_stdout(Exact("bar\n"))
+    .run()
+}
+
+#[test]
+fn short_lambda_immediate_call() -> Result {
+  Test::new()?
+    .program("println((fn(foo, bar) => foo + bar)(2, 3))")
+    .expected_stdout(Exact("5\n"))
+    .run()
+}
+
+#[test]
+fn short_lambda_in_list() -> Result {
+  Test::new()?
+    .program(indoc! {
+      "
+      foo = [fn(bar) => [bar, bar + 1]]
+      println(foo[0](2)[1])
+      "
+    })
+    .expected_stdout(Exact("3\n"))
+    .run()
+}
+
+#[test]
+fn short_lambda_observes_outer_scope_changes() -> Result {
+  Test::new()?
+    .program(indoc! {
+      "
+      foo = 2
+      bar = fn() => foo
+      foo = 3
+      println(bar())
+      "
+    })
+    .expected_stdout(Exact("3\n"))
+    .run()
+}
+
+#[test]
+fn short_lambda_returns_closure() -> Result {
+  Test::new()?
+    .program(indoc! {
+      "
+      foo = fn(bar) => fn(baz) => bar + baz
+      println(foo(2)(3))
+      "
+    })
+    .expected_stdout(Exact("5\n"))
+    .run()
+}
+
+#[test]
+fn short_lambda_returns_null() -> Result {
+  Test::new()?
+    .program(indoc! {
+      "
+      foo = fn() => null
+      println(foo())
+      "
+    })
+    .expected_stdout(Exact("null\n"))
+    .run()
+}
+
 #[test]
 fn simple_break() -> Result {
   Test::new()?