Refactor BlockLayoutEncoder APIs for std430.

This splits HLSL SSBO access into two steps.

First we compute a mapping from the collected SSBO variable names to
TField pointers. Then during tree traversal we use a block encoding
visitor class that uses the shader names to store BlockMemberInfo
structures for the structures and variables. Each nested structure
is traversed separately so that the BlockMemberInfo offsets are
relative to the structure start rather than the enclosing block. The
array stride for a structure is the size of the struct after all the
alignment is included.

This gives the correct results for the SSBO access chain in the HLSL
code. It also will allow us to use the same encoding and visiting logic
for SSBOs on the API side.

Bug: angleproject:3024
Change-Id: I42b1db0e7547782ae77fe5f64a797f803f203f45
Reviewed-on: https://chromium-review.googlesource.com/c/1352731
Commit-Queue: Jamie Madill <jmadill@chromium.org>
Reviewed-by: Yuly Novikov <ynovikov@chromium.org>
diff --git a/src/compiler/translator/OutputHLSL.cpp b/src/compiler/translator/OutputHLSL.cpp
index d660b9a..cd15daa 100644
--- a/src/compiler/translator/OutputHLSL.cpp
+++ b/src/compiler/translator/OutputHLSL.cpp
@@ -235,7 +235,8 @@
                        ShCompileOptions compileOptions,
                        sh::WorkGroupSize workGroupSize,
                        TSymbolTable *symbolTable,
-                       PerformanceDiagnostics *perfDiagnostics)
+                       PerformanceDiagnostics *perfDiagnostics,
+                       const std::vector<InterfaceBlock> &shaderStorageBlocks)
     : TIntermTraverser(true, true, true, symbolTable),
       mShaderType(shaderType),
       mShaderVersion(shaderVersion),
@@ -303,7 +304,8 @@
     // Reserve registers for the default uniform block and driver constants
     mResourcesHLSL->reserveUniformBlockRegisters(2);
 
-    mSSBOOutputHLSL = new ShaderStorageBlockOutputHLSL(this, symbolTable, mResourcesHLSL);
+    mSSBOOutputHLSL =
+        new ShaderStorageBlockOutputHLSL(this, symbolTable, mResourcesHLSL, shaderStorageBlocks);
 }
 
 OutputHLSL::~OutputHLSL()
diff --git a/src/compiler/translator/OutputHLSL.h b/src/compiler/translator/OutputHLSL.h
index 92fbb6f..da1974d 100644
--- a/src/compiler/translator/OutputHLSL.h
+++ b/src/compiler/translator/OutputHLSL.h
@@ -47,7 +47,8 @@
                ShCompileOptions compileOptions,
                sh::WorkGroupSize workGroupSize,
                TSymbolTable *symbolTable,
-               PerformanceDiagnostics *perfDiagnostics);
+               PerformanceDiagnostics *perfDiagnostics,
+               const std::vector<InterfaceBlock> &shaderStorageBlocks);
 
     ~OutputHLSL();
 
diff --git a/src/compiler/translator/ShaderStorageBlockOutputHLSL.cpp b/src/compiler/translator/ShaderStorageBlockOutputHLSL.cpp
index 7dc570b..f3eb53b 100644
--- a/src/compiler/translator/ShaderStorageBlockOutputHLSL.cpp
+++ b/src/compiler/translator/ShaderStorageBlockOutputHLSL.cpp
@@ -132,149 +132,160 @@
     return nullptr;
 }
 
-void GetShaderStorageBlockFieldMemberInfo(const TFieldList &fields,
-                                          sh::BlockLayoutEncoder *encoder,
-                                          TLayoutBlockStorage storage,
-                                          bool rowMajor,
-                                          bool isSSBOFieldMember,
-                                          BlockMemberInfoMap *blockInfoOut);
-
-size_t GetBlockFieldMemberInfoAndReturnBlockSize(const TFieldList &fields,
-                                                 TLayoutBlockStorage storage,
-                                                 bool rowMajor,
-                                                 BlockMemberInfoMap *blockInfoOut,
-                                                 int *structureBaseAlignment)
+const InterfaceBlock *FindInterfaceBlock(const TInterfaceBlock *needle,
+                                         const std::vector<InterfaceBlock> &haystack)
 {
-    sh::Std140BlockEncoder std140Encoder;
-    sh::Std430BlockEncoder std430Encoder;
-    sh::HLSLBlockEncoder hlslEncoder(sh::HLSLBlockEncoder::ENCODE_PACKED, false);
-    sh::BlockLayoutEncoder *structureEncoder = nullptr;
-
-    if (storage == EbsStd140)
+    for (const InterfaceBlock &block : haystack)
     {
-        structureEncoder = &std140Encoder;
-    }
-    else if (storage == EbsStd430)
-    {
-        structureEncoder = &std430Encoder;
-    }
-    else
-    {
-        structureEncoder = &hlslEncoder;
+        if (strcmp(block.name.c_str(), needle->name().data()) == 0)
+        {
+            ASSERT(block.fields.size() == needle->fields().size());
+            return &block;
+        }
     }
 
-    GetShaderStorageBlockFieldMemberInfo(fields, structureEncoder, storage, rowMajor, false,
-                                         blockInfoOut);
-    structureEncoder->exitAggregateType();
-    *structureBaseAlignment = static_cast<int>(structureEncoder->getStructureBaseAlignment());
-    return structureEncoder->getBlockSize();
+    UNREACHABLE();
+    return nullptr;
 }
 
