9 files changed, 188 insertions, 101 deletions

diff --git a/include/amber/shader_info.h b/include/amber/shader_info.h
index 4f1fc31..9dc5553 100644
--- a/include/amber/shader_info.h
+++ b/include/amber/shader_info.h
@@ -35,6 +35,7 @@ enum ShaderType {
   kShaderTypeVertex,
   kShaderTypeTessellationControl,
   kShaderTypeTessellationEvaluation,
+  kShaderTypeMulti,
 };
 
 struct ShaderInfo {
diff --git a/src/amberscript/parser.cc b/src/amberscript/parser.cc
index c397b85..3ef125a 100644
--- a/src/amberscript/parser.cc
+++ b/src/amberscript/parser.cc
@@ -78,6 +78,8 @@ Result Parser::ToShaderType(const std::string& str, ShaderType* type) {
     *type = kShaderTypeTessellationControl;
   else if (str == "compute")
     *type = kShaderTypeCompute;
+  else if (str == "multi")
+    *type = kShaderTypeMulti;
   else
     return Result("unknown shader type: " + str);
   return {};
@@ -301,8 +303,6 @@ Result Parser::ParsePipelineBlock() {
       break;
     } else if (tok == "ATTACH") {
       r = ParsePipelineAttach(pipeline.get());
-    } else if (tok == "ENTRY_POINT") {
-      r = ParsePipelineEntryPoint(pipeline.get());
     } else if (tok == "SHADER_OPTIMIZATION") {
       r = ParsePipelineShaderOptimizations(pipeline.get());
     } else {
@@ -335,31 +335,58 @@ Result Parser::ParsePipelineAttach(Pipeline* pipeline) {
   if (!shader)
     return Result("unknown shader in ATTACH command");
 
-  Result r = pipeline->AddShader(shader);
-  if (!r.IsSuccess())
-    return r;
-
-  return ValidateEndOfStatement("ATTACH command");
-}
+  token = tokenizer_->NextToken();
+  if (token->IsEOL() || token->IsEOS()) {
+    if (shader->GetType() == kShaderTypeMulti)
+      return Result("multi shader ATTACH requires TYPE");
 
-Result Parser::ParsePipelineEntryPoint(Pipeline* pipeline) {
-  auto token = tokenizer_->NextToken();
+    Result r = pipeline->AddShader(shader, shader->GetType());
+    if (!r.IsSuccess())
+      return r;
+    return {};
+  }
   if (!token->IsString())
-    return Result("missing shader name in ENTRY_POINT command");
+    return Result("Invalid token after ATTACH");
 
-  auto* shader = script_->GetShader(token->AsString());
-  if (!shader)
-    return Result("unknown shader in ENTRY_POINT command");
+  bool set_shader_type = false;
+  ShaderType shader_type = shader->GetType();
+  auto type = token->AsString();
+  if (type == "TYPE") {
+    token = tokenizer_->NextToken();
+    if (!token->IsString())
+      return Result("invalid type in ATTACH");
+
+    Result r = ToShaderType(token->AsString(), &shader_type);
+    if (!r.IsSuccess())
+      return r;
+
+    set_shader_type = true;
+
+    token = tokenizer_->NextToken();
+    if (!token->IsString())
+      return Result("ATTACH TYPE requires an ENTRY_POINT");
+
+    type = token->AsString();
+  }
+  if (type != "ENTRY_POINT")
+    return Result("Unknown ATTACH parameter: " + type);
+
+  if (shader->GetType() == ShaderType::kShaderTypeMulti && !set_shader_type)
+    return Result("ATTACH missing TYPE for multi shader");
+
+  Result r = pipeline->AddShader(shader, shader_type);
+  if (!r.IsSuccess())
+    return r;
 
   token = tokenizer_->NextToken();
   if (!token->IsString())
-    return Result("invalid value in ENTRY_POINT command");
+    return Result("missing shader name in ATTACH ENTRY_POINT command");
 
-  Result r = pipeline->SetShaderEntryPoint(shader, token->AsString());
+  r = pipeline->SetShaderEntryPoint(shader, token->AsString());
   if (!r.IsSuccess())
     return r;
 
-  return ValidateEndOfStatement("ENTRY_POINT command");
+  return ValidateEndOfStatement("ATTACH command");
 }
 
 Result Parser::ParsePipelineShaderOptimizations(Pipeline* pipeline) {
diff --git a/src/amberscript/parser.h b/src/amberscript/parser.h
index e64924a..23046b9 100644
--- a/src/amberscript/parser.h
+++ b/src/amberscript/parser.h
@@ -55,7 +55,6 @@ class Parser : public amber::Parser {
   Result ParseShaderBlock();
   Result ParsePipelineBlock();
   Result ParsePipelineAttach(Pipeline*);
-  Result ParsePipelineEntryPoint(Pipeline*);
   Result ParsePipelineShaderOptimizations(Pipeline*);
 
   std::unique_ptr<Tokenizer> tokenizer_;
diff --git a/src/amberscript/parser_test.cc b/src/amberscript/parser_test.cc
index 56bdd90..c72972c 100644
--- a/src/amberscript/parser_test.cc
+++ b/src/amberscript/parser_test.cc
@@ -150,7 +150,8 @@ INSTANTIATE_TEST_CASE_P(
                     NameData{"geometry"},
                     NameData{"tessellation_evaluation"},
                     NameData{"tessellation_control"},
-                    NameData{"compute"}), );  // NOLINT(whitespace/parens)
+                    NameData{"compute"},
+                    NameData{"multi"}), );  // NOLINT(whitespace/parens)
 
 TEST_F(AmberScriptParserTest, ShaderPassThroughUnknownShaderType) {
   std::string in = "SHADER UNKNOWN my_shader PASSTHROUGH";
@@ -331,8 +332,9 @@ INSTANTIATE_TEST_CASE_P(
         ShaderTypeData{"tessellation_evaluation",
                        kShaderTypeTessellationEvaluation},
         ShaderTypeData{"tessellation_control", kShaderTypeTessellationControl},
-        ShaderTypeData{"compute",
-                       kShaderTypeCompute}), );  // NOLINT(whitespace/parens)
+        ShaderTypeData{"compute", kShaderTypeCompute},
+        ShaderTypeData{"multi",
+                       kShaderTypeMulti}), );  // NOLINT(whitespace/parens)
 
 using AmberScriptParserShaderFormatTest =
     testing::TestWithParam<ShaderFormatData>;
