defers

compiler/NubLang/Generation/QBE/QBEGenerator.cs

@@ -498,6 +498,9 @@ public class QBEGenerator
                 EmitScopeCleanup();
                 _writer.Indented($"jmp {_continueLabels.Peek()}");
                 break;
+            case DeferNode defer:
+                Scope.Defers.Push(defer);
+                break;
             case IfNode ifStatement:
                 EmitIf(ifStatement);
                 break;
@@ -523,6 +526,11 @@ public class QBEGenerator
 
     private void EmitScopeCleanup()
     {
+        while (Scope.Defers.TryPop(out var defer))
+        {
+            EmitStatement(defer.Statement);
+        }
+
         while (Scope.Variables.TryPop(out var variable))
         {
             if (variable.Type is NubStructType structType)
@@ -1283,6 +1291,7 @@ public class QBEGenerator
 public class Scope
 {
     public readonly Stack<Variable> Variables = [];
+    public readonly Stack<DeferNode> Defers = [];
 }
 
 public record Variable(string Name, NubType Type);

compiler/NubLang/TypeChecking/Node/StatementNode.cs

@@ -23,3 +23,5 @@ public record ContinueNode : TerminalStatementNode;
 public record BreakNode : TerminalStatementNode;
 
 public record WhileNode(ExpressionNode Condition, BlockNode Body) : StatementNode;
+
+public record DeferNode(StatementNode Statement) : StatementNode;

compiler/NubLang/TypeChecking/TypeChecker.cs

@@ -1,5 +1,4 @@
 using System.Diagnostics;
-using System.Diagnostics.CodeAnalysis;
 using System.Globalization;
 using NubLang.Diagnostics;
 using NubLang.Modules;
@@ -764,12 +763,6 @@ public sealed class TypeChecker
         BeginScope(false);
 
         foreach (var statement in node.Statements)
-        {
-            if (statement is DeferSyntax deferStmt)
-            {
-                CurrentScope.PushDefer(deferStmt);
-            }
-            else
         {
             var checkedStatement = CheckStatement(statement);
 
@@ -777,21 +770,6 @@ public sealed class TypeChecker
             {
                 if (checkedStatement is TerminalStatementNode)
                 {
-                        while (CurrentScope.TryPopDefer(out var defer))
-                        {
-                            var checkedDeferredStatement = CheckStatement(defer.Statement);
-                            if (checkedDeferredStatement is TerminalStatementNode)
-                            {
-                                throw new TypeCheckerException(Diagnostic
-                                    .Error("Cannot defer terminal statements")
-                                    .WithHelp("Avoid using defer with terminal statements such as 'return', 'break', 'continue' etc.")
-                                    .At(defer)
-                                    .Build());
-                            }
-
-                            statements.Add(checkedDeferredStatement);
-                        }
-
                     reachable = false;
                 }
 
@@ -806,22 +784,6 @@ public sealed class TypeChecker
                 }
             }
         }
-        }
-
-        while (CurrentScope.TryPopDefer(out var defer))
-        {
-            var checkedDeferredStatement = CheckStatement(defer.Statement);
-            if (checkedDeferredStatement is TerminalStatementNode)
-            {
-                throw new TypeCheckerException(Diagnostic
-                    .Error("Cannot defer terminal statements")
-                    .WithHelp("Avoid using defer with terminal statements such as 'return', 'break', 'continue' etc.")
-                    .At(defer)
-                    .Build());
-            }
-
-            statements.Add(checkedDeferredStatement);
-        }
 
         EndScope();
 
@@ -841,7 +803,7 @@ public sealed class TypeChecker
             StatementExpressionSyntax stmtExpr => CheckStatementExpression(stmtExpr),
             VariableDeclarationSyntax varDeclStmt => CheckVariableDeclaration(varDeclStmt),
             WhileSyntax whileStmt => CheckWhile(whileStmt),
-            DeferSyntax => throw new InvalidOperationException($"Compiler bug: defer statement should not have reached {nameof(CheckStatement)}"),
+            DeferSyntax defer => new DeferNode(CheckStatement(defer.Statement)),
             _ => throw new ArgumentOutOfRangeException(nameof(statement))
         };
     }
@@ -957,7 +919,6 @@ public record Variable(string Name, NubType Type, VariableKind Kind);
 
 public class Scope(Scope? parent = null)
 {
-    private readonly Stack<DeferSyntax> _defers = [];
     private readonly List<Variable> _variables = [];
 
     public Variable? LookupVariable(string name)
@@ -976,16 +937,6 @@ public class Scope(Scope? parent = null)
         _variables.Add(variable);
     }
 
-    public void PushDefer(DeferSyntax defer)
-    {
-        _defers.Push(defer);
-    }
-
-    public bool TryPopDefer([NotNullWhen(true)] out DeferSyntax? defer)
-    {
-        return _defers.TryPop(out defer);
-    }
-
     public Scope SubScope()
     {
         return new Scope(this);