-void GetShaderStorageBlockFieldMemberInfo(const TFieldList &fields,
-                                          sh::BlockLayoutEncoder *encoder,
-                                          TLayoutBlockStorage storage,
-                                          bool rowMajor,
-                                          bool isSSBOFieldMember,
-                                          BlockMemberInfoMap *blockInfoOut)
+std::string StripArrayIndices(const std::string &nameIn)
 {
-    for (const TField *field : fields)
+    std::string name = nameIn;
+    size_t pos       = name.find('[');
+    while (pos != std::string::npos)
     {
-        const TType &fieldType = *field->type();
-        bool isRowMajorLayout  = rowMajor;
-        if (isSSBOFieldMember)
-        {
-            isRowMajorLayout = (fieldType.getLayoutQualifier().matrixPacking == EmpRowMajor);
-        }
-        if (fieldType.getStruct())
-        {
-            int structureBaseAlignment = 0;
-            // This is to set structure member offset and array stride using a new encoder to ensure
-            // that the first field member offset in structure is always zero.
-            size_t structureStride = GetBlockFieldMemberInfoAndReturnBlockSize(
-                fieldType.getStruct()->fields(), storage, isRowMajorLayout, blockInfoOut,
-                &structureBaseAlignment);
-            // According to OpenGL ES 3.1 spec, session 7.6.2.2 Standard Uniform Block Layout. In
-            // rule 9, if the member is a structure, the base alignment of the structure is N, where
-            // N is the largest base alignment value of any of its members. When using the std430
-            // storage layout, the base alignment and stride of structures in rule 9 are not rounded
-            // up a multiple of the base alignment of a vec4. So we must set structure base
-            // alignment before enterAggregateType.
-            encoder->setStructureBaseAlignment(structureBaseAlignment);
-            encoder->enterAggregateType();
-            const BlockMemberInfo memberInfo(static_cast<int>(encoder->getBlockSize()),
-                                             static_cast<int>(structureStride), 0, false);
-            (*blockInfoOut)[field] = memberInfo;
+        size_t closePos = name.find(']', pos);
+        ASSERT(closePos != std::string::npos && closePos > pos);
+        name.erase(pos, closePos - pos + 1);
+        pos = name.find('[');
+    }
+    ASSERT(name.find(']') == std::string::npos);
+    return name;
+}
 
-            // Below if-else is in order to get correct offset for the field members after structure
-            // field.
-            if (fieldType.isArray())
-            {
-                size_t size = fieldType.getArraySizeProduct() * structureStride;
-                encoder->increaseCurrentOffset(size);
-            }
-            else
-            {
-                encoder->increaseCurrentOffset(structureStride);
-            }
-        }
-        else if (fieldType.isArrayOfArrays())
-        {
-            size_t beginSize                        = encoder->getBlockSize();
-            const TVector<unsigned int> &arraySizes = *fieldType.getArraySizes();
-            // arraySizes[0] stores the innermost array's size.
-            std::vector<unsigned int> innermostArraySize(1u, arraySizes[0]);
-            const BlockMemberInfo &memberInfo =
-                encoder->encodeType(GLVariableType(fieldType), innermostArraySize,
-                                    isRowMajorLayout && fieldType.isMatrix());
-            (*blockInfoOut)[field] = memberInfo;
-            size_t endSize         = encoder->getBlockSize();
+// Does not include any array indices.
+void MapVariableToField(const ShaderVariable &variable,
+                        const TField *field,
+                        std::string currentName,
+                        ShaderVarToFieldMap *shaderVarToFieldMap)
+{
+    ASSERT((field->type()->getStruct() == nullptr) == variable.fields.empty());
+    (*shaderVarToFieldMap)[currentName] = field;
 
-            // The total size of array of arrays is memberInfo.arrayStride *
-            // fieldType.getArraySizeProduct(). However, encoder->encodeType will change the current
-            // offset of encoder. So the final increase size will be total size of arrays of arrays
-            // minus the increased sized by encoder->encodeType.
-            size_t arrayOfArraysSize = memberInfo.arrayStride * fieldType.getArraySizeProduct();
-            size_t increaseSize      = arrayOfArraysSize - (endSize - beginSize);
-            encoder->increaseCurrentOffset(increaseSize);
-        }
-        else
+    if (!variable.fields.empty())
+    {
+        const TStructure *subStruct = field->type()->getStruct();
+        ASSERT(variable.fields.size() == subStruct->fields().size());
+
+        for (size_t index = 0; index < variable.fields.size(); ++index)
         {
-            std::vector<unsigned int> fieldArraySizes;
-            if (auto *arraySizes = fieldType.getArraySizes())
-            {
-                fieldArraySizes.assign(arraySizes->begin(), arraySizes->end());
-            }
-            const BlockMemberInfo &memberInfo =
-                encoder->encodeType(GLVariableType(fieldType), fieldArraySizes,
-                                    isRowMajorLayout && fieldType.isMatrix());
-            (*blockInfoOut)[field] = memberInfo;
+            const TField *subField            = subStruct->fields()[index];
+            const ShaderVariable &subVariable = variable.fields[index];
+            std::string subName               = currentName + "." + subVariable.name;
+            MapVariableToField(subVariable, subField, subName, shaderVarToFieldMap);
         }
     }
 }
 
