Update MetalArgumentBuffer to support buffers

libs/filamat/src/MetalArgumentBuffer.cpp

@@ -18,6 +18,7 @@
 
 #include <sstream>
 #include <utility>
+#include <variant>
 
 namespace filamat {
 
@@ -128,9 +129,6 @@ std::ostream& MetalArgumentBuffer::Builder::BufferArgument::write(std::ostream&
 MetalArgumentBuffer::MetalArgumentBuffer(Builder& builder) {
     mName = builder.mName;
 
-    std::stringstream ss;
-    ss << "struct " << mName << " {" << std::endl;
-
     auto& args = builder.mArguments;
 
     // Sort the arguments by index.
@@ -145,6 +143,18 @@ MetalArgumentBuffer::MetalArgumentBuffer(Builder& builder) {
                         [](auto const& x, auto const& y) { return x.index == y.index; }, lhs, rhs);
             }) == args.end());
 
+    std::stringstream ss;
+
+    // Add forward declarations of buffers.
+    for (const auto& a : builder.mArguments) {
+        if (std::holds_alternative<Builder::BufferArgument>(a)) {
+            const auto& bufferArg = std::get<Builder::BufferArgument>(a);
+            ss << "struct " << bufferArg.type << ";" << std::endl;
+        }
+    }
+
+    ss << "struct " << mName << " {" << std::endl;
+
     for (const auto& a : builder.mArguments) {
         std::visit([&](auto&& arg) {
             arg.write(ss);

libs/filamat/tests/test_argBufferFixup.cpp

@@ -97,6 +97,7 @@ TEST(ArgBufferFixup, Buffer) {
     auto argBufferStr = argBuffer->getMsl();
 
     const std::string expected =
+            "struct FrameUniforms;\n"
             "struct myArgumentBuffer {\n"
             "constant FrameUniforms* frameUniforms [[id(0)]];\n"
             "}";
@@ -106,6 +107,28 @@ TEST(ArgBufferFixup, Buffer) {
     MetalArgumentBuffer::destroy(&argBuffer);
 }
 
+TEST(ArgBufferFixup, MultipleBuffers) {
+    auto argBuffer =
+            MetalArgumentBuffer::Builder()
+                    .name("myArgumentBuffer")
+                    .buffer(0, "FrameUniforms", "frameUniforms")
+                    .buffer(1, "ObjectUniforms", "objectUniforms")
+                    .build();
+    auto argBufferStr = argBuffer->getMsl();
+
+    const std::string expected =
+            "struct FrameUniforms;\n"
+            "struct ObjectUniforms;\n"
+            "struct myArgumentBuffer {\n"
+            "constant FrameUniforms* frameUniforms [[id(0)]];\n"
+            "constant ObjectUniforms* objectUniforms [[id(1)]];\n"
+            "}";
+
+    EXPECT_EQ(argBuffer->getMsl(), expected);
+
+    MetalArgumentBuffer::destroy(&argBuffer);
+}
+
 TEST(ArgBufferFixup, TextureAndSamplerMS) {
     auto argBuffer =
             MetalArgumentBuffer::Builder()
@@ -139,6 +162,7 @@ TEST(ArgBufferFixup, Sorted) {
     auto argBufferStr = argBuffer->getMsl();
 
     const std::string expected =
+            "struct FrameUniforms;\n"
             "struct myArgumentBuffer {\n"
             "texture2d<float> textureA [[id(0)]];\n"
             "sampler samplerA [[id(1)]];\n"