Store location value into `sem::Function`.

This CL adds an optional location value to the `sem::Function`
which will store the resolved `@location` value.

Bug: tint:1633
Change-Id: I95130858d8a1cecae1389be74120da29fec2b448
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/101063
Reviewed-by: Ben Clayton <bclayton@google.com>
Commit-Queue: Dan Sinclair <dsinclair@chromium.org>
Kokoro: Kokoro <noreply+kokoro@google.com>

src/tint/resolver/resolver.cc

@@ -906,6 +906,19 @@ sem::Function* Resolver::Function(const ast::Function* decl) {
         return_type = builder_->create<sem::Void>();
     }
 
+    // Determine if the return type has a location
+    std::optional<uint32_t> return_location;
+    for (auto* attr : decl->return_type_attributes) {
+        Mark(attr);
+
+        if (auto* a = attr->As<ast::LocationAttribute>()) {
+            return_location = a->value;
+        }
+    }
+    if (!validator_.NoDuplicateAttributes(decl->attributes)) {
+        return nullptr;
+    }
+
     if (auto* str = return_type->As<sem::Struct>()) {
         if (!ApplyStorageClassUsageToType(ast::StorageClass::kNone, str, decl->source)) {
             AddNote(
@@ -929,7 +942,8 @@ sem::Function* Resolver::Function(const ast::Function* decl) {
         }
     }
 
-    auto* func = builder_->create<sem::Function>(decl, return_type, std::move(parameters));
+    auto* func =
+        builder_->create<sem::Function>(decl, return_type, return_location, std::move(parameters));
     builder_->Sem().Add(decl, func);
 
     TINT_SCOPED_ASSIGNMENT(current_function_, func);
@@ -968,13 +982,7 @@ sem::Function* Resolver::Function(const ast::Function* decl) {
     for (auto* attr : decl->attributes) {
         Mark(attr);
     }
-    if (!validator_.NoDuplicateAttributes(decl->attributes)) {
-        return nullptr;
-    }
 
-    for (auto* attr : decl->return_type_attributes) {
-        Mark(attr);
-    }
     if (!validator_.NoDuplicateAttributes(decl->return_type_attributes)) {
         return nullptr;
     }

src/tint/resolver/resolver_test.cc

@@ -802,6 +802,45 @@ TEST_F(ResolverTest, Function_RegisterInputOutputVariables) {
     EXPECT_EQ(vars[2]->Declaration(), priv_var);
 }
 
+TEST_F(ResolverTest, Function_ReturnType_Location) {
+    auto* func = Func("my_func", utils::Empty, ty.f32(),
+                      utils::Vector{
+                          Return(1_f),
+                      },
+                      utils::Vector{
+                          Stage(ast::PipelineStage::kFragment),
+                      },
+                      utils::Vector{
+                          Location(2),
+                      });
+
+    EXPECT_TRUE(r()->Resolve()) << r()->error();
+
+    auto* sem = Sem().Get(func);
+    ASSERT_NE(nullptr, sem);
+    EXPECT_EQ(2u, sem->ReturnLocation());
+}
+
+TEST_F(ResolverTest, Function_ReturnType_NoLocation) {
+    GlobalVar("my_vec", ty.vec4<f32>(), ast::StorageClass::kPrivate);
+    auto* func = Func("my_func", utils::Empty, ty.vec4<f32>(),
+                      utils::Vector{
+                          Return("my_vec"),
+                      },
+                      utils::Vector{
+                          Stage(ast::PipelineStage::kVertex),
+                      },
+                      utils::Vector{
+                          Builtin(ast::BuiltinValue::kPosition),
+                      });
+
+    EXPECT_TRUE(r()->Resolve()) << r()->error();
+
+    auto* sem = Sem().Get(func);
+    ASSERT_NE(nullptr, sem);
+    EXPECT_FALSE(sem->ReturnLocation());
+}
+
 TEST_F(ResolverTest, Function_RegisterInputOutputVariables_SubFunction) {
     auto* s = Structure("S", utils::Vector{Member("m", ty.u32())});
 

src/tint/sem/function.cc

@@ -40,10 +40,12 @@ utils::VectorRef<const Parameter*> SetOwner(utils::VectorRef<Parameter*> paramet
 
 Function::Function(const ast::Function* declaration,
                    Type* return_type,
+                   std::optional<uint32_t> return_location,
                    utils::VectorRef<Parameter*> parameters)
     : Base(return_type, SetOwner(std::move(parameters), this), EvaluationStage::kRuntime),
       declaration_(declaration),
-      workgroup_size_{WorkgroupDimension{1}, WorkgroupDimension{1}, WorkgroupDimension{1}} {}
+      workgroup_size_{WorkgroupDimension{1}, WorkgroupDimension{1}, WorkgroupDimension{1}},
+      return_location_(return_location) {}
 
 Function::~Function() = default;
 

src/tint/sem/function.h

@@ -16,6 +16,7 @@
 #define SRC_TINT_SEM_FUNCTION_H_
 
 #include <array>
+#include <optional>
 #include <utility>
 #include <vector>
 
@@ -60,9 +61,11 @@ class Function final : public Castable<Function, CallTarget> {
     /// Constructor
     /// @param declaration the ast::Function
     /// @param return_type the return type of the function
+    /// @param return_location the location value for the return, if provided
     /// @param parameters the parameters to the function
     Function(const ast::Function* declaration,
              Type* return_type,
+             std::optional<uint32_t> return_location,
              utils::VectorRef<Parameter*> parameters);
 
     /// Destructor
@@ -254,6 +257,9 @@ class Function final : public Castable<Function, CallTarget> {
     /// @return the behaviors of this function
     sem::Behaviors& Behaviors() { return behaviors_; }
 
+    /// @return the location for the return, if provided
+    std::optional<uint32_t> ReturnLocation() const { return return_location_; }
+
   private:
     Function(const Function&) = delete;
     Function(Function&&) = delete;
@@ -274,6 +280,8 @@ class Function final : public Castable<Function, CallTarget> {
     std::vector<const Function*> ancestor_entry_points_;
     bool has_discard_ = false;
     sem::Behaviors behaviors_{sem::Behavior::kNext};
+
+    std::optional<uint32_t> return_location_;
 };
 
 }  // namespace tint::sem