+class BlockInfoVisitor final : public BlockEncoderVisitor
+{
+  public:
+    BlockInfoVisitor(const std::string &prefix,
+                     TLayoutBlockStorage storage,
+                     const ShaderVarToFieldMap &shaderVarToFieldMap,
+                     BlockMemberInfoMap *blockInfoOut)
+        : BlockEncoderVisitor(prefix, "", getEncoder(storage)),
+          mShaderVarToFieldMap(shaderVarToFieldMap),
+          mBlockInfoOut(blockInfoOut),
+          mHLSLEncoder(HLSLBlockEncoder::ENCODE_PACKED, false),
+          mStorage(storage)
+    {}
+
+    BlockLayoutEncoder *getEncoder(TLayoutBlockStorage storage)
+    {
+        switch (storage)
+        {
+            case EbsStd140:
+                return &mStd140Encoder;
+            case EbsStd430:
+                return &mStd430Encoder;
+            default:
+                return &mHLSLEncoder;
+        }
+    }
+
+    void enterStructAccess(const ShaderVariable &structVar, bool isRowMajor) override
+    {
+        BlockEncoderVisitor::enterStructAccess(structVar, isRowMajor);
+
+        std::string variableName = StripArrayIndices(collapseNameStack());
+
+        // Remove the trailing "."
+        variableName.pop_back();
+
+        BlockInfoVisitor childVisitor(variableName, mStorage, mShaderVarToFieldMap, mBlockInfoOut);
+        childVisitor.getEncoder(mStorage)->enterAggregateType(structVar);
+        TraverseShaderVariables(structVar.fields, isRowMajor, &childVisitor);
+        childVisitor.getEncoder(mStorage)->exitAggregateType(structVar);
+
+        int offset      = getEncoder(mStorage)->getCurrentOffset();
+        int arrayStride = childVisitor.getEncoder(mStorage)->getCurrentOffset();
+
+        auto iter = mShaderVarToFieldMap.find(variableName);
+        if (iter == mShaderVarToFieldMap.end())
+            return;
+
+        const TField *structField = iter->second;
+        if (mBlockInfoOut->count(structField) == 0)
+        {
+            mBlockInfoOut->emplace(structField, BlockMemberInfo(offset, arrayStride, -1, false));
+        }
+    }
+
+    void encodeVariable(const ShaderVariable &variable,
+                        const BlockMemberInfo &variableInfo,
+                        const std::string &name,
+                        const std::string &mappedName) override
+    {
+        auto iter = mShaderVarToFieldMap.find(StripArrayIndices(name));
+        if (iter == mShaderVarToFieldMap.end())
+            return;
+
+        const TField *field = iter->second;
+        if (mBlockInfoOut->count(field) == 0)
+        {
+            mBlockInfoOut->emplace(field, variableInfo);
+        }
+    }
+
+  private:
+    const ShaderVarToFieldMap &mShaderVarToFieldMap;
+    BlockMemberInfoMap *mBlockInfoOut;
+    Std140BlockEncoder mStd140Encoder;
+    Std430BlockEncoder mStd430Encoder;
+    HLSLBlockEncoder mHLSLEncoder;
+    TLayoutBlockStorage mStorage;
+};
+
 void GetShaderStorageBlockMembersInfo(const TInterfaceBlock *interfaceBlock,
+                                      const std::vector<InterfaceBlock> &shaderStorageBlocks,
                                       BlockMemberInfoMap *blockInfoOut)
 {
-    sh::Std140BlockEncoder std140Encoder;
-    sh::Std430BlockEncoder std430Encoder;
-    sh::HLSLBlockEncoder hlslEncoder(sh::HLSLBlockEncoder::ENCODE_PACKED, false);
-    sh::BlockLayoutEncoder *encoder = nullptr;
+    // Find the sh::InterfaceBlock.
+    const InterfaceBlock *block = FindInterfaceBlock(interfaceBlock, shaderStorageBlocks);
+    ASSERT(block);
 
-    if (interfaceBlock->blockStorage() == EbsStd140)
+    // Map ShaderVariable to TField.
+    ShaderVarToFieldMap shaderVarToFieldMap;
+    for (size_t index = 0; index < block->fields.size(); ++index)
     {
-        encoder = &std140Encoder;
-    }
-    else if (interfaceBlock->blockStorage() == EbsStd430)
-    {
-        encoder = &std430Encoder;
-    }
-    else
-    {
-        encoder = &hlslEncoder;
+        const TField *field            = interfaceBlock->fields()[index];
+        const ShaderVariable &variable = block->fields[index];
+        MapVariableToField(variable, field, variable.name, &shaderVarToFieldMap);
     }
 
-    GetShaderStorageBlockFieldMemberInfo(interfaceBlock->fields(), encoder,
-                                         interfaceBlock->blockStorage(), false, true, blockInfoOut);
+    BlockInfoVisitor visitor("", interfaceBlock->blockStorage(), shaderVarToFieldMap, blockInfoOut);
+    TraverseShaderVariables(block->fields, false, &visitor);
 }
 
 bool IsInArrayOfArraysChain(TIntermTyped *node)
@@ -290,18 +301,20 @@
 
     return false;
 }
-
 }  // anonymous namespace
 
-ShaderStorageBlockOutputHLSL::ShaderStorageBlockOutputHLSL(OutputHLSL *outputHLSL,
-                                                           TSymbolTable *symbolTable,
-                                                           ResourcesHLSL *resourcesHLSL)
+ShaderStorageBlockOutputHLSL::ShaderStorageBlockOutputHLSL(
+    OutputHLSL *outputHLSL,
+    TSymbolTable *symbolTable,
+    ResourcesHLSL *resourcesHLSL,
+    const std::vector<InterfaceBlock> &shaderStorageBlocks)
     : TIntermTraverser(true, true, true, symbolTable),
       mMatrixStride(0),
       mRowMajor(false),
       mLocationAsTheLastArgument(false),
       mOutputHLSL(outputHLSL),
