writer/wgsl: Handle for-loops

Bug: tint:952
Change-Id: I712a0b448f187f02ef97f7f9b3887e00f674716a
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/56763
Kokoro: Kokoro <noreply+kokoro@google.com>
Auto-Submit: Ben Clayton <bclayton@google.com>
Reviewed-by: David Neto <dneto@google.com>

src/writer/wgsl/generator_impl.cc

@@ -75,9 +75,11 @@ bool GeneratorImpl::Generate() {
         return false;
       }
     } else if (auto* var = decl->As<ast::Variable>()) {
+      make_indent();
       if (!EmitVariable(var)) {
         return false;
       }
+      out_ << ";" << std::endl;
     } else {
       TINT_UNREACHABLE(Writer, diagnostics_);
       return false;
@@ -575,8 +577,6 @@ bool GeneratorImpl::EmitStructType(const ast::Struct* str) {
 }
 
 bool GeneratorImpl::EmitVariable(ast::Variable* var) {
-  make_indent();
-
   if (!var->decorations().empty()) {
     if (!EmitDecorations(var->decorations())) {
       return false;
@@ -617,7 +617,6 @@ bool GeneratorImpl::EmitVariable(ast::Variable* var) {
       return false;
     }
   }
-  out_ << ";" << std::endl;
 
   return true;
 }
@@ -819,16 +818,6 @@ bool GeneratorImpl::EmitBlock(const ast::BlockStatement* stmt) {
   return true;
 }
 
-bool GeneratorImpl::EmitIndentedBlockAndNewline(
-    const ast::BlockStatement* stmt) {
-  make_indent();
-  const bool result = EmitBlock(stmt);
-  if (result) {
-    out_ << std::endl;
-  }
-  return result;
-}
-
 bool GeneratorImpl::EmitBlockAndNewline(const ast::BlockStatement* stmt) {
   const bool result = EmitBlock(stmt);
   if (result) {
@@ -838,22 +827,32 @@ bool GeneratorImpl::EmitBlockAndNewline(const ast::BlockStatement* stmt) {
 }
 
 bool GeneratorImpl::EmitStatement(ast::Statement* stmt) {
+  make_indent();
+
+  if (!EmitRawStatement(stmt)) {
+    return false;
+  }
+
+  if (!stmt->IsAnyOf<ast::BlockStatement, ast::IfStatement,
+                     ast::SwitchStatement, ast::LoopStatement,
+                     ast::ForLoopStatement>()) {
+    out_ << ";" << std::endl;
+  }
+  return true;
+}
+
+bool GeneratorImpl::EmitRawStatement(ast::Statement* stmt) {
   if (auto* a = stmt->As<ast::AssignmentStatement>()) {
     return EmitAssign(a);
   }
   if (auto* b = stmt->As<ast::BlockStatement>()) {
-    return EmitIndentedBlockAndNewline(b);
+    return EmitBlockAndNewline(b);
   }
   if (auto* b = stmt->As<ast::BreakStatement>()) {
     return EmitBreak(b);
   }
   if (auto* c = stmt->As<ast::CallStatement>()) {
-    make_indent();
+    return EmitCall(c->expr());
-    if (!EmitCall(c->expr())) {
-      return false;
-    }
-    out_ << ";" << std::endl;
-    return true;
   }
   if (auto* c = stmt->As<ast::ContinueStatement>()) {
     return EmitContinue(c);
@@ -870,6 +869,9 @@ bool GeneratorImpl::EmitStatement(ast::Statement* stmt) {
   if (auto* l = stmt->As<ast::LoopStatement>()) {
     return EmitLoop(l);
   }
+  if (auto* l = stmt->As<ast::ForLoopStatement>()) {
+    return EmitForLoop(l);
+  }
   if (auto* r = stmt->As<ast::ReturnStatement>()) {
     return EmitReturn(r);
   }
@@ -886,8 +888,6 @@ bool GeneratorImpl::EmitStatement(ast::Statement* stmt) {
 }
 
 bool GeneratorImpl::EmitAssign(ast::AssignmentStatement* stmt) {
-  make_indent();
-
   if (!EmitExpression(stmt->lhs())) {
     return false;
   }
@@ -898,14 +898,11 @@ bool GeneratorImpl::EmitAssign(ast::AssignmentStatement* stmt) {
     return false;
   }
 
-  out_ << ";" << std::endl;
-
   return true;
 }
 
 bool GeneratorImpl::EmitBreak(ast::BreakStatement*) {
-  make_indent();
+  out_ << "break";
-  out_ << "break;" << std::endl;
   return true;
 }
 
@@ -935,8 +932,7 @@ bool GeneratorImpl::EmitCase(ast::CaseStatement* stmt) {
 }
 
 bool GeneratorImpl::EmitContinue(ast::ContinueStatement*) {
-  make_indent();
+  out_ << "continue";
-  out_ << "continue;" << std::endl;
   return true;
 }
 
@@ -955,14 +951,11 @@ bool GeneratorImpl::EmitElse(ast::ElseStatement* stmt) {
 }
 
 bool GeneratorImpl::EmitFallthrough(ast::FallthroughStatement*) {
-  make_indent();
+  out_ << "fallthrough";
-  out_ << "fallthrough;" << std::endl;
   return true;
 }
 
 bool GeneratorImpl::EmitIf(ast::IfStatement* stmt) {
-  make_indent();
-
   out_ << "if (";
   if (!EmitExpression(stmt->condition())) {
     return false;
@@ -984,14 +977,11 @@ bool GeneratorImpl::EmitIf(ast::IfStatement* stmt) {
 }
 
 bool GeneratorImpl::EmitDiscard(ast::DiscardStatement*) {
-  make_indent();
+  out_ << "discard";
-  out_ << "discard;" << std::endl;
   return true;
 }
 
 bool GeneratorImpl::EmitLoop(ast::LoopStatement* stmt) {
-  make_indent();
-
   out_ << "loop {" << std::endl;
   increment_indent();
 
@@ -1019,9 +1009,50 @@ bool GeneratorImpl::EmitLoop(ast::LoopStatement* stmt) {
   return true;
 }
 
-bool GeneratorImpl::EmitReturn(ast::ReturnStatement* stmt) {
+bool GeneratorImpl::EmitForLoop(ast::ForLoopStatement* stmt) {
-  make_indent();
+  out_ << "for";
+  {
+    ScopedParen sp(out_);
+    if (auto* init = stmt->initializer()) {
+      if (!EmitRawStatement(init)) {
+        return false;
+      }
+    }
 
+    out_ << "; ";
+
+    if (auto* cond = stmt->condition()) {
+      if (!EmitExpression(cond)) {
+        return false;
+      }
+    }
+
+    out_ << "; ";
+
+    if (auto* cont = stmt->continuing()) {
+      if (!EmitRawStatement(cont)) {
+        return false;
+      }
+    }
+  }
+  out_ << " {" << std::endl;
+
+  {
+    ScopedIndent si(this);
+    for (auto* s : stmt->body()->statements()) {
+      if (!EmitStatement(s)) {
+        return false;
+      }
+    }
+  }
+
+  make_indent();
+  out_ << "}" << std::endl;
+
+  return true;
+}
+
+bool GeneratorImpl::EmitReturn(ast::ReturnStatement* stmt) {
   out_ << "return";
   if (stmt->has_value()) {
     out_ << " ";
@@ -1029,13 +1060,10 @@ bool GeneratorImpl::EmitReturn(ast::ReturnStatement* stmt) {
       return false;
     }
   }
-  out_ << ";" << std::endl;
   return true;
 }
 
 bool GeneratorImpl::EmitSwitch(ast::SwitchStatement* stmt) {
-  make_indent();
-
   out_ << "switch(";
   if (!EmitExpression(stmt->condition())) {
     return false;

src/writer/wgsl/generator_impl.h

@@ -25,6 +25,7 @@
 #include "src/ast/continue_statement.h"
 #include "src/ast/discard_statement.h"
 #include "src/ast/fallthrough_statement.h"
+#include "src/ast/for_loop_statement.h"
 #include "src/ast/if_statement.h"
 #include "src/ast/loop_statement.h"
 #include "src/ast/member_accessor_expression.h"
@@ -81,10 +82,6 @@ class GeneratorImpl : public TextGenerator {
   /// Handles a block statement with a newline at the end
   /// @param stmt the statement to emit
   /// @returns true if the statement was emitted successfully
-  bool EmitIndentedBlockAndNewline(const ast::BlockStatement* stmt);
-  /// Handles a block statement with a newline at the end
-  /// @param stmt the statement to emit
-  /// @returns true if the statement was emitted successfully
   bool EmitBlockAndNewline(const ast::BlockStatement* stmt);
   /// Handles a break statement
   /// @param stmt the statement to emit
@@ -146,6 +143,10 @@ class GeneratorImpl : public TextGenerator {
   /// @param stmt the statement to emit
   /// @returns true if the statement was emtited
   bool EmitLoop(ast::LoopStatement* stmt);
+  /// Handles a for-loop statement
+  /// @param stmt the statement to emit
+  /// @returns true if the statement was emtited
+  bool EmitForLoop(ast::ForLoopStatement* stmt);
   /// Handles a member accessor expression
   /// @param expr the member accessor expression
   /// @returns true if the member accessor was emitted
@@ -158,6 +159,11 @@ class GeneratorImpl : public TextGenerator {
   /// @param stmt the statement to emit
   /// @returns true if the statement was emitted
   bool EmitStatement(ast::Statement* stmt);
+  /// Emits a statement without an indentation or trailing semi-colon and
+  /// newline
+  /// @param stmt the statement to emit
+  /// @returns true if the statement was emitted
+  bool EmitRawStatement(ast::Statement* stmt);
   /// Handles generating a switch statement
   /// @param stmt the statement to emit
   /// @returns true if the statement was emitted

src/writer/wgsl/generator_impl_variable_test.cc

@@ -29,8 +29,7 @@ TEST_F(WgslGeneratorImplTest, EmitVariable) {
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitVariable(v)) << gen.error();
-  EXPECT_EQ(gen.result(), R"(var<private> a : f32;
+  EXPECT_EQ(gen.result(), R"(var<private> a : f32)");
-)");
 }
 
 TEST_F(WgslGeneratorImplTest, EmitVariable_StorageClass) {
@@ -39,8 +38,7 @@ TEST_F(WgslGeneratorImplTest, EmitVariable_StorageClass) {
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitVariable(v)) << gen.error();
-  EXPECT_EQ(gen.result(), R"(var<private> a : f32;
+  EXPECT_EQ(gen.result(), R"(var<private> a : f32)");
-)");
 }
 
 TEST_F(WgslGeneratorImplTest, EmitVariable_Access_Read) {
@@ -56,8 +54,8 @@ TEST_F(WgslGeneratorImplTest, EmitVariable_Access_Read) {
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitVariable(v)) << gen.error();
-  EXPECT_EQ(gen.result(), R"([[binding(0), group(0)]] var<storage, read> a : S;
+  EXPECT_EQ(gen.result(),
-)");
+            R"([[binding(0), group(0)]] var<storage, read> a : S)");
 }
 
 TEST_F(WgslGeneratorImplTest, EmitVariable_Access_Write) {
@@ -73,8 +71,8 @@ TEST_F(WgslGeneratorImplTest, EmitVariable_Access_Write) {
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitVariable(v)) << gen.error();
-  EXPECT_EQ(gen.result(), R"([[binding(0), group(0)]] var<storage, write> a : S;
+  EXPECT_EQ(gen.result(),
-)");
+            R"([[binding(0), group(0)]] var<storage, write> a : S)");
 }
 
 TEST_F(WgslGeneratorImplTest, EmitVariable_Access_ReadWrite) {
@@ -91,8 +89,7 @@ TEST_F(WgslGeneratorImplTest, EmitVariable_Access_ReadWrite) {
 
   ASSERT_TRUE(gen.EmitVariable(v)) << gen.error();
   EXPECT_EQ(gen.result(),
-            R"([[binding(0), group(0)]] var<storage, read_write> a : S;
+            R"([[binding(0), group(0)]] var<storage, read_write> a : S)");
-)");
 }
 
 TEST_F(WgslGeneratorImplTest, EmitVariable_Decorated) {
@@ -106,9 +103,7 @@ TEST_F(WgslGeneratorImplTest, EmitVariable_Decorated) {
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitVariable(v)) << gen.error();
-  EXPECT_EQ(gen.result(),
+  EXPECT_EQ(gen.result(), R"([[group(1), binding(2)]] var a : sampler)");
-            R"([[group(1), binding(2)]] var a : sampler;
-)");
 }
 
 TEST_F(WgslGeneratorImplTest, EmitVariable_Constructor) {
@@ -117,8 +112,7 @@ TEST_F(WgslGeneratorImplTest, EmitVariable_Constructor) {
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitVariable(v)) << gen.error();
-  EXPECT_EQ(gen.result(), R"(var<private> a : f32 = 1.0;
+  EXPECT_EQ(gen.result(), R"(var<private> a : f32 = 1.0)");
-)");
 }
 
 TEST_F(WgslGeneratorImplTest, EmitVariable_Const) {
@@ -128,8 +122,7 @@ TEST_F(WgslGeneratorImplTest, EmitVariable_Const) {
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitVariable(v)) << gen.error();
-  EXPECT_EQ(gen.result(), R"(let a : f32 = 1.0;
+  EXPECT_EQ(gen.result(), R"(let a : f32 = 1.0)");
-)");
 }
 
 }  // namespace