[spirv-writer] Add preliminary support for GLSL methods

This CL adds the type determination and builder code to support
outputting the GLSL Round call.

Bug: tint:5
Change-Id: I84dadebaf19aee3361fb13b5f32ce1a9f1b0c421
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/19923
Reviewed-by: David Neto <dneto@google.com>
diff --git a/BUILD.gn b/BUILD.gn
index 8cd94fc..49a445b 100644
--- a/BUILD.gn
+++ b/BUILD.gn
@@ -550,6 +550,7 @@
     "src/writer/spirv/binary_writer_test.cc",
     "src/writer/spirv/builder_assign_test.cc",
     "src/writer/spirv/builder_binary_expression_test.cc",
+    "src/writer/spirv/builder_call_test.cc",
     "src/writer/spirv/builder_constructor_expression_test.cc",
     "src/writer/spirv/builder_entry_point_test.cc",
     "src/writer/spirv/builder_function_test.cc",
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index cb677ab..6e2d82f 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -423,6 +423,7 @@
     writer/spirv/binary_writer_test.cc
     writer/spirv/builder_assign_test.cc
     writer/spirv/builder_binary_expression_test.cc
+    writer/spirv/builder_call_test.cc
     writer/spirv/builder_constructor_expression_test.cc
     writer/spirv/builder_entry_point_test.cc
     writer/spirv/builder_function_test.cc
diff --git a/src/ast/import.h b/src/ast/import.h
index c3d6b82..9c73f3e 100644
--- a/src/ast/import.h
+++ b/src/ast/import.h
@@ -18,6 +18,7 @@
 #include <memory>
 #include <ostream>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 #include "src/ast/node.h"
@@ -31,13 +32,13 @@
   /// Create a new empty import statement
   Import() = default;
   /// Create a new import statement
-  /// @param path The import path e.g. GLSL.std.430
-  /// @param name The import reference name e.g. std::
+  /// @param path The import path e.g. GLSL.std.450
+  /// @param name The import reference name e.g. std
   Import(const std::string& path, const std::string& name);
   /// Create a new import statement
   /// @param source The input source for the import statement
   /// @param path The import path e.g. GLSL.std.430
-  /// @param name The import reference name e.g. std::
+  /// @param name The import reference name e.g. std
   Import(const Source& source,
          const std::string& path,
          const std::string& name);
@@ -57,6 +58,24 @@
   /// @returns the import name
   const std::string& name() const { return name_; }
 
+  /// Add the given |name| to map to the given |id|
+  /// @param name the name to map
+  /// @param id the id to map too
+  void AddMethodId(const std::string& name, uint32_t id) {
+    method_to_id_map_[name] = id;
+  }
+
+  /// Retrieves the id for a given name
+  /// @param name the name to lookup
+  /// @returns the id for the given name or 0 on failure
+  uint32_t GetIdForMethod(const std::string& name) const {
+    auto val = method_to_id_map_.find(name);
+    if (val == method_to_id_map_.end()) {
+      return 0;
+    }
+    return val->second;
+  }
+
   /// @returns true if the name and path are both present
   bool IsValid() const override;
 
@@ -70,6 +89,7 @@
 
   std::string path_;
   std::string name_;
+  std::unordered_map<std::string, uint32_t> method_to_id_map_;
 };
 
 /// A list of unique imports
diff --git a/src/ast/module.cc b/src/ast/module.cc
index 79e566b..88d97d8 100644
--- a/src/ast/module.cc
+++ b/src/ast/module.cc
@@ -25,7 +25,7 @@
 
 Module::~Module() = default;
 
