perf: fix O(n²) descendant traversal + close #33 (disallow Luau keywords)

Loom.Core/Diagnostics/InternalCodes.cs

@@ -2,7 +2,7 @@
 
 public static class InternalCodes
 {
     public const string Unknown = "L???";
     public const string NotImplemented = "L000";
     public const string CompilerError = "L001";
 
@@ -74,6 +74,7 @@
     public const string InvalidLuauNameAttribute = "L350";
     public const string AnonymousEventDisconnect = "L351";
     public const string UnresolvedEventDisconnect = "L352";
-
+    public const string ReservedLuauKeyword = "L353";
+
     public const string SimplifiableCode = "L400";
 }

Loom.Core/FlowAnalysis/FlowAnalyzer.cs

@@ -12,7 +12,11 @@
     private readonly Dictionary<Node, FlowState> _states = [];
 
     public FlowState GetState(Node node) => _states.TryGetValue(node, out var existingState) ? existingState : FlowState.Empty;
-    public FlowAnalyzerResult Analyze() => new(BindState(semanticModel.Tree, AnalyzeStatements(semanticModel.Tree.Statements, new FlowState())), _diagnostics);
+    public FlowAnalyzerResult Analyze()
+    {
+        BindState(semanticModel.Tree, AnalyzeStatements(semanticModel.Tree.Statements, new FlowState()));
+        return new FlowAnalyzerResult(_diagnostics);
+    }
 
     private FlowState AnalyzeStatements(IReadOnlyList<Statement> statements, FlowState state) =>
         statements.Aggregate(state, (current, statement) => AnalyzeStatement(statement, current));
@@ -26,6 +30,7 @@
             FunctionDeclaration functionDeclaration => AnalyzeFunctionDeclaration(functionDeclaration, state),
             EventDeclaration eventDeclaration => AnalyzeEventDeclaration(eventDeclaration, state),
             InterfaceDeclaration interfaceDeclaration => AnalyzeInterfaceDeclaration(interfaceDeclaration, state),
+            EnumDeclaration enumDeclaration => AnalyzeEnumDeclaration(enumDeclaration, state),
             Implement implement => AnalyzeImplement(implement, state),
             Return @return => AnalyzeReturn(@return, state),
             Break @break => AnalyzeBreak(@break, state),
@@ -101,6 +106,14 @@
                 : state
         );
 
+    private FlowState AnalyzeEnumDeclaration(EnumDeclaration enumDeclaration, FlowState state) =>
+        BindState(
+            enumDeclaration,
+            semanticModel.GetDeclarationSymbol(enumDeclaration, SymbolKind.Variable) is { } symbol
+                ? state.WithInitialized(symbol)
+                : state
+        );
+
     private FlowState AnalyzeFunctionDeclaration(FunctionDeclaration functionDeclaration, FlowState state)
     {
         var newState = semanticModel.GetDeclarationSymbol(functionDeclaration) is { } symbol
@@ -251,7 +264,7 @@
     private static FlowState ComputeLoopExitState(FlowState entryState, FlowState bodyState, List<FlowState> breakStates)
     {
         var definitely = entryState.DefinitelyInitialized;
         foreach (var exit in breakStates)
             definitely = definitely.Intersect(exit.DefinitelyInitialized);
 
         var maybeBuilder = entryState.MaybeInitialized.ToBuilder();

Loom.Core/FlowAnalysis/FlowAnalyzerResult.cs

@@ -2,5 +2,5 @@
 
 namespace Loom.Core.FlowAnalysis;
 
-public sealed record FlowAnalyzerResult(FlowState FlowState, DiagnosticBag Diagnostics)
+public sealed record FlowAnalyzerResult(DiagnosticBag Diagnostics)
     : DiagnosedResult(Diagnostics);

Loom.Core/Generation/LuauGenerator.Types.cs

@@ -1,7 +1,6 @@
 using Loom.Core.Diagnostics;
 using Loom.Core.Parsing.AST;
 using Loom.Core.Resolving;
-using Loom.Core.TypeChecking.Types;
 using Loom.Luau;
 using Loom.Luau.AST;
 using ArrayType = Loom.Core.Parsing.AST.ArrayType;

Loom.Core/Generation/Macros/MacroExpander.cs

@@ -4,7 +4,6 @@
 using Loom.Core.Parsing.AST;
 using Loom.Core.Resolving;
 using Loom.Core.TypeChecking;
-using Loom.Core.TypeChecking.Types;
 using Loom.Luau.AST;
 using ElementAccess = Loom.Core.Parsing.AST.ElementAccess;
 using Identifier = Loom.Core.Parsing.AST.Identifier;

Loom.Core/Lexing/Lexer.cs

@@ -24,7 +24,7 @@ public LexerResult Tokenize()
             significantTokens.Add(token);
         }
 
-        return new LexerResult(file, significantTokens, allTokens, _diagnostics);
+        return new LexerResult(significantTokens, allTokens, _diagnostics);
     }
 
     private IEnumerable<Token> GetTokens()