@@ -557,7 +559,7 @@ END)";
   Parser parser;
   Result r = parser.Parse(in);
   ASSERT_FALSE(r.IsSuccess());
-  EXPECT_EQ("5: can not add duplicate shader to pipeline", r.Error());
+  EXPECT_EQ("6: can not add duplicate shader to pipeline", r.Error());
 }
 
 TEST_F(AmberScriptParserTest, AttachInvalidToken) {
@@ -581,7 +583,7 @@ END)";
   Parser parser;
   Result r = parser.Parse(in);
   ASSERT_FALSE(r.IsSuccess());
-  EXPECT_EQ("4: extra parameters after ATTACH command", r.Error());
+  EXPECT_EQ("4: Unknown ATTACH parameter: INVALID", r.Error());
 }
 
 TEST_F(AmberScriptParserTest, AttachMissingValue) {
@@ -611,7 +613,7 @@ END)";
   Parser parser;
   Result r = parser.Parse(in);
   ASSERT_FALSE(r.IsSuccess()) << r.Error();
-  EXPECT_EQ("8: can not add a compute shader to a graphics pipeline",
+  EXPECT_EQ("9: can not add a compute shader to a graphics pipeline",
             r.Error());
 }
 
@@ -633,7 +635,7 @@ END)";
   Parser parser;
   Result r = parser.Parse(in);
   ASSERT_FALSE(r.IsSuccess()) << r.Error();
