2 files changed, 29 insertions, 15 deletions

diff --git a/src/vkscript/command_parser.cc b/src/vkscript/command_parser.cc
index b60993f..5123991 100644
--- a/src/vkscript/command_parser.cc
+++ b/src/vkscript/command_parser.cc
@@ -594,10 +594,12 @@ Result CommandParser::ProcessSSBO() {
       return Result("Invalid type provided: " + token->AsString());
 
     auto* buf = cmd->GetBuffer();
-    if (buf->FormatIsDefault() || !buf->GetFormat())
-      buf->SetFormat(std::move(fmt));
-    else if (!buf->GetFormat()->Equal(fmt.get()))
+    if (buf->FormatIsDefault() || !buf->GetFormat()) {
+      buf->SetFormat(fmt.get());
+      script_->RegisterFormat(std::move(fmt));
+    } else if (!buf->GetFormat()->Equal(fmt.get())) {
       return Result("probe ssbo format does not match buffer format");
+    }
 
     token = tokenizer_->NextToken();
     if (!token->IsInteger()) {
@@ -641,7 +643,8 @@ Result CommandParser::ProcessSSBO() {
     if (!buf->GetFormat()) {
       FormatParser fp;
       auto fmt = fp.Parse("R8_SINT");
-      buf->SetFormat(std::move(fmt));
+      buf->SetFormat(fmt.get());
+      script_->RegisterFormat(std::move(fmt));
       // This has to come after the SetFormat() call because SetFormat() resets
       // the value back to false.
       buf->SetFormatIsDefault(true);
@@ -746,10 +749,12 @@ Result CommandParser::ProcessUniform() {
     fmt->SetLayout(Format::Layout::kStd140);
 
   auto* buf = cmd->GetBuffer();
-  if (buf->FormatIsDefault() || !buf->GetFormat())
-    buf->SetFormat(std::move(fmt));
-  else if (!buf->GetFormat()->Equal(fmt.get()))
+  if (buf->FormatIsDefault() || !buf->GetFormat()) {
+    buf->SetFormat(fmt.get());
+    script_->RegisterFormat(std::move(fmt));
+  } else if (!buf->GetFormat()->Equal(fmt.get())) {
     return Result("probe ssbo format does not match buffer format");
+  }
 
   token = tokenizer_->NextToken();
   if (!token->IsInteger()) {
@@ -2056,18 +2061,21 @@ Result CommandParser::ProcessProbeSSBO() {
                   std::to_string(binding));
   }
 
-  if (buffer->FormatIsDefault() || !buffer->GetFormat())
-    buffer->SetFormat(MakeUnique<Format>(*fmt));
-  else if (buffer->GetFormat() && !buffer->GetFormat()->Equal(fmt.get()))
+  if (buffer->FormatIsDefault() || !buffer->GetFormat()) {
+    buffer->SetFormat(fmt.get());
+  } else if (buffer->GetFormat() && !buffer->GetFormat()->Equal(fmt.get())) {
     return Result("probe format does not match buffer format");
+  }
 
   auto cmd = MakeUnique<ProbeSSBOCommand>(buffer);
   cmd->SetLine(cur_line);
   cmd->SetTolerances(current_tolerances_);
-  cmd->SetFormat(std::move(fmt));
+  cmd->SetFormat(fmt.get());
   cmd->SetDescriptorSet(set);
   cmd->SetBinding(binding);
 
+  script_->RegisterFormat(std::move(fmt));
+
   if (!token->IsInteger())
     return Result("Invalid offset for probe ssbo command: " +
                   token->ToOriginalString());
diff --git a/src/vkscript/parser.cc b/src/vkscript/parser.cc
index f3ad45f..6ae7684 100644
--- a/src/vkscript/parser.cc
+++ b/src/vkscript/parser.cc
@@ -178,7 +178,8 @@ Result Parser::ProcessRequireBlock(const SectionParser::Section& section) {
       }
       script_->GetPipeline(kDefaultPipelineName)
           ->GetColorAttachments()[0]
-          .buffer->SetFormat(std::move(fmt));
+          .buffer->SetFormat(fmt.get());
+      script_->RegisterFormat(std::move(fmt));
 
     } else if (str == "depthstencil") {
       token = tokenizer.NextToken();
@@ -199,7 +200,9 @@ Result Parser::ProcessRequireBlock(const SectionParser::Section& section) {
 
       // Generate and add a depth buffer
       auto depth_buf = pipeline->GenerateDefaultDepthAttachmentBuffer();
-      depth_buf->SetFormat(std::move(fmt));
+      depth_buf->SetFormat(fmt.get());
+      script_->RegisterFormat(std::move(fmt));
+
       Result r = pipeline->SetDepthBuffer(depth_buf.get());
       if (!r.IsSuccess())
         return r;
@@ -292,8 +295,10 @@ Result Parser::ProcessIndicesBlock(const SectionParser::Section& section) {
     auto b = MakeUnique<Buffer>(BufferType::kIndex);
     auto* buf = b.get();
     b->SetName("indices");
-    b->SetFormat(std::move(fmt));
+    b->SetFormat(fmt.get());
     b->SetData(std::move(indices));
+    script_->RegisterFormat(std::move(fmt));
+
     Result r = script_->AddBuffer(std::move(b));
     if (!r.IsSuccess())
       return r;
@@ -422,9 +427,10 @@ Result Parser::ProcessVertexDataBlock(const SectionParser::Section& section) {
     auto buffer = MakeUnique<Buffer>(BufferType::kVertex);
     auto* buf = buffer.get();
     buffer->SetName("Vertices" + std::to_string(i));
-    buffer->SetFormat(std::move(headers[i].format));
+    buffer->SetFormat(headers[i].format.get());
     buffer->SetData(std::move(values[i]));
     script_->AddBuffer(std::move(buffer));
+    script_->RegisterFormat(std::move(headers[i].format));
 
     pipeline->AddVertexBuffer(buf, headers[i].location);
   }