-      mResourcesHLSL(resourcesHLSL)
+      mResourcesHLSL(resourcesHLSL),
+      mShaderStorageBlocks(shaderStorageBlocks)
 {
     mSSBOFunctionHLSL = new ShaderStorageBlockFunctionHLSL;
 }
@@ -432,7 +445,7 @@
     mSSBOFunctionHLSL->shaderStorageBlockFunctionHeader(out);
 }
 
-// Check if the current node is the end of the sssbo access chain. If true, we should output ')' for
+// Check if the current node is the end of the SSBO access chain. If true, we should output ')' for
 // Load method.
 bool ShaderStorageBlockOutputHLSL::isEndOfSSBOAccessChain()
 {
@@ -484,7 +497,8 @@
             }
             mReferencedShaderStorageBlocks[interfaceBlock->uniqueId().get()] =
                 new TReferencedBlock(interfaceBlock, instanceVariable);
-            GetShaderStorageBlockMembersInfo(interfaceBlock, &mBlockMemberInfoMap);
+            GetShaderStorageBlockMembersInfo(interfaceBlock, mShaderStorageBlocks,
+                                             &mBlockMemberInfoMap);
         }
         if (variableType.isInterfaceBlock())
         {
@@ -572,7 +586,8 @@
                     {
                         mReferencedShaderStorageBlocks[interfaceBlock->uniqueId().get()] =
                             new TReferencedBlock(interfaceBlock, &instanceArraySymbol->variable());
-                        GetShaderStorageBlockMembersInfo(interfaceBlock, &mBlockMemberInfoMap);
+                        GetShaderStorageBlockMembersInfo(interfaceBlock, mShaderStorageBlocks,
+                                                         &mBlockMemberInfoMap);
                     }
 
                     const int arrayIndex = node->getRight()->getAsConstantUnion()->getIConst(0);
diff --git a/src/compiler/translator/ShaderStorageBlockOutputHLSL.h b/src/compiler/translator/ShaderStorageBlockOutputHLSL.h
index 340e709..caceca3 100644
--- a/src/compiler/translator/ShaderStorageBlockOutputHLSL.h
+++ b/src/compiler/translator/ShaderStorageBlockOutputHLSL.h
@@ -34,12 +34,15 @@
 // Used to save shader storage block field member information.
 using BlockMemberInfoMap = std::map<const TField *, BlockMemberInfo>;
 
+using ShaderVarToFieldMap = std::map<std::string, const TField *>;
+
 class ShaderStorageBlockOutputHLSL : public TIntermTraverser
 {
   public:
     ShaderStorageBlockOutputHLSL(OutputHLSL *outputHLSL,
                                  TSymbolTable *symbolTable,
-                                 ResourcesHLSL *resourcesHLSL);
+                                 ResourcesHLSL *resourcesHLSL,
+                                 const std::vector<InterfaceBlock> &shaderStorageBlocks);
 
     ~ShaderStorageBlockOutputHLSL();
 
@@ -82,6 +85,7 @@
     ReferencedInterfaceBlocks mReferencedShaderStorageBlocks;
 
     BlockMemberInfoMap mBlockMemberInfoMap;
+    const std::vector<InterfaceBlock> &mShaderStorageBlocks;
 };
 }  // namespace sh
 
diff --git a/src/compiler/translator/TranslatorHLSL.cpp b/src/compiler/translator/TranslatorHLSL.cpp
index d72348b..a8cf7f7 100644
--- a/src/compiler/translator/TranslatorHLSL.cpp
+++ b/src/compiler/translator/TranslatorHLSL.cpp
@@ -138,7 +138,7 @@
     sh::OutputHLSL outputHLSL(getShaderType(), getShaderVersion(), getExtensionBehavior(),
                               getSourcePath(), getOutputType(), numRenderTargets, getUniforms(),
                               compileOptions, getComputeShaderLocalSize(), &getSymbolTable(),
-                              perfDiagnostics);
+                              perfDiagnostics, mShaderStorageBlocks);
 
     outputHLSL.output(root, getInfoSink().obj);
 
diff --git a/src/compiler/translator/blocklayout.cpp b/src/compiler/translator/blocklayout.cpp
index 56c0521..d754eaf 100644
--- a/src/compiler/translator/blocklayout.cpp
+++ b/src/compiler/translator/blocklayout.cpp
@@ -46,10 +46,6 @@
                            bool inRowMajorLayout,
                            BlockLayoutMap *blockInfoOut)
 {
-    // TODO(jiajia.qin@intel.com):we need to set the right structure base alignment before
-    // enterAggregateType for std430 layout just like GetShaderStorageBlockFieldMemberInfo did in
-    // ShaderStorageBlockOutputHLSL.cpp. http://anglebug.com/1920
-
     BlockLayoutMapVisitor visitor(blockInfoOut, prefix, encoder);
     TraverseShaderVariables(fields, inRowMajorLayout, &visitor);
 }