-  EXPECT_EQ("8: only compute shaders allowed in a compute pipeline", r.Error());
+  EXPECT_EQ("9: only compute shaders allowed in a compute pipeline", r.Error());
 }
 INSTANTIATE_TEST_CASE_P(
     AmberScriptParserPipelineAttachTests,
@@ -656,10 +658,8 @@ SHADER fragment my_fragment GLSL
 END
 
 PIPELINE graphics my_pipeline
-  ATTACH my_shader
+  ATTACH my_shader ENTRY_POINT green
   ATTACH my_fragment
-
-  ENTRY_POINT my_shader green
 END
 )";
 
@@ -690,14 +690,13 @@ SHADER compute my_compute GLSL
 # Compute Shader
 END
 PIPELINE compute my_pipeline
-  ATTACH my_compute
-  ENTRY_POINT my_compute 1234
+  ATTACH my_compute ENTRY_POINT 1234
 END)";
 
   Parser parser;
   Result r = parser.Parse(in);
-  ASSERT_FALSE(r.IsSuccess()) << r.Error();
-  EXPECT_EQ("7: invalid value in ENTRY_POINT command", r.Error());
+  ASSERT_FALSE(r.IsSuccess());
+  EXPECT_EQ("6: missing shader name in ATTACH ENTRY_POINT command", r.Error());
 }
 
 TEST_F(AmberScriptParserTest, PipelineEntryPointMissingValue) {
@@ -706,79 +705,117 @@ SHADER compute my_compute GLSL
 # Compute Shader
 END
 PIPELINE compute my_pipeline
-  ATTACH my_compute
-  ENTRY_POINT
+  ATTACH my_compute ENTRY_POINT
 END)";
 
   Parser parser;
   Result r = parser.Parse(in);
-  ASSERT_FALSE(r.IsSuccess()) << r.Error();
-  EXPECT_EQ("8: missing shader name in ENTRY_POINT command", r.Error());
+  ASSERT_FALSE(r.IsSuccess());
+  EXPECT_EQ("7: missing shader name in ATTACH ENTRY_POINT command", r.Error());
 }
 
-TEST_F(AmberScriptParserTest, PipelineEntryPointMissingEntryPointName) {
+TEST_F(AmberScriptParserTest, PipelineEntryPointExtraParameter) {
   std::string in = R"(
 SHADER compute my_compute GLSL
 # Compute Shader
 END
 PIPELINE compute my_pipeline
-  ATTACH my_compute
-  ENTRY_POINT my_compute
+  ATTACH my_compute ENTRY_POINT green INVALID
 END)";
 
   Parser parser;
   Result r = parser.Parse(in);
-  ASSERT_FALSE(r.IsSuccess()) << r.Error();
-  EXPECT_EQ("8: invalid value in ENTRY_POINT command", r.Error());
+  ASSERT_FALSE(r.IsSuccess());
+  EXPECT_EQ("6: extra parameters after ATTACH command", r.Error());
 }
 
-TEST_F(AmberScriptParserTest, PipelineEntryPointExtraParameter) {
+TEST_F(AmberScriptParserTest, PiplineMultiShaderAttach) {
   std::string in = R"(
-SHADER compute my_compute GLSL
-# Compute Shader
+SHADER multi my_shader GLSL
+# shaders
 END
 PIPELINE compute my_pipeline
-  ATTACH my_compute
-  ENTRY_POINT my_compute green INVALID
+  ATTACH my_shader TYPE compute ENTRY_POINT my_entry_point
 END)";
 
   Parser parser;
   Result r = parser.Parse(in);
-  ASSERT_FALSE(r.IsSuccess()) << r.Error();
-  EXPECT_EQ("7: extra parameters after ENTRY_POINT command", r.Error());
+  ASSERT_TRUE(r.IsSuccess()) << r.Error();
+
+  auto script = parser.GetScript();
+  const auto& pipelines = script->GetPipelines();
+  ASSERT_EQ(1U, pipelines.size());
+
+  const auto* pipeline = pipelines[0].get();
+  const auto& shaders = pipeline->GetShaders();
+  ASSERT_EQ(1U, shaders.size());
+
+  ASSERT_TRUE(shaders[0].GetShader() != nullptr);
+  EXPECT_EQ(kShaderTypeMulti, shaders[0].GetShader()->GetType());
+  EXPECT_EQ(kShaderTypeCompute, shaders[0].GetShaderType());
+  EXPECT_EQ("my_entry_point", shaders[0].GetEntryPoint());
 }
 
