From be5555a18a78706b05d5f0017ffb5d2dd03f6325 Mon Sep 17 00:00:00 2001
From: DragonJoker <dragonjoker59@hotmail.com>
Date: Wed, 23 Mar 2022 21:17:44 +0100
Subject: [PATCH] Moved texture/image access calls result to an SSA alias.

---
 source/CompilerSpirV/SpirvExprVisitor.cpp  | 38 +++++++++-------------
 source/ShaderAST/Visitors/TransformSSA.cpp | 17 ++++++++--
 test/ShaderWriter/TestWriterShader.cpp     | 28 ++++++++++++++++
 3 files changed, 58 insertions(+), 25 deletions(-)

diff --git a/source/CompilerSpirV/SpirvExprVisitor.cpp b/source/CompilerSpirV/SpirvExprVisitor.cpp
index 76b9b687..ad16d1b3 100644
--- a/source/CompilerSpirV/SpirvExprVisitor.cpp
+++ b/source/CompilerSpirV/SpirvExprVisitor.cpp
@@ -684,31 +684,23 @@ namespace spirv
 				shuffle.emplace_back( loadedLhsId );
 				shuffle.emplace_back( rhsId );
 				ast::expr::SwizzleKind rhsSwizzleKind;
+				auto rhsCount = getComponentCount( expr->getRHS()->getType()->getKind() );
+				assert( rhsCount <= 4u );
 
-				if ( expr->getRHS()->getKind() == ast::expr::Kind::eSwizzle
-					&& static_cast< ast::expr::Swizzle const & >( *expr->getRHS() ).getOuterExpr()->getKind() != ast::expr::Kind::eIntrinsicCall
-					&& static_cast< ast::expr::Swizzle const & >( *expr->getRHS() ).getOuterExpr()->getKind() != ast::expr::Kind::eCombinedImageAccessCall
-					&& static_cast< ast::expr::Swizzle const & >( *expr->getRHS() ).getOuterExpr()->getKind() != ast::expr::Kind::eImageAccessCall )
+				switch ( rhsCount )
 				{
-					rhsSwizzleKind = static_cast< ast::expr::Swizzle const & >( *expr->getRHS() ).getSwizzle();
-				}
-				else
-				{
-					auto rhsCount = getComponentCount( expr->getRHS()->getType()->getKind() );
-					assert( rhsCount >= 2u && rhsCount <= 4u );
-
-					switch ( rhsCount )
-					{
-					case 2u:
-						rhsSwizzleKind = ast::expr::SwizzleKind::e01;
-						break;
-					case 3u:
-						rhsSwizzleKind = ast::expr::SwizzleKind::e012;
-						break;
-					default:
-						rhsSwizzleKind = ast::expr::SwizzleKind::e0123;
-						break;
-					}
+				case 1u:
+					rhsSwizzleKind = ast::expr::SwizzleKind::e0;
+					break;
+				case 2u:
+					rhsSwizzleKind = ast::expr::SwizzleKind::e01;
+					break;
+				case 3u:
+					rhsSwizzleKind = ast::expr::SwizzleKind::e012;
+					break;
+				default:
+					rhsSwizzleKind = ast::expr::SwizzleKind::e0123;
+					break;
 				}
 
 				auto swizzleComponents = convert( getSwizzleComponents( lhsSwizzleKind
diff --git a/source/ShaderAST/Visitors/TransformSSA.cpp b/source/ShaderAST/Visitors/TransformSSA.cpp
index 7688adc2..c3c3cba8 100644
--- a/source/ShaderAST/Visitors/TransformSSA.cpp
+++ b/source/ShaderAST/Visitors/TransformSSA.cpp
@@ -1174,6 +1174,14 @@ namespace ast
 							, ast::expr::SwizzleKind::e012 );
 					}
 				}
+
+				if ( srcType->getKind() != ast::type::Kind::eVoid )
+				{
+					auto type = m_result->getType();
+					auto alias = doCreateAliasVar( type
+						, std::move( m_result ) );
+					m_result = ast::expr::makeIdentifier( m_cache, alias );
+				}
 			}
 
 			void visitInitExpr( expr::Init * expr )override
@@ -1322,8 +1330,8 @@ namespace ast
 				}
 
 				m_result = ast::expr::makeCombinedImageAccessCall( returnType
-					, kind
-					, std::move( args ) );
+						, kind
+						, std::move( args ) );
 
 				if ( returnComponentsCount != InvalidComponentCount && returnComponentsCount != count )
 				{
@@ -1347,6 +1355,11 @@ namespace ast
 
 					m_result = ast::expr::makeSwizzle( std::move( m_result ), swizzleKind );
 				}
+
+				auto type = m_result->getType();
+				auto alias = doCreateAliasVar( type
+					, std::move( m_result ) );
+				m_result = ast::expr::makeIdentifier( m_cache, alias );
 			}
 
 		private:
diff --git a/test/ShaderWriter/TestWriterShader.cpp b/test/ShaderWriter/TestWriterShader.cpp
index 583b589c..eba05884 100644
--- a/test/ShaderWriter/TestWriterShader.cpp
+++ b/test/ShaderWriter/TestWriterShader.cpp
@@ -1632,6 +1632,33 @@ namespace
 			, CurrentCompilers );
 		testEnd();
 	}
+
+	void arraySsboTextureLookup( test::sdw_test::TestCounts & testCounts )
+	{
+		testBegin( "arraySsboTextureLookup" );
+		using namespace sdw;
+		{
+			FragmentWriter writer;
+			auto fragUvw = writer.declInput< Vec3 >( "fragUvw", 0u );
+			auto outColor = writer.declOutput< Vec4 >( "outColor", 0u );
+
+			sdw::Ssbo colorsSsbo = writer.declStorageBuffer( "colorsBuffer", 0u, 1u );
+			auto colors = colorsSsbo.declMemberArray< Vec4 >( "colors" );
+			colorsSsbo.end();
+
+			auto volumeTexture = writer.declCombinedImg< UCombinedImage3DR32 >( "volumeTexture", 1u, 1u );
+
+			writer.implementMain(
+				[&]( sdw::FragmentIn in, sdw::FragmentOut out )
+				{
+					outColor = colors[volumeTexture.sample( fragUvw )];
+				} );
+			test::writeShader( writer
+				, testCounts
+				, CurrentCompilers );
+		}
+		testEnd();
+	}
 }
 
 sdwTestSuiteMain( TestWriterShader )
@@ -1652,6 +1679,7 @@ sdwTestSuiteMain( TestWriterShader )
 	basicPipeline( testCounts );
 	voxelPipeline( testCounts );
 	tessellationPipeline( testCounts );
+	arraySsboTextureLookup( testCounts );
 	sdwTestSuiteEnd();
 }