From d1684ed28776ed525775aa7a7a41d8e02866a4ca Mon Sep 17 00:00:00 2001
From: dan sinclair <dsinclair@google.com>
Date: Wed, 15 Jul 2020 18:26:17 +0000
Subject: [PATCH] [msl-writer] Refactor some entry point variable code.

This CL cleans up some of the entry point variable code.

Bug: tint:8
Change-Id: Ic89ff3608d34af26eeb9fe9408cdbd5c7de22dda
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/24782
Reviewed-by: David Neto <dneto@google.com>
---
 src/writer/msl/generator_impl.cc | 59 ++++++++++++++++++++------------
 src/writer/msl/generator_impl.h  |  3 ++
 2 files changed, 40 insertions(+), 22 deletions(-)

diff --git a/src/writer/msl/generator_impl.cc b/src/writer/msl/generator_impl.cc
index defabf5627..260d6fcd79 100644
--- a/src/writer/msl/generator_impl.cc
+++ b/src/writer/msl/generator_impl.cc
@@ -278,6 +278,27 @@ bool GeneratorImpl::EmitBreak(ast::BreakStatement*) {
   return true;
 }
 
+std::string GeneratorImpl::current_ep_var_name(VarType type) {
+  std::string name = "";
+  switch (type) {
+    case VarType::kIn: {
+      auto in_it = ep_name_to_in_data_.find(current_ep_name_);
+      if (in_it != ep_name_to_in_data_.end()) {
+        name = in_it->second.var_name;
+      }
+      break;
+    }
+    case VarType::kOut: {
+      auto out_it = ep_name_to_out_data_.find(current_ep_name_);
+      if (out_it != ep_name_to_out_data_.end()) {
+        name = out_it->second.var_name;
+      }
+      break;
+    }
+  }
+  return name;
+}
+
 bool GeneratorImpl::EmitCall(ast::CallExpression* expr) {
   if (!expr->func()->IsIdentifier()) {
     error_ = "invalid function name";
@@ -301,19 +322,18 @@ bool GeneratorImpl::EmitCall(ast::CallExpression* expr) {
     out_ << name << "(";
 
     bool first = true;
-
-    auto in_it = ep_name_to_in_data_.find(current_ep_name_);
-    if (in_it != ep_name_to_in_data_.end()) {
-      out_ << in_it->second.var_name;
+    auto var_name = current_ep_var_name(VarType::kIn);
+    if (!var_name.empty()) {
+      out_ << var_name;
       first = false;
     }
 
-    auto out_it = ep_name_to_out_data_.find(current_ep_name_);
-    if (out_it != ep_name_to_out_data_.end()) {
+    var_name = current_ep_var_name(VarType::kOut);
+    if (!var_name.empty()) {
       if (!first) {
         out_ << ", ";
       }
-      out_ << out_it->second.var_name;
+      out_ << var_name;
       first = false;
     }
 
@@ -832,23 +852,18 @@ bool GeneratorImpl::EmitIdentifier(ast::IdentifierExpression* expr) {
 
   ast::Variable* var = nullptr;
   if (global_variables_.get(ident->name(), &var)) {
-    if (var->storage_class() == ast::StorageClass::kInput &&
-        var->IsDecorated() && var->AsDecorated()->HasLocationDecoration()) {
-      auto it = ep_name_to_in_data_.find(current_ep_name_);
-      if (it == ep_name_to_in_data_.end()) {
-        error_ = "unable to find entry point data for input";
+    if (var->IsDecorated() && var->AsDecorated()->HasLocationDecoration() &&
+        (var->storage_class() == ast::StorageClass::kInput ||
+         var->storage_class() == ast::StorageClass::kOutput)) {
+      auto var_type = var->storage_class() == ast::StorageClass::kInput
+                          ? VarType::kIn
+                          : VarType::kOut;
+      auto name = current_ep_var_name(var_type);
+      if (name.empty()) {
+        error_ = "unable to find entry point data for variable";
         return false;
       }
-      out_ << it->second.var_name << ".";
-    } else if (var->storage_class() == ast::StorageClass::kOutput &&
-               var->IsDecorated() &&
-               var->AsDecorated()->HasLocationDecoration()) {
-      auto it = ep_name_to_out_data_.find(current_ep_name_);
-      if (it == ep_name_to_out_data_.end()) {
-        error_ = "unable to find entry point data for output";
-        return false;
-      }
-      out_ << it->second.var_name << ".";
+      out_ << name << ".";
     }
   }
   out_ << namer_.NameFor(ident->name());
diff --git a/src/writer/msl/generator_impl.h b/src/writer/msl/generator_impl.h
index c851ee76c3..6f7f6e5fad 100644
--- a/src/writer/msl/generator_impl.h
+++ b/src/writer/msl/generator_impl.h
@@ -202,6 +202,9 @@ class GeneratorImpl : public TextGenerator {
   Namer* namer_for_testing() { return &namer_; }
 
  private:
+  enum class VarType { kIn, kOut };
+  std::string current_ep_var_name(VarType type);
+
   Namer namer_;
   ScopeStack<ast::Variable*> global_variables_;
   std::string current_ep_name_;