writer/wgsl: Fix printing of for-loops

Fix various issue with formatting for loop. Add tests.

Bug: tint:952
Change-Id: I704341a15f0050ebf82df219d0c7d068a3a63c26
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/58064
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: James Price <jrprice@google.com>

src/writer/msl/generator_impl_loop_test.cc

@@ -181,7 +181,7 @@ TEST_F(MslGeneratorImplTest, Emit_ForLoopWithSimpleInit) {
 
 TEST_F(MslGeneratorImplTest, Emit_ForLoopWithMultiStmtInit) {
   // var<workgroup> a : atomic<i32>;
-  // for(var b = atomicCompareExchangeWeak(&a, 1, 2); ; ) {
+  // for({ignore(1); ignore(2);}; ; ) {
   //   return;
   // }
   Global("a", ty.atomic<i32>(), ast::StorageClass::kWorkgroup);

src/writer/text_generator.cc

@@ -108,11 +108,11 @@ void TextGenerator::TextBuffer::Insert(const TextBuffer& tb,
   }
 }
 
-std::string TextGenerator::TextBuffer::String() const {
+std::string TextGenerator::TextBuffer::String(uint32_t indent /* = 0 */) const {
   std::stringstream ss;
   for (auto& line : lines) {
     if (!line.content.empty()) {
-      for (uint32_t i = 0; i < line.indent; i++) {
+      for (uint32_t i = 0; i < indent + line.indent; i++) {
         ss << " ";
       }
       ss << line.content;

src/writer/text_generator.h

@@ -106,7 +106,8 @@ class TextGenerator {
     void Insert(const TextBuffer& tb, size_t before, uint32_t indent);
 
     /// @returns the buffer's content as a single string
-    std::string String() const;
+    /// @param indent additional indentation to apply to each line
+    std::string String(uint32_t indent = 0) const;
 
     /// The current indentation of the TextBuffer. Lines appended to the
     /// TextBuffer will use this indentation.

src/writer/wgsl/generator_impl.cc

@@ -1034,17 +1034,20 @@ bool GeneratorImpl::EmitForLoop(ast::ForLoopStatement* stmt) {
       ScopedParen sp(out);
       switch (init_buf.lines.size()) {
         case 0:  // No initializer
-          out << ";";
           break;
         case 1:  // Single line initializer statement
-          out << init_buf.lines[0].content;
+          out << TrimSuffix(init_buf.lines[0].content, ";");
           break;
         default:  // Block initializer statement
-          current_buffer_->Append(init_buf);
+          for (size_t i = 1; i < init_buf.lines.size(); i++) {
+            // Indent all by the first line
+            init_buf.lines[i].indent += current_buffer_->current_indent;
+          }
+          out << TrimSuffix(init_buf.String(), "\n");
           break;
       }
 
-      out << " ";
+      out << "; ";
 
       if (auto* cond = stmt->condition()) {
         if (!EmitExpression(out, cond)) {
@@ -1056,13 +1059,16 @@ bool GeneratorImpl::EmitForLoop(ast::ForLoopStatement* stmt) {
 
       switch (cont_buf.lines.size()) {
         case 0:  // No continuing
-          out << ";";
           break;
         case 1:  // Single line continuing statement
           out << TrimSuffix(cont_buf.lines[0].content, ";");
           break;
         default:  // Block continuing statement
-          current_buffer_->Append(cont_buf);
+          for (size_t i = 1; i < cont_buf.lines.size(); i++) {
+            // Indent all by the first line
+            cont_buf.lines[i].indent += current_buffer_->current_indent;
+          }
+          out << TrimSuffix(cont_buf.String(), "\n");
           break;
       }
     }

src/writer/wgsl/generator_impl_loop_test.cc

@@ -61,6 +61,142 @@ TEST_F(WgslGeneratorImplTest, Emit_LoopWithContinuing) {
 )");
 }
 
+TEST_F(WgslGeneratorImplTest, Emit_ForLoopWithMultiStmtInit) {
+  // var<workgroup> a : atomic<i32>;
+  // for({ignore(1); ignore(2);}; ; ) {
+  //   return;
+  // }
+  Global("a", ty.atomic<i32>(), ast::StorageClass::kWorkgroup);
+  auto* multi_stmt = Block(Ignore(1), Ignore(2));
+  auto* f = For(multi_stmt, nullptr, nullptr, Block(Return()));
+  WrapInFunction(f);
+
+  GeneratorImpl& gen = Build();
+
+  gen.increment_indent();
+
+  ASSERT_TRUE(gen.EmitStatement(f)) << gen.error();
+  EXPECT_EQ(gen.result(), R"(  for({
+    ignore(1);
+    ignore(2);
+  }; ; ) {
+    return;
+  }
+)");
+}
+
+TEST_F(WgslGeneratorImplTest, Emit_ForLoopWithSimpleCond) {
+  // for(; true; ) {
+  //   return;
+  // }
+
+  auto* f = For(nullptr, true, nullptr, Block(Return()));
+  WrapInFunction(f);
+
+  GeneratorImpl& gen = Build();
+
+  gen.increment_indent();
+
+  ASSERT_TRUE(gen.EmitStatement(f)) << gen.error();
+  EXPECT_EQ(gen.result(), R"(  for(; true; ) {
+    return;
+  }
+)");
+}
+
+TEST_F(WgslGeneratorImplTest, Emit_ForLoopWithSimpleCont) {
+  // for(; ; i = i + 1) {
+  //   return;
+  // }
+
+  auto* v = Decl(Var("i", ty.i32()));
+  auto* f = For(nullptr, nullptr, Assign("i", Add("i", 1)), Block(Return()));
+  WrapInFunction(v, f);
+
+  GeneratorImpl& gen = Build();
+
+  gen.increment_indent();
+
+  ASSERT_TRUE(gen.EmitStatement(f)) << gen.error();
+  EXPECT_EQ(gen.result(), R"(  for(; ; i = (i + 1)) {
+    return;
+  }
+)");
+}
+
+TEST_F(WgslGeneratorImplTest, Emit_ForLoopWithMultiStmtCont) {
+  // var<workgroup> a : atomic<i32>;
+  // for(; ; { ignore(1); ignore(2); }) {
+  //   return;
+  // }
+
+  Global("a", ty.atomic<i32>(), ast::StorageClass::kWorkgroup);
+  auto* multi_stmt = Block(Ignore(1), Ignore(2));
+  auto* f = For(nullptr, nullptr, multi_stmt, Block(Return()));
+  WrapInFunction(f);
+
+  GeneratorImpl& gen = Build();
+
+  gen.increment_indent();
+
+  ASSERT_TRUE(gen.EmitStatement(f)) << gen.error();
+  EXPECT_EQ(gen.result(), R"(  for(; ; {
+    ignore(1);
+    ignore(2);
+  }) {
+    return;
+  }
+)");
+}
+
+TEST_F(WgslGeneratorImplTest, Emit_ForLoopWithSimpleInitCondCont) {
+  // for(var i : i32; true; i = i + 1) {
+  //   return;
+  // }
+
+  auto* f = For(Decl(Var("i", ty.i32())), true, Assign("i", Add("i", 1)),
+                Block(Return()));
+  WrapInFunction(f);
+
+  GeneratorImpl& gen = Build();
+
+  gen.increment_indent();
+
+  ASSERT_TRUE(gen.EmitStatement(f)) << gen.error();
+  EXPECT_EQ(gen.result(), R"(  for(var i : i32; true; i = (i + 1)) {
+    return;
+  }
+)");
+}
+
+TEST_F(WgslGeneratorImplTest, Emit_ForLoopWithMultiStmtInitCondCont) {
+  // var<workgroup> a : atomic<i32>;
+  // for({ ignore(1); ignore(2); }; true; { ignore(3); ignore(4); }) {
+  //   return;
+  // }
+  Global("a", ty.atomic<i32>(), ast::StorageClass::kWorkgroup);
+  auto* multi_stmt_a = Block(Ignore(1), Ignore(2));
+  auto* multi_stmt_b = Block(Ignore(3), Ignore(4));
+  auto* f = For(multi_stmt_a, Expr(true), multi_stmt_b, Block(Return()));
+  WrapInFunction(f);
+
+  GeneratorImpl& gen = Build();
+
+  gen.increment_indent();
+
+  ASSERT_TRUE(gen.EmitStatement(f)) << gen.error();
+  EXPECT_EQ(gen.result(), R"(  for({
+    ignore(1);
+    ignore(2);
+  }; true; {
+    ignore(3);
+    ignore(4);
+  }) {
+    return;
+  }
+)");
+}
+
 }  // namespace
 }  // namespace wgsl
 }  // namespace writer