@@ -60,13 +56,12 @@
 {
     const std::vector<ShaderVariable> &fields = variable.fields;
 
-    visitor->enterStructAccess(variable);
+    visitor->enterStructAccess(variable, isRowMajorLayout);
     TraverseShaderVariables(fields, isRowMajorLayout, visitor);
-    visitor->exitStructAccess(variable);
+    visitor->exitStructAccess(variable, isRowMajorLayout);
 }
 
 void TraverseStructArrayVariable(const ShaderVariable &variable,
-                                 unsigned int arrayNestingIndex,
                                  bool inRowMajorLayout,
                                  ShaderVariableVisitor *visitor)
 {
@@ -74,7 +69,7 @@
 
     // Nested arrays are processed starting from outermost (arrayNestingIndex 0u) and ending at the
     // innermost. We make a special case for unsized arrays.
-    const unsigned int currentArraySize = variable.getNestedArraySize(arrayNestingIndex);
+    const unsigned int currentArraySize = variable.getNestedArraySize(0);
     unsigned int count                  = std::max(currentArraySize, 1u);
     for (unsigned int arrayElement = 0u; arrayElement < count; ++arrayElement)
     {
@@ -83,9 +78,9 @@
         ShaderVariable elementVar = variable;
         elementVar.indexIntoArray(arrayElement);
 
-        if (arrayNestingIndex + 1u < variable.arraySizes.size())
+        if (variable.arraySizes.size() > 1u)
         {
-            TraverseStructArrayVariable(elementVar, arrayNestingIndex, inRowMajorLayout, visitor);
+            TraverseStructArrayVariable(elementVar, inRowMajorLayout, visitor);
         }
         else
         {
@@ -145,9 +140,34 @@
     }
     return strstr.str();
 }
+
+size_t GetStd430BaseAlignment(GLenum variableType, bool isRowMajor)
+{
+    GLenum flippedType   = isRowMajor ? variableType : gl::TransposeMatrixType(variableType);
+    size_t numComponents = static_cast<size_t>(gl::VariableColumnCount(flippedType));
+    return ComponentAlignment(numComponents);
+}
+
+class BaseAlignmentVisitor : public ShaderVariableVisitor
+{
+  public:
+    BaseAlignmentVisitor() = default;
+    void visitVariable(const ShaderVariable &variable, bool isRowMajor) override
+    {
+        size_t baseAlignment = GetStd430BaseAlignment(variable.type, isRowMajor);
+        mCurrentAlignment    = std::max(mCurrentAlignment, baseAlignment);
+    }
+
+    // This is in components rather than bytes.
+    size_t getBaseAlignment() const { return mCurrentAlignment; }
+
+  private:
+    size_t mCurrentAlignment = 0;
+};
 }  // anonymous namespace
 
-BlockLayoutEncoder::BlockLayoutEncoder() : mCurrentOffset(0), mStructureBaseAlignment(0) {}
+// BlockLayoutEncoder implementation.
+BlockLayoutEncoder::BlockLayoutEncoder() : mCurrentOffset(0) {}
 
 BlockMemberInfo BlockLayoutEncoder::encodeType(GLenum type,
                                                const std::vector<unsigned int> &arraySizes,
@@ -168,16 +188,6 @@
     return memberInfo;
 }
 
-void BlockLayoutEncoder::increaseCurrentOffset(size_t offsetInBytes)
-{
-    mCurrentOffset += (offsetInBytes / kBytesPerComponent);
-}
-
-void BlockLayoutEncoder::setStructureBaseAlignment(size_t baseAlignment)
-{
-    mStructureBaseAlignment = baseAlignment;
-}
-
 // static
 size_t BlockLayoutEncoder::GetBlockRegister(const BlockMemberInfo &info)
 {
@@ -190,21 +200,22 @@
     return (info.offset / kBytesPerComponent) % kComponentsPerRegister;
 }
 
-void BlockLayoutEncoder::nextRegister()
+void BlockLayoutEncoder::align(size_t baseAlignment)
 {
-    mCurrentOffset = rx::roundUp<size_t>(mCurrentOffset, kComponentsPerRegister);
+    mCurrentOffset = rx::roundUp<size_t>(mCurrentOffset, baseAlignment);
 }
 
+// Std140BlockEncoder implementation.
 Std140BlockEncoder::Std140BlockEncoder() {}
 
-void Std140BlockEncoder::enterAggregateType()
+void Std140BlockEncoder::enterAggregateType(const ShaderVariable &structVar)
 {
-    nextRegister();
+    align(getBaseAlignment(structVar));
 }
 
-void Std140BlockEncoder::exitAggregateType()
+void Std140BlockEncoder::exitAggregateType(const ShaderVariable &structVar)
 {
-    nextRegister();
+    align(getBaseAlignment(structVar));
 }
 
 void Std140BlockEncoder::getBlockLayoutInfo(GLenum type,
@@ -222,24 +233,24 @@
 
     if (gl::IsMatrixType(type))
     {
-        baseAlignment = kComponentsPerRegister;
-        matrixStride  = kComponentsPerRegister;
+        baseAlignment = getTypeBaseAlignment(type, isRowMajorMatrix);
+        matrixStride  = getTypeBaseAlignment(type, isRowMajorMatrix);
 
         if (!arraySizes.empty())
         {
             const int numRegisters = gl::MatrixRegisterCount(type, isRowMajorMatrix);
-            arrayStride            = kComponentsPerRegister * numRegisters;
+            arrayStride            = getTypeBaseAlignment(type, isRowMajorMatrix) * numRegisters;
         }
     }
     else if (!arraySizes.empty())
     {
-        baseAlignment = kComponentsPerRegister;
-        arrayStride   = kComponentsPerRegister;
+        baseAlignment = getTypeBaseAlignment(type, false);
+        arrayStride   = getTypeBaseAlignment(type, false);
     }
     else
     {
-        const int numComponents = gl::VariableComponentCount(type);
-        baseAlignment           = (numComponents == 3 ? 4u : static_cast<size_t>(numComponents));
+        const size_t numComponents = static_cast<size_t>(gl::VariableComponentCount(type));
+        baseAlignment              = ComponentAlignment(numComponents);
     }
 
     mCurrentOffset = rx::roundUp(mCurrentOffset, baseAlignment);
@@ -269,52 +280,24 @@
     }
 }
 
