[resolver] Track pipeline stage usages for structs

This will be used by the generators to determine how to handle
location decorations.

Change-Id: Ib0e0ce852a5da3819781b402c5625a440c4c9544
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/46400
Auto-Submit: James Price <jrprice@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>
Commit-Queue: Ben Clayton <bclayton@google.com>

src/resolver/resolver.cc

@@ -293,6 +293,54 @@ bool Resolver::Function(ast::Function* func) {
   variable_stack_.push_scope();
   for (auto* param : func->params()) {
     variable_stack_.set(param->symbol(), CreateVariableInfo(param));
+
+    if (auto* str =
+            param->declared_type()->UnwrapAliasIfNeeded()->As<type::Struct>()) {
+      auto* info = Structure(str);
+      if (!info) {
+        return false;
+      }
+      switch (func->pipeline_stage()) {
+        case ast::PipelineStage::kVertex:
+          info->pipeline_stage_uses.emplace(
+              semantic::PipelineStageUsage::kVertexInput);
+          break;
+        case ast::PipelineStage::kFragment:
+          info->pipeline_stage_uses.emplace(
+              semantic::PipelineStageUsage::kFragmentInput);
+          break;
+        case ast::PipelineStage::kCompute:
+          info->pipeline_stage_uses.emplace(
+              semantic::PipelineStageUsage::kComputeInput);
+          break;
+        case ast::PipelineStage::kNone:
+          break;
+      }
+    }
+  }
+
+  if (auto* str =
+          func->return_type()->UnwrapAliasIfNeeded()->As<type::Struct>()) {
+    auto* info = Structure(str);
+    if (!info) {
+      return false;
+    }
+    switch (func->pipeline_stage()) {
+      case ast::PipelineStage::kVertex:
+        info->pipeline_stage_uses.emplace(
+            semantic::PipelineStageUsage::kVertexOutput);
+        break;
+      case ast::PipelineStage::kFragment:
+        info->pipeline_stage_uses.emplace(
+            semantic::PipelineStageUsage::kFragmentOutput);
+        break;
+      case ast::PipelineStage::kCompute:
+        info->pipeline_stage_uses.emplace(
+            semantic::PipelineStageUsage::kComputeOutput);
+        break;
+      case ast::PipelineStage::kNone:
+        break;
+    }
   }
 
   if (!BlockStatement(func->body())) {
@@ -1359,7 +1407,8 @@ void Resolver::CreateSemanticNodes() const {
     builder_->Sem().Add(
         str, builder_->create<semantic::Struct>(
                  str, std::move(info->members), info->align, info->size,
-                 info->size_no_padding, info->storage_class_usage));
+                 info->size_no_padding, info->storage_class_usage,
+                 info->pipeline_stage_uses));
   }
 }
 

src/resolver/resolver.h

@@ -24,6 +24,7 @@
 #include "src/intrinsic_table.h"
 #include "src/program_builder.h"
 #include "src/scope_stack.h"
