[spirv-writer] Add function variables

This Cl adds function variables to the SPIR-V output. This requires some
refactoring to split function instructions and variables apart in the
builder.

Bug: tint:5
Change-Id: I4d0045f5a02311cf9a2803929c66c648278e3734
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/18600
Reviewed-by: David Neto <dneto@google.com>
diff --git a/samples/main.cc b/samples/main.cc
index 2937852..b4af202 100644
--- a/samples/main.cc
+++ b/samples/main.cc
@@ -224,9 +224,11 @@
   tools.SetMessageConsumer(msg_consumer);
 
   std::string result;
-  tools.Disassemble(data, &result,
-                    SPV_BINARY_TO_TEXT_OPTION_INDENT |
-                        SPV_BINARY_TO_TEXT_OPTION_FRIENDLY_NAMES);
+  if (!tools.Disassemble(data, &result,
+                         SPV_BINARY_TO_TEXT_OPTION_INDENT |
+                             SPV_BINARY_TO_TEXT_OPTION_FRIENDLY_NAMES)) {
+    std::cerr << spv_errors << std::endl;
+  }
   return result;
 }
 #endif  // TINT_BUILD_SPV_WRITER
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index 1c0d736..8ad3e2c 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -231,6 +231,8 @@
     writer/spirv/binary_writer.h
     writer/spirv/builder.cc
     writer/spirv/builder.h
+    writer/spirv/function.cc
+    writer/spirv/function.h
     writer/spirv/generator.cc
     writer/spirv/generator.h
     writer/spirv/instruction.cc
@@ -416,6 +418,7 @@
     writer/spirv/builder_constructor_expression_test.cc
     writer/spirv/builder_entry_point_test.cc
     writer/spirv/builder_function_test.cc
+    writer/spirv/builder_function_variable_test.cc
     writer/spirv/builder_global_variable_test.cc
     writer/spirv/builder_literal_test.cc
     writer/spirv/builder_return_test.cc
diff --git a/src/writer/spirv/builder.cc b/src/writer/spirv/builder.cc
index 35017b9..5a794f7 100644
--- a/src/writer/spirv/builder.cc
+++ b/src/writer/spirv/builder.cc
@@ -39,6 +39,7 @@
 #include "src/ast/type/vector_type.h"
 #include "src/ast/type_constructor_expression.h"
 #include "src/ast/uint_literal.h"
+#include "src/ast/variable_decl_statement.h"
 
 namespace tint {
 namespace writer {
@@ -125,7 +126,9 @@
   size += size_of(debug_);
   size += size_of(annotations_);
   size += size_of(types_);
-  size += size_of(instructions_);
+  for (const auto& func : functions_) {
+    size += func.word_length();
+  }
 
   return size;
 }
@@ -143,8 +146,8 @@
   for (const auto& inst : types_) {
     cb(inst);
   }
-  for (const auto& inst : instructions_) {
-    cb(inst);
+  for (const auto& func : functions_) {
+    func.iterate(cb);
   }
 }
 
@@ -199,10 +202,13 @@
   }
 
   // TODO(dsinclair): Handle parameters
-  push_inst(spv::Op::OpFunction, {Operand::Int(ret_id), func_op,
-                                  Operand::Int(SpvFunctionControlMaskNone),
-                                  Operand::Int(func_type_id)});
-  push_inst(spv::Op::OpLabel, {result_op()});
+
+  auto definition_inst = Instruction{
+      spv::Op::OpFunction,
+      {Operand::Int(ret_id), func_op, Operand::Int(SpvFunctionControlMaskNone),
+       Operand::Int(func_type_id)}};
+  std::vector<Instruction> params;
+  push_function(Function{definition_inst, result_op(), std::move(params)});
 
   for (const auto& stmt : func->body()) {
     if (!GenerateStatement(stmt.get())) {
@@ -210,8 +216,6 @@
     }
   }
 
-  push_inst(spv::Op::OpFunctionEnd, {});
-
   func_name_to_id_[func->name()] = func_id;
   return true;
 }
@@ -237,6 +241,52 @@
   return func_type_id;
 }
 
