Swizzle result of buffer load based on element type

When lowering `buf[i]` to `texelFetch(..., i)` we need to deal with the case where the type of `b` might be `Buffer<float>` in which case we want to add a `.x` swizzle to the end of the fetch.

2 files changed, 48 insertions, 2 deletions

diff --git a/source/slang/emit.cpp b/source/slang/emit.cpp
index 76dc9c75f..8660b5f93 100644
--- a/source/slang/emit.cpp
+++ b/source/slang/emit.cpp
@@ -1943,6 +1943,52 @@ struct EmitVisitor
                                 }
                                 break;
 
+                            case 'z':
+                                // If we are calling a D3D texturing operation in the form t.Foo(s, ...),
+                                // where `t` is a `Texture*<T>`, then this is the step where we try to
+                                // properly swizzle the output of the equivalent GLSL call into the right
+                                // shape.
+                                assert(argCount > 0);
+                                if (auto memberExpr = callExpr->FunctionExpr.As<MemberExpressionSyntaxNode>())
+                                {
+                                    auto base = memberExpr->BaseExpression;
+                                    if (auto baseTextureType = base->Type->As<TextureType>())
+                                    {
+                                        auto elementType = baseTextureType->elementType;
+                                        if (auto basicType = elementType->As<BasicExpressionType>())
+                                        {
+                                            // A scalar result is expected
+                                            Emit(".x");
+                                        }
+                                        else if (auto vectorType = elementType->As<VectorExpressionType>())
+                                        {
+                                            // A vector result is expected
+                                            auto elementCount = GetIntVal(vectorType->elementCount);
+
+                                            if (elementCount < 4)
+                                            {
+                                                char const* swiz[] = { "", ".x", ".xy", ".xyz", "" };
+                                                Emit(swiz[elementCount]);
+                                            }
+                                        }
+                                        else
+                                        {
+                                            // What other cases are possible?
+                                        }
+                                    }
+                                    else
+                                    {
+                                        assert(!"unexpected");
+                                    }
+
+                                }
+                                else
+                                {
+                                    assert(!"unexpected");
+                                }
+                                break;
+
+
                             default:
                                 assert(!"unexpected");
                                 break;
diff --git a/source/slang/slang-stdlib.cpp b/source/slang/slang-stdlib.cpp
index 8aa0d0e4c..a62db693c 100644
--- a/source/slang/slang-stdlib.cpp
+++ b/source/slang/slang-stdlib.cpp
@@ -2024,12 +2024,12 @@ namespace Slang
 
             sb << "__intrinsic void GetDimensions(out uint dim);\n";
 
-            sb << "__intrinsic(glsl, \"texelFetch($P, $0)\")\n";
+            sb << "__intrinsic(glsl, \"texelFetch($P, $0)$z\")\n";
             sb << "__intrinsic T Load(int location);\n";
 
             sb << "__intrinsic T Load(int location, out uint status);\n";
 
-            sb << "__intrinsic(glsl, \"texelFetch($P, int($0))\")\n";
+            sb << "__intrinsic(glsl, \"texelFetch($P, int($0))$z\")\n";
             sb << "__intrinsic __subscript(uint index) -> T";
 
             if (kBaseBufferAccessLevels[aa].access != SLANG_RESOURCE_ACCESS_READ)