-TEST_F(AmberScriptParserTest, PipelineMultipleEntryPointsForOneShader) {
+TEST_F(AmberScriptParserTest,
+       PipelineMultiShaderMismatchPipelineAndShaderType) {
   std::string in = R"(
-SHADER compute my_compute GLSL
-# Compute Shader
+SHADER multi my_shader GLSL
+# shaders
 END
-PIPELINE compute my_pipeline
-  ATTACH my_compute
-  ENTRY_POINT my_compute green
-  ENTRY_POINT my_compute red
+PIPELINE graphics my_pipeline
+  ATTACH my_shader TYPE compute ENTRY_POINT my_entry_point
 END)";
 
   Parser parser;
   Result r = parser.Parse(in);
-  ASSERT_FALSE(r.IsSuccess()) << r.Error();
-  EXPECT_EQ("8: multiple entry points given for the same shader", r.Error());
+  ASSERT_FALSE(r.IsSuccess());
+  EXPECT_EQ("6: can not add a compute shader to a graphics pipeline",
+            r.Error());
 }
 
-TEST_F(AmberScriptParserTest, PipelineEntryPointForInvalidShader) {
+TEST_F(AmberScriptParserTest, PipelineMultiShaderMissingEntryPoint) {
   std::string in = R"(
-SHADER compute my_compute GLSL
-# Compute Shader
+SHADER multi my_shader GLSL
+# shaders
 END
-PIPELINE compute my_pipeline
-  ATTACH my_compute
-  ENTRY_POINT INVALID green
+PIPELINE graphics my_pipeline
+  ATTACH my_shader TYPE fragment
 END)";
 
   Parser parser;
   Result r = parser.Parse(in);
-  ASSERT_FALSE(r.IsSuccess()) << r.Error();
-  EXPECT_EQ("7: unknown shader in ENTRY_POINT command", r.Error());
+  ASSERT_FALSE(r.IsSuccess());
+  EXPECT_EQ("7: ATTACH TYPE requires an ENTRY_POINT", r.Error());
+}
+
+TEST_F(AmberScriptParserTest, PipelineMultiShaderMissingType) {
+  std::string in = R"(
+SHADER multi my_shader GLSL
+# shaders
+END
+PIPELINE graphics my_pipeline
+  ATTACH my_shader
+END)";
+
+  Parser parser;
+  Result r = parser.Parse(in);
+  ASSERT_FALSE(r.IsSuccess());
+  EXPECT_EQ("7: multi shader ATTACH requires TYPE", r.Error());
+}
+
+TEST_F(AmberScriptParserTest, PipelineMultiShaderMissingTypeWithEntryPoint) {
+  std::string in = R"(
+SHADER multi my_shader GLSL
+# shaders
+END
+PIPELINE graphics my_pipeline
+  ATTACH my_shader ENTRY_POINT my_ep
+END)";
+
+  Parser parser;
+  Result r = parser.Parse(in);
+  ASSERT_FALSE(r.IsSuccess());
+  EXPECT_EQ("6: ATTACH missing TYPE for multi shader", r.Error());
 }
 
 TEST_F(AmberScriptParserTest, PipelineShaderOptimization) {
diff --git a/src/pipeline.cc b/src/pipeline.cc
index cde8329..04e8e5d 100644
--- a/src/pipeline.cc
+++ b/src/pipeline.cc
@@ -19,8 +19,8 @@
 
 namespace amber {
 
-Pipeline::ShaderInfo::ShaderInfo(const Shader* shader)
-    : shader_(shader), entry_point_("main") {}
+Pipeline::ShaderInfo::ShaderInfo(const Shader* shader, ShaderType type)
+    : shader_(shader), shader_type_(type), entry_point_("main") {}
 
 Pipeline::ShaderInfo::ShaderInfo(const ShaderInfo&) = default;
 
@@ -30,16 +30,16 @@ Pipeline::Pipeline(PipelineType type) : pipeline_type_(type) {}
 
 Pipeline::~Pipeline() = default;
 
-Result Pipeline::AddShader(const Shader* shader) {
+Result Pipeline::AddShader(const Shader* shader, ShaderType shader_type) {
   if (!shader)
     return Result("shader can not be null when attached to pipeline");
 
   if (pipeline_type_ == PipelineType::kCompute &&
-      shader->GetType() != kShaderTypeCompute) {
+      shader_type != kShaderTypeCompute) {
     return Result("only compute shaders allowed in a compute pipeline");
   }
   if (pipeline_type_ == PipelineType::kGraphics &&
-      shader->GetType() == kShaderTypeCompute) {
+      shader_type == kShaderTypeCompute) {
     return Result("can not add a compute shader to a graphics pipeline");
   }
 
@@ -47,11 +47,11 @@ Result Pipeline::AddShader(const Shader* shader) {
     const auto* is = info.GetShader();
     if (is == shader)
       return Result("can not add duplicate shader to pipeline");
-    if (is->GetType() == shader->GetType())
+    if (is->GetType() == shader_type)
       return Result("can not add duplicate shader type to pipeline");
   }
 
-  shaders_.emplace_back(shader);
+  shaders_.emplace_back(shader, shader_type);
   return {};
 }
 
@@ -101,6 +101,21 @@ Result Pipeline::SetShaderEntryPoint(const Shader* shader,
                 shader->GetName());
 }
 
+Result Pipeline::SetShaderType(const Shader* shader, ShaderType type) {
+  if (!shader)
+    return Result("invalid shader specified for shader type");
+
+  for (auto& info : shaders_) {
+    if (info.GetShader() == shader) {
+      info.SetShaderType(type);
+      return {};
+    }
+  }
+
+  return Result("unknown shader specified for shader type: " +
+                shader->GetName());
+}
+
 Result Pipeline::Validate() const {
   if (pipeline_type_ == PipelineType::kGraphics)
     return ValidateGraphics();
diff --git a/src/pipeline.h b/src/pipeline.h
index 881ecc6..0c6ed55 100644
--- a/src/pipeline.h
+++ b/src/pipeline.h
@@ -29,7 +29,7 @@ class Pipeline {
  public:
   class ShaderInfo {
    public:
-    explicit ShaderInfo(const Shader*);
+    ShaderInfo(const Shader*, ShaderType type);
     ShaderInfo(const ShaderInfo&);
     ~ShaderInfo();
 
@@ -45,8 +45,12 @@ class Pipeline {
     void SetEntryPoint(const std::string& ep) { entry_point_ = ep; }
     std::string GetEntryPoint() const { return entry_point_; }
 
+    void SetShaderType(ShaderType type) { shader_type_ = type; }
+    ShaderType GetShaderType() const { return shader_type_; }
+
    private:
     const Shader* shader_ = nullptr;
+    ShaderType shader_type_;
     std::vector<std::string> shader_optimizations_;
     std::string entry_point_;
   };
@@ -59,9 +63,10 @@ class Pipeline {
   void SetName(const std::string& name) { name_ = name; }
   const std::string& GetName() const { return name_; }
 
-  Result AddShader(const Shader*);
+  Result AddShader(const Shader*, ShaderType);
   const std::vector<ShaderInfo>& GetShaders() const { return shaders_; }
 
+  Result SetShaderType(const Shader* shader, ShaderType type);
   Result SetShaderEntryPoint(const Shader* shader, const std::string& name);
   Result SetShaderOptimizations(const Shader* shader,
                                 const std::vector<std::string>& opts);
diff --git a/src/pipeline_test.cc b/src/pipeline_test.cc
index 635d2e8..7a54a45 100644
--- a/src/pipeline_test.cc
+++ b/src/pipeline_test.cc
@@ -31,10 +31,10 @@ TEST_F(AmberScriptPipelineTest, AddShader) {
   Shader f(kShaderTypeFragment);
 
   Pipeline p(PipelineType::kGraphics);
-  Result r = p.AddShader(&v);
+  Result r = p.AddShader(&v, kShaderTypeVertex);
   ASSERT_TRUE(r.IsSuccess()) << r.Error();
 
-  r = p.AddShader(&f);
+  r = p.AddShader(&f, kShaderTypeFragment);
   ASSERT_TRUE(r.IsSuccess()) << r.Error();
 
   const auto& shaders = p.GetShaders();
@@ -46,7 +46,7 @@ TEST_F(AmberScriptPipelineTest, AddShader) {
 
 TEST_F(AmberScriptPipelineTest, MissingShader) {
   Pipeline p(PipelineType::kGraphics);
-  Result r = p.AddShader(nullptr);
+  Result r = p.AddShader(nullptr, kShaderTypeVertex);
   ASSERT_FALSE(r.IsSuccess());
   EXPECT_EQ("shader can not be null when attached to pipeline", r.Error());
 }
@@ -56,13 +56,13 @@ TEST_F(AmberScriptPipelineTest, DuplicateShaders) {
   Shader f(kShaderTypeFragment);
 
   Pipeline p(PipelineType::kGraphics);
-  Result r = p.AddShader(&v);
+  Result r = p.AddShader(&v, kShaderTypeVertex);
   ASSERT_TRUE(r.IsSuccess()) << r.Error();
 
-  r = p.AddShader(&f);
+  r = p.AddShader(&f, kShaderTypeFragment);
   ASSERT_TRUE(r.IsSuccess()) << r.Error();
 
-  r = p.AddShader(&v);
+  r = p.AddShader(&v, kShaderTypeVertex);
   ASSERT_FALSE(r.IsSuccess());
   EXPECT_EQ("can not add duplicate shader to pipeline", r.Error());
 }
@@ -72,10 +72,10 @@ TEST_F(AmberScriptPipelineTest, DuplicateShaderType) {
   Shader f(kShaderTypeVertex);
 
   Pipeline p(PipelineType::kGraphics);
-  Result r = p.AddShader(&v);
+  Result r = p.AddShader(&v, kShaderTypeVertex);
   ASSERT_TRUE(r.IsSuccess()) << r.Error();
 
-  r = p.AddShader(&f);
+  r = p.AddShader(&f, kShaderTypeVertex);
   ASSERT_FALSE(r.IsSuccess());
   EXPECT_EQ("can not add duplicate shader type to pipeline", r.Error());
 }
@@ -89,7 +89,7 @@ TEST_P(AmberScriptPipelineComputePipelineTest,
   Shader s(test_data.type);
 
   Pipeline p(PipelineType::kCompute);
-  Result r = p.AddShader(&s);
+  Result r = p.AddShader(&s, test_data.type);
   ASSERT_FALSE(r.IsSuccess());
   EXPECT_EQ("only compute shaders allowed in a compute pipeline", r.Error());
 }
@@ -108,7 +108,7 @@ TEST_F(AmberScriptPipelineTest, SettingComputeShaderToGraphicsPipeline) {
   Shader c(kShaderTypeCompute);
 
   Pipeline p(PipelineType::kGraphics);
-  Result r = p.AddShader(&c);
+  Result r = p.AddShader(&c, kShaderTypeCompute);
   ASSERT_FALSE(r.IsSuccess());
   EXPECT_EQ("can not add a compute shader to a graphics pipeline", r.Error());
 }
@@ -118,10 +118,10 @@ TEST_F(AmberScriptPipelineTest, SetShaderOptimizations) {
   Shader f(kShaderTypeFragment);
 
   Pipeline p(PipelineType::kGraphics);
-  Result r = p.AddShader(&v);
+  Result r = p.AddShader(&v, kShaderTypeVertex);
   ASSERT_TRUE(r.IsSuccess()) << r.Error();
 
-  r = p.AddShader(&f);
+  r = p.AddShader(&f, kShaderTypeFragment);
   ASSERT_TRUE(r.IsSuccess()) << r.Error();
 
   std::vector<std::string> first = {"First", "Second"};
@@ -143,7 +143,7 @@ TEST_F(AmberScriptPipelineTest, DuplicateShaderOptimizations) {
   Shader v(kShaderTypeVertex);
 
   Pipeline p(PipelineType::kGraphics);
-  Result r = p.AddShader(&v);
+  Result r = p.AddShader(&v, kShaderTypeVertex);
   ASSERT_TRUE(r.IsSuccess()) << r.Error();
 
   std::vector<std::string> data = {"One", "One"};
@@ -176,13 +176,13 @@ TEST_F(AmberScriptPipelineTest,
   Shader g(kShaderTypeGeometry);
 
   Pipeline p(PipelineType::kGraphics);
-  Result r = p.AddShader(&v);
+  Result r = p.AddShader(&v, kShaderTypeVertex);
   EXPECT_TRUE(r.IsSuccess()) << r.Error();
 
-  r = p.AddShader(&g);
+  r = p.AddShader(&g, kShaderTypeGeometry);
   EXPECT_TRUE(r.IsSuccess()) << r.Error();
 
-  r = p.AddShader(&f);
+  r = p.AddShader(&f, kShaderTypeFragment);
   EXPECT_TRUE(r.IsSuccess()) << r.Error();
 
   r = p.Validate();
@@ -194,10 +194,10 @@ TEST_F(AmberScriptPipelineTest, GraphicsPipelineMissingFragmentShader) {
   Shader g(kShaderTypeGeometry);
 
   Pipeline p(PipelineType::kGraphics);
-  Result r = p.AddShader(&v);
+  Result r = p.AddShader(&v, kShaderTypeVertex);
   EXPECT_TRUE(r.IsSuccess()) << r.Error();
 
-  r = p.AddShader(&g);
+  r = p.AddShader(&g, kShaderTypeGeometry);
   EXPECT_TRUE(r.IsSuccess()) << r.Error();
 
   r = p.Validate();
@@ -210,10 +210,10 @@ TEST_F(AmberScriptPipelineTest, GraphicsPipelineMissingVertexShader) {
   Shader g(kShaderTypeGeometry);
 
   Pipeline p(PipelineType::kGraphics);
-  Result r = p.AddShader(&g);
+  Result r = p.AddShader(&g, kShaderTypeGeometry);
   EXPECT_TRUE(r.IsSuccess()) << r.Error();
 
-  r = p.AddShader(&f);
+  r = p.AddShader(&f, kShaderTypeFragment);
   EXPECT_TRUE(r.IsSuccess()) << r.Error();
 
   r = p.Validate();
@@ -226,7 +226,7 @@ TEST_F(AmberScriptPipelineTest,
   Shader g(kShaderTypeGeometry);
 
   Pipeline p(PipelineType::kGraphics);
-  Result r = p.AddShader(&g);
+  Result r = p.AddShader(&g, kShaderTypeGeometry);
   EXPECT_TRUE(r.IsSuccess()) << r.Error();
 
   r = p.Validate();
@@ -247,7 +247,7 @@ TEST_F(AmberScriptPipelineTest, ComputePipelineRequiresComputeShader) {
   Shader c(kShaderTypeCompute);
 
   Pipeline p(PipelineType::kCompute);
-  Result r = p.AddShader(&c);
+  Result r = p.AddShader(&c, kShaderTypeCompute);
   EXPECT_TRUE(r.IsSuccess()) << r.Error();
 
   r = p.Validate();
@@ -281,7 +281,7 @@ TEST_F(AmberScriptPipelineTest, SetEntryPointForNullShader) {
 TEST_F(AmberScriptPipelineTest, SetBlankEntryPoint) {
   Shader c(kShaderTypeCompute);
   Pipeline p(PipelineType::kCompute);
-  Result r = p.AddShader(&c);
+  Result r = p.AddShader(&c, kShaderTypeCompute);
   ASSERT_TRUE(r.IsSuccess()) << r.Error();
 
   r = p.SetShaderEntryPoint(&c, "");
@@ -292,7 +292,7 @@ TEST_F(AmberScriptPipelineTest, SetBlankEntryPoint) {
 TEST_F(AmberScriptPipelineTest, ShaderDefaultEntryPoint) {
   Shader c(kShaderTypeCompute);
   Pipeline p(PipelineType::kCompute);
-  Result r = p.AddShader(&c);
+  Result r = p.AddShader(&c, kShaderTypeCompute);
   ASSERT_TRUE(r.IsSuccess()) << r.Error();
 
   const auto& shaders = p.GetShaders();
@@ -303,7 +303,7 @@ TEST_F(AmberScriptPipelineTest, ShaderDefaultEntryPoint) {
 TEST_F(AmberScriptPipelineTest, SetShaderEntryPoint) {
   Shader c(kShaderTypeCompute);
   Pipeline p(PipelineType::kCompute);
-  Result r = p.AddShader(&c);
+  Result r = p.AddShader(&c, kShaderTypeCompute);
   ASSERT_TRUE(r.IsSuccess()) << r.Error();
 
   r = p.SetShaderEntryPoint(&c, "my_main");
@@ -317,7 +317,7 @@ TEST_F(AmberScriptPipelineTest, SetShaderEntryPoint) {
 TEST_F(AmberScriptPipelineTest, SetEntryPointMulitpleTimes) {
   Shader c(kShaderTypeCompute);
   Pipeline p(PipelineType::kCompute);
-  Result r = p.AddShader(&c);
+  Result r = p.AddShader(&c, kShaderTypeCompute);
   ASSERT_TRUE(r.IsSuccess()) << r.Error();
 
   r = p.SetShaderEntryPoint(&c, "my_main");
diff --git a/src/script.cc b/src/script.cc
index 012e6aa..ce9f3aa 100644
--- a/src/script.cc
+++ b/src/script.cc
@@ -145,7 +145,7 @@ std::vector<ShaderInfo> Script::GetShaderInfo() const {
     // pipelines everywhere
 
     // TODO(dsinclair): The optimization passes should be retrieved from the
-    // pipeline and returned here instead of an empy array.
+    // pipeline and returned here instead of an empty array.
     ret.emplace_back(ShaderInfo{shader->GetFormat(),
                                 shader->GetType(),
                                 shader->GetName(),
diff --git a/src/vulkan/engine_vulkan.cc b/src/vulkan/engine_vulkan.cc
index ce3f31c..c22ae3d 100644
--- a/src/vulkan/engine_vulkan.cc
+++ b/src/vulkan/engine_vulkan.cc
@@ -45,6 +45,9 @@ VkShaderStageFlagBits ToVkShaderStage(ShaderType type) {
       return VK_SHADER_STAGE_TESSELLATION_EVALUATION_BIT;
     case kShaderTypeCompute:
       return VK_SHADER_STAGE_COMPUTE_BIT;
+    case kShaderTypeMulti:
+      // It's an error if this arrives here ...
+      break;
   }
 
   assert(false && "Vulkan::Unknown shader stage");