+bool Builder::GenerateFunctionVariable(ast::Variable* var) {
+  uint32_t init_id = 0;
+  if (var->has_constructor()) {
+    init_id = GenerateExpression(var->constructor());
+    if (init_id == 0) {
+      return false;
+    }
+  }
+
+  if (var->is_const()) {
+    if (!var->has_constructor()) {
+      error_ = "missing constructor for constant";
+      return false;
+    }
+
+    // TODO(dsinclair): Store variable name to id
+    return true;
+  }
+
+  auto result = result_op();
+  auto var_id = result.to_i();
+  auto sc = ast::StorageClass::kFunction;
+  ast::type::PointerType pt(var->type(), sc);
+  auto type_id = GenerateTypeIfNeeded(&pt);
+  if (type_id == 0) {
+    return false;
+  }
+
+  push_debug(spv::Op::OpName,
+             {Operand::Int(var_id), Operand::String(var->name())});
+
+  // TODO(dsinclair) We could detect if the constructor is fully const and emit
+  // an initializer value for the variable instead of doing the OpLoad.
+
+  push_function_var(
+      {Operand::Int(type_id), result, Operand::Int(ConvertStorageClass(sc))});
+  if (var->has_constructor()) {
+    push_function_inst(spv::Op::OpStore,
+                       {Operand::Int(var_id), Operand::Int(init_id)});
+  }
+
+  // TODO(dsinclair) Mapping of variable name to id
+
+  return true;
+}
+
 bool Builder::GenerateGlobalVariable(ast::Variable* var) {
   uint32_t init_id = 0;
   if (var->has_constructor()) {
@@ -344,10 +394,14 @@
     out << "__const";
 
     std::vector<Operand> ops;
+    bool constructor_is_const = true;
     for (const auto& e : init->values()) {
-      if (is_global_init && !e->IsConstructor()) {
-        error_ = "constructor must be a constant expression";
-        return 0;
+      if (!e->IsConstructor()) {
+        if (is_global_init) {
+          error_ = "constructor must be a constant expression";
+          return 0;
+        }
+        constructor_is_const = false;
       }
       auto id =
           GenerateConstructorExpression(e->AsConstructor(), is_global_init);
@@ -371,9 +425,11 @@
 
     const_to_id_[str] = result.to_i();
 
-    // TODO(dsinclair) For non-global constant's this should be
-    // in the instructions and ben an OpCompositeConstruct call.
-    push_type(spv::Op::OpConstantComposite, ops);
+    if (constructor_is_const) {
+      push_type(spv::Op::OpConstantComposite, ops);
+    } else {
+      push_function_inst(spv::Op::OpCompositeConstruct, ops);
+    }
     return result.to_i();
   }
 
@@ -425,9 +481,9 @@
     if (val_id == 0) {
       return false;
     }
-    push_inst(spv::Op::OpReturnValue, {Operand::Int(val_id)});
+    push_function_inst(spv::Op::OpReturnValue, {Operand::Int(val_id)});
   } else {
-    push_inst(spv::Op::OpReturn, {});
+    push_function_inst(spv::Op::OpReturn, {});
   }
 
   return true;
@@ -437,11 +493,18 @@
   if (stmt->IsReturn()) {
     return GenerateReturnStatement(stmt->AsReturn());
   }
+  if (stmt->IsVariableDecl()) {
+    return GenerateVariableDeclStatement(stmt->AsVariableDecl());
+  }
 
   error_ = "Unknown statement";
   return false;
 }
 
+bool Builder::GenerateVariableDeclStatement(ast::VariableDeclStatement* stmt) {
+  return GenerateFunctionVariable(stmt->variable());
+}
+
 uint32_t Builder::GenerateTypeIfNeeded(ast::type::Type* type) {
   if (type->IsAlias()) {
     return GenerateTypeIfNeeded(type->AsAlias()->type());
diff --git a/src/writer/spirv/builder.h b/src/writer/spirv/builder.h
index 79082d7..e6f0d17 100644
--- a/src/writer/spirv/builder.h
+++ b/src/writer/spirv/builder.h
@@ -25,6 +25,7 @@
 #include "src/ast/literal.h"
 #include "src/ast/module.h"
 #include "src/ast/struct_member.h"
+#include "src/writer/spirv/function.h"
 #include "src/writer/spirv/instruction.h"
 
 namespace tint {
@@ -107,14 +108,6 @@
   }
   /// @returns the type instructions
   const std::vector<Instruction>& types() const { return types_; }
-  /// Adds an instruction to the instruction list
-  /// @param op the op to set
-  /// @param operands the operands for the instruction
-  void push_inst(spv::Op op, const std::vector<Operand>& operands) {
-    instructions_.push_back(Instruction{op, operands});
-  }
-  /// @returns the instruction list
-  const std::vector<Instruction>& instructions() const { return instructions_; }
   /// Adds an instruction to the annotations
   /// @param op the op to set
   /// @param operands the operands for the instruction
@@ -124,6 +117,23 @@
   /// @returns the annotations
   const std::vector<Instruction>& annots() const { return annotations_; }
 
+  /// Adds a function to the builder
+  /// @param func the function to add
+  void push_function(const Function& func) { functions_.push_back(func); }
+  /// @returns the functions
+  const std::vector<Function>& functions() const { return functions_; }
+  /// Pushes an instruction to the current function
+  /// @param op the operation
+  /// @param operands the operands
+  void push_function_inst(spv::Op op, const std::vector<Operand>& operands) {
+    functions_.back().push_inst(op, operands);
+  }
+  /// Pushes a variable to the current function
+  /// @param operands the variable operands
+  void push_function_var(const std::vector<Operand>& operands) {
+    functions_.back().push_var(operands);
+  }
+
   /// Converts a storage class to a SPIR-V storage class.
   /// @param klass the storage class to convert
   /// @returns the SPIR-V storage class or SpvStorageClassMax on error.
@@ -149,6 +159,10 @@
   /// @param func the function to generate for
   /// @returns the ID to use for the function type. Returns 0 on failure.
   uint32_t GenerateFunctionTypeIfNeeded(ast::Function* func);
+  /// Generates a function variable
+  /// @param var the variable
+  /// @returns true if the variable was generated
+  bool GenerateFunctionVariable(ast::Variable* var);
   /// Generates a global variable
   /// @param var the variable to generate
   /// @returns true if the variable is emited.
@@ -199,11 +213,6 @@
   /// @returns true if the vector was successfully generated
   bool GenerateStructType(ast::type::StructType* struct_type,
                           const Operand& result);
-  /// Generates a vector type declaration
-  /// @param vec the vector to generate
-  /// @param result the result operand
-  /// @returns true if the vector was successfully generated
-  bool GenerateVectorType(ast::type::VectorType* vec, const Operand& result);
   /// Generates a struct member
   /// @param struct_id the id of the parent structure
   /// @param idx the index of the member
@@ -212,6 +221,15 @@
   uint32_t GenerateStructMember(uint32_t struct_id,
                                 uint32_t idx,
                                 ast::StructMember* member);
+  /// Generates a variable declaration statement
+  /// @param stmt the statement to generate
+  /// @returns true on successfull generation
+  bool GenerateVariableDeclStatement(ast::VariableDeclStatement* stmt);
+  /// Generates a vector type declaration
+  /// @param vec the vector to generate
+  /// @param result the result operand
+  /// @returns true if the vector was successfully generated
+  bool GenerateVectorType(ast::type::VectorType* vec, const Operand& result);
 
  private:
   /// @returns an Operand with a new result ID in it. Increments the next_id_
@@ -223,8 +241,8 @@
   std::vector<Instruction> preamble_;
   std::vector<Instruction> debug_;
   std::vector<Instruction> types_;
-  std::vector<Instruction> instructions_;
   std::vector<Instruction> annotations_;
+  std::vector<Function> functions_;
 
   std::unordered_map<std::string, uint32_t> import_name_to_id_;
   std::unordered_map<std::string, uint32_t> func_name_to_id_;
diff --git a/src/writer/spirv/builder_function_test.cc b/src/writer/spirv/builder_function_test.cc
index 7022e0d..d5ee73e 100644
--- a/src/writer/spirv/builder_function_test.cc
+++ b/src/writer/spirv/builder_function_test.cc
@@ -41,9 +41,10 @@
   EXPECT_EQ(DumpInstructions(b.types()), R"(%2 = OpTypeVoid
 %1 = OpTypeFunction %2
 )");
-  EXPECT_EQ(DumpInstructions(b.instructions()), R"(%3 = OpFunction %2 None %1
-%4 = OpLabel
-OpFunctionEnd
+
+  ASSERT_GE(b.functions().size(), 1);
+  const auto& ret = b.functions()[0];
+  EXPECT_EQ(DumpInstruction(ret.declaration()), R"(%3 = OpFunction %2 None %1
 )");
 }
 
diff --git a/src/writer/spirv/builder_function_variable_test.cc b/src/writer/spirv/builder_function_variable_test.cc
new file mode 100644
index 0000000..e97f0ef
--- /dev/null
+++ b/src/writer/spirv/builder_function_variable_test.cc
@@ -0,0 +1,180 @@
+// 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/binding_decoration.h"
+#include "src/ast/builtin.h"
+#include "src/ast/builtin_decoration.h"
+#include "src/ast/decorated_variable.h"
+#include "src/ast/float_literal.h"
+#include "src/ast/location_decoration.h"
+#include "src/ast/relational_expression.h"
+#include "src/ast/scalar_constructor_expression.h"
+#include "src/ast/set_decoration.h"
+#include "src/ast/storage_class.h"
+#include "src/ast/type/f32_type.h"
+#include "src/ast/type/vector_type.h"
+#include "src/ast/type_constructor_expression.h"
+#include "src/ast/variable.h"
+#include "src/ast/variable_decoration.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, FunctionVar_NoStorageClass) {
+  ast::type::F32Type f32;
+  ast::Variable v("var", ast::StorageClass::kNone, &f32);
+
+  Builder b;
+  b.push_function(Function{});
+  EXPECT_TRUE(b.GenerateFunctionVariable(&v)) << b.error();
+  EXPECT_EQ(DumpInstructions(b.debug()), R"(OpName %1 "var"
+)");
+  EXPECT_EQ(DumpInstructions(b.types()), R"(%3 = OpTypeFloat 32
+%2 = OpTypePointer Function %3
+)");
+
+  const auto& func = b.functions()[0];
+  EXPECT_EQ(DumpInstructions(func.variables()), R"(%1 = OpVariable %2 Function
+)");
+}
+
+TEST_F(BuilderTest, FunctionVar_WithConstantConstructor) {
+  ast::type::F32Type f32;
+  ast::type::VectorType vec(&f32, 3);
+
+  std::vector<std::unique_ptr<ast::Expression>> 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)));
+
+  auto init =
+      std::make_unique<ast::TypeConstructorExpression>(&vec, std::move(vals));
+
+  ast::Variable v("var", ast::StorageClass::kOutput, &f32);
+  v.set_constructor(std::move(init));
+
+  Builder b;
+  b.push_function(Function{});
+  EXPECT_TRUE(b.GenerateFunctionVariable(&v)) << b.error();
+  ASSERT_FALSE(b.has_error()) << b.error();
+
+  EXPECT_EQ(DumpInstructions(b.debug()), R"(OpName %6 "var"
+)");
+  EXPECT_EQ(DumpInstructions(b.types()), R"(%2 = OpTypeFloat 32
+%1 = OpTypeVector %2 3
+%3 = OpConstant %2 1
+%4 = OpConstant %2 3
+%5 = OpConstantComposite %1 %3 %3 %4
+%7 = OpTypePointer Function %2
+)");
+  EXPECT_EQ(DumpInstructions(b.functions()[0].variables()),
+            R"(%6 = OpVariable %7 Function
+)");
+  EXPECT_EQ(DumpInstructions(b.functions()[0].instructions()), R"(OpStore %6 %5
+)");
+}
+
+// DISABLED until we have RelationalExpression Output
+TEST_F(BuilderTest, DISABLED_FunctionVar_WithNonConstantConstructor) {
+  ast::type::F32Type f32;
+  ast::type::VectorType vec(&f32, 2);
+
+  auto rel = std::make_unique<ast::RelationalExpression>(
+      ast::Relation::kAdd,
+      std::make_unique<ast::ScalarConstructorExpression>(
+          std::make_unique<ast::FloatLiteral>(&f32, 3.0f)),
+      std::make_unique<ast::ScalarConstructorExpression>(
+          std::make_unique<ast::FloatLiteral>(&f32, 3.0f)));
+
+  std::vector<std::unique_ptr<ast::Expression>> vals;
+  vals.push_back(std::make_unique<ast::ScalarConstructorExpression>(
+      std::make_unique<ast::FloatLiteral>(&f32, 1.0f)));
+  vals.push_back(std::move(rel));
+
+  auto init =
+      std::make_unique<ast::TypeConstructorExpression>(&vec, std::move(vals));
+
+  ast::Variable v("var", ast::StorageClass::kOutput, &f32);
+  v.set_constructor(std::move(init));
+
+  Builder b;
+  b.push_function(Function{});
+  EXPECT_TRUE(b.GenerateFunctionVariable(&v)) << b.error();
+  ASSERT_FALSE(b.has_error()) << b.error();
+
+  EXPECT_EQ(DumpInstructions(b.debug()), R"(OpName %6 "var"
+)");
+  EXPECT_EQ(DumpInstructions(b.types()), R"(%2 = OpTypeFloat 32
+%1 = OpTypeVector %2 2
+%3 = OpConstant %2 1
+%4 = OpConstant %2 3
+%7 = OpTypePointer Output %2
+)");
+  EXPECT_EQ(DumpInstructions(b.functions()[0].variables()),
+            R"(%6 = OpVariable %7 Output %5
+)");
+  EXPECT_EQ(DumpInstructions(b.functions()[0].instructions()),
+            R"(%10 = OpIAdd %3 %3
+%9 = OpCompositeConstruct %1 %3 %10
+OpStore %6 %9
+)");
+}
+
+TEST_F(BuilderTest, FunctionVar_Const) {
+  ast::type::F32Type f32;
+  ast::type::VectorType vec(&f32, 3);
+
+  std::vector<std::unique_ptr<ast::Expression>> 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)));
+
+  auto init =
+      std::make_unique<ast::TypeConstructorExpression>(&vec, std::move(vals));
+
+  ast::Variable v("var", ast::StorageClass::kOutput, &f32);
+  v.set_constructor(std::move(init));
+  v.set_is_const(true);
+
+  Builder b;
+  EXPECT_TRUE(b.GenerateFunctionVariable(&v)) << b.error();
+  ASSERT_FALSE(b.has_error()) << b.error();
+
+  EXPECT_EQ(DumpInstructions(b.types()), R"(%2 = OpTypeFloat 32
+%1 = OpTypeVector %2 3
+%3 = OpConstant %2 1
+%4 = OpConstant %2 3
+%5 = OpConstantComposite %1 %3 %3 %4
+)");
+}
+
+}  // namespace
+}  // namespace spirv
+}  // namespace writer
+}  // namespace tint
diff --git a/src/writer/spirv/builder_return_test.cc b/src/writer/spirv/builder_return_test.cc
index f61805b..2792876 100644
--- a/src/writer/spirv/builder_return_test.cc
+++ b/src/writer/spirv/builder_return_test.cc
@@ -35,10 +35,11 @@
   ast::ReturnStatement ret;
 
   Builder b;