+// Std430BlockEncoder implementation.
 Std430BlockEncoder::Std430BlockEncoder() {}
 
-void Std430BlockEncoder::nextRegister()
+size_t Std430BlockEncoder::getBaseAlignment(const ShaderVariable &shaderVar) const
 {
-    mCurrentOffset = rx::roundUp<size_t>(mCurrentOffset, mStructureBaseAlignment);
+    if (shaderVar.isStruct())
+    {
+        BaseAlignmentVisitor visitor;
+        TraverseShaderVariables(shaderVar.fields, false, &visitor);
+        return visitor.getBaseAlignment();
+    }
+
+    return GetStd430BaseAlignment(shaderVar.type, shaderVar.isRowMajorLayout);
 }
 
-void Std430BlockEncoder::getBlockLayoutInfo(GLenum type,
-                                            const std::vector<unsigned int> &arraySizes,
-                                            bool isRowMajorMatrix,
-                                            int *arrayStrideOut,
-                                            int *matrixStrideOut)
+size_t Std430BlockEncoder::getTypeBaseAlignment(GLenum type, bool isRowMajorMatrix) const
 {
-    // We assume we are only dealing with 4 byte components (no doubles or half-words currently)
-    ASSERT(gl::VariableComponentSize(gl::VariableComponentType(type)) == kBytesPerComponent);
-
-    size_t baseAlignment = 0;
-    int matrixStride     = 0;
-    int arrayStride      = 0;
-
-    if (gl::IsMatrixType(type))
-    {
-        const int numComponents = gl::MatrixComponentCount(type, isRowMajorMatrix);
-        baseAlignment           = (numComponents == 3 ? 4u : static_cast<size_t>(numComponents));
-        matrixStride            = baseAlignment;
-
-        if (!arraySizes.empty())
-        {
-            const int numRegisters = gl::MatrixRegisterCount(type, isRowMajorMatrix);
-            arrayStride            = matrixStride * numRegisters;
-        }
-    }
-    else
-    {
-        const int numComponents = gl::VariableComponentCount(type);
-        baseAlignment           = (numComponents == 3 ? 4u : static_cast<size_t>(numComponents));
-        if (!arraySizes.empty())
-        {
-            arrayStride = baseAlignment;
-        }
-    }
-    mStructureBaseAlignment = std::max(baseAlignment, mStructureBaseAlignment);
-    mCurrentOffset          = rx::roundUp(mCurrentOffset, baseAlignment);
-
-    *matrixStrideOut = matrixStride;
-    *arrayStrideOut  = arrayStride;
+    return GetStd430BaseAlignment(type, isRowMajorMatrix);
 }
 
 void GetInterfaceBlockInfo(const std::vector<InterfaceBlockField> &fields,
@@ -366,13 +349,13 @@
     mMappedNameStack.pop_back();
 }
 
-void VariableNameVisitor::enterStructAccess(const ShaderVariable &structVar)
+void VariableNameVisitor::enterStructAccess(const ShaderVariable &structVar, bool isRowMajor)
 {
     mNameStack.push_back(".");
     mMappedNameStack.push_back(".");
 }
 
