This CL allows passing the pipeline name to the -B flag. (#733)

Previously we could only extract buffers from the first pipeline in the
file. This Cl extends the -B flag to allow providing the pipeline name
as a prefix. So, my_pipeline:0:2 would extract from the pipeline named
my_pipeline, descriptor set 0, binding 2.

Fixes #732

6 files changed, 52 insertions, 15 deletions

diff --git a/samples/amber.cc b/samples/amber.cc
index d508da8..19e5182 100644
--- a/samples/amber.cc
+++ b/samples/amber.cc
@@ -89,9 +89,9 @@ const char kUsage[] = R"(Usage: amber [options] SCRIPT [SCRIPTS...]
                                or as a PPM image otherwise.
   -I <buffername>           -- Name of framebuffer to dump. Defaults to 'framebuffer'.
   -b <filename>             -- Write contents of a UBO or SSBO to <filename>.
-  -B [<desc set>:]<binding> -- Descriptor set and binding of buffer to write.
-                               Default is [0:]0.
-  -w <filename>             -- Write shader assembly to |filename|                               
+  -B [<pipeline name>:][<desc set>:]<binding> -- Identifier of buffer to write.
+                               Default is [first pipeline:][0:]0.
+  -w <filename>             -- Write shader assembly to |filename|
   -e <engine>               -- Specify graphics engine: vulkan, dawn. Default is vulkan.
   -v <engine version>       -- Engine version (eg, 1.1 for Vulkan). Default 1.0.
   -V, --version             -- Output version information for Amber and libraries.
diff --git a/src/amber.cc b/src/amber.cc
index 598b4a1..43d358e 100644
--- a/src/amber.cc
+++ b/src/amber.cc
@@ -180,9 +180,6 @@ amber::Result Amber::ExecuteWithShaderData(const amber::Recipe* recipe,
     return {};
   }
 
-  // TODO(dsinclair): Figure out how extractions work with multiple pipelines.
-  auto* pipeline = script->GetPipelines()[0].get();
-
   // Try to perform each extraction, copying the buffer data into |buffer_info|.
   // We do not overwrite |executor_result| if extraction fails.
   for (BufferInfo& buffer_info : opts->extractions) {
@@ -197,14 +194,21 @@ amber::Result Amber::ExecuteWithShaderData(const amber::Recipe* recipe,
       continue;
     }
 
-    DescriptorSetAndBindingParser desc_set_and_binding_parser;
-    r = desc_set_and_binding_parser.Parse(buffer_info.buffer_name);
+    DescriptorSetAndBindingParser p;
+    r = p.Parse(buffer_info.buffer_name);
     if (!r.IsSuccess())
       continue;
 
-    const auto* buffer = pipeline->GetBufferForBinding(
-        desc_set_and_binding_parser.GetDescriptorSet(),
-        desc_set_and_binding_parser.GetBinding());
+    // Extract the named pipeline from the request, otherwise use the
+    // first pipeline which was parsed.
+    Pipeline* pipeline = nullptr;
+    if (p.HasPipelineName())
+      pipeline = script->GetPipeline(p.PipelineName());
+    else
+      pipeline = script->GetPipelines()[0].get();
+
+    const auto* buffer =
+        pipeline->GetBufferForBinding(p.GetDescriptorSet(), p.GetBinding());
     if (!buffer)
       continue;
 
diff --git a/src/descriptor_set_and_binding_parser.cc b/src/descriptor_set_and_binding_parser.cc
index d563573..c9a1a03 100644
--- a/src/descriptor_set_and_binding_parser.cc
+++ b/src/descriptor_set_and_binding_parser.cc
@@ -14,6 +14,7 @@
 
 #include "src/descriptor_set_and_binding_parser.h"
 
+#include <cctype>
 #include <iostream>
 #include <limits>
 
@@ -26,7 +27,24 @@ DescriptorSetAndBindingParser::DescriptorSetAndBindingParser() = default;
 DescriptorSetAndBindingParser::~DescriptorSetAndBindingParser() = default;
 
 Result DescriptorSetAndBindingParser::Parse(const std::string& buffer_id) {
-  Tokenizer t(buffer_id);
+  size_t idx = 0;
+
+  // Pipeline name
+  if (!std::isdigit(buffer_id[idx]) && std::isalpha(buffer_id[idx]) &&
+      buffer_id[idx] != ':' && buffer_id[idx] != '-') {
+    idx++;
+    while (idx < buffer_id.size() && buffer_id[idx] != ':')
+      idx++;
+
+    pipeline_name_ = buffer_id.substr(0, idx);
+
+    // Move past the :
+    idx += 1;
+  }
+  if (idx >= buffer_id.size())
+    return Result("Invalid buffer id: " + buffer_id);
+
+  Tokenizer t(buffer_id.substr(idx));
   auto token = t.NextToken();
   if (token->IsInteger()) {
     if (token->AsInt32() < 0) {
@@ -45,8 +63,6 @@ Result DescriptorSetAndBindingParser::Parse(const std::string& buffer_id) {
     }
 
     descriptor_set_ = val;
-  } else {
-    descriptor_set_ = 0;
   }
 
   if (!token->IsString())
diff --git a/src/descriptor_set_and_binding_parser.h b/src/descriptor_set_and_binding_parser.h
index 44a412a..145aa42 100644
--- a/src/descriptor_set_and_binding_parser.h
+++ b/src/descriptor_set_and_binding_parser.h
@@ -32,9 +32,13 @@ class DescriptorSetAndBindingParser {
   /// be a single non-negative integer or two those integers separated by
   /// ':'. For example, ":0", "1", and "2:3" are valid strings for |buffer_id|,
   /// but "", "-4", ":-5", ":", "a", and "b:c" are invalid strings for
-  /// |buffer_id|.
+  /// |buffer_id|. The |buffer_id| can be prefixed by the name of a pipeline.
   Result Parse(const std::string& buffer_id);
 
+  /// Returns true if a pipeline name was specified
+  bool HasPipelineName() const { return !pipeline_name_.empty(); }
+  std::string PipelineName() const { return pipeline_name_; }
+
   /// Return descriptor set that is the result of Parse().
   uint32_t GetDescriptorSet() const { return descriptor_set_; }
 
@@ -42,6 +46,7 @@ class DescriptorSetAndBindingParser {
   uint32_t GetBinding() const { return binding_; }
 
  private:
+  std::string pipeline_name_;
   uint32_t descriptor_set_ = 0;
   uint32_t binding_ = 0;
 };
diff --git a/src/descriptor_set_and_binding_parser_test.cc b/src/descriptor_set_and_binding_parser_test.cc
index 6010130..ce13f10 100644
--- a/src/descriptor_set_and_binding_parser_test.cc
+++ b/src/descriptor_set_and_binding_parser_test.cc
@@ -25,6 +25,7 @@ TEST_F(DescriptorSetAndBindingParserTest, CommaAndBinding) {
   Result r = parser.Parse(":1234");
   ASSERT_TRUE(r.IsSuccess()) << r.Error();
 
+  EXPECT_FALSE(parser.HasPipelineName());
   EXPECT_EQ(0, parser.GetDescriptorSet());
   EXPECT_EQ(1234, parser.GetBinding());
 }
@@ -107,4 +108,15 @@ TEST_F(DescriptorSetAndBindingParserTest, DescSetAndNegativeBinding) {
       r.Error());
 }
 
+TEST_F(DescriptorSetAndBindingParserTest, WithPipelineName) {
+  DescriptorSetAndBindingParser parser;
+  Result r = parser.Parse("pipeline1:123:234");
+  ASSERT_TRUE(r.IsSuccess()) << r.Error();
+
+  EXPECT_TRUE(parser.HasPipelineName());
+  EXPECT_EQ("pipeline1", parser.PipelineName());
+  EXPECT_EQ(123, parser.GetDescriptorSet());
+  EXPECT_EQ(234, parser.GetBinding());
+}
+
 }  // namespace amber
diff --git a/tests/cases/position_to_ssbo.amber b/tests/cases/position_to_ssbo.vkscript
index 39f7aae..39f7aae 100644
--- a/tests/cases/position_to_ssbo.amber
+++ b/tests/cases/position_to_ssbo.vkscript