Add defer statement (#6619)

1 files changed, 93 insertions, 3 deletions

diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index e3c4ddf05..e6ec68660 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -19,6 +19,7 @@
 #include "slang-ir-insert-debug-value-store.h"
 #include "slang-ir-insts.h"
 #include "slang-ir-loop-inversion.h"
+#include "slang-ir-lower-defer.h"
 #include "slang-ir-lower-error-handling.h"
 #include "slang-ir-lower-expand-type.h"
 #include "slang-ir-missing-return.h"
@@ -593,6 +594,9 @@ struct IRGenContext
     // The element index if we are inside an `expand` expression.
     IRInst* expandIndex = nullptr;
 
+    // The current scope end for use with `defer`.
+    IRBlock* scopeEndBlock = nullptr;
+
     // Callback function to call when after lowering a type.
     std::function<IRType*(IRGenContext* context, Type* type, IRType* irType)> lowerTypeCallback =
         nullptr;
@@ -6021,6 +6025,53 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
         startBlock();
     }
 
+    /// Create a new scope end block and return the previous one.
+    ///
+    /// This is needed for `defer` to be aware of scopes. `preallocated` can
+    /// be specified if you already have a block at the end of the scope, like
+    /// in `for` loops.
+    IRBlock* pushScopeBlock(IRBlock* preallocated = nullptr)
+    {
+        IRBlock* prevScopeEndBlock = context->scopeEndBlock;
+
+        auto builder = getBuilder();
+        context->scopeEndBlock = preallocated ? preallocated : builder->createBlock();
+        return prevScopeEndBlock;
+    }
+
+    /// Pop the current scope end block and restore the previous one.
+    ///
+    /// This is needed for `defer` to be aware of scopes. `previous` should be
+    /// the block returned from the corresponding pushScopeBlock. `preallocated`
+    /// should be true if the corresponding pushScopeBlock was given a block
+    /// as a parameter.
+    void popScopeBlock(IRBlock* previous, bool preallocated)
+    {
+        if (!preallocated)
+        {
+            // If pushScopeBlock actually created the block, we have to insert
+            // or deallocate it here. Otherwise, we assume that the caller
+            // handles the end block.
+            auto builder = getBuilder();
+            if (context->scopeEndBlock->hasUses())
+            {
+                // The end of the scope was referenced, so we need to actually
+                // keep it around and jump through it.
+                // Move the terminator to the scope end block.
+                emitBranchIfNeeded(context->scopeEndBlock);
+                builder->insertBlock(context->scopeEndBlock);
+                builder->setInsertInto(context->scopeEndBlock);
+            }
+            else
+            {
+                // Scope end block was left unused, so we may as well delete it.
+                context->scopeEndBlock->removeAndDeallocate();
+            }
+        }
+
+        context->scopeEndBlock = previous;
+    }
+
     void visitIfStmt(IfStmt* stmt)
     {
         auto builder = getBuilder();
@@ -6043,11 +6094,13 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
             ifInst = builder->emitIfElse(irCond, thenBlock, elseBlock, afterBlock);
 
             insertBlock(thenBlock);
+            IRBlock* prevScopeEndBlock = pushScopeBlock(afterBlock);
             lowerStmt(context, thenStmt);
             emitBranchIfNeeded(afterBlock);
 
             insertBlock(elseBlock);
             lowerStmt(context, elseStmt);
+            popScopeBlock(prevScopeEndBlock, true);
 
             insertBlock(afterBlock);
         }
@@ -6059,7 +6112,10 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
             ifInst = builder->emitIf(irCond, thenBlock, afterBlock);
 
             insertBlock(thenBlock);
+
+            IRBlock* prevScopeEndBlock = pushScopeBlock(afterBlock);
             lowerStmt(context, thenStmt);
+            popScopeBlock(prevScopeEndBlock, true);
 
             insertBlock(afterBlock);
         }
@@ -6150,7 +6206,9 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
 
         // Emit the body of the loop
         insertBlock(bodyLabel);
+        IRBlock* prevScopeEndBlock = pushScopeBlock(continueLabel);
         lowerStmt(context, stmt->statement);
+        popScopeBlock(prevScopeEndBlock, true);
 
         if (auto inferredMaxIters = stmt->findModifier<InferredMaxItersAttribute>())
         {
@@ -6256,7 +6314,9 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
 
         // Emit the body of the loop
         insertBlock(bodyLabel);
+        IRBlock* prevScopeEndBlock = pushScopeBlock(continueLabel);
         lowerStmt(context, stmt->statement);
+        popScopeBlock(prevScopeEndBlock, true);
 
         // At the end of the body we need to jump back to the top.
         emitBranchIfNeeded(loopHead);
@@ -6300,7 +6360,9 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
         insertBlock(loopHead);
 
         // Emit the body of the loop
+        IRBlock* prevScopeEndBlock = pushScopeBlock(continueLabel);
         lowerStmt(context, stmt->statement);
+        popScopeBlock(prevScopeEndBlock, true);
 
         insertBlock(testLabel);
 
@@ -6429,10 +6491,12 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
 
     void visitBlockStmt(BlockStmt* stmt)
     {
-        // To lower a block (scope) statement,
-        // just lower its body. The IR doesn't
-        // need to reflect the scoping of the AST.
+        IRBlock* prevScopeEndBlock = pushScopeBlock(nullptr);
+
+        // To lower a block (scope) statement, just lower its body.
         lowerStmt(context, stmt->body);
+
+        popScopeBlock(prevScopeEndBlock, false);
     }
 
     void visitReturnStmt(ReturnStmt* stmt)
@@ -6523,6 +6587,29 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
         }
     }
 
+    void visitDeferStmt(DeferStmt* stmt)
+    {
+        auto builder = getBuilder();
+        startBlockIfNeeded(stmt);
+
+        IRBlock* deferBlock = builder->createBlock();
+        IRBlock* mergeBlock = builder->createBlock();
+
+        builder->emitDefer(deferBlock, mergeBlock, context->scopeEndBlock);
+
+        builder->insertBlock(deferBlock);
+        builder->setInsertInto(deferBlock);
+
+        IRBlock* prevScopeEndBlock = pushScopeBlock(mergeBlock);
+        lowerStmt(context, stmt->statement);
+        popScopeBlock(prevScopeEndBlock, true);
+
+        builder->emitBranch(mergeBlock);
+
+        builder->insertBlock(mergeBlock);
+        builder->setInsertInto(mergeBlock);
+    }
+
     void visitDiscardStmt(DiscardStmt* stmt)
     {
         startBlockIfNeeded(stmt);
@@ -11714,6 +11801,9 @@ RefPtr<IRModule> generateIRForTranslationUnit(
     // normal `call` + `ifElse`, etc.
     lowerErrorHandling(module, compileRequest->getSink());
 
+    // Lower `defer` so that later passes need not be aware of it.
+    lowerDefer(module, compileRequest->getSink());
+
     // Synthesize some code we want to make sure is inlined and simplified
     synthesizeBitFieldAccessors(module);