-void VariableNameVisitor::exitStructAccess(const ShaderVariable &structVar)
+void VariableNameVisitor::exitStructAccess(const ShaderVariable &structVar, bool isRowMajor)
 {
     mNameStack.pop_back();
     mMappedNameStack.pop_back();
@@ -380,7 +363,7 @@
 
 void VariableNameVisitor::enterArray(const ShaderVariable &arrayVar)
 {
-    if (!arrayVar.hasParentArrayIndex())
+    if (!arrayVar.hasParentArrayIndex() && !arrayVar.isStruct())
     {
         mNameStack.push_back(arrayVar.name);
         mMappedNameStack.push_back(arrayVar.mappedName);
@@ -389,7 +372,7 @@
 
 void VariableNameVisitor::exitArray(const ShaderVariable &arrayVar)
 {
-    if (!arrayVar.hasParentArrayIndex())
+    if (!arrayVar.hasParentArrayIndex() && !arrayVar.isStruct())
     {
         mNameStack.pop_back();
         mMappedNameStack.pop_back();
@@ -465,16 +448,16 @@
 
 BlockEncoderVisitor::~BlockEncoderVisitor() = default;
 
-void BlockEncoderVisitor::enterStructAccess(const ShaderVariable &structVar)
+void BlockEncoderVisitor::enterStructAccess(const ShaderVariable &structVar, bool isRowMajor)
 {
-    VariableNameVisitor::enterStructAccess(structVar);
-    mEncoder->enterAggregateType();
+    VariableNameVisitor::enterStructAccess(structVar, isRowMajor);
+    mEncoder->enterAggregateType(structVar);
 }
 
-void BlockEncoderVisitor::exitStructAccess(const ShaderVariable &structVar)
+void BlockEncoderVisitor::exitStructAccess(const ShaderVariable &structVar, bool isRowMajor)
 {
-    mEncoder->exitAggregateType();
-    VariableNameVisitor::exitStructAccess(structVar);
+    mEncoder->exitAggregateType(structVar);
+    VariableNameVisitor::exitStructAccess(structVar, isRowMajor);
 }
 
 void BlockEncoderVisitor::visitNamedVariable(const ShaderVariable &variable,
@@ -502,16 +485,16 @@
 
     if (variable.isStruct())
     {
+        visitor->enterStruct(variable);
         if (variable.isArray())
         {
-            TraverseStructArrayVariable(variable, 0u, rowMajorLayout, visitor);
+            TraverseStructArrayVariable(variable, rowMajorLayout, visitor);
         }
         else
         {
-            visitor->enterStruct(variable);
             TraverseStructVariable(variable, rowMajorLayout, visitor);
-            visitor->exitStruct(variable);
         }
+        visitor->exitStruct(variable);
     }
     else if (variable.isArrayOfArrays())
     {
diff --git a/src/compiler/translator/blocklayout.h b/src/compiler/translator/blocklayout.h
index af9ab0c..bd11b95 100644
--- a/src/compiler/translator/blocklayout.h
+++ b/src/compiler/translator/blocklayout.h
@@ -66,6 +66,11 @@
     int topLevelArrayStride = -1;
 };
 
+constexpr size_t ComponentAlignment(size_t numComponents)
+{
+    return (numComponents == 3u ? 4u : numComponents);
+}
+
 constexpr BlockMemberInfo kDefaultBlockMemberInfo;
 
 class BlockLayoutEncoder
@@ -78,13 +83,11 @@
                                const std::vector<unsigned int> &arraySizes,
                                bool isRowMajorMatrix);
 
-    size_t getBlockSize() const { return mCurrentOffset * kBytesPerComponent; }
-    size_t getStructureBaseAlignment() const { return mStructureBaseAlignment; }
-    void increaseCurrentOffset(size_t offsetInBytes);
-    void setStructureBaseAlignment(size_t baseAlignment);
+    size_t getCurrentOffset() const { return mCurrentOffset * kBytesPerComponent; }
 
-    virtual void enterAggregateType() = 0;
-    virtual void exitAggregateType()  = 0;
+    // Called when entering/exiting a structure variable.
+    virtual void enterAggregateType(const ShaderVariable &structVar) = 0;
+    virtual void exitAggregateType(const ShaderVariable &structVar)  = 0;
 
     static constexpr size_t kBytesPerComponent           = 4u;
     static constexpr unsigned int kComponentsPerRegister = 4u;
@@ -93,10 +96,7 @@
     static size_t GetBlockRegisterElement(const BlockMemberInfo &info);
 
   protected:
-    size_t mCurrentOffset;
-    size_t mStructureBaseAlignment;
-
-    virtual void nextRegister();
+    void align(size_t baseAlignment);
 
     virtual void getBlockLayoutInfo(GLenum type,
                                     const std::vector<unsigned int> &arraySizes,
@@ -108,6 +108,8 @@
                                bool isRowMajorMatrix,
                                int arrayStride,
                                int matrixStride)          = 0;
+
+    size_t mCurrentOffset;
 };
 
 // Will return default values for everything.
@@ -116,8 +118,8 @@
   public:
     DummyBlockEncoder() = default;
 
-    void enterAggregateType() override {}
-    void exitAggregateType() override {}
+    void enterAggregateType(const ShaderVariable &structVar) override {}
+    void exitAggregateType(const ShaderVariable &structVar) override {}
 
   protected:
     void getBlockLayoutInfo(GLenum type,
@@ -146,8 +148,8 @@
   public:
     Std140BlockEncoder();
 
-    void enterAggregateType() override;
-    void exitAggregateType() override;
+    void enterAggregateType(const ShaderVariable &structVar) override;
+    void exitAggregateType(const ShaderVariable &structVar) override;
 
   protected:
     void getBlockLayoutInfo(GLenum type,
@@ -160,6 +162,16 @@
                        bool isRowMajorMatrix,
                        int arrayStride,
                        int matrixStride) override;
+
+    virtual size_t getBaseAlignment(const ShaderVariable &variable) const
+    {
+        return kComponentsPerRegister;
+    }
+
+    virtual size_t getTypeBaseAlignment(GLenum type, bool isRowMajorMatrix) const
+    {
+        return kComponentsPerRegister;
+    }
 };
 
 class Std430BlockEncoder : public Std140BlockEncoder
@@ -168,12 +180,8 @@
     Std430BlockEncoder();
 
   protected:
-    void nextRegister() override;
-    void getBlockLayoutInfo(GLenum type,
-                            const std::vector<unsigned int> &arraySizes,
-                            bool isRowMajorMatrix,
-                            int *arrayStrideOut,
-                            int *matrixStrideOut) override;
+    size_t getBaseAlignment(const ShaderVariable &variable) const override;
+    size_t getTypeBaseAlignment(GLenum type, bool isRowMajorMatrix) const override;
 };
 
 using BlockLayoutMap = std::map<std::string, BlockMemberInfo>;
@@ -197,8 +205,8 @@
     virtual void enterStruct(const ShaderVariable &structVar) {}
     virtual void exitStruct(const ShaderVariable &structVar) {}
 
-    virtual void enterStructAccess(const ShaderVariable &structVar) {}
-    virtual void exitStructAccess(const ShaderVariable &structVar) {}
+    virtual void enterStructAccess(const ShaderVariable &structVar, bool isRowMajor) {}
+    virtual void exitStructAccess(const ShaderVariable &structVar, bool isRowMajor) {}
 
     virtual void enterArray(const ShaderVariable &arrayVar) {}
     virtual void exitArray(const ShaderVariable &arrayVar) {}
@@ -222,8 +230,8 @@
 
     void enterStruct(const ShaderVariable &structVar) override;
     void exitStruct(const ShaderVariable &structVar) override;
-    void enterStructAccess(const ShaderVariable &structVar) override;
-    void exitStructAccess(const ShaderVariable &structVar) override;
+    void enterStructAccess(const ShaderVariable &structVar, bool isRowMajor) override;
+    void exitStructAccess(const ShaderVariable &structVar, bool isRowMajor) override;
     void enterArray(const ShaderVariable &arrayVar) override;
     void exitArray(const ShaderVariable &arrayVar) override;
     void enterArrayElement(const ShaderVariable &arrayVar, unsigned int arrayElement) override;
@@ -243,11 +251,12 @@
                                     const std::string &name,
                                     const std::string &mappedName) = 0;
 
+    std::string collapseNameStack() const;
+    std::string collapseMappedNameStack() const;
+
   private:
     void visitSampler(const sh::ShaderVariable &sampler) final;
     void visitVariable(const ShaderVariable &variable, bool isRowMajor) final;
-    std::string collapseNameStack() const;
-    std::string collapseMappedNameStack() const;
 
     std::vector<std::string> mNameStack;
     std::vector<std::string> mMappedNameStack;
@@ -261,8 +270,8 @@
                         BlockLayoutEncoder *encoder);
     ~BlockEncoderVisitor();
 
-    void enterStructAccess(const ShaderVariable &structVar) override;
-    void exitStructAccess(const ShaderVariable &structVar) override;
+    void enterStructAccess(const ShaderVariable &structVar, bool isRowMajor) override;
+    void exitStructAccess(const ShaderVariable &structVar, bool isRowMajor) override;
 
     void visitNamedVariable(const ShaderVariable &variable,
                             bool isRowMajor,
diff --git a/src/compiler/translator/blocklayoutHLSL.cpp b/src/compiler/translator/blocklayoutHLSL.cpp
index 66556c4..b94f645 100644
--- a/src/compiler/translator/blocklayoutHLSL.cpp
+++ b/src/compiler/translator/blocklayoutHLSL.cpp
@@ -19,12 +19,12 @@
     : mEncoderStrategy(strategy), mTransposeMatrices(transposeMatrices)
 {}
 
-void HLSLBlockEncoder::enterAggregateType()
+void HLSLBlockEncoder::enterAggregateType(const ShaderVariable &structVar)
 {
-    nextRegister();
+    align(kComponentsPerRegister);
 }
 
-void HLSLBlockEncoder::exitAggregateType() {}
+void HLSLBlockEncoder::exitAggregateType(const ShaderVariable &structVar) {}
 
 void HLSLBlockEncoder::getBlockLayoutInfo(GLenum typeIn,
                                           const std::vector<unsigned int> &arraySizes,
@@ -45,7 +45,7 @@
     // register
     if (!isPacked() || gl::IsMatrixType(type) || !arraySizes.empty())
     {
-        nextRegister();
+        align(kComponentsPerRegister);
     }
 
     if (gl::IsMatrixType(type))
@@ -67,7 +67,7 @@
         int numComponents = gl::VariableComponentCount(type);
         if ((numComponents + (mCurrentOffset % kComponentsPerRegister)) > kComponentsPerRegister)
         {
-            nextRegister();
+            align(kComponentsPerRegister);
         }
     }
 
@@ -134,14 +134,14 @@
     {
         for (size_t arrayElement = 0; arrayElement < variable.getArraySizeProduct(); arrayElement++)
         {
-            encoder->enterAggregateType();
+            encoder->enterAggregateType(variable);
 
-            for (size_t fieldIndex = 0; fieldIndex < variable.fields.size(); fieldIndex++)
+            for (const ShaderVariable &field : variable.fields)
             {
-                HLSLVariableRegisterCount(variable.fields[fieldIndex], encoder);
+                HLSLVariableRegisterCount(field, encoder);
             }
 
-            encoder->exitAggregateType();
+            encoder->exitAggregateType(variable);
         }
     }
     else
@@ -157,7 +157,7 @@
     HLSLVariableRegisterCount(variable, &encoder);
 
     const size_t registerBytes = (encoder.kBytesPerComponent * encoder.kComponentsPerRegister);
-    return static_cast<unsigned int>(rx::roundUp<size_t>(encoder.getBlockSize(), registerBytes) /
-                                     registerBytes);
+    return static_cast<unsigned int>(
+        rx::roundUp<size_t>(encoder.getCurrentOffset(), registerBytes) / registerBytes);
 }
 }  // namespace sh
diff --git a/src/compiler/translator/blocklayoutHLSL.h b/src/compiler/translator/blocklayoutHLSL.h
index a310d2f..98d93fc 100644
--- a/src/compiler/translator/blocklayoutHLSL.h
+++ b/src/compiler/translator/blocklayoutHLSL.h
@@ -35,8 +35,8 @@
 
     HLSLBlockEncoder(HLSLBlockEncoderStrategy strategy, bool transposeMatrices);
 
-    void enterAggregateType() override;
-    void exitAggregateType() override;
+    void enterAggregateType(const ShaderVariable &structVar) override;
+    void exitAggregateType(const ShaderVariable &structVar) override;
     void skipRegisters(unsigned int numRegisters);
 
     bool isPacked() const { return mEncoderStrategy == ENCODE_PACKED; }