+  b.push_function(Function{});
   EXPECT_TRUE(b.GenerateReturnStatement(&ret));
   ASSERT_FALSE(b.has_error()) << b.error();
 
-  EXPECT_EQ(DumpInstructions(b.instructions()), R"(OpReturn
+  EXPECT_EQ(DumpInstructions(b.functions()[0].instructions()), R"(OpReturn
 )");
 }
 
@@ -60,6 +61,7 @@
   ast::ReturnStatement ret(std::move(val));
 
   Builder b;
+  b.push_function(Function{});
   EXPECT_TRUE(b.GenerateReturnStatement(&ret));
   ASSERT_FALSE(b.has_error()) << b.error();
 
@@ -69,7 +71,8 @@
 %4 = OpConstant %2 3
 %5 = OpConstantComposite %1 %3 %3 %4
 )");
-  EXPECT_EQ(DumpInstructions(b.instructions()), R"(OpReturnValue %5
+  EXPECT_EQ(DumpInstructions(b.functions()[0].instructions()),
+            R"(OpReturnValue %5
 )");
 }
 
diff --git a/src/writer/spirv/function.cc b/src/writer/spirv/function.cc
new file mode 100644
index 0000000..2b4dae2
--- /dev/null
+++ b/src/writer/spirv/function.cc
@@ -0,0 +1,53 @@
+// 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 "src/writer/spirv/function.h"
+
+namespace tint {
+namespace writer {
+namespace spirv {
+
+Function::Function()
+    : declaration_(Instruction{spv::Op::OpNop, {}}),
+      label_op_(Operand::Int(0)) {}
+
+Function::Function(const Instruction& declaration,
+                   const Operand& label_op,
+                   const std::vector<Instruction>& params)
+    : declaration_(declaration), label_op_(label_op), params_(params) {}
+
+Function::~Function() = default;
+
+void Function::iterate(std::function<void(const Instruction&)> cb) const {
+  cb(declaration_);
+
+  for (const auto& param : params_) {
+    cb(param);
+  }
+
+  cb(Instruction{spv::Op::OpLabel, {label_op_}});
+
+  for (const auto& var : vars_) {
+    cb(var);
+  }
+  for (const auto& inst : instructions_) {
+    cb(inst);
+  }
+
+  cb(Instruction{spv::Op::OpFunctionEnd, {}});
+}
+
+}  // namespace spirv
+}  // namespace writer
+}  // namespace tint
diff --git a/src/writer/spirv/function.h b/src/writer/spirv/function.h
new file mode 100644
index 0000000..e15377f
--- /dev/null
+++ b/src/writer/spirv/function.h
@@ -0,0 +1,100 @@
+// 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.
+
+#ifndef SRC_WRITER_SPIRV_FUNCTION_H_
+#define SRC_WRITER_SPIRV_FUNCTION_H_
+
+#include <vector>
+
+#include "spirv/unified1/spirv.hpp11"
+#include "src/writer/spirv/instruction.h"
+#include "src/writer/spirv/operand.h"
+
+namespace tint {
+namespace writer {
+namespace spirv {
+
+/// A SPIR-V function
+class Function {
+ public:
+  /// Constructor for testing purposes
+  /// This creates a bad declaration, so won't generate correct SPIR-V
+  Function();
+
+  /// Constructor
+  /// @param declaration the function declaration
+  /// @param label_op the operand for the initial function label
+  /// @param params the function parameters
+  Function(const Instruction& declaration,
+           const Operand& label_op,
+           const std::vector<Instruction>& params);
+  /// Copy constructor
+  /// @param other the function to copy
+  Function(const Function& other) = default;
+  ~Function();
+
+  /// Iterates over the function call the cb on each instruction
+  /// @param cb the callback to call
+  void iterate(std::function<void(const Instruction&)> cb) const;
+
+  /// @returns the declaration
+  const Instruction& declaration() const { return declaration_; }
+
+  /// Adds an instruction to the instruction list
+  /// @param op the op to set
+  /// @param operands the operands for the instruction
+  void push_inst(spv::Op op, const std::vector<Operand>& operands) {
+    instructions_.push_back(Instruction{op, operands});
+  }
+  /// @returns the instruction list
+  const std::vector<Instruction>& instructions() const { return instructions_; }
+
+  /// Adds a variable to the variable list
+  /// @param operands the operands for the variable
+  void push_var(const std::vector<Operand>& operands) {
+    vars_.push_back(Instruction{spv::Op::OpVariable, operands});
+  }
+  /// @returns the variable list
+  const std::vector<Instruction>& variables() const { return vars_; }
+
+  /// @returns the word length of the function
+  uint32_t word_length() const {
+    // 1 for the Label and 1 for the FunctionEnd
+    uint32_t size = 2 + declaration_.word_length();
+
+    for (const auto& param : params_) {
+      size += param.word_length();
+    }
+    for (const auto& var : vars_) {
+      size += var.word_length();
+    }
+    for (const auto& inst : instructions_) {
+      size += inst.word_length();
+    }
+    return size;
+  }
+
+ private:
+  Instruction declaration_;
+  Operand label_op_;
+  std::vector<Instruction> params_;
+  std::vector<Instruction> vars_;
+  std::vector<Instruction> instructions_;
+};
+
+}  // namespace spirv
+}  // namespace writer
+}  // namespace tint
+
+#endif  // SRC_WRITER_SPIRV_FUNCTION_H_