test/bug/tint/990.wgsl.expected.wgsl

@@ -1,5 +1,5 @@
 fn f() {
   var i : i32;
-  for(let p = &(i); ; ;) {
+  for(let p = &(i); ; ) {
   }
 }

test/statements/for/condition.wgsl.expected.wgsl

@@ -1,5 +1,5 @@
 fn f() {
   var i : i32;
-  for(; (i < 4); ;) {
+  for(; (i < 4); ) {
   }
 }

test/statements/for/empty.wgsl.expected.wgsl

@@ -1,4 +1,4 @@
 fn f() {
-  for(; ; ;) {
+  for(; ; ) {
   }
 }

test/statements/for/initializer.wgsl.expected.wgsl

@@ -1,4 +1,4 @@
 fn f() {
-  for(var i : i32 = 0; ; ;) {
+  for(var i : i32 = 0; ; ) {
   }
 }

test/statements/for/scoping.wgsl.expected.wgsl

@@ -1,5 +1,5 @@
 fn f() {
-  for(var must_not_collide : i32 = 0; ; ;) {
+  for(var must_not_collide : i32 = 0; ; ) {
   }
   var must_not_collide : i32;
 }

test/test.wgsl.expected.hlsl

@@ -0,0 +1,11 @@
+[numthreads(1, 1, 1)]
+void unused_entry_point() {
+  return;
+}
+
+void f() {
+  {
+    for(; ; ) {
+    }
+  }
+}

test/test.wgsl.expected.msl

@@ -0,0 +1,8 @@
+#include <metal_stdlib>
+
+using namespace metal;
+void f() {
+  for(; ; ) {
+  }
+}
+

test/test.wgsl.expected.spvasm

@@ -0,0 +1,30 @@
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 11
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %unused_entry_point "unused_entry_point"
+               OpExecutionMode %unused_entry_point LocalSize 1 1 1
+               OpName %unused_entry_point "unused_entry_point"
+               OpName %f "f"
+       %void = OpTypeVoid
+          %1 = OpTypeFunction %void
+%unused_entry_point = OpFunction %void None %1
+          %4 = OpLabel
+               OpReturn
+               OpFunctionEnd
+          %f = OpFunction %void None %1
+          %6 = OpLabel
+               OpBranch %7
+          %7 = OpLabel
+               OpLoopMerge %8 %9 None
+               OpBranch %10
+         %10 = OpLabel
+               OpBranch %9
+          %9 = OpLabel
+               OpBranch %7
+          %8 = OpLabel
+               OpReturn
+               OpFunctionEnd

test/test.wgsl.expected.wgsl

@@ -0,0 +1,4 @@
+fn f() {
+  for(; ; ) {
+  }
+}