Add type nodes

Note that names are placeholders

src/LuauAST/impl/create.ts

@@ -45,8 +45,8 @@ let lastTempId = 0;
 /**
  * Creates a new temporary identifier for a node.
  */
-export function tempId(name = "") {
-	return luau.create(luau.SyntaxKind.TemporaryIdentifier, { name, id: lastTempId++ });
+export function tempId(name = "", annotation: luau.TypeExpression | undefined = undefined) {
+	return luau.create(luau.SyntaxKind.TemporaryIdentifier, { name, id: lastTempId++, annotation });
 }
 
 /**
@@ -102,8 +102,8 @@ export function string(value: string) {
  * Creates a new identifier node.
  * @param name The name of the identifier.
  */
-export function id(name: string) {
-	return luau.create(luau.SyntaxKind.Identifier, { name });
+export function id(name: string, annotation: luau.TypeExpression | undefined = undefined) {
+	return luau.create(luau.SyntaxKind.Identifier, { name, annotation });
 }
 
 /**
@@ -182,3 +182,24 @@ export function call(expression: luau.IndexableExpression, args: Array<luau.Expr
 		args: luau.list.make(...args),
 	});
 }
+
+export function typeId(name: string) {
+	return luau.create(luau.SyntaxKind.TypeIdentifier, {
+		module: undefined,
+		name,
+	});
+}
+
+export function any() {
+	return luau.create(luau.SyntaxKind.TypeIdentifier, {
+		module: undefined,
+		name: "any",
+	});
+}
+
+export function cast(expression: luau.Expression, type: luau.TypeExpression) {
+	return luau.create(luau.SyntaxKind.TypeCast, {
+		expression,
+		type,
+	});
+}

src/LuauAST/impl/enums.ts

@@ -43,6 +43,17 @@ export enum SyntaxKind {
 	ReturnStatement,
 	Comment,
 
+	// types
+	TypeIdentifier,
+	TypeMixedTable,
+	TypeMixedTableField,
+	TypeMixedTableIndexedField,
+	TypeFunction,
+	TypeParameter,
+	TypeTypeOf,
+	TypeStatement,
+	TypeCast,
+
 	// fields
 	MapField,
 	InterpolatedStringPart,

src/LuauAST/types/mapping.ts

@@ -27,6 +27,7 @@ export interface ExpressionByKind extends IndexableExpressionByKind {
 	[luau.SyntaxKind.Map]: luau.Map;
 	[luau.SyntaxKind.Set]: luau.Set;
 	[luau.SyntaxKind.MixedTable]: luau.MixedTable;
+	[luau.SyntaxKind.TypeCast]: luau.TypeCast;
 }
 
 export interface StatementByKind {
@@ -45,14 +46,26 @@ export interface StatementByKind {
 	[luau.SyntaxKind.VariableDeclaration]: luau.VariableDeclaration;
 	[luau.SyntaxKind.ReturnStatement]: luau.ReturnStatement;
 	[luau.SyntaxKind.Comment]: luau.Comment;
+	[luau.SyntaxKind.TypeStatement]: luau.TypeStatement;
 }
 
 export interface FieldByKind {
 	[luau.SyntaxKind.MapField]: luau.MapField;
 	[luau.SyntaxKind.InterpolatedStringPart]: luau.InterpolatedStringPart;
 }
 
-export interface NodeByKind extends luau.ExpressionByKind, luau.StatementByKind, luau.FieldByKind {}
+export interface TypeByKind {
+	[luau.SyntaxKind.TypeIdentifier]: luau.TypeIdentifier;
+	[luau.SyntaxKind.TypeMixedTableField]: luau.TypeMixedTableField;
+	[luau.SyntaxKind.TypeFunction]: luau.TypeFunction;
+	[luau.SyntaxKind.TypeParameter]: luau.TypeParameter;
+	[luau.SyntaxKind.TypeMixedTable]: luau.TypeMixedTable;
+	[luau.SyntaxKind.TypeMixedTableIndexedField]: luau.TypeMixedTableIndexedField;
+	[luau.SyntaxKind.TypeTypeOf]: luau.TypeTypeOf;
+	[luau.SyntaxKind.None]: luau.None;
+}
+
+export interface NodeByKind extends luau.ExpressionByKind, luau.StatementByKind, luau.FieldByKind, TypeByKind {}
 
 export type IndexableExpression<T extends luau.SyntaxKind = luau.SyntaxKind> = {
 	[K in keyof IndexableExpressionByKind]: IndexableExpressionByKind[K]["kind"] extends T

src/LuauAST/types/nodes.ts

@@ -16,6 +16,8 @@ export interface BaseExpression<T extends keyof luau.ExpressionByKind = keyof lu
 export interface BaseStatement<T extends keyof luau.StatementByKind = keyof luau.StatementByKind>
 	extends luau.BaseNode<T> {}
 
+export interface BaseType<T extends keyof luau.TypeByKind = keyof luau.TypeByKind> extends luau.BaseNode<T> {}
+
 export interface BaseField<T extends keyof luau.FieldByKind = keyof luau.FieldByKind> extends luau.BaseNode<T> {}
 
 export interface HasParameters {
@@ -25,6 +27,13 @@ export interface HasParameters {
 
 export type AnyIdentifier = luau.Identifier | luau.TemporaryIdentifier;
 
+export type TypeExpression =
+	| luau.TypeIdentifier
+	| luau.TypeMixedTable
+	| luau.TypeFunction
+	| luau.TypeTypeOf
+	| luau.None;
+
 export type WritableExpression = luau.AnyIdentifier | luau.PropertyAccessExpression | luau.ComputedIndexExpression;
 
 export type SimpleTypes =
@@ -63,11 +72,13 @@ export interface FunctionExpression extends luau.BaseExpression<luau.SyntaxKind.
 
 export interface Identifier extends luau.BaseExpression<luau.SyntaxKind.Identifier> {
 	name: string;
+	annotation?: TypeExpression;
 }
 
 export interface TemporaryIdentifier extends luau.BaseExpression<luau.SyntaxKind.TemporaryIdentifier> {
 	name: string;
 	id: number;
+	annotation?: TypeExpression;
 }
 
 export interface ComputedIndexExpression extends luau.BaseExpression<luau.SyntaxKind.ComputedIndexExpression> {
@@ -185,6 +196,7 @@ export interface FunctionDeclaration extends luau.BaseStatement<luau.SyntaxKind.
 	localize: boolean;
 	name: luau.AnyIdentifier | luau.PropertyAccessExpression;
 	statements: luau.List<luau.Statement>;
+	returnType?: TypeExpression;
 }
 
 export interface MethodDeclaration extends luau.BaseStatement<luau.SyntaxKind.MethodDeclaration>, HasParameters {
@@ -215,3 +227,48 @@ export interface MapField extends luau.BaseField<luau.SyntaxKind.MapField> {
 export interface InterpolatedStringPart extends luau.BaseField<luau.SyntaxKind.InterpolatedStringPart> {
 	text: string;
 }
+
+// types
+export interface TypeIdentifier extends luau.BaseType<luau.SyntaxKind.TypeIdentifier> {
+	module?: string;
+	name: string;
+}
+
+export interface TypeMixedTable extends luau.BaseType<luau.SyntaxKind.TypeMixedTable> {
+	fields: luau.List<luau.TypeMixedTableField | luau.TypeMixedTableIndexedField>;
+}
+
+export interface TypeMixedTableField extends luau.BaseType<luau.SyntaxKind.TypeMixedTableField> {
+	index: luau.TypeIdentifier;
+	value: luau.TypeExpression;
+}
+
+export interface TypeMixedTableIndexedField extends luau.BaseType<luau.SyntaxKind.TypeMixedTableIndexedField> {
+	index: luau.TypeIdentifier;
+	value: luau.TypeExpression;
+}
+
+export interface TypeFunction extends luau.BaseType<luau.SyntaxKind.TypeFunction> {
+	parameters?: luau.List<TypeParameter>;
+	dotDotDot?: TypeExpression;
+	returnType?: TypeExpression;
+}
+
+export interface TypeParameter extends luau.BaseType<luau.SyntaxKind.TypeParameter> {
+	name?: string;
+	value: luau.TypeExpression;
+}
+
+export interface TypeTypeOf extends luau.BaseType<luau.SyntaxKind.TypeTypeOf> {
+	expression: luau.Expression;
+}
+
+export interface TypeStatement extends luau.BaseStatement<luau.SyntaxKind.TypeStatement> {
+	identifier: luau.TypeIdentifier;
+	expression: TypeExpression;
+}
+
+export interface TypeCast extends luau.BaseExpression<luau.SyntaxKind.TypeCast> {
+	expression: luau.Expression;
+	type: TypeExpression;
+}

src/LuauRenderer/nodes/expressions/indexable/renderIdentifier.ts

@@ -1,8 +1,16 @@
 import luau from "LuauAST";
+import { isNone } from "LuauAST/bundle";
 import { assert } from "LuauAST/util/assert";
-import { RenderState } from "LuauRenderer";
+import { render, RenderState } from "LuauRenderer";
 
 export function renderIdentifier(state: RenderState, node: luau.Identifier) {
 	assert(luau.isValidIdentifier(node.name), `Invalid Luau Identifier: "${node.name}"`);
 	return node.name;
 }
+
+export function renderIdentifierWithType(state: RenderState, node: luau.Identifier) {
+	if (node.annotation && !isNone(node)) {
+		return `${node.name}: ${render(state, node.annotation)}`;
+	}
+	return node.name;
+}

src/LuauRenderer/nodes/expressions/indexable/renderTemporaryIdentifier.ts

@@ -1,9 +1,19 @@
 import luau from "LuauAST";
+import { isNode } from "LuauAST/bundle";
 import { assert } from "LuauAST/util/assert";
-import { RenderState } from "LuauRenderer";
+import { render, RenderState } from "LuauRenderer";
 
 export function renderTemporaryIdentifier(state: RenderState, node: luau.TemporaryIdentifier) {
 	const name = state.getTempName(node);
 	assert(luau.isValidIdentifier(name), `Invalid Temporary Identifier: "${name}"`);
 	return name;
 }
+
+export function renderTemporaryIdentifierWithType(state: RenderState, node: luau.TemporaryIdentifier) {
+	const name = state.getTempName(node);
+	assert(luau.isValidIdentifier(name), `Invalid Temporary Identifier: "${name}"`);
+	if (node.annotation && !isNode(node.annotation)) {
+		return `${name}: ${render(state, node.annotation)}`;
+	}
+	return name;
+}

src/LuauRenderer/nodes/statements/renderFunctionDeclaration.ts

@@ -12,7 +12,9 @@ export function renderFunctionDeclaration(state: RenderState, node: luau.Functio
 	const paramStr = renderParameters(state, node);
 
 	let result = "";
-	result += state.line(`${node.localize ? "local " : ""}function ${nameStr}(${paramStr})`);
+	result += state.line(
+		`${node.localize ? "local " : ""}function ${nameStr}(${paramStr})${node.returnType ? `: ${render(state, node.returnType)}` : ""}`,
+	);
 	result += state.block(() => renderStatements(state, node.statements));
 	result += state.line(`end`);
 	return result;

src/LuauRenderer/nodes/statements/renderTypeDeclaration.ts

@@ -0,0 +1,6 @@
+import luau from "LuauAST";
+import { render, RenderState } from "LuauRenderer";
+
+export function renderTypeDeclaration(state: RenderState, node: luau.TypeStatement) {
+	return state.line(`type ${render(state, node.identifier)} = ${render(state, node.expression)}`);
+}

src/LuauRenderer/nodes/statements/renderVariableDeclaration.ts

@@ -1,14 +1,14 @@
 import luau from "LuauAST";
 import { assert } from "LuauAST/util/assert";
-import { render, RenderState } from "LuauRenderer";
+import { render, RenderState, renderTyped } from "LuauRenderer";
 
 export function renderVariableDeclaration(state: RenderState, node: luau.VariableDeclaration) {
 	let leftStr: string;
 	if (luau.list.isList(node.left)) {
 		assert(!luau.list.isEmpty(node.left));
-		leftStr = luau.list.mapToArray(node.left, id => render(state, id)).join(", ");
+		leftStr = luau.list.mapToArray(node.left, id => renderTyped(state, id)).join(", ");
 	} else {
-		leftStr = render(state, node.left);
+		leftStr = renderTyped(state, node.left);
 	}
 
 	if (node.right) {

src/LuauRenderer/nodes/types/renderMixedTable.ts

@@ -0,0 +1,21 @@
+import luau from "LuauAST";
+import { render, RenderState } from "LuauRenderer";
+
+export function renderMixedTableType(state: RenderState, node: luau.TypeMixedTable) {
+	const inner: Array<string> = [];
+	for (const field of luau.list.toArray(node.fields)) {
+		inner.push(render(state, field));
+	}
+	return `{ ${inner.join(", ")} }`;
+}
+
+export function renderMixedTableFieldType(state: RenderState, node: luau.TypeMixedTableField) {
+	if (!luau.isValidIdentifier(node.index.name)) {
+		return `["${node.index.name}"]: ${render(state, node.value)}`;
+	}
+	return `${render(state, node.index)}: ${render(state, node.value)}`;
+}
+
+export function renderMixedTableIndexedFieldType(state: RenderState, node: luau.TypeMixedTableIndexedField) {
+	return `[${render(state, node.index)}]: ${render(state, node.value)}`;
+}

src/LuauRenderer/nodes/types/renderTypeCast.ts

@@ -0,0 +1,6 @@
+import luau from "LuauAST";
+import { render, RenderState } from "LuauRenderer";
+
+export function renderTypeCast(state: RenderState, node: luau.TypeCast) {
+	return `(${render(state, node.expression)} :: ${render(state, node.type)})`;
+}

src/LuauRenderer/nodes/types/renderTypeFunction.ts

@@ -0,0 +1,22 @@
+import luau from "LuauAST";
+import { render, RenderState } from "LuauRenderer";
+
+export function renderTypeFunction(state: RenderState, node: luau.TypeFunction) {
+	let params: string;
+
+	if (node.parameters !== undefined) {
+		const rendered = luau.list.mapToArray(node.parameters, expr => render(state, expr));
+		if (node.dotDotDot) {
+			rendered.push(`...${render(state, node.dotDotDot)}`);
+		}
+		params = rendered.join(", ");
+	} else {
+		if (node.dotDotDot) {
+			params = `...${render(state, node.dotDotDot)}`;
+		} else {
+			params = "...any";
+		}
+	}
+
+	return `(${params}) -> ${node.returnType ? render(state, node.returnType) : "()"}`;
+}

src/LuauRenderer/nodes/types/renderTypeIdentifier.ts

@@ -0,0 +1,9 @@
+import assert from "assert";
+import luau from "LuauAST";
+import { RenderState } from "LuauRenderer";
+
+export function renderTypeIdentifier(state: RenderState, node: luau.TypeIdentifier) {
+	assert(luau.isValidIdentifier(node.name), `Invalid Luau Identifier: "${node.name}"`);
+	if (node.module) return `${node.module}.${node.name}`;
+	return node.name;
+}

src/LuauRenderer/nodes/types/renderTypeOf.ts

@@ -0,0 +1,6 @@
+import luau from "LuauAST";
+import { render, RenderState } from "LuauRenderer";
+
+export function renderTypeOf(state: RenderState, node: luau.TypeTypeOf) {
+	return `typeof(${render(state, node.expression)})`;
+}

src/LuauRenderer/nodes/types/renderTypeParameter.ts

@@ -0,0 +1,12 @@
+import luau from "LuauAST";
+import { assert } from "LuauAST/util/assert";
+import { render, RenderState } from "LuauRenderer";
+
+export function renderTypeParameter(state: RenderState, node: luau.TypeParameter) {
+	if (node.name) {
+		assert(luau.isValidIdentifier(node.name), `Invalid Luau Identifier: "${node.name}"`);
+		return `${node.name}: ${render(state, node.value)}`;
+	}
+
+	return render(state, node.value);
+}