Luau: add support for user-defined type functions (#334)

Closes #330

CHANGELOG.md

@@ -6,6 +6,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- Luau: added support for user defined type functions (rfc: https://github.com/luau-lang/rfcs/blob/master/docs/user-defined-type-functions.md)
+
 ## [1.1.2] - 2024-11-17
 
 ### Fixed

full-moon/src/ast/luau.rs

@@ -762,6 +762,137 @@ impl ExportedTypeDeclaration {
     }
 }
 
+/// A user defined type function, such as `type function foo() ... end`.
+///
+/// See more: https://github.com/luau-lang/rfcs/blob/master/docs/user-defined-type-functions.md
+#[derive(Clone, Debug, Display, PartialEq, Node, Visit)]
+#[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
+#[display("{}{}{}{}", type_token, function_token, function_name, function_body)]
+pub struct TypeFunction {
+    pub(crate) type_token: TokenReference,
+    pub(crate) function_token: TokenReference,
+    pub(crate) function_name: TokenReference,
+    pub(crate) function_body: FunctionBody,
+}
+
+impl TypeFunction {
+    /// Creates a new TypeFunction from the given function name and body
+    pub fn new(function_name: TokenReference, function_body: FunctionBody) -> Self {
+        Self {
+            type_token: TokenReference::new(
+                Vec::new(),
+                Token::new(TokenType::Identifier {
+                    identifier: "type".into(),
+                }),
+                vec![Token::new(TokenType::spaces(1))],
+            ),
+            function_token: TokenReference::basic_symbol("function "),
+            function_name,
+            function_body,
+        }
+    }
+
+    /// The token `type`.
+    pub fn type_token(&self) -> &TokenReference {
+        &self.type_token
+    }
+
+    /// The token `function`.
+    pub fn function_token(&self) -> &TokenReference {
+        &self.function_token
+    }
+
+    /// The name of the type function, `Pairs` in `type function Pairs() ... end`.
+    pub fn function_name(&self) -> &TokenReference {
+        &self.function_name
+    }
+
+    /// The body of the type function.
+    pub fn function_body(&self) -> &FunctionBody {
+        &self.function_body
+    }
+
+    /// Returns a new TypeFunction with the given `type` token
+    pub fn with_type_token(self, type_token: TokenReference) -> Self {
+        Self { type_token, ..self }
+    }
+
+    /// Returns a new TypeFunction with the given function name
+    pub fn with_function_token(self, function_token: TokenReference) -> Self {
+        Self {
+            function_token,
+            ..self
+        }
+    }
+
+    /// Returns a new TypeFunction with the given function name
+    pub fn with_function_name(self, function_name: TokenReference) -> Self {
+        Self {
+            function_name,
+            ..self
+        }
+    }
+
+    /// Returns a new TypeFunction with the given function body
+    pub fn with_function_body(self, function_body: FunctionBody) -> Self {
+        Self {
+            function_body,
+            ..self
+        }
+    }
+}
+
+/// An exported type function, such as `export type function Pairs() ... end`
+#[derive(Clone, Debug, Display, PartialEq, Node, Visit)]
+#[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
+#[display("{export_token}{type_function}")]
+pub struct ExportedTypeFunction {
+    pub(crate) export_token: TokenReference,
+    pub(crate) type_function: TypeFunction,
+}
+
+impl ExportedTypeFunction {
+    /// Creates a new ExportedTypeFunction with the given type function
+    pub fn new(type_function: TypeFunction) -> Self {
+        Self {
+            export_token: TokenReference::new(
+                vec![],
+                Token::new(TokenType::Identifier {
+                    identifier: ShortString::new("export"),
+                }),
+                vec![Token::new(TokenType::spaces(1))],
+            ),
+            type_function,
+        }
+    }
+
+    /// The token `export`.
+    pub fn export_token(&self) -> &TokenReference {
+        &self.export_token
+    }
+
+    /// The type function, `type function Pairs() ... end`.
+    pub fn type_function(&self) -> &TypeFunction {
+        &self.type_function
+    }
+
+    /// Returns a new ExportedTypeFunction with the `export` token
+    pub fn with_export_token(self, export_token: TokenReference) -> Self {
+        Self {
+            export_token,
+            ..self
+        }
+    }
+
+    /// Returns a new ExportedTypeFunction with the given type function
+    pub fn with_type_function(self, type_function: TypeFunction) -> Self {
+        Self {
+            type_function,
+            ..self
+        }
+    }
+}
+
 #[derive(Clone, Debug, Display, PartialEq, Eq, Node, Visit)]
 #[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
 #[non_exhaustive]

full-moon/src/ast/mod.rs

@@ -402,6 +402,14 @@ pub enum Stmt {
     /// Only available when the "luau" feature flag is enabled.
     #[cfg(feature = "luau")]
     TypeDeclaration(TypeDeclaration),
+    /// An exported type function, such as `export type function Pairs(...) end`
+    /// Only available when the "luau" feature flag is enabled.
+    #[cfg(feature = "luau")]
+    ExportedTypeFunction(ExportedTypeFunction),
+    /// A type function, such as `type function Pairs(...) end`
+    /// Only available when the "luau" feature flag is enabled.
+    #[cfg(feature = "luau")]
+    TypeFunction(TypeFunction),
 
     /// A goto statement, such as `goto label`
     /// Only available when the "lua52" or "luajit" feature flag is enabled.

full-moon/src/ast/parsers.rs

@@ -351,6 +351,25 @@ fn parse_stmt(state: &mut ParserState) -> ParserResult<StmtVariant> {
                                         Err(()) => return ParserResult::LexerMoved,
                                     };
 
+                                    if let Some(function_token) = state.consume_if(Symbol::Function)
+                                    {
+                                        return ParserResult::Value(StmtVariant::Stmt(
+                                            ast::Stmt::ExportedTypeFunction(
+                                                ast::ExportedTypeFunction {
+                                                    export_token,
+                                                    type_function: match expect_type_function(
+                                                        state,
+                                                        type_token,
+                                                        function_token,
+                                                    ) {
+                                                        Ok(type_function) => type_function,
+                                                        Err(()) => return ParserResult::LexerMoved,
+                                                    },
+                                                },
+                                            ),
+                                        ));
+                                    }
+
                                     return ParserResult::Value(StmtVariant::Stmt(
                                         ast::Stmt::ExportedTypeDeclaration(
                                             ast::ExportedTypeDeclaration {
@@ -370,6 +389,22 @@ fn parse_stmt(state: &mut ParserState) -> ParserResult<StmtVariant> {
                                 {
                                     let type_token = token;
 
+                                    if let Some(function_token) = state.consume_if(Symbol::Function)
+                                    {
+                                        return ParserResult::Value(StmtVariant::Stmt(
+                                            ast::Stmt::TypeFunction(
+                                                match expect_type_function(
+                                                    state,
+                                                    type_token,
+                                                    function_token,
+                                                ) {
+                                                    Ok(type_function) => type_function,
+                                                    Err(()) => return ParserResult::LexerMoved,
+                                                },
+                                            ),
+                                        ));
+                                    }
+
                                     return ParserResult::Value(StmtVariant::Stmt(
                                         ast::Stmt::TypeDeclaration(
                                             match expect_type_declaration(state, type_token) {
@@ -1332,6 +1367,40 @@ fn expect_type_declaration(
     })
 }
 
+#[cfg(feature = "luau")]
+fn expect_type_function(
+    state: &mut ParserState,
+    type_token: TokenReference,
+    function_token: TokenReference,
+) -> Result<ast::TypeFunction, ()> {
+    let function_name = match state.current() {
+        Ok(token) if token.token_kind() == TokenKind::Identifier => state.consume().unwrap(),
+
+        Ok(token) => {
+            state.token_error(token.clone(), "expected a type function name");
+            return Err(());
+        }
+
+        Err(()) => return Err(()),
+    };
+
+    let function_body = match parse_function_body(state) {
+        ParserResult::Value(body) => body,
+        ParserResult::LexerMoved => return Err(()),
+        ParserResult::NotFound => {
+            state.token_error(function_token, "expected a type function body");
+            return Err(());
+        }
+    };
+
+    Ok(ast::TypeFunction {
+        type_token,
+        function_token,
+        function_name,
+        function_body,
+    })
+}
+
 fn parse_prefix(state: &mut ParserState) -> ParserResult<ast::Prefix> {
     let current_token = match state.current() {
         Ok(token) => token,

full-moon/src/visitors.rs

@@ -282,6 +282,7 @@ create_visitor!(ast: {
         visit_compound_op => CompoundOp,
         visit_else_if_expression => ElseIfExpression,
         visit_exported_type_declaration => ExportedTypeDeclaration,
+        visit_exported_type_function => ExportedTypeFunction,
         visit_generic_declaration => GenericDeclaration,
         visit_generic_declaration_parameter => GenericDeclarationParameter,
         visit_generic_parameter_info => GenericParameterInfo,
@@ -293,6 +294,7 @@ create_visitor!(ast: {
         visit_type_declaration => TypeDeclaration,
         visit_type_field => TypeField,
         visit_type_field_key => TypeFieldKey,
+        visit_type_function => TypeFunction,
         visit_type_info => TypeInfo,
         visit_type_intersection => TypeIntersection,
         visit_type_specifier => TypeSpecifier,