[spirv-writer] Add function calls

This CL adds calls to functions to the SPIR-V writer.

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

src/type_determiner.cc

@@ -172,10 +172,6 @@ bool TypeDeterminer::Determine() {
     }
   }
 
-  for (const auto& func : mod_->functions()) {
-    name_to_function_[func->name()] = func.get();
-  }
-
   if (!DetermineFunctions(mod_->functions())) {
     return false;
   }
@@ -192,7 +188,13 @@ bool TypeDeterminer::DetermineFunctions(const ast::FunctionList& funcs) {
 }
 
 bool TypeDeterminer::DetermineFunction(ast::Function* func) {
+  name_to_function_[func->name()] = func;
+
   variable_stack_.push_scope();
+  for (const auto& param : func->params()) {
+    variable_stack_.set(param->name(), param.get());
+  }
+
   if (!DetermineStatements(func->body())) {
     return false;
   }

src/writer/spirv/builder.cc

@@ -1334,51 +1334,64 @@ uint32_t Builder::GenerateCallExpression(ast::CallExpression* expr) {
     return GenerateIntrinsic(ident->name(), expr);
   }
 
-  // TODO(dsinclair): Support regular function calls
-  if (!ident->has_path()) {
-    error_ = "function calls not supported yet.";
-    return 0;
-  }
-
   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)};
+  spv::Op op = spv::Op::OpNop;
+  std::vector<Operand> ops = {Operand::Int(type_id), result};
+
+  // Handle regular function calls
+  if (!ident->has_path()) {
+    auto func_id = func_name_to_id_[ident->name()];
+    if (func_id == 0) {
+      error_ = "unable to find called function: " + ident->name();
+      return 0;
+    }
+    ops.push_back(Operand::Int(func_id));
+
+    op = spv::Op::OpFunctionCall;
+  } else {
+    // Imported function call
+    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;
+    }
+
+    ops.push_back(Operand::Int(set_id));
+    ops.push_back(Operand::Int(inst_id));
+
+    op = spv::Op::OpExtInst;
+  }
 
   for (const auto& param : expr->params()) {
     auto id = GenerateExpression(param.get());
     if (id == 0) {
       return 0;
     }
-    ops.push_back(Operand::Int(GenerateLoadIfNeeded(param->result_type(), id)));
+    id = GenerateLoadIfNeeded(param->result_type(), id);
+    ops.push_back(Operand::Int(id));
   }
 
-  push_function_inst(spv::Op::OpExtInst, std::move(ops));
+  push_function_inst(op, std::move(ops));
 
   return result_id;
 }

src/writer/spirv/builder_call_test.cc

@@ -15,11 +15,15 @@
 #include <memory>
 
 #include "gtest/gtest.h"
+#include "src/ast/binary_expression.h"
 #include "src/ast/call_expression.h"
 #include "src/ast/float_literal.h"
 #include "src/ast/identifier_expression.h"
+#include "src/ast/return_statement.h"
 #include "src/ast/scalar_constructor_expression.h"
+#include "src/ast/sint_literal.h"
 #include "src/ast/type/f32_type.h"
+#include "src/ast/type/i32_type.h"
 #include "src/ast/type/void_type.h"
 #include "src/context.h"
 #include "src/type_determiner.h"
@@ -125,6 +129,74 @@ OpFunctionEnd
 )");
 }
 
+TEST_F(BuilderTest, Call) {
+  ast::type::F32Type f32;
+  ast::type::VoidType void_type;
+
+  ast::VariableList func_params;
+  func_params.push_back(
+      std::make_unique<ast::Variable>("a", ast::StorageClass::kFunction, &f32));
+  func_params.push_back(
+      std::make_unique<ast::Variable>("b", ast::StorageClass::kFunction, &f32));
+
+  ast::Function a_func("a_func", std::move(func_params), &f32);
+
+  ast::StatementList body;
+  body.push_back(std::make_unique<ast::ReturnStatement>(
+      std::make_unique<ast::BinaryExpression>(
+          ast::BinaryOp::kAdd, std::make_unique<ast::IdentifierExpression>("a"),
+          std::make_unique<ast::IdentifierExpression>("b"))));
+  a_func.set_body(std::move(body));
+
+  ast::Function func("main", {}, &void_type);
+
+  ast::ExpressionList call_params;
+  call_params.push_back(std::make_unique<ast::ScalarConstructorExpression>(
+      std::make_unique<ast::FloatLiteral>(&f32, 1.f)));
+  call_params.push_back(std::make_unique<ast::ScalarConstructorExpression>(
+      std::make_unique<ast::FloatLiteral>(&f32, 1.f)));
+
+  ast::CallExpression expr(
+      std::make_unique<ast::IdentifierExpression>("a_func"),
+      std::move(call_params));
+
+  Context ctx;
+  ast::Module mod;
+  TypeDeterminer td(&ctx, &mod);
+  ASSERT_TRUE(td.DetermineFunction(&func)) << td.error();
+  ASSERT_TRUE(td.DetermineFunction(&a_func)) << td.error();
+  ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
+
+  Builder b(&mod);
+  ASSERT_TRUE(b.GenerateFunction(&a_func)) << b.error();
+  ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
+
+  EXPECT_EQ(b.GenerateCallExpression(&expr), 14u) << b.error();
+  EXPECT_EQ(DumpBuilder(b), R"(OpName %3 "a_func"
+OpName %4 "a"
+OpName %5 "b"
+OpName %12 "main"
+%2 = OpTypeFloat 32
+%1 = OpTypeFunction %2 %2 %2
+%11 = OpTypeVoid
+%10 = OpTypeFunction %11
+%15 = OpConstant %2 1
+%3 = OpFunction %2 None %1
+%4 = OpFunctionParameter %2
+%5 = OpFunctionParameter %2
+%6 = OpLabel
+%7 = OpLoad %2 %4
+%8 = OpLoad %2 %5
+%9 = OpFAdd %2 %7 %8
+OpReturnValue %9
+OpFunctionEnd
+%12 = OpFunction %11 None %10
+%13 = OpLabel
+%14 = OpFunctionCall %2 %3 %15 %15
+OpFunctionEnd
+)");
+}
+
 }  // namespace
 }  // namespace spirv
 }  // namespace writer