+#include "src/semantic/struct.h"
 #include "src/utils/unique_vector.h"
 
 namespace tint {
@@ -139,6 +140,7 @@ class Resolver {
     uint32_t size = 0;
     uint32_t size_no_padding = 0;
     std::unordered_set<ast::StorageClass> storage_class_usage;
+    std::unordered_set<semantic::PipelineStageUsage> pipeline_stage_uses;
   };
 
   /// Structure holding semantic information about a block (i.e. scope), such as

src/resolver/struct_pipeline_stage_use_test.cc

@@ -0,0 +1,204 @@
+// Copyright 2021 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/resolver/resolver.h"
+
+#include "gmock/gmock.h"
+#include "src/ast/stage_decoration.h"
+#include "src/resolver/resolver_test_helper.h"
+#include "src/semantic/struct.h"
+
+using ::testing::UnorderedElementsAre;
+
+namespace tint {
+namespace resolver {
+namespace {
+
+using ResolverPipelineStageUseTest = ResolverTest;
+
+TEST_F(ResolverPipelineStageUseTest, UnusedStruct) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_TRUE(sem->PipelineStageUses().empty());
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsNonEntryPointParam) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  Func("foo", {Const("param", s)}, ty.void_(), {}, {});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_TRUE(sem->PipelineStageUses().empty());
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsNonEntryPointReturnType) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  Func("foo", {}, s, {Return(Construct(s, Expr(0.f)))}, {});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_TRUE(sem->PipelineStageUses().empty());
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsVertexShaderParam) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  Func("main", {Const("param", s)}, ty.void_(), {},
+       {create<ast::StageDecoration>(ast::PipelineStage::kVertex)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(sem->PipelineStageUses(),
+              UnorderedElementsAre(semantic::PipelineStageUsage::kVertexInput));
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsVertexShaderReturnType) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  Func("main", {}, s, {Return(Construct(s, Expr(0.f)))},
+       {create<ast::StageDecoration>(ast::PipelineStage::kVertex)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(
+      sem->PipelineStageUses(),
+      UnorderedElementsAre(semantic::PipelineStageUsage::kVertexOutput));
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsFragmentShaderParam) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  Func("main", {Const("param", s)}, ty.void_(), {},
+       {create<ast::StageDecoration>(ast::PipelineStage::kFragment)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(
+      sem->PipelineStageUses(),
+      UnorderedElementsAre(semantic::PipelineStageUsage::kFragmentInput));
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsFragmentShaderReturnType) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  Func("main", {}, s, {Return(Construct(s, Expr(0.f)))},
+       {create<ast::StageDecoration>(ast::PipelineStage::kFragment)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(
+      sem->PipelineStageUses(),
+      UnorderedElementsAre(semantic::PipelineStageUsage::kFragmentOutput));
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsComputeShaderParam) {
+  auto* s = Structure("S", {Member("a", ty.u32(),
+                                   {create<ast::BuiltinDecoration>(
+                                       ast::Builtin::kLocalInvocationIndex)})});
+
+  Func("main", {Const("param", s)}, ty.void_(), {},
+       {create<ast::StageDecoration>(ast::PipelineStage::kCompute)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(
+      sem->PipelineStageUses(),
+      UnorderedElementsAre(semantic::PipelineStageUsage::kComputeInput));
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedMultipleStages) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  Func("vert_main", {Const("param", s)}, s, {Return(Construct(s, Expr(0.f)))},
+       {create<ast::StageDecoration>(ast::PipelineStage::kVertex)});
+
+  Func("frag_main", {Const("param", s)}, ty.void_(), {},
+       {create<ast::StageDecoration>(ast::PipelineStage::kFragment)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(
+      sem->PipelineStageUses(),
+      UnorderedElementsAre(semantic::PipelineStageUsage::kVertexInput,
+                           semantic::PipelineStageUsage::kVertexOutput,
+                           semantic::PipelineStageUsage::kFragmentInput));
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsShaderParamViaAlias) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+  auto* s_alias = ty.alias("S_alias", s);
+
+  Func("main", {Const("param", s_alias)}, ty.void_(), {},
+       {create<ast::StageDecoration>(ast::PipelineStage::kFragment)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(
+      sem->PipelineStageUses(),
+      UnorderedElementsAre(semantic::PipelineStageUsage::kFragmentInput));
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsShaderReturnTypeViaAlias) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+  auto* s_alias = ty.alias("S_alias", s);
+
+  Func("main", {}, s_alias, {Return(Construct(s_alias, Expr(0.f)))},
+       {create<ast::StageDecoration>(ast::PipelineStage::kFragment)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(
+      sem->PipelineStageUses(),
+      UnorderedElementsAre(semantic::PipelineStageUsage::kFragmentOutput));
+}
+
+}  // namespace
+}  // namespace resolver
+}  // namespace tint