// Copyright (c) 2024 Epic Games, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "struct_packing_pass.h"

#include <algorithm>

#include "source/opt/instruction.h"
#include "source/opt/ir_context.h"

namespace spvtools {
namespace opt {

/*
Std140 packing rules from the original GLSL 140 specification (see
https://registry.khronos.org/OpenGL/extensions/ARB/ARB_uniform_buffer_object.txt)

When using the "std140" storage layout, structures will be laid out in
buffer storage with its members stored in monotonically increasing order
based on their location in the declaration. A structure and each
structure member have a base offset and a base alignment, from which an
aligned offset is computed by rounding the base offset up to a multiple of
the base alignment. The base offset of the first member of a structure is
taken from the aligned offset of the structure itself. The base offset of
all other structure members is derived by taking the offset of the last
basic machine unit consumed by the previous member and adding one. Each
structure member is stored in memory at its aligned offset. The members
of a top-level uniform block are laid out in buffer storage by treating
the uniform block as a structure with a base offset of zero.

(1) If the member is a scalar consuming <N> basic machine units, the
    base alignment is <N>.

(2) If the member is a two- or four-component vector with components
    consuming <N> basic machine units, the base alignment is 2<N> or
    4<N>, respectively.

(3) If the member is a three-component vector with components consuming
    <N> basic machine units, the base alignment is 4<N>.

(4) If the member is an array of scalars or vectors, the base alignment
    and array stride are set to match the base alignment of a single
    array element, according to rules (1), (2), and (3), and rounded up
    to the base alignment of a vec4. The array may have padding at the
    end; the base offset of the member following the array is rounded up
    to the next multiple of the base alignment.

(5) If the member is a column-major matrix with <C> columns and <R>
    rows, the matrix is stored identically to an array of <C> column
    vectors with <R> components each, according to rule (4).

(6) If the member is an array of <S> column-major matrices with <C>
    columns and <R> rows, the matrix is stored identically to a row of
    <S>*<C> column vectors with <R> components each, according to rule
    (4).

(7) If the member is a row-major matrix with <C> columns and <R> rows,
    the matrix is stored identically to an array of <R> row vectors
    with <C> components each, according to rule (4).

(8) If the member is an array of <S> row-major matrices with <C> columns
    and <R> rows, the matrix is stored identically to a row of <S>*<R>
    row vectors with <C> components each, according to rule (4).

(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, and rounded up to the base alignment of a vec4. The
    individual members of this sub-structure are then assigned offsets
    by applying this set of rules recursively, where the base offset of
    the first member of the sub-structure is equal to the aligned offset
    of the structure. The structure may have padding at the end; the
    base offset of the member following the sub-structure is rounded up
    to the next multiple of the base alignment of the structure.

(10) If the member is an array of <S> structures, the <S> elements of
    the array are laid out in order, according to rule (9).
*/

static bool isPackingVec4Padded(StructPackingPass::PackingRules rules) {
  switch (rules) {
    case StructPackingPass::PackingRules::Std140:
    case StructPackingPass::PackingRules::Std140EnhancedLayout:
    case StructPackingPass::PackingRules::HlslCbuffer:
    case StructPackingPass::PackingRules::HlslCbufferPackOffset:
      return true;
    default:
      return false;
  }
}

static bool isPackingScalar(StructPackingPass::PackingRules rules) {
  switch (rules) {
    case StructPackingPass::PackingRules::Scalar:
    case StructPackingPass::PackingRules::ScalarEnhancedLayout:
      return true;
    default:
      return false;
  }
}

static bool isPackingHlsl(StructPackingPass::PackingRules rules) {
  switch (rules) {
    case StructPackingPass::PackingRules::HlslCbuffer:
    case StructPackingPass::PackingRules::HlslCbufferPackOffset:
      return true;
    default:
      return false;
  }
}

static uint32_t getPackedBaseSize(const analysis::Type& type) {
  switch (type.kind()) {
    case analysis::Type::kBool:
      return 1;
    case analysis::Type::kInteger:
      return type.AsInteger()->width() / 8;
    case analysis::Type::kFloat:
      return type.AsFloat()->width() / 8;
    case analysis::Type::kVector:
      return getPackedBaseSize(*type.AsVector()->element_type());
    case analysis::Type::kMatrix:
      return getPackedBaseSize(*type.AsMatrix()->element_type());
    default:
      break;  // we only expect bool, int, float, vec, and mat here
  }
  assert(0 && "Unrecognized type to get base size");
  return 0;
}

static uint32_t getScalarElementCount(const analysis::Type& type) {
  switch (type.kind()) {
    case analysis::Type::kVector:
      return type.AsVector()->element_count();
    case analysis::Type::kMatrix:
      return getScalarElementCount(*type.AsMatrix()->element_type());
    case analysis::Type::kStruct:
      assert(0 && "getScalarElementCount() does not recognized struct types");
      return 0;
    default:
      return 1;
  }
}

// Aligns the specified value to a multiple of alignment, whereas the
// alignment must be a power-of-two.
static uint32_t alignPow2(uint32_t value, uint32_t alignment) {
  return (value + alignment - 1) & ~(alignment - 1);
}

void StructPackingPass::buildConstantsMap() {
  constantsMap_.clear();
  for (Instruction* instr : context()->module()->GetConstants()) {
    constantsMap_[instr->result_id()] = instr;
  }
}

uint32_t StructPackingPass::getPackedAlignment(
    const analysis::Type& type) const {
  switch (type.kind()) {
    case analysis::Type::kArray: {
      // Get alignment of base type and round up to minimum alignment
      const uint32_t minAlignment = isPackingVec4Padded(packingRules_) ? 16 : 1;
      return std::max<uint32_t>(
          minAlignment, getPackedAlignment(*type.AsArray()->element_type()));
    }
    case analysis::Type::kStruct: {
      // Rule 9. Struct alignment is maximum alignmnet of its members
      uint32_t alignment = 1;

      for (const analysis::Type* elementType :
           type.AsStruct()->element_types()) {
        alignment =
            std::max<uint32_t>(alignment, getPackedAlignment(*elementType));
      }

      if (isPackingVec4Padded(packingRules_))
        alignment = std::max<uint32_t>(alignment, 16u);

      return alignment;
    }
    default: {
      const uint32_t baseAlignment = getPackedBaseSize(type);

      // Scalar block layout always uses alignment for the most basic component
      if (isPackingScalar(packingRules_)) return baseAlignment;

      if (const analysis::Matrix* matrixType = type.AsMatrix()) {
        // Rule 5/7
        if (isPackingVec4Padded(packingRules_) ||
            matrixType->element_count() == 3)
          return baseAlignment * 4;
        else
          return baseAlignment * matrixType->element_count();
      } else if (const analysis::Vector* vectorType = type.AsVector()) {
        // Rule 1
        if (vectorType->element_count() == 1) return baseAlignment;

        // Rule 2
        if (vectorType->element_count() == 2 ||
            vectorType->element_count() == 4)
          return baseAlignment * vectorType->element_count();

        // Rule 3
        if (vectorType->element_count() == 3) return baseAlignment * 4;
      } else {
        // Rule 1
        return baseAlignment;
      }
    }
  }
  assert(0 && "Unrecognized type to get packed alignment");
  return 0;
}

static uint32_t getPadAlignment(const analysis::Type& type,
                                uint32_t packedAlignment) {
  // The next member following a struct member is aligned to the base alignment
  // of a previous struct member.
  return type.kind() == analysis::Type::kStruct ? packedAlignment : 1;
}

uint32_t StructPackingPass::getPackedSize(const analysis::Type& type) const {
  switch (type.kind()) {
    case analysis::Type::kArray: {
      if (const analysis::Array* arrayType = type.AsArray()) {
        uint32_t size =
            getPackedArrayStride(*arrayType) * getArrayLength(*arrayType);

        // For arrays of vector and matrices in HLSL, the last element has a
        // size depending on its vector/matrix size to allow packing other
        // vectors in the last element.
        const analysis::Type* arraySubType = arrayType->element_type();
        if (isPackingHlsl(packingRules_) &&
            arraySubType->kind() != analysis::Type::kStruct) {
          size -= (4 - getScalarElementCount(*arraySubType)) *
                  getPackedBaseSize(*arraySubType);
        }
        return size;
      }
      break;
    }
    case analysis::Type::kStruct: {
      uint32_t size = 0;
      uint32_t padAlignment = 1;
      for (const analysis::Type* memberType :
           type.AsStruct()->element_types()) {
        const uint32_t packedAlignment = getPackedAlignment(*memberType);
        const uint32_t alignment =
            std::max<uint32_t>(packedAlignment, padAlignment);
        padAlignment = getPadAlignment(*memberType, packedAlignment);
        size = alignPow2(size, alignment);
        size += getPackedSize(*memberType);
      }
      return size;
    }
    default: {
      const uint32_t baseAlignment = getPackedBaseSize(type);
      if (isPackingScalar(packingRules_)) {
        return getScalarElementCount(type) * baseAlignment;
      } else {
        uint32_t size = 0;
        if (const analysis::Matrix* matrixType = type.AsMatrix()) {
          const analysis::Vector* matrixSubType =
              matrixType->element_type()->AsVector();
          assert(matrixSubType != nullptr &&
                 "Matrix sub-type is expected to be a vector type");
          if (isPackingVec4Padded(packingRules_) ||
              matrixType->element_count() == 3)
            size = matrixSubType->element_count() * baseAlignment * 4;
          else
            size = matrixSubType->element_count() * baseAlignment *
                   matrixType->element_count();

          // For matrices in HLSL, the last element has a size depending on its
          // vector size to allow packing other vectors in the last element.
          if (isPackingHlsl(packingRules_)) {
            size -= (4 - matrixSubType->element_count()) *
                    getPackedBaseSize(*matrixSubType);
          }
        } else if (const analysis::Vector* vectorType = type.AsVector()) {
          size = vectorType->element_count() * baseAlignment;
        } else {
          size = baseAlignment;
        }
        return size;
      }
    }
  }
  assert(0 && "Unrecognized type to get packed size");
  return 0;
}

uint32_t StructPackingPass::getPackedArrayStride(
    const analysis::Array& arrayType) const {
  // Array stride is equal to aligned size of element type
  const uint32_t elementSize = getPackedSize(*arrayType.element_type());
  const uint32_t alignment = getPackedAlignment(arrayType);
  return alignPow2(elementSize, alignment);
}

uint32_t StructPackingPass::getArrayLength(
    const analysis::Array& arrayType) const {
  return getConstantInt(arrayType.LengthId());
}

uint32_t StructPackingPass::getConstantInt(spv::Id id) const {
  auto it = constantsMap_.find(id);
  assert(it != constantsMap_.end() &&
         "Failed to map SPIR-V instruction ID to constant value");
  [[maybe_unused]] const analysis::Type* constType =
      context()->get_type_mgr()->GetType(it->second->type_id());
  assert(constType != nullptr &&
         "Failed to map SPIR-V instruction result type to definition");
  assert(constType->kind() == analysis::Type::kInteger &&
         "Failed to map SPIR-V instruction result type to integer type");
  return it->second->GetOperand(2).words[0];
}

StructPackingPass::PackingRules StructPackingPass::ParsePackingRuleFromString(
    const std::string& s) {
  if (s == "std140") return PackingRules::Std140;
  if (s == "std140EnhancedLayout") return PackingRules::Std140EnhancedLayout;
  if (s == "std430") return PackingRules::Std430;
  if (s == "std430EnhancedLayout") return PackingRules::Std430EnhancedLayout;
  if (s == "hlslCbuffer") return PackingRules::HlslCbuffer;
  if (s == "hlslCbufferPackOffset") return PackingRules::HlslCbufferPackOffset;
  if (s == "scalar") return PackingRules::Scalar;
  if (s == "scalarEnhancedLayout") return PackingRules::ScalarEnhancedLayout;
  return PackingRules::Undefined;
}

StructPackingPass::StructPackingPass(const char* structToPack,
                                     PackingRules rules)
    : structToPack_{structToPack != nullptr ? structToPack : ""},
      packingRules_{rules} {}

Pass::Status StructPackingPass::Process() {
  if (packingRules_ == PackingRules::Undefined) {
    if (consumer()) {
      consumer()(SPV_MSG_ERROR, "", {0, 0, 0},
                 "Cannot pack struct with undefined rule");
    }
    return Status::Failure;
  }

  // Build Id-to-instruction map for easier access
  buildConstantsMap();

  // Find structure of interest
  const uint32_t structIdToPack = findStructIdByName(structToPack_.c_str());

  const Instruction* structDef =
      context()->get_def_use_mgr()->GetDef(structIdToPack);
  if (structDef == nullptr || structDef->opcode() != spv::Op::OpTypeStruct) {
    if (consumer()) {
      const std::string message =
          "Failed to find struct with name " + structToPack_;
      consumer()(SPV_MSG_ERROR, "", {0, 0, 0}, message.c_str());
    }
    return Status::Failure;
  }

  // Find all struct member types
  std::vector<const analysis::Type*> structMemberTypes =
      findStructMemberTypes(*structDef);

  return assignStructMemberOffsets(structIdToPack, structMemberTypes);
}

uint32_t StructPackingPass::findStructIdByName(const char* structName) const {
  for (Instruction& instr : context()->module()->debugs2()) {
    if (instr.opcode() == spv::Op::OpName &&
        instr.GetOperand(1).AsString() == structName) {
      return instr.GetOperand(0).AsId();
    }
  }
  return 0;
}

std::vector<const analysis::Type*> StructPackingPass::findStructMemberTypes(
    const Instruction& structDef) const {
  // Found struct type to pack, now collect all types of its members
  assert(structDef.NumOperands() > 0 &&
         "Number of operands in OpTypeStruct instruction must not be zero");
  const uint32_t numMembers = structDef.NumOperands() - 1;
  std::vector<const analysis::Type*> structMemberTypes;
  structMemberTypes.resize(numMembers);
  for (uint32_t i = 0; i < numMembers; ++i) {
    const spv::Id memberTypeId = structDef.GetOperand(1 + i).AsId();
    if (const analysis::Type* memberType =
            context()->get_type_mgr()->GetType(memberTypeId)) {
      structMemberTypes[i] = memberType;
    }
  }
  return structMemberTypes;
}

Pass::Status StructPackingPass::assignStructMemberOffsets(
    uint32_t structIdToPack,
    const std::vector<const analysis::Type*>& structMemberTypes) {
  // Returns true if the specified instruction is a OpMemberDecorate for the
  // struct we're looking for with an offset decoration
  auto isMemberOffsetDecoration =
      [structIdToPack](const Instruction& instr) -> bool {
    return instr.opcode() == spv::Op::OpMemberDecorate &&
           instr.GetOperand(0).AsId() == structIdToPack &&
           static_cast<spv::Decoration>(instr.GetOperand(2).words[0]) ==
               spv::Decoration::Offset;
  };

  bool modified = false;

  // Find and re-assign all member offset decorations
  for (auto it = context()->module()->annotation_begin(),
            itEnd = context()->module()->annotation_end();
       it != itEnd; ++it) {
    if (isMemberOffsetDecoration(*it)) {
      // Found first member decoration with offset, we expect all other
      // offsets right after the first one
      uint32_t prevMemberIndex = 0;
      uint32_t currentOffset = 0;
      uint32_t padAlignment = 1;
      do {
        const uint32_t memberIndex = it->GetOperand(1).words[0];
        if (memberIndex < prevMemberIndex) {
          // Failure: we expect all members to appear in consecutive order
          return Status::Failure;
        }

        // Apply alignment rules to current offset
        const analysis::Type& memberType = *structMemberTypes[memberIndex];
        uint32_t packedAlignment = getPackedAlignment(memberType);
        uint32_t packedSize = getPackedSize(memberType);

        if (isPackingHlsl(packingRules_)) {
          // If a member crosses vec4 boundaries, alignment is size of vec4
          if (currentOffset / 16 != (currentOffset + packedSize - 1) / 16)
            packedAlignment = std::max<uint32_t>(packedAlignment, 16u);
        }

        const uint32_t alignment =
            std::max<uint32_t>(packedAlignment, padAlignment);
        currentOffset = alignPow2(currentOffset, alignment);
        padAlignment = getPadAlignment(memberType, packedAlignment);

        // Override packed offset in instruction
        if (it->GetOperand(3).words[0] < currentOffset) {
          // Failure: packing resulted in higher offset for member than
          // previously generated
          return Status::Failure;
        }

        it->GetOperand(3).words[0] = currentOffset;
        modified = true;

        // Move to next member
        ++it;
        prevMemberIndex = memberIndex;
        currentOffset += packedSize;
      } while (it != itEnd && isMemberOffsetDecoration(*it));

      // We're done with all decorations for the struct of interest
      break;
    }
  }

  return modified ? Status::SuccessWithChange : Status::SuccessWithoutChange;
}

}  // namespace opt
}  // namespace spvtools