Implement EnumerateDF to visit all sub-statements

Improve Switch-Case

Author: sor

backend/Core/Stmt.cs

@@ -29,6 +29,14 @@ public bool IsAttrib( string key, string value )
 			&& attribs.TryGetValue( key, out Strings values )
 			&& values.Contains( value );
 
+		// Enumerate (depth first) through all contained Stmt and itself
+		// Only overloaded in Stmt which contain more Stmt itself
+		// Filter results with e.g. EnumerateDF.OfType<ReturnStmt>()
+		[Pure]
+		public virtual IEnumerable<Stmt> EnumerateDF {
+			get { yield return this; }
+		}
+
 		public virtual void AssignAttribs( Attribs inAttribs )
 		{
 			attribs = inAttribs;
@@ -189,6 +197,21 @@ public class IfStmt : Stmt
 		public List<CondThen> ifThens;
 		public Stmt?          elseStmt;
 
+		[Pure]
+		public override IEnumerable<Stmt> EnumerateDF {
+			get {
+				foreach( CondThen ifThen in ifThens )
+					foreach( Stmt subStmt in ifThen.thenStmt.EnumerateDF )
+						yield return subStmt;
+
+				if( elseStmt != null )
+					foreach( Stmt subStmt in elseStmt.EnumerateDF )
+						yield return subStmt;
+
+				yield return this;
+			}
+		}
+
 		public override Strings Gen( int level )
 		{
 			Strings ret    = new();
@@ -211,10 +234,10 @@ public override Strings Gen( int level )
 		}
 	}
 
-	public class CaseStmt : Stmt
+	public struct CaseStmt
 	{
-		public List<Expr> caseExprs;
-		public List<Stmt> bodyStmts;
+		public List<Expr> caseExprs; // can have multiple ORed conditions
+		public MultiStmt  bodyStmt;
 		public bool       autoBreak;
 	}
 
@@ -224,6 +247,22 @@ public class SwitchStmt : Stmt
 		public List<CaseStmt> caseStmts;
 		public List<Stmt>?    elseStmts;
 
+		[Pure]
+		public override IEnumerable<Stmt> EnumerateDF {
+			get {
+				foreach( CaseStmt caseStmt in caseStmts )
+					foreach( Stmt subStmt in caseStmt.bodyStmt.EnumerateDF )
+						yield return subStmt;
+
+				if( elseStmts != null )
+					foreach( Stmt elseStmt in elseStmts )
+						foreach( Stmt subStmt in elseStmt.EnumerateDF )
+							yield return subStmt;
+
+				yield return this;
+			}
+		}
+
 		public override Strings Gen( int level )
 		{
 			throw new NotImplementedException();
@@ -235,6 +274,16 @@ public class LoopStmt : Stmt
 	{
 		public Stmt bodyStmt;
 
+		[Pure]
+		public override IEnumerable<Stmt> EnumerateDF {
+			get {
+				foreach( Stmt subStmt in bodyStmt.EnumerateDF )
+					yield return subStmt;
+
+				yield return this;
+			}
+		}
+
 		public override Strings Gen( int level )
 		{
 			Strings ret    = new();
@@ -255,6 +304,21 @@ public class ForStmt : LoopStmt
 		public Expr  iterExpr;
 		public Stmt? elseStmt;
 
+		[Pure]
+		public override IEnumerable<Stmt> EnumerateDF {
+			get {
+				foreach( Stmt subStmt in initStmt.EnumerateDF )
+					yield return subStmt;
+
+				if( elseStmt != null )
+					foreach( Stmt subStmt in elseStmt.EnumerateDF )
+						yield return subStmt;
+
+				foreach( Stmt baseStmt in base.EnumerateDF )
+					yield return baseStmt;
+			}
+		}
+
 		public override Strings Gen( int level )
 		{
 			if( elseStmt != null )
@@ -280,6 +344,18 @@ public class WhileStmt : LoopStmt
 		public Expr  condExpr;
 		public Stmt? elseStmt;
 
+		[Pure]
+		public override IEnumerable<Stmt> EnumerateDF {
+			get {
+				if( elseStmt != null )
+					foreach( Stmt subStmt in elseStmt.EnumerateDF )
+						yield return subStmt;
+
+				foreach( Stmt baseStmt in base.EnumerateDF )
+					yield return baseStmt;
+			}
+		}
+
 		public override Strings Gen( int level )
 		{
 			if( elseStmt != null )
@@ -347,6 +423,17 @@ public class MultiStmt : Stmt
 		// TODO merge this with Block?
 		// public bool       IsBlock { get; init; }
 
+		[Pure]
+		public override IEnumerable<Stmt> EnumerateDF {
+			get {
+				foreach( Stmt stmt in stmts )
+					foreach( Stmt subStmt in stmt.EnumerateDF )
+						yield return subStmt;
+
+				yield return this;
+			}
+		}
+
 		public override void AssignAttribs( Attribs inAttribs )
 		{
 			stmts.ForEach( v => v.AssignAttribs( inAttribs ) );
@@ -366,6 +453,17 @@ public class Block : Stmt
 	{
 		public List<Stmt> stmts;
 
+		[Pure]
+		public override IEnumerable<Stmt> EnumerateDF {
+			get {
+				foreach( Stmt stmt in stmts )
+					foreach( Stmt subStmt in stmt.EnumerateDF )
+						yield return subStmt;
+
+				yield return this;
+			}
+		}
+
 		// only overriden here
 		public override Strings GenWithoutCurly( int level )
 		{

backend/Visitor/VStmt.cs

@@ -166,10 +166,21 @@ public override Stmt VisitIfStmt( IfStmtContext c )
 
 		public new CaseStmt VisitCaseStmt( CaseStmtContext c )
 		{
+			MultiStmt multiStmt = new() {
+				srcPos = c.ToSrcPos(),
+				stmts  = c.levStmt().Select( Visit ).ToList(),
+			};
+			bool isFall = c.FALL() != null;
+			if( !isFall )
+				multiStmt.stmts.Append( new BreakStmt { depth = 1 } );
+			// TODO: else C++ [[fallthrough]]
+			//else
+			//	multiStmt.stmts.Append( new FallStmt() );
+			//	multiStmt.stmts.Append( new FreetextStmt( "[[fallthrough]]" ) );
+
 			CaseStmt ret = new() {
-				srcPos    = c.ToSrcPos(),
 				caseExprs = c.expr().Select( q => q.Visit() ).ToList(),
-				bodyStmts = c.levStmt().Select( Visit ).ToList(),
+				bodyStmt  = multiStmt,
 				autoBreak = c.FALL() != null,
 			};
 			return ret;