Add helpers for referenced variables.

This CL adds a helper to get referenced builtins and referenced
locations from the ast::Function.

Change-Id: I95cf7efd6b0fe9569c5e514d3245112e78147ffe
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/24783
Reviewed-by: David Neto <dneto@google.com>

src/ast/function.cc

@@ -16,6 +16,8 @@
 
 #include <sstream>
 
+#include "src/ast/decorated_variable.h"
+
 namespace tint {
 namespace ast {
 
@@ -51,6 +53,42 @@ void Function::add_referenced_module_variable(Variable* var) {
   referenced_module_vars_.push_back(var);
 }
 
+const std::vector<std::pair<Variable*, LocationDecoration*>>
+Function::referenced_location_variables() const {
+  std::vector<std::pair<Variable*, LocationDecoration*>> ret;
+
+  for (auto* var : referenced_module_variables()) {
+    if (!var->IsDecorated()) {
+      continue;
+    }
+    for (auto& deco : var->AsDecorated()->decorations()) {
+      if (deco->IsLocation()) {
+        ret.push_back({var, deco.get()->AsLocation()});
+        break;
+      }
+    }
+  }
+  return ret;
+}
+
+const std::vector<std::pair<Variable*, BuiltinDecoration*>>
+Function::referenced_builtin_variables() const {
+  std::vector<std::pair<Variable*, BuiltinDecoration*>> ret;
+
+  for (auto* var : referenced_module_variables()) {
+    if (!var->IsDecorated()) {
+      continue;
+    }
+    for (auto& deco : var->AsDecorated()->decorations()) {
+      if (deco->IsBuiltin()) {
+        ret.push_back({var, deco.get()->AsBuiltin()});
+        break;
+      }
+    }
+  }
+  return ret;
+}
+
 void Function::add_ancestor_entry_point(const std::string& ep) {
   for (const auto& point : ancestor_entry_points_) {
     if (point == ep) {

src/ast/function.h

@@ -21,7 +21,9 @@
 #include <utility>
 #include <vector>
 
+#include "src/ast/builtin_decoration.h"
 #include "src/ast/expression.h"
+#include "src/ast/location_decoration.h"
 #include "src/ast/node.h"
 #include "src/ast/statement.h"
 #include "src/ast/type/type.h"
@@ -76,6 +78,14 @@ class Function : public Node {
   const std::vector<Variable*>& referenced_module_variables() const {
     return referenced_module_vars_;
   }
+  /// Retrieves any referenced location variables
+  /// @returns the <variable, decoration> pair.
+  const std::vector<std::pair<Variable*, LocationDecoration*>>
+  referenced_location_variables() const;
+  /// Retrieves any referenced builtin variables
+  /// @returns the <variable, decoration> pair.
+  const std::vector<std::pair<Variable*, BuiltinDecoration*>>
+  referenced_builtin_variables() const;
 
   /// Adds an ancestor entry point
   /// @param ep the entry point ancestor

src/ast/function_test.cc

@@ -15,7 +15,10 @@
 #include "src/ast/function.h"
 
 #include "gtest/gtest.h"
+#include "src/ast/builtin_decoration.h"
+#include "src/ast/decorated_variable.h"
 #include "src/ast/kill_statement.h"
+#include "src/ast/location_decoration.h"
 #include "src/ast/pipeline_stage.h"
 #include "src/ast/type/f32_type.h"
 #include "src/ast/type/i32_type.h"
@@ -78,6 +81,92 @@ TEST_F(FunctionTest, AddDuplicateReferencedVariables) {
   EXPECT_EQ(f.referenced_module_variables()[1], &v2);
 }
 
+TEST_F(FunctionTest, GetReferenceLocations) {
+  type::VoidType void_type;
+  type::I32Type i32;
+
+  VariableDecorationList decos;
+  DecoratedVariable loc1(
+      std::make_unique<ast::Variable>("loc1", StorageClass::kInput, &i32));
+  decos.push_back(std::make_unique<ast::LocationDecoration>(0));
+  loc1.set_decorations(std::move(decos));
+
+  DecoratedVariable loc2(
+      std::make_unique<ast::Variable>("loc2", StorageClass::kInput, &i32));
+  decos.push_back(std::make_unique<ast::LocationDecoration>(1));
+  loc2.set_decorations(std::move(decos));
+
+  DecoratedVariable builtin1(
+      std::make_unique<ast::Variable>("builtin1", StorageClass::kInput, &i32));
+  decos.push_back(
+      std::make_unique<ast::BuiltinDecoration>(ast::Builtin::kPosition));
+  builtin1.set_decorations(std::move(decos));
+
+  DecoratedVariable builtin2(
+      std::make_unique<ast::Variable>("builtin2", StorageClass::kInput, &i32));
+  decos.push_back(
+      std::make_unique<ast::BuiltinDecoration>(ast::Builtin::kFragDepth));
+  builtin2.set_decorations(std::move(decos));
+
+  Function f("func", VariableList{}, &void_type);
+
+  f.add_referenced_module_variable(&loc1);
+  f.add_referenced_module_variable(&builtin1);
+  f.add_referenced_module_variable(&loc2);
+  f.add_referenced_module_variable(&builtin2);
+  ASSERT_EQ(f.referenced_module_variables().size(), 4u);
+
+  auto ref_locs = f.referenced_location_variables();
+  ASSERT_EQ(ref_locs.size(), 2u);
+  EXPECT_EQ(ref_locs[0].first, &loc1);
+  EXPECT_EQ(ref_locs[0].second->value(), 0);
+  EXPECT_EQ(ref_locs[1].first, &loc2);
+  EXPECT_EQ(ref_locs[1].second->value(), 1);
+}
+
+TEST_F(FunctionTest, GetReferenceBuiltins) {
+  type::VoidType void_type;
+  type::I32Type i32;
+
+  VariableDecorationList decos;
+  DecoratedVariable loc1(
+      std::make_unique<ast::Variable>("loc1", StorageClass::kInput, &i32));
+  decos.push_back(std::make_unique<ast::LocationDecoration>(0));
+  loc1.set_decorations(std::move(decos));
+
+  DecoratedVariable loc2(
+      std::make_unique<ast::Variable>("loc2", StorageClass::kInput, &i32));
+  decos.push_back(std::make_unique<ast::LocationDecoration>(1));
+  loc2.set_decorations(std::move(decos));
+
+  DecoratedVariable builtin1(
+      std::make_unique<ast::Variable>("builtin1", StorageClass::kInput, &i32));
+  decos.push_back(
+      std::make_unique<ast::BuiltinDecoration>(ast::Builtin::kPosition));
+  builtin1.set_decorations(std::move(decos));
+
+  DecoratedVariable builtin2(
+      std::make_unique<ast::Variable>("builtin2", StorageClass::kInput, &i32));
+  decos.push_back(
+      std::make_unique<ast::BuiltinDecoration>(ast::Builtin::kFragDepth));
+  builtin2.set_decorations(std::move(decos));
+
+  Function f("func", VariableList{}, &void_type);
+
+  f.add_referenced_module_variable(&loc1);
+  f.add_referenced_module_variable(&builtin1);
+  f.add_referenced_module_variable(&loc2);
+  f.add_referenced_module_variable(&builtin2);
+  ASSERT_EQ(f.referenced_module_variables().size(), 4u);
+
+  auto ref_locs = f.referenced_builtin_variables();
+  ASSERT_EQ(ref_locs.size(), 2u);
+  EXPECT_EQ(ref_locs[0].first, &builtin1);
+  EXPECT_EQ(ref_locs[0].second->value(), ast::Builtin::kPosition);
+  EXPECT_EQ(ref_locs[1].first, &builtin2);
+  EXPECT_EQ(ref_locs[1].second->value(), ast::Builtin::kFragDepth);
+}
+
 TEST_F(FunctionTest, AddDuplicateEntryPoints) {
   ast::type::VoidType void_type;
   Function f("func", VariableList{}, &void_type);

src/writer/msl/generator_impl.cc

@@ -518,21 +518,9 @@ bool GeneratorImpl::EmitEntryPointData(ast::EntryPoint* ep) {
 
   std::vector<std::pair<ast::Variable*, uint32_t>> in_locations;
   std::vector<std::pair<ast::Variable*, uint32_t>> out_locations;
-  for (auto* var : func->referenced_module_variables()) {
-    if (!var->IsDecorated()) {
-      continue;
-    }
-    auto* decorated = var->AsDecorated();
-    ast::LocationDecoration* locn_deco = nullptr;
-    for (auto& deco : decorated->decorations()) {
-      if (deco->IsLocation()) {
-        locn_deco = deco.get()->AsLocation();
-        break;
-      }
-    }
-    if (locn_deco == nullptr) {
-      continue;
-    }
+  for (auto data : func->referenced_location_variables()) {
+    auto var = data.first;
+    auto locn_deco = data.second;
 
     if (var->storage_class() == ast::StorageClass::kInput) {
       in_locations.push_back({var, locn_deco->value()});
@@ -753,6 +741,8 @@ bool GeneratorImpl::EmitFunctionInternal(ast::Function* func,
 
   // TODO(dsinclair): Handle any entry point builtin params used here
 
+  // TODO(dsinclair): Binding/Set inputs
+
   for (const auto& v : func->params()) {
     if (!first) {
       out_ << ", ";
@@ -816,6 +806,9 @@ bool GeneratorImpl::EmitEntryPointFunction(ast::EntryPoint* ep) {
   }
 
   // TODO(dsinclair): Output other builtin inputs
+
+  // TODO(dsinclair): Binding/Set inputs
+
   out_ << ") {" << std::endl;
 
   increment_indent();