Loom.Core/Lexing/LexerResult.cs

@@ -3,5 +3,5 @@
 
 namespace Loom.Core.Lexing;
 
-public sealed record LexerResult(SourceFile File, List<Token> Tokens, List<Token> TokensWithTrivia, DiagnosticBag Diagnostics)
+public sealed record LexerResult(List<Token> Tokens, List<Token> TokensWithTrivia, DiagnosticBag Diagnostics)
     : DiagnosedResult(Diagnostics);

Loom.Core/Parsing/AST/EventDeclaration.cs

@@ -4,7 +4,7 @@
 
 namespace Loom.Core.Parsing.AST;
 
-public class EventDeclaration(Token keyword, Token name, TypeParameters? typeParameters, Parameters? parameters, Attributes? attributes = null)
+public class EventDeclaration(Token keyword, Token name, TypeParameters? typeParameters, Parameters? parameters, Attributes? attributes)
     : GenericNamedDeclaration([], keyword, name, typeParameters, attributes),
       IWithAttributes
 {

Loom.Core/Parsing/AST/Node.cs

@@ -10,7 +10,7 @@
     public NodeId Id { get; }
     public IReadOnlyList<Node> Children { get; }
     public IReadOnlyList<Token> Tokens { get; }
     public TextSpan Span { get; }
     public LocationSpan LocationSpan => new(new Location(File, Span.Position), new Location(File, Span.End));
     public SourceFile File => field ??= Tokens.Count == 0 ? SourceFile.Empty : Tokens[0].File;
     [MaybeNull] public Node Parent { get; private set; }
@@ -29,7 +29,21 @@
     public abstract T Accept<T>(Visitor<T> visitor);
     public override string ToString() => LocationSpan.GetText().ToString();
     public IReadOnlyList<T> GetDescendants<T>() where T : Node => GetDescendants().OfType<T>().ToArray();
-    public IReadOnlyList<Node> GetDescendants() => [..Children, ..Children.SelectMany(c => c.GetDescendants())];
+
+    public IReadOnlyList<Node> GetDescendants()
+    {
+        var result = new List<Node>();
+        var queue = new Queue<Node>(Children);
+        while (queue.Count > 0)
+        {
+            var node = queue.Dequeue();
+            result.Add(node);
+            foreach (var child in node.Children)
+                queue.Enqueue(child);
+        }
+
+        return result;
+    }
     public bool IsDescendantOf<T>() where T : Node => FirstAncestorOfType<T>() != null;
 
     public T? FirstAncestorOfType<T>()

Loom.Core/Resolving/Resolver.cs

@@ -4,6 +4,7 @@
 using Loom.Core.Parsing.AST;
 using Loom.Core.Text;
 using Loom.Core.TypeChecking;
+using Loom.Luau;
 using Attribute = Loom.Core.Parsing.AST.Attribute;
 
 namespace Loom.Core.Resolving;
@@ -689,6 +690,15 @@ private void DeclareSymbol(Symbol symbol)
     {
         AddToLookup(symbol);
         AddDeclaration(symbol);
+        if (LuauFactory.Keywords.Contains(symbol.Name))
+        {
+            _diagnostics.Error(
+                symbol.Declaration,
+                InternalCodes.ReservedLuauKeyword,
+                $"'{symbol.Name}' is a reserved Luau keyword and cannot be used as a declaration name."
+            );
+        }
+
         if (TypeChecker.EmitDebugDiagnostics)
             _diagnostics.Debug(symbol.Declaration, $"Declared symbol: {symbol}");
 

Loom.Core/Resolving/SemanticModel.cs

@@ -11,16 +11,6 @@ namespace Loom.Core.Resolving;
 public sealed record SemanticModel(Tree Tree, DiagnosticBag Diagnostics, SymbolTable Declarations, SymbolTable References)
     : DiagnosedResult(Diagnostics)
 {
-    internal int RuntimeReferences = 0;
-
-    /// <summary>
-    /// Node IDs of reference-site nodes that originate from a non-intrinsic source file.
-    /// Populated by the resolver as references are recorded, so
-    /// <see cref="MustImportRuntimeLibrary"/> can filter references by the file of the
-    /// referencing node without holding on to the node instances themselves.
-    /// </summary>
-    internal HashSet<NodeId> NonIntrinsicReferenceNodes { get; } = [];
-
     public bool DisableRuntimeLibraryImport { get; set; }
     public bool MustImportRuntimeLibrary =>
         !DisableRuntimeLibraryImport
@@ -32,7 +22,16 @@ public sealed record SemanticModel(Tree Tree, DiagnosticBag Diagnostics, SymbolT
             )
         );
 
+    /// <summary>
+    /// Node IDs of reference-site nodes that originate from a non-intrinsic source file.
+    /// Populated by the resolver as references are recorded, so
+    /// <see cref="MustImportRuntimeLibrary"/> can filter references by the file of the
+    /// referencing node without holding on to the node instances themselves.
+    /// </summary>
+    internal HashSet<NodeId> NonIntrinsicReferenceNodes { get; } = [];
+    internal int RuntimeReferences = 0;
     internal TypeSolver TypeSolver { get; } = new(new DiagnosticBag());
+    private SymbolLookup DeclarationsByName => field ??= Declarations.Values.SelectMany(s => s).GroupBy(s => s.Name).ToDictionary(g => g.Key, g => g.ToList());
 
     public bool IsCompileTimeConstant(Expression expression) =>
         expression is Literal or NameOf
@@ -43,7 +42,7 @@ expression is Literal or NameOf
     public object? GetConstantValue(Expression expression) =>
         expression switch
         {
-            QualifiedName qn when GetType(qn.Identifier) is TypeChecking.Types.ObjectType objectType
+            QualifiedName qn when GetType(qn.Identifier) is ObjectType objectType
                 && objectType.GetProperty(qn.Names.First().Name.Text) is { ValueType: TypeChecking.Types.LiteralType literalType } =>
                 literalType.Value,
             _ when GetType(expression) is TypeChecking.Types.LiteralType literalType => literalType.Value,
@@ -114,12 +113,6 @@ public bool TryGetIntrinsicAttribute(Expression expression, string name, [MaybeN
     public Type GetType(Node node) => TypeSolver.GetType(node);
     public Type? GetDeclarationType(Node node) => GetSymbol(node) is { } symbol ? TypeSolver.GetType(symbol.Declaration) : null;
     public T? FindIntrinsicDeclarationSymbol<T>(string name) where T : Symbol => FindDeclarationSymbol<T>(name, s => s.IsIntrinsic);
-
-    private Dictionary<string, List<Symbol>>? _declarationsByName;
-
-    private Dictionary<string, List<Symbol>> DeclarationsByName =>
-        _declarationsByName ??= Declarations.Values.SelectMany(s => s).GroupBy(s => s.Name).ToDictionary(g => g.Key, g => g.ToList());
-
     private T? FindDeclarationSymbol<T>(string name) where T : Symbol => FindDeclarationSymbol<T>(name, static _ => true);
 
     private T? FindDeclarationSymbol<T>(string name, Func<T, bool> predicate) where T : Symbol =>

Loom.Core/TypeChecking/BinaryOperatorBinder.cs

@@ -1,5 +1,4 @@
 using Loom.Core.Parsing.AST;
-using Loom.Core.Resolving;
 using Loom.Core.Text;
 using PrimitiveType = Loom.Core.TypeChecking.Types.PrimitiveType;
 using Type = Loom.Core.TypeChecking.Types.Type;

Loom.Core/TypeChecking/Intrinsic/generated/None.loom

@@ -9448,15 +9448,15 @@ declare interface DragDetector: ClickDetector {
 	[luau_name("WorldSecondaryAxis")]
 	mut world_secondary_axis: Vector3;
 	[luau_name("AddConstraintFunction"), luau_method]
-	add_constraint_function: fn(priority: number, function: fn: void): ScriptConnection;
+	add_constraint_function: fn(priority: number, callback: fn: void): ScriptConnection;
 	[luau_name("GetReferenceFrame"), luau_method]
 	get_reference_frame: fn: CFrame;
 	[luau_name("RestartDrag"), luau_method]
 	restart_drag: fn: void;
 	[luau_name("SetDragStyleFunction"), luau_method]
-	set_drag_style_function: fn(function: fn: void): void;
+	set_drag_style_function: fn(callback: fn: void): void;
 	[luau_name("SetPermissionPolicyFunction"), luau_method]
-	set_permission_policy_function: fn(function: fn: void): void;
+	set_permission_policy_function: fn(callback: fn: void): void;
 	[luau_name("DragContinue")]
 	event drag_continue(playerWhoDragged: Player, cursorRay: Ray, viewFrame: CFrame, vrInputFrame: CFrame?, isModeSwitchKeyDown: bool);
 	[luau_name("DragEnd")]
@@ -14077,9 +14077,9 @@ declare interface Model: PVInstance {
 #: A logical grouping of datamodel nodes which can enable scripts to run in parallel. :#
 declare interface Actor: Model {
 	[luau_name("BindToMessage"), luau_method]
-	bind_to_message: fn(topic: string, function: fn: void): ScriptConnection;
+	bind_to_message: fn(topic: string, callback: fn: void): ScriptConnection;
 	[luau_name("BindToMessageParallel"), luau_method]
-	bind_to_message_parallel: fn(topic: string, function: fn: void): ScriptConnection;
+	bind_to_message_parallel: fn(topic: string, callback: fn: void): ScriptConnection;
 	[luau_name("SendMessage"), luau_method]
 	send_message: fn(topic: string, message: unknown): void;
 }
@@ -15199,9 +15199,9 @@ declare interface RunService: Instance {
 	[luau_name("FrameNumber")]
 	frame_number: number;
 	[luau_name("BindToRenderStep"), luau_method]
-	bind_to_render_step: fn(name: string, priority: number, function: fn: void): void;
+	bind_to_render_step: fn(name: string, priority: number, callback: fn: void): void;
 	[luau_name("BindToSimulation"), luau_method]
-	bind_to_simulation: fn(function: fn: void, frequency: Enum["StepFrequency"]?, priority: number?): ScriptConnection;
+	bind_to_simulation: fn(callback: fn: void, frequency: Enum["StepFrequency"]?, priority: number?): ScriptConnection;
 	[luau_name("GetPredictionStatus"), luau_method]
 	get_prediction_status: fn(context: Instance): Enum["PredictionStatus"];
 	[luau_name("IsClient"), luau_method]
@@ -15466,7 +15466,7 @@ declare interface DataModel: ServiceProvider {
 	[luau_name("RunService")]
 	run_service: RunService?;
 	[luau_name("BindToClose"), luau_method]
-	bind_to_close: fn(function: fn: void): void;
+	bind_to_close: fn(callback: fn: void): void;
 	[luau_name("GetMessage"), luau_method]
 	get_message: fn: string;
 	[luau_name("GetRemoteBuildMode"), luau_method]
@@ -16866,13 +16866,13 @@ declare interface UIDragDetector: UIComponent {
 	[luau_name("UIDragSpeedAxisMapping")]
 	mut ui_drag_speed_axis_mapping: Enum["UIDragSpeedAxisMapping"];
 	[luau_name("AddConstraintFunction"), luau_method]
-	add_constraint_function: fn(priority: number, function: fn: void): ScriptConnection;
+	add_constraint_function: fn(priority: number, callback: fn: void): ScriptConnection;
 	[luau_name("GetReferencePosition"), luau_method]
 	get_reference_position: fn: UDim2;
 	[luau_name("GetReferenceRotation"), luau_method]
 	get_reference_rotation: fn: number;
 	[luau_name("SetDragStyleFunction"), luau_method]
-	set_drag_style_function: fn(function: fn: void): void;
+	set_drag_style_function: fn(callback: fn: void): void;
 	[luau_name("DragContinue")]
 	event drag_continue(inputPosition: Vector2);
 	[luau_name("DragEnd")]
@@ -16928,7 +16928,7 @@ declare interface UIGridStyleLayout: UILayout {
 	[luau_name("ApplyLayout"), luau_method]
 	apply_layout: fn: void;
 	[luau_name("SetCustomSortFunction"), luau_method]
-	set_custom_sort_function: fn(function: fn: void?): void;
+	set_custom_sort_function: fn(callback: fn: void?): void;
 }
 
 declare interface UIGridLayout: UIGridStyleLayout {

Loom.Core/TypeChecking/Intrinsic/generated/PluginSecurity.loom

@@ -8294,7 +8294,7 @@ declare interface PluginGui: LayerCollector {
 	[luau_name("Title")]
 	mut title: string;
 	[luau_name("BindToClose"), luau_method]
-	bind_to_close: fn(function: fn: void?): void;
+	bind_to_close: fn(callback: fn: void?): void;
 	[luau_name("GetRelativeMousePosition"), luau_method]
 	get_relative_mouse_position: fn: Vector2;
 	[luau_name("PluginDragDropped")]

Loom.Core/TypeChecking/TypeChecker.cs

@@ -26,7 +26,7 @@
 public sealed partial class TypeChecker
     : Visitor<Type>
 {
     public static bool EmitDebugDiagnostics { get; set; }
 
     private readonly DiagnosticBag _diagnostics = new();
     private readonly Dictionary<Node, FlowState> _exitStates = [];
@@ -415,7 +415,7 @@
         return BindType(invocation, substitutedReturnType);
     }
 
     private Type CheckEventInvocation(Invocation invocation, InstantiatedType eventType)
     {
         var argumentList = invocation.Arguments.ArgumentList;
         var argumentTypes = argumentList.ConvertAll(Visit);
@@ -1089,17 +1089,32 @@
 
         var possibleReturnTypes = functionDeclaration.Body is ExpressionBody body
             ? [Visit(body)]
-            : functionDeclaration.Body
-                .GetDescendants<Return>()
-                .Where(returnStatement => returnStatement.FirstAncestorOfType<FunctionDeclaration>() == functionDeclaration)
-                .Select(Visit)
-                .ToList();
+            : GetOwnReturnStatements(functionDeclaration.Body).Select(Visit).ToList();
 
         return possibleReturnTypes.Count == 0
             ? Types.PrimitiveType.Void
             : TypeSimplifier.Simplify(new Types.UnionType(possibleReturnTypes));
     }
 
+    private static List<Return> GetOwnReturnStatements(Node node)
+    {
+        var result = new List<Return>();
+        CollectOwnReturnStatements(node, result);
+        return result;
+    }
+
+    private static void CollectOwnReturnStatements(Node node, List<Return> result)
+    {
+        foreach (var child in node.Children)
+        {
+            if (child is Return returnStatement)
+                result.Add(returnStatement);
+
+            if (child is not FunctionDeclaration)
+                CollectOwnReturnStatements(child, result);
+        }
+    }
+
     private Type ResolveHoistedType(Symbol symbol)
     {
         var type = GetTypeFromSymbol(symbol);
@@ -1178,7 +1193,7 @@
         return type;
     }
 
     private Type ReportCannotUseToIndex(Node node, Type objectType, Type indexType, string? cannotFindReason = "")
     {
         _diagnostics.Error(node, InternalCodes.InvalidAccess, $"Expression of type '{indexType}' cannot be used to index type '{objectType}'.{cannotFindReason}");
         return BindType(node, Types.PrimitiveType.Never);

Loom.Core/TypeChecking/TypeNarrower.cs

@@ -4,7 +4,6 @@
 using Loom.Core.Parsing.AST;
 using Loom.Core.Resolving;
 using Loom.Core.Text;
-using Loom.Core.TypeChecking.Types;
 using LiteralType = Loom.Core.TypeChecking.Types.LiteralType;
 using PrimitiveType = Loom.Core.TypeChecking.Types.PrimitiveType;
 using Type = Loom.Core.TypeChecking.Types.Type;

Loom.Luau/AST/PropertyAccess.cs

@@ -6,6 +6,21 @@ public class PropertyAccess(LuauExpression target, List<string> names) : LuauExp
     public LuauExpression Target { get; } = target;
     public List<string> Names { get; } = names;
 
-    public override string Render(RenderState state) =>
-        state.ParenthesizeIfNeeded(Target) + (Names.Count > 1 ? '.' : "") + string.Join('.', Names[..^1]) + Operator + Names.Last();
+    public override string Render(RenderState state)
+    {
+        var result = state.ParenthesizeIfNeeded(Target);
+        for (var i = 0; i < Names.Count; i++)
+        {
+            var name = Names[i];
+            if (LuauFactory.Keywords.Contains(name))
+            {
+                result += $"[{RenderState.StringDelimiter}{RenderState.Escape(name)}{RenderState.StringDelimiter}]";
+                continue;
+            }
+
+            result += (i == Names.Count - 1 ? Operator : '.') + name;
+        }
+
+        return result;
+    }
 }

Loom.Luau/LuauFactory.cs

@@ -8,6 +8,13 @@
 
     public static readonly Identifier Self = new("self");
 
+    public static readonly HashSet<string> Keywords =
+    [
+        "and", "break", "do", "else", "elseif", "end", "false", "for", "function",
+        "if", "in", "local", "nil", "not", "or", "repeat", "return", "then",
+        "true", "until", "while", "continue"
+    ];
+
     public static QualifiedTypeName QualifyRuntimeType(TypeName typeName) => new([RuntimeImportName], typeName);
 
     public static Call Bit32Call(string name, List<LuauExpression> arguments) => LibraryCall("bit32", [name], arguments);
@@ -16,7 +23,7 @@
     public static Call TableCall(string name, List<LuauExpression> arguments) => LibraryCall("table", [name], arguments);
     public static Call StringCall(string name, List<LuauExpression> arguments) => LibraryCall("string", [name], arguments);
     public static Call TaskCall(string name, List<LuauExpression> arguments) => LibraryCall("task", [name], arguments);
     public static Call RequireCall(string path) => new(new Identifier("require"), [new StringLiteral(path)]);
     public static Call SetMetatableCall(Table main, LuauExpression meta) => new(new Identifier("setmetatable"), [main, meta]);
 
     public static Call RuntimeLibraryCall(List<string> path, List<LuauExpression> arguments) =>

Loom.Testing/AstInspectorTest.cs

@@ -158,7 +158,7 @@
     [Fact]
     public void InspectTree_RendersStatements()
     {
-        var lexerResult = new LexerResult(SourceFile.Empty, [], [], new DiagnosticBag());
+        var lexerResult = new LexerResult([], [], new DiagnosticBag());
         var tree = new Tree(lexerResult, [new InspectorTestStatement { Name = "first" }, new InspectorTestStatement { Name = "second" }]);
 
         var inspectionResult = InspectTree(tree);
@@ -171,11 +171,11 @@
     private sealed class InspectorTestNode()
         : Node([], [])
     {
         public string? Name { get; set; }
         public Token? MyToken { get; set; }
         public InspectorTestNode? Child { get; set; }
         public IReadOnlyList<object>? Items { get; set; }
         public int Number { get; set; }
 
         public override T Accept<T>(Visitor<T> visitor) => throw new NotImplementedException();
     }
@@ -183,7 +183,7 @@
     private sealed class InspectorTestType(string typeName)
         : TypeExpression([], [])
     {
         public string TypeName { get; } = typeName;
 
         public override T Accept<T>(Visitor<T> visitor) => throw new NotImplementedException();
     }
@@ -191,7 +191,7 @@
     private sealed class InspectorTestStatement()
         : Statement([], [])
     {
         public string? Name { get; set; }
 
         public override T Accept<T>(Visitor<T> visitor) => throw new NotImplementedException();
     }

Loom.Testing/CompilationUnitTest.cs

@@ -24,8 +24,9 @@ public void Compiles_Project_NoEmit()
         var path = config.Files.OutputDirectory;
         Directory.Delete(path, true);
         Directory.CreateDirectory(path);
+        File.Create(Path.Combine(path, ".gitkeep"));
 
-        var luauFiles = Directory.EnumerateFiles(path);
+        var luauFiles = Directory.EnumerateFiles(path, "*.luau", SearchOption.TopDirectoryOnly);
         Assert.Empty(luauFiles);
     }