-const Import* Module::FindImportByName(const std::string& name) {
+Import* Module::FindImportByName(const std::string& name) {
   for (const auto& import : imports_) {
     if (import->name() == name)
       return import.get();
diff --git a/src/ast/module.h b/src/ast/module.h
index 5484a81..26e9da1 100644
--- a/src/ast/module.h
+++ b/src/ast/module.h
@@ -47,7 +47,7 @@
   /// Find the import of the given name
   /// @param name The import name to search for
   /// @returns the import with the given name if found, nullptr otherwise.
-  const Import* FindImportByName(const std::string& name);
+  Import* FindImportByName(const std::string& name);
 
   /// Add a global variable to the module
   /// @param var the variable to add
diff --git a/src/type_determiner.cc b/src/type_determiner.cc
index 4b8c44b..7736255 100644
--- a/src/type_determiner.cc
+++ b/src/type_determiner.cc
@@ -16,6 +16,7 @@
 
 #include <memory>
 
+#include "spirv/unified1/GLSL.std.450.h"
 #include "src/ast/array_accessor_expression.h"
 #include "src/ast/as_expression.h"
 #include "src/ast/assignment_statement.h"
@@ -214,6 +215,15 @@
   return false;
 }
 
+bool TypeDeterminer::DetermineResultType(const ast::ExpressionList& list) {
+  for (const auto& expr : list) {
+    if (!DetermineResultType(expr.get())) {
+      return false;
+    }
+  }
+  return true;
+}
+
 bool TypeDeterminer::DetermineResultType(ast::Expression* expr) {
   // This is blindly called above, so in some cases the expression won't exist.
   if (!expr) {
@@ -285,14 +295,43 @@
 }
 
 bool TypeDeterminer::DetermineCall(ast::CallExpression* expr) {
-  for (const auto& param : expr->params()) {
-    if (!DetermineResultType(param.get())) {
-      return false;
-    }
+  if (!DetermineResultType(expr->params())) {
+    return false;
   }
 
-  if (!DetermineResultType(expr->func())) {
-    return false;
+  // The expression has to be an identifier as you can't store function pointers
+  // but, if it isn't we'll just use the normal result determination to be on
+  // the safe side.
+  if (expr->func()->IsIdentifier()) {
+    auto* ident = expr->func()->AsIdentifier();
+
+    if (ident->has_path()) {
+      auto* imp = mod_->FindImportByName(ident->path());
+      if (imp == nullptr) {
+        error_ = "Unable to find import for " + ident->name();
+        return false;
+      }
+
+      uint32_t ext_id = 0;
+      auto* result_type =
+          GetImportData(imp->path(), ident->name(), expr->params(), &ext_id);
+      if (result_type == nullptr) {
+        return false;
+      }
+
+      imp->AddMethodId(ident->name(), ext_id);
+      expr->func()->set_result_type(result_type);
+    } else {
+      // An identifier with a single name is a function call, not an import
+      // lookup which we can handle with the regular identifier lookup.
+      if (!DetermineResultType(ident)) {
+        return false;
+      }
+    }
+  } else {
+    if (!DetermineResultType(expr->func())) {
+      return false;
+    }
   }
   expr->set_result_type(expr->func()->result_type());
   return true;
@@ -314,8 +353,8 @@
 
 bool TypeDeterminer::DetermineIdentifier(ast::IdentifierExpression* expr) {
   if (expr->has_path()) {
-    // TODO(dsinclair): Handle imports
-    set_error(expr->source(), "imports not handled in type determination");
+    set_error(expr->source(),
+              "determine identifier should not be called with imports");
     return false;
   }
 
@@ -459,10 +498,8 @@
 }
 
 bool TypeDeterminer::DetermineUnaryMethod(ast::UnaryMethodExpression* expr) {
-  for (const auto& param : expr->params()) {
-    if (!DetermineResultType(param.get())) {
-      return false;
-    }
+  if (!DetermineResultType(expr->params())) {
+    return false;
   }
 
   switch (expr->op()) {
@@ -530,4 +567,32 @@
   return true;
 }
 
+ast::type::Type* TypeDeterminer::GetImportData(
+    const std::string& path,
+    const std::string& name,
+    const ast::ExpressionList& params,
+    uint32_t* id) {
+  if (path != "GLSL.std.450") {
+    return nullptr;
+  }
+
+  if (name == "round") {
+    if (params.size() != 1) {
+      error_ = "incorrect number of parameters for round. Expected 1 got " +
+               std::to_string(params.size());
+      return nullptr;
+    }
+    if (!params[0]->result_type()->is_float_scalar_or_vector()) {
+      error_ =
+          "incorrect type for round. Requires a float scalar or a float vector";
+      return nullptr;
+    }
+
+    *id = GLSLstd450Round;
+    return params[0]->result_type();
+  }
+
+  return nullptr;
+}
+
 }  // namespace tint
diff --git a/src/type_determiner.h b/src/type_determiner.h
index 4ad39b1..ca73c6c 100644
--- a/src/type_determiner.h
+++ b/src/type_determiner.h
@@ -71,6 +71,10 @@
   /// @param stmt the statement to check
   /// @returns true if the determination was successful
   bool DetermineResultType(ast::Statement* stmt);
+  /// Determines type information for an expression list
+  /// @param list the expression list to check
+  /// @returns true if the determination was successful
+  bool DetermineResultType(const ast::ExpressionList& list);
   /// Determines type information for an expression
   /// @param expr the expression to check
   /// @returns true if the determination was successful
@@ -87,6 +91,17 @@
     variable_stack_.set(var->name(), var);
   }
 
+  /// Retrieves information for the requested import.
+  /// @param path the import path
+  /// @param name the method name to get information on
+  /// @param params the parameters to the method call
+  /// @param id out parameter for the external call ID. Must not be a nullptr.
+  /// @returns the return type of |name| in |path| or nullptr on error.
+  ast::type::Type* GetImportData(const std::string& path,
+                                 const std::string& name,
+                                 const ast::ExpressionList& params,
+                                 uint32_t* id);
+
  private:
   void set_error(const Source& src, const std::string& msg);
 
diff --git a/src/type_determiner_test.cc b/src/type_determiner_test.cc
index 1fab3b3..7cf3eae 100644
--- a/src/type_determiner_test.cc
+++ b/src/type_determiner_test.cc
@@ -16,8 +16,10 @@
 
 #include <memory>
 #include <utility>
+#include <vector>
 
 #include "gtest/gtest.h"
+#include "spirv/unified1/GLSL.std.450.h"
 #include "src/ast/array_accessor_expression.h"
 #include "src/ast/as_expression.h"
 #include "src/ast/assignment_statement.h"
@@ -570,6 +572,27 @@
   EXPECT_TRUE(param_ptr->result_type()->IsF32());
 }
 
+TEST_F(TypeDeterminerTest, Expr_Call_GLSLImport) {
+  ast::type::F32Type f32;
+
+  mod()->AddImport(std::make_unique<ast::Import>("GLSL.std.450", "std"));
+
+  // Register the function
+  EXPECT_TRUE(td()->Determine());
+
+  ast::ExpressionList call_params;
+  call_params.push_back(std::make_unique<ast::ScalarConstructorExpression>(
+      std::make_unique<ast::FloatLiteral>(&f32, 2.4)));
+
+  std::vector<std::string> name{"std", "round"};
+  ast::CallExpression call(std::make_unique<ast::IdentifierExpression>(name),
+                           std::move(call_params));
+
+  EXPECT_TRUE(td()->DetermineResultType(&call));
+  ASSERT_NE(call.result_type(), nullptr);
+  EXPECT_TRUE(call.result_type()->IsF32());
+}
+
 TEST_F(TypeDeterminerTest, Expr_Cast) {
   ast::type::F32Type f32;
   ast::CastExpression cast(&f32,
@@ -1456,5 +1479,92 @@
             "function variable has a non-function storage class");
 }
 
+TEST_F(TypeDeterminerTest, ImportData_Round_Scalar) {
+  ast::type::F32Type f32;
+
+  ast::ExpressionList params;
+  params.push_back(std::make_unique<ast::ScalarConstructorExpression>(
+      std::make_unique<ast::IntLiteral>(&f32, 1.f)));
+
+  ASSERT_TRUE(td()->DetermineResultType(params)) << td()->error();
+
+  uint32_t id = 0;
+  auto* type = td()->GetImportData("GLSL.std.450", "round", params, &id);
+  ASSERT_NE(type, nullptr);
+  EXPECT_TRUE(type->is_float_scalar());
+  EXPECT_EQ(id, GLSLstd450Round);
+}
+
+TEST_F(TypeDeterminerTest, ImportData_Round_Vector) {
+  ast::type::F32Type f32;
+  ast::type::VectorType vec(&f32, 3);
+
+  ast::ExpressionList vals;
+  vals.push_back(std::make_unique<ast::ScalarConstructorExpression>(
+      std::make_unique<ast::FloatLiteral>(&f32, 1.0f)));
+  vals.push_back(std::make_unique<ast::ScalarConstructorExpression>(
+      std::make_unique<ast::FloatLiteral>(&f32, 1.0f)));
+  vals.push_back(std::make_unique<ast::ScalarConstructorExpression>(
+      std::make_unique<ast::FloatLiteral>(&f32, 3.0f)));
+
+  ast::ExpressionList params;
+  params.push_back(
+      std::make_unique<ast::TypeConstructorExpression>(&vec, std::move(vals)));
+
+  ASSERT_TRUE(td()->DetermineResultType(params)) << td()->error();
+
+  uint32_t id = 0;
+  auto* type = td()->GetImportData("GLSL.std.450", "round", params, &id);
+  ASSERT_NE(type, nullptr);
+  EXPECT_TRUE(type->is_float_vector());
+  EXPECT_EQ(type->AsVector()->size(), 3);
+  EXPECT_EQ(id, GLSLstd450Round);
+}
+
+TEST_F(TypeDeterminerTest, ImportData_Round_Error_Integer) {
+  ast::type::I32Type i32;
+
+  ast::ExpressionList params;
+  params.push_back(std::make_unique<ast::ScalarConstructorExpression>(
+      std::make_unique<ast::IntLiteral>(&i32, 1)));
+
+  ASSERT_TRUE(td()->DetermineResultType(params)) << td()->error();
+
+  uint32_t id = 0;
+  auto* type = td()->GetImportData("GLSL.std.450", "round", params, &id);
+  ASSERT_EQ(type, nullptr);
+  EXPECT_EQ(
+      td()->error(),
+      "incorrect type for round. Requires a float scalar or a float vector");
+}
+
+TEST_F(TypeDeterminerTest, ImportData_Round_Error_NoParams) {
+  ast::ExpressionList params;
+  uint32_t id = 0;
+  auto* type = td()->GetImportData("GLSL.std.450", "round", params, &id);
+  ASSERT_EQ(type, nullptr);
+  EXPECT_EQ(td()->error(),
+            "incorrect number of parameters for round. Expected 1 got 0");
+}
+
+TEST_F(TypeDeterminerTest, ImportData_Round_Error_MultipleParams) {
+  ast::type::F32Type f32;
+  ast::ExpressionList params;
+  params.push_back(std::make_unique<ast::ScalarConstructorExpression>(
+      std::make_unique<ast::FloatLiteral>(&f32, 1.f)));
+  params.push_back(std::make_unique<ast::ScalarConstructorExpression>(
+      std::make_unique<ast::FloatLiteral>(&f32, 1.f)));
+  params.push_back(std::make_unique<ast::ScalarConstructorExpression>(
+      std::make_unique<ast::FloatLiteral>(&f32, 1.f)));
+
+  ASSERT_TRUE(td()->DetermineResultType(params)) << td()->error();
+
+  uint32_t id = 0;
+  auto* type = td()->GetImportData("GLSL.std.450", "round", params, &id);
+  ASSERT_EQ(type, nullptr);
+  EXPECT_EQ(td()->error(),
+            "incorrect number of parameters for round. Expected 1 got 3");
+}
+
 }  // namespace
 }  // namespace tint
diff --git a/src/writer/spirv/builder.cc b/src/writer/spirv/builder.cc
index 4b5f67b..66136a2 100644
--- a/src/writer/spirv/builder.cc
+++ b/src/writer/spirv/builder.cc
@@ -22,6 +22,7 @@
 #include "src/ast/binding_decoration.h"
 #include "src/ast/bool_literal.h"
 #include "src/ast/builtin_decoration.h"
+#include "src/ast/call_expression.h"
 #include "src/ast/constructor_expression.h"
 #include "src/ast/decorated_variable.h"
 #include "src/ast/else_statement.h"
@@ -430,10 +431,18 @@
 
 uint32_t Builder::GenerateIdentifierExpression(
     ast::IdentifierExpression* expr) {
-  // TODO(dsinclair): handle names with namespaces in them ...
-
   uint32_t val = 0;
-  if (!scope_stack_.get(expr->name(), &val)) {
+  if (expr->has_path()) {
+    auto* imp = mod_->FindImportByName(expr->path());
+    if (imp == nullptr) {
+      error_ = "unable to find import for " + expr->path();
+      return 0;
+    }
+    val = imp->GetIdForMethod(expr->name());
+    if (val == 0) {
+      error_ = "unable to lookup: " + expr->name() + " in " + expr->path();
+    }
+  } else if (!scope_stack_.get(expr->name(), &val)) {
     error_ = "unable to find name for identifier: " + expr->name();
     return 0;
   }
@@ -691,6 +700,59 @@
   return result_id;
 }
 
+uint32_t Builder::GenerateCallExpression(ast::CallExpression* expr) {
+  // TODO(dsinclair): Support regular function calls
+  if (!expr->func()->IsIdentifier() ||
+      !expr->func()->AsIdentifier()->has_path()) {
+    error_ = "function calls not supported yet.";
+    return 0;
+  }
+
+  auto* ident = expr->func()->AsIdentifier();
+
+  auto type_id = GenerateTypeIfNeeded(expr->func()->result_type());
+  if (type_id == 0) {
+    return 0;
+  }
+
+  auto set_iter = import_name_to_id_.find(ident->path());
+  if (set_iter == import_name_to_id_.end()) {
+    error_ = "unknown import " + ident->path();
+    return 0;
+  }
+  auto set_id = set_iter->second;
+
+  auto* imp = mod_->FindImportByName(ident->path());
+  if (imp == nullptr) {
+    error_ = "unknown import " + ident->path();
+    return 0;
+  }
+
+  auto inst_id = imp->GetIdForMethod(ident->name());
+  if (inst_id == 0) {
+    error_ = "unknown method " + ident->name();
+    return 0;
+  }
+
+  auto result = result_op();
+  auto result_id = result.to_i();
+
+  std::vector<Operand> ops{Operand::Int(type_id), result, Operand::Int(set_id),
+                           Operand::Int(inst_id)};
+
+  for (const auto& param : expr->params()) {
+    auto id = GenerateExpression(param.get());
+    if (id == 0) {
+      return 0;
+    }
+    ops.push_back(Operand::Int(id));
+  }
+
+  push_function_inst(spv::Op::OpExtInst, std::move(ops));
+
+  return result_id;
+}
+
 bool Builder::GenerateConditionalBlock(
     ast::Expression* cond,
     const ast::StatementList& true_body,
diff --git a/src/writer/spirv/builder.h b/src/writer/spirv/builder.h
index 36f53d0..ae8377d 100644
--- a/src/writer/spirv/builder.h
+++ b/src/writer/spirv/builder.h
@@ -202,6 +202,10 @@
   /// @param expr the expression to generate
   /// @returns the expression ID on success or 0 otherwise
   uint32_t GenerateBinaryExpression(ast::BinaryExpression* expr);
+  /// Generates a call expression
+  /// @param expr the expression to generate
+  /// @returns the expression ID on success or 0 otherwise
+  uint32_t GenerateCallExpression(ast::CallExpression* expr);
   /// Generates a loop statement
   /// @param stmt the statement to generate
   /// @returns true on successful generation
diff --git a/src/writer/spirv/builder_call_test.cc b/src/writer/spirv/builder_call_test.cc
new file mode 100644
index 0000000..7fef65a
--- /dev/null
+++ b/src/writer/spirv/builder_call_test.cc
@@ -0,0 +1,81 @@
+// Copyright 2020 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 <memory>
+
+#include "gtest/gtest.h"
+#include "src/ast/call_expression.h"
+#include "src/ast/float_literal.h"
+#include "src/ast/identifier_expression.h"
+#include "src/ast/scalar_constructor_expression.h"
+#include "src/ast/type/f32_type.h"
+#include "src/ast/type/void_type.h"
+#include "src/context.h"
+#include "src/type_determiner.h"
+#include "src/writer/spirv/builder.h"
+#include "src/writer/spirv/spv_dump.h"
+
+namespace tint {
+namespace writer {
+namespace spirv {
+namespace {
+
+using BuilderTest = testing::Test;
+
+TEST_F(BuilderTest, Call_GLSLMethod) {
+  ast::type::F32Type f32;
+  ast::type::VoidType void_type;
+
+  ast::ExpressionList params;
+  params.push_back(std::make_unique<ast::ScalarConstructorExpression>(
+      std::make_unique<ast::FloatLiteral>(&f32, 1.f)));
+
+  ast::CallExpression expr(std::make_unique<ast::IdentifierExpression>(
+                               std::vector<std::string>{"std", "round"}),
+                           std::move(params));
+
+  Context ctx;
+  ast::Module mod;
+  TypeDeterminer td(&ctx, &mod);
+
+  auto imp = std::make_unique<ast::Import>("GLSL.std.450", "std");
+  auto* glsl = imp.get();
+  mod.AddImport(std::move(imp));
+
+  ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
+
+  ast::Function func("a_func", {}, &void_type);
+
+  Builder b(&mod);
+  b.GenerateImport(glsl);
+  ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
+
+  EXPECT_EQ(b.GenerateCallExpression(&expr), 7u) << b.error();
+  EXPECT_EQ(DumpBuilder(b), R"(%1 = OpExtInstImport "GLSL.std.450"
+OpName %4 "a_func"
+%3 = OpTypeVoid
+%2 = OpTypeFunction %3
+%6 = OpTypeFloat 32
+%8 = OpConstant %6 1
+%4 = OpFunction %3 None %2
+%5 = OpLabel
+%7 = OpExtInst %6 %1 Round %8
+OpFunctionEnd
+)");
+}
+
+}  // namespace
+}  // namespace spirv
+}  // namespace writer
+}  // namespace tint
diff --git a/src/writer/spirv/builder_ident_expression_test.cc b/src/writer/spirv/builder_ident_expression_test.cc
index 05222eb..c692894 100644
--- a/src/writer/spirv/builder_ident_expression_test.cc
+++ b/src/writer/spirv/builder_ident_expression_test.cc
@@ -217,6 +217,31 @@
 )");
 }
 
+TEST_F(BuilderTest, IdentifierExpression_ImportMethod) {
+  auto imp = std::make_unique<ast::Import>("GLSL.std.450", "std");
+  imp->AddMethodId("round", 42u);
+
+  ast::Module mod;
+  mod.AddImport(std::move(imp));
+  Builder b(&mod);
+
+  ast::IdentifierExpression expr(std::vector<std::string>({"std", "round"}));
+  EXPECT_EQ(b.GenerateIdentifierExpression(&expr), 42u) << b.error();
+}
+
+TEST_F(BuilderTest, IdentifierExpression_ImportMethod_NotFound) {
+  auto imp = std::make_unique<ast::Import>("GLSL.std.450", "std");
+
+  ast::Module mod;
+  mod.AddImport(std::move(imp));
+  Builder b(&mod);
+
+  ast::IdentifierExpression expr(std::vector<std::string>({"std", "ceil"}));
+  EXPECT_EQ(b.GenerateIdentifierExpression(&expr), 0u);
+  ASSERT_TRUE(b.has_error());
+  EXPECT_EQ(b.error(), "unable to lookup: ceil in std");
+}
+
 }  // namespace
 }  // namespace spirv
 }  // namespace writer