Add support for type function statements

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+* add support for type functions ([#333](https://github.com/seaofvoices/darklua/pull/333))
 * add support for property modifiers in table types (like `read` in a type like `{ read name: string }`) ([#332](https://github.com/seaofvoices/darklua/pull/332))
 
 ## 0.17.3

src/ast_converter.rs

@@ -852,6 +852,36 @@ impl<'a> AstConverter<'a> {
                             .into(),
                     );
                 }
+                ConvertWork::MakeTypeFunctionStatement {
+                    statement,
+                    export_token,
+                } => {
+                    let builder = self.convert_function_body_attributes(
+                        statement.function_body(),
+                        self.convert_token(statement.function_token())?,
+                    )?;
+
+                    let mut name = Identifier::new(statement.function_name().token().to_string());
+                    let mut type_token = None;
+                    let mut export = None;
+
+                    if self.hold_token_data {
+                        name.set_token(self.convert_token(statement.function_name())?);
+                        type_token = Some(self.convert_token(statement.type_token())?);
+                        export = export_token
+                            .map(|token| self.convert_token(token))
+                            .transpose()?;
+                    }
+
+                    let mut type_function_statement =
+                        builder.into_type_function_statement(name, type_token, export);
+
+                    if export_token.is_some() {
+                        type_function_statement.set_exported();
+                    }
+
+                    self.statements.push(type_function_statement.into());
+                }
                 ConvertWork::MakeLocalAssignStatement { statement } => {
                     let variables = statement
                         .names()
@@ -1547,6 +1577,22 @@ impl<'a> AstConverter<'a> {
                     });
                 self.push_function_body_work(local_function.body());
             }
+            ast::Stmt::TypeFunction(type_function) => {
+                self.work_stack
+                    .push(ConvertWork::MakeTypeFunctionStatement {
+                        statement: type_function,
+                        export_token: None,
+                    });
+                self.push_function_body_work(type_function.function_body());
+            }
+            ast::Stmt::ExportedTypeFunction(type_function) => {
+                self.work_stack
+                    .push(ConvertWork::MakeTypeFunctionStatement {
+                        statement: type_function.type_function(),
+                        export_token: Some(type_function.export_token()),
+                    });
+                self.push_function_body_work(type_function.type_function().function_body());
+            }
             ast::Stmt::NumericFor(numeric_for) => {
                 self.work_stack.push(ConvertWork::MakeNumericForStatement {
                     statement: numeric_for,
@@ -2866,6 +2912,10 @@ enum ConvertWork<'a> {
         type_declaration: &'a ast::luau::TypeDeclaration,
         export_token: Option<&'a tokenizer::TokenReference>,
     },
+    MakeTypeFunctionStatement {
+        statement: &'a ast::luau::TypeFunction,
+        export_token: Option<&'a tokenizer::TokenReference>,
+    },
     MakePrefixFromExpression {
         prefix: &'a ast::Prefix,
     },

src/generator/dense.rs

@@ -637,6 +637,42 @@ impl LuaGenerator for DenseLuaGenerator {
         self.write_type(statement.get_type());
     }
 
+    fn write_type_function_statement(&mut self, function: &nodes::TypeFunctionStatement) {
+        if function.is_exported() {
+            self.push_str("export");
+        }
+        self.push_str("type");
+        self.push_str("function");
+
+        self.write_identifier(function.get_identifier());
+
+        if let Some(generics) = function.get_generic_parameters() {
+            self.write_function_generics(generics);
+        }
+
+        self.push_char('(');
+
+        let parameters = function.get_parameters();
+        self.write_function_parameters(
+            parameters,
+            function.is_variadic(),
+            function.get_variadic_type(),
+        );
+        self.push_char(')');
+
+        if let Some(return_type) = function.get_return_type() {
+            self.push_char(':');
+            self.write_function_return_type(return_type);
+        }
+
+        let block = function.get_block();
+
+        if !block.is_empty() {
+            self.write_block(block);
+        }
+        self.push_str("end");
+    }
+
     fn write_false_expression(&mut self, _token: &Option<nodes::Token>) {
         self.push_str("false");
     }

src/generator/mod.rs

@@ -37,6 +37,7 @@ pub trait LuaGenerator {
             Repeat(statement) => self.write_repeat_statement(statement),
             While(statement) => self.write_while_statement(statement),
             TypeDeclaration(statement) => self.write_type_declaration_statement(statement),
+            TypeFunction(statement) => self.write_type_function_statement(statement),
         }
     }
 
@@ -53,6 +54,7 @@ pub trait LuaGenerator {
     fn write_repeat_statement(&mut self, repeat: &nodes::RepeatStatement);
     fn write_while_statement(&mut self, while_statement: &nodes::WhileStatement);
     fn write_type_declaration_statement(&mut self, statement: &nodes::TypeDeclarationStatement);
+    fn write_type_function_statement(&mut self, statement: &nodes::TypeFunctionStatement);
 
     fn write_variable(&mut self, variable: &nodes::Variable) {
         use nodes::Variable::*;
@@ -883,6 +885,18 @@ mod $mod_name {
                     ),
         ));
 
+        snapshot_node!($mod_name, $generator, type_function, write_type_function_statement => (
+            empty => TypeFunctionStatement::from_name("nothing", Block::default()),
+            empty_exported => TypeFunctionStatement::from_name("nothing", Block::default())
+                .export(),
+            empty_with_parameter => TypeFunctionStatement::from_name("nothing", Block::default())
+                .with_parameter("param"),
+            empty_with_parameters_and_return_type => TypeFunctionStatement::from_name("nothing", Block::default())
+                .with_parameter("param")
+                .with_return_type(Type::from(true)),
+        ));
+
+
         snapshot_node!($mod_name, $generator, if_statement, write_statement => (
             empty => IfStatement::create(false, Block::default()),
             empty_with_empty_else => IfStatement::create(false, Block::default())

src/generator/readable.rs

@@ -16,6 +16,7 @@ enum StatementType {
     Repeat,
     While,
     TypeDeclaration,
+    TypeFunction,
     Return,
     Break,
     Continue,
@@ -38,6 +39,7 @@ impl From<&nodes::Statement> for StatementType {
             Repeat(_) => Self::Repeat,
             While(_) => Self::While,
             TypeDeclaration(_) => Self::TypeDeclaration,
+            TypeFunction(_) => Self::TypeFunction,
         }
     }
 }
@@ -871,6 +873,45 @@ impl LuaGenerator for ReadableLuaGenerator {
         self.write_type(statement.get_type());
     }
 
+    fn write_type_function_statement(&mut self, function: &nodes::TypeFunctionStatement) {
+        self.push_can_add_new_line(false);
+        if function.is_exported() {
+            self.push_str("export");
+        }
+        self.push_str("type function ");
+        self.write_identifier(function.get_identifier());
+
+        if let Some(generics) = function.get_generic_parameters() {
+            self.write_function_generics(generics);
+        }
+
+        self.raw_push_char('(');
+
+        self.pop_can_add_new_line();
+
+        let parameters = function.get_parameters();
+        self.write_function_parameters(
+            parameters,
+            function.is_variadic(),
+            function.get_variadic_type(),
+        );
+        self.raw_push_char(')');
+
+        if let Some(return_type) = function.get_return_type() {
+            self.write_function_return_type_suffix(return_type);
+        }
+
+        let block = function.get_block();
+
+        if block.is_empty() {
+            self.raw_push_str(" end");
+        } else {
+            self.push_new_line();
+            self.indent_and_write_block(block);
+            self.push_str("end");
+        }
+    }
+
     fn write_false_expression(&mut self, _token: &Option<nodes::Token>) {
         self.push_str("false");
     }

src/generator/snapshots/darklua_core__generator__test__dense__snapshots__type_function__dense_type_function_empty.snap

@@ -0,0 +1,5 @@
+---
+source: src/generator/mod.rs
+expression: generator.into_string()
+---
+type function nothing()end

src/generator/snapshots/darklua_core__generator__test__dense__snapshots__type_function__dense_type_function_empty_exported.snap

@@ -0,0 +1,5 @@
+---
+source: src/generator/mod.rs
+expression: generator.into_string()
+---
+export type function nothing()end

src/generator/snapshots/darklua_core__generator__test__dense__snapshots__type_function__dense_type_function_empty_with_parameter.snap

@@ -0,0 +1,5 @@
+---
+source: src/generator/mod.rs
+expression: generator.into_string()
+---
+type function nothing(param)end

src/generator/snapshots/darklua_core__generator__test__dense__snapshots__type_function__dense_type_function_empty_with_parameters_and_return_type.snap

@@ -0,0 +1,5 @@
+---
+source: src/generator/mod.rs
+expression: generator.into_string()
+---
+type function nothing(param):true end

src/generator/snapshots/darklua_core__generator__test__readable__snapshots__type_function__readable_type_function_empty.snap

@@ -0,0 +1,5 @@
+---
+source: src/generator/mod.rs
+expression: generator.into_string()
+---
+type function nothing() end

src/generator/snapshots/darklua_core__generator__test__readable__snapshots__type_function__readable_type_function_empty_exported.snap

@@ -0,0 +1,5 @@
+---
+source: src/generator/mod.rs
+expression: generator.into_string()
+---
+export type function nothing() end

src/generator/snapshots/darklua_core__generator__test__readable__snapshots__type_function__readable_type_function_empty_with_parameter.snap

@@ -0,0 +1,5 @@
+---
+source: src/generator/mod.rs
+expression: generator.into_string()
+---
+type function nothing(param) end

src/generator/snapshots/darklua_core__generator__test__readable__snapshots__type_function__readable_type_function_empty_with_parameters_and_return_type.snap

@@ -0,0 +1,5 @@
+---
+source: src/generator/mod.rs
+expression: generator.into_string()
+---
+type function nothing(param): true end

src/generator/snapshots/darklua_core__generator__test__token_based__snapshots__type_function__token_based_type_function_empty.snap

@@ -0,0 +1,5 @@
+---
+source: src/generator/mod.rs
+expression: generator.into_string()
+---
+type function nothing()end

src/generator/snapshots/darklua_core__generator__test__token_based__snapshots__type_function__token_based_type_function_empty_exported.snap

@@ -0,0 +1,5 @@
+---
+source: src/generator/mod.rs
+expression: generator.into_string()
+---
+export type function nothing()end

src/generator/snapshots/darklua_core__generator__test__token_based__snapshots__type_function__token_based_type_function_empty_with_parameter.snap

@@ -0,0 +1,5 @@
+---
+source: src/generator/mod.rs
+expression: generator.into_string()
+---
+type function nothing(param)end

src/generator/snapshots/darklua_core__generator__test__token_based__snapshots__type_function__token_based_type_function_empty_with_parameters_and_return_type.snap

@@ -0,0 +1,5 @@
+---
+source: src/generator/mod.rs
+expression: generator.into_string()
+---
+type function nothing(param):true end

src/generator/token_based.rs

@@ -687,6 +687,35 @@ impl<'a> TokenBasedLuaGenerator<'a> {
         self.write_type(statement.get_type());
     }
 
+    fn write_type_function_with_tokens(
+        &mut self,
+        function: &TypeFunctionStatement,
+        tokens: &TypeFunctionStatementTokens,
+    ) {
+        if function.is_exported() {
+            if let Some(export_token) = &tokens.export {
+                self.write_token(export_token);
+            } else {
+                self.write_symbol("export");
+            }
+        }
+        self.write_token(&tokens.r#type);
+        self.write_token(&tokens.function_body.function);
+
+        self.write_identifier(function.get_identifier());
+
+        self.write_function_attributes(
+            tokens,
+            function.get_generic_parameters(),
+            function.parameters_count(),
+            function.iter_parameters(),
+            function.is_variadic(),
+            function.get_variadic_type(),
+            function.get_return_type(),
+            function.get_block(),
+        );
+    }
+
     fn write_generic_parameters_with_default_with_tokens(
         &mut self,
         generic_parameters: &GenericParametersWithDefaults,
@@ -1397,10 +1426,37 @@ impl<'a> TokenBasedLuaGenerator<'a> {
         TypeDeclarationTokens {
             r#type: Token::from_content("type"),
             equal: Token::from_content("="),
-            export: if statement.is_exported() {
-                Some(Token::from_content("export"))
-            } else {
-                None
+            export: statement
+                .is_exported()
+                .then(|| Token::from_content("export")),
+        }
+    }
+
+    fn generate_type_function_tokens(
+        &self,
+        statement: &TypeFunctionStatement,
+    ) -> TypeFunctionStatementTokens {
+        TypeFunctionStatementTokens {
+            r#type: Token::from_content("type"),
+            export: statement
+                .is_exported()
+                .then(|| Token::from_content("export")),
+            function_body: FunctionBodyTokens {
+                function: Token::from_content("function"),
+                opening_parenthese: Token::from_content("("),
+                closing_parenthese: Token::from_content(")"),
+                end: Token::from_content("end"),
+                parameter_commas: intersect_with_token(
+                    comma_token(),
+                    statement.parameters_count() + usize::from(statement.is_variadic()),
+                ),
+                variable_arguments: statement.is_variadic().then(|| Token::from_content("...")),
+                variable_arguments_colon: statement
+                    .has_variadic_type()
+                    .then(|| Token::from_content(":")),
+                return_type_colon: statement
+                    .has_return_type()
+                    .then(|| Token::from_content(":")),
             },
         }
     }
@@ -1896,6 +1952,17 @@ impl LuaGenerator for TokenBasedLuaGenerator<'_> {
         }
     }
 
+    fn write_type_function_statement(&mut self, statement: &TypeFunctionStatement) {
+        if let Some(tokens) = statement.get_tokens() {
+            self.write_type_function_with_tokens(statement, tokens);
+        } else {
+            self.write_type_function_with_tokens(
+                statement,
+                &self.generate_type_function_tokens(statement),
+            );
+        }
+    }
+
     fn write_false_expression(&mut self, token: &Option<Token>) {
         if let Some(token) = token {
             self.write_token(token);