Add syntax for multi-level break. (#2431)

* Add syntax for multi-level break.

* Fix.

* Fix.

Co-authored-by: Yong He <yhe@nvidia.com>

8 files changed, 119 insertions, 20 deletions

diff --git a/source/slang/slang-ast-iterator.h b/source/slang/slang-ast-iterator.h
index a145d759b..257a910ea 100644
--- a/source/slang/slang-ast-iterator.h
+++ b/source/slang/slang-ast-iterator.h
@@ -297,6 +297,12 @@ struct ASTIterator
                 dispatchIfNotNull(stmt);
         }
 
+        void visitLabelStmt(LabelStmt* stmt)
+        {
+            iterator->maybeDispatchCallback(stmt);
+            dispatchIfNotNull(stmt->innerStmt);
+        }
+
         void visitBreakStmt(BreakStmt* stmt) { iterator->maybeDispatchCallback(stmt); }
 
         void visitContinueStmt(ContinueStmt* stmt) { iterator->maybeDispatchCallback(stmt); }
diff --git a/source/slang/slang-ast-stmt.h b/source/slang/slang-ast-stmt.h
index 2edc844c6..b18eb077d 100644
--- a/source/slang/slang-ast-stmt.h
+++ b/source/slang/slang-ast-stmt.h
@@ -23,6 +23,15 @@ class SeqStmt : public Stmt
     List<Stmt*> stmts;
 };
 
+// A statement with a label.
+class LabelStmt : public Stmt
+{
+    SLANG_AST_CLASS(LabelStmt)
+
+    Token label;
+    Stmt* innerStmt;
+};
+
 // The simplest kind of scope statement: just a `{...}` block
 class BlockStmt : public ScopeStmt 
 {
@@ -196,6 +205,8 @@ class JumpStmt : public ChildStmt
 class BreakStmt : public JumpStmt 
 {
     SLANG_AST_CLASS(BreakStmt)
+
+    Token targetLabel;
 };
 
 class ContinueStmt : public JumpStmt 
diff --git a/source/slang/slang-check-impl.h b/source/slang/slang-check-impl.h
index b926907e2..433c6ca70 100644
--- a/source/slang/slang-check-impl.h
+++ b/source/slang/slang-check-impl.h
@@ -1849,12 +1849,16 @@ namespace Slang
         template<typename T>
         T* FindOuterStmt();
 
+        Stmt* findOuterStmtWithLabel(Name* label);
+
         void visitDeclStmt(DeclStmt* stmt);
 
         void visitBlockStmt(BlockStmt* stmt);
 
         void visitSeqStmt(SeqStmt* stmt);
 
+        void visitLabelStmt(LabelStmt* stmt);
+
         void visitBreakStmt(BreakStmt *stmt);
 
         void visitContinueStmt(ContinueStmt *stmt);
diff --git a/source/slang/slang-check-stmt.cpp b/source/slang/slang-check-stmt.cpp
index a25a8683d..0f450340f 100644
--- a/source/slang/slang-check-stmt.cpp
+++ b/source/slang/slang-check-stmt.cpp
@@ -67,6 +67,12 @@ namespace Slang
         }
     }
 
+    void SemanticsStmtVisitor::visitLabelStmt(LabelStmt* stmt)
+    {
+        WithOuterStmt subContext(this, stmt);
+        subContext.checkStmt(stmt->innerStmt);
+    }
+
     void SemanticsStmtVisitor::checkStmt(Stmt* stmt)
     {
         SemanticsVisitor::checkStmt(stmt, *this);
@@ -85,14 +91,51 @@ namespace Slang
         return nullptr;
     }
 
+    Stmt* SemanticsStmtVisitor::findOuterStmtWithLabel(Name* label)
+    {
+        for (auto outerStmtInfo = m_outerStmts; outerStmtInfo; outerStmtInfo = outerStmtInfo->next)
+        {
+            auto outerStmt = outerStmtInfo->stmt;
+            auto found = as<LabelStmt>(outerStmt);
+            if (found)
+            {
+                if (found->label.getName() == label)
+                {
+                    return found->innerStmt;
+                }
+            }
+        }
+        return nullptr;
+    }
+
     void SemanticsStmtVisitor::visitBreakStmt(BreakStmt *stmt)
     {
-        auto outer = FindOuterStmt<BreakableStmt>();
-        if (!outer)
+        Stmt* targetStmt = nullptr;
+        if (stmt->targetLabel.type == TokenType::Identifier)
         {
-            getSink()->diagnose(stmt, Diagnostics::breakOutsideLoop);
+            // This is a break statement with an explicit target label.
+            // Try to find the outer stmt with the label.
+            targetStmt = findOuterStmtWithLabel(stmt->targetLabel.getName());
+            if (!targetStmt)
+            {
+                getSink()->diagnose(stmt, Diagnostics::breakLabelNotFound, stmt->targetLabel.getName());
+            }
+            if (!as<BreakableStmt>(targetStmt))
+            {
+                getSink()->diagnose(stmt, Diagnostics::targetLabelDoesNotMarkBreakableStmt, stmt->targetLabel.getName());
+            }
         }
-        stmt->parentStmt = outer;
+        else
+        {
+            // For `break` statements without an explicit target,
+            // find the inner most breakable stmt.
+            targetStmt = FindOuterStmt<BreakableStmt>();
+            if (!targetStmt)
+            {
+                getSink()->diagnose(stmt, Diagnostics::breakOutsideLoop);
+            }
+        }
+        stmt->parentStmt = targetStmt;
     }
 
     void SemanticsStmtVisitor::visitContinueStmt(ContinueStmt *stmt)
diff --git a/source/slang/slang-diagnostic-defs.h b/source/slang/slang-diagnostic-defs.h
index 9abb4941a..d7e56309a 100644
--- a/source/slang/slang-diagnostic-defs.h
+++ b/source/slang/slang-diagnostic-defs.h
@@ -235,7 +235,7 @@ DIAGNOSTIC(20014, Error, classIsReservedKeyword, "'class' is a reserved keyword
 // 3xxxx - Semantic analysis
 //
 DIAGNOSTIC(30002, Error, divideByZero, "divide by zero")
-DIAGNOSTIC(30003, Error, breakOutsideLoop, "'break' must appear inside loop constructs.")
+DIAGNOSTIC(30003, Error, breakOutsideLoop, "'break' must appear inside loop or switch constructs.")
 DIAGNOSTIC(30004, Error, continueOutsideLoop, "'continue' must appear inside loop constructs.")
 DIAGNOSTIC(30005, Error, whilePredicateTypeError, "'while': expression must evaluate to int.")
 DIAGNOSTIC(30006, Error, ifPredicateTypeError,  "'if': expression must evaluate to int.")
@@ -272,6 +272,8 @@ DIAGNOSTIC(30050, Error,  mutatingMethodOnImmutableValue, "mutating method '$0'
 
 DIAGNOSTIC(30051, Error, invalidValueForArgument, "invalid value for argument '$0'")
 DIAGNOSTIC(30052, Error, invalidSwizzleExpr, "invalid swizzle pattern '$0' on type '$1'")
+DIAGNOSTIC(30053, Error, breakLabelNotFound, "label '$0' used as break target is not found.")
+DIAGNOSTIC(30054, Error, targetLabelDoesNotMarkBreakableStmt, "invalid break target: statement labeled '$0' is not breakable.")
 DIAGNOSTIC(30043, Error, getStringHashRequiresStringLiteral, "getStringHash parameter can only accept a string literal")
 
 DIAGNOSTIC(30060, Error, expectedAType, "expected a type, got a '$0'")
diff --git a/source/slang/slang-language-server-ast-lookup.cpp b/source/slang/slang-language-server-ast-lookup.cpp
index a44c7a7d9..cd211b0f5 100644
--- a/source/slang/slang-language-server-ast-lookup.cpp
+++ b/source/slang/slang-language-server-ast-lookup.cpp
@@ -451,6 +451,13 @@ struct ASTLookupStmtVisitor : public StmtVisitor<ASTLookupStmtVisitor, bool>
         return false;
     }
 
+    bool visitLabelStmt(LabelStmt* stmt)
+    {
+        if (_isLocInRange(context, stmt->label.loc, stmt->label.getContent().getLength()))
+            return true;
+        return dispatchIfNotNull(stmt->innerStmt);
+    }
+
     bool visitBreakStmt(BreakStmt*) { return false; }
 
     bool visitContinueStmt(ContinueStmt*) { return false; }
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 753e0d6e9..e2ca9a711 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -4508,6 +4508,11 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
         SLANG_UNEXPECTED("`case` or `default` not under `switch`");
     }
 
+    void visitLabelStmt(LabelStmt* stmt)
+    {
+        lowerStmt(context, stmt->innerStmt);
+    }
+
     void visitCompileTimeForStmt(CompileTimeForStmt* stmt)
     {
         // The user is asking us to emit code for the loop
diff --git a/source/slang/slang-parser.cpp b/source/slang/slang-parser.cpp
index 55b6afd6a..c0c035211 100644
--- a/source/slang/slang-parser.cpp
+++ b/source/slang/slang-parser.cpp
@@ -170,21 +170,22 @@ namespace Slang
         bool LookAheadToken(TokenType type, int offset);
         bool LookAheadToken(const char* string, int offset);
 
-        void                                        parseSourceFile(ModuleDecl* program);
-        Decl*					ParseStruct();
-        ClassDecl*					    ParseClass();
-        Stmt*					ParseStatement();
-        Stmt*			        parseBlockStatement();
-        DeclStmt*			parseVarDeclrStatement(Modifiers modifiers);
-        IfStmt*				parseIfStatement();
-        ForStmt*				ParseForStatement();
-        WhileStmt*			ParseWhileStatement();
-        DoWhileStmt*			ParseDoWhileStatement();
-        BreakStmt*			ParseBreakStatement();
-        ContinueStmt*			ParseContinueStatement();
-        ReturnStmt*			ParseReturnStatement();
-        ExpressionStmt*		ParseExpressionStatement();
-        Expr*				ParseExpression(Precedence level = Precedence::Comma);
+        void parseSourceFile(ModuleDecl* program);
+        Decl* ParseStruct();
+        ClassDecl* ParseClass();
+        Stmt* ParseStatement();
+        Stmt* parseBlockStatement();
+        Stmt* parseLabelStatement();
+        DeclStmt* parseVarDeclrStatement(Modifiers modifiers);
+        IfStmt* parseIfStatement();
+        ForStmt* ParseForStatement();
+        WhileStmt* ParseWhileStatement();
+        DoWhileStmt* ParseDoWhileStatement();
+        BreakStmt* ParseBreakStatement();
+        ContinueStmt* ParseContinueStatement();
+        ReturnStmt* ParseReturnStatement();
+        ExpressionStmt* ParseExpressionStatement();
+        Expr* ParseExpression(Precedence level = Precedence::Comma);
 
         // Parse an expression that might be used in an initializer or argument context, so we should avoid operator-comma
         inline Expr*			ParseInitExpr() { return ParseExpression(Precedence::Assignment); }
@@ -4313,6 +4314,12 @@ namespace Slang
         }
         else if (LookAheadToken(TokenType::Identifier))
         {
+            if (LookAheadToken(TokenType::Colon, 1))
+            {
+                // An identifier followed by an ":" is a label.
+                return parseLabelStatement();
+            }
+
             // We might be looking at a local declaration, or an
             // expression statement, and we need to figure out which.
             //
@@ -4484,6 +4491,16 @@ namespace Slang
         return blockStatement;
     }
 
+    Stmt* Parser::parseLabelStatement()
+    {
+        LabelStmt* stmt = astBuilder->create<LabelStmt>();
+        FillPosition(stmt);
+        stmt->label = ReadToken(TokenType::Identifier);
+        ReadToken(TokenType::Colon);
+        stmt->innerStmt = ParseStatement();
+        return stmt;
+    }
+
     DeclStmt* Parser::parseVarDeclrStatement(
         Modifiers modifiers)
     {
@@ -4604,6 +4621,10 @@ namespace Slang
         BreakStmt* breakStatement = astBuilder->create<BreakStmt>();
         FillPosition(breakStatement);
         ReadToken("break");
+        if (LookAheadToken(TokenType::Identifier))
+        {
+            breakStatement->targetLabel = ReadToken();
+        }
         ReadToken(TokenType::Semicolon);
         return breakStatement;
     }