//
// Copyright 2024 The ANGLE Project Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
//

#ifndef COMPILER_TRANSLATOR_WGSL_OUTPUT_UNIFORM_BLOCKS_H_
#define COMPILER_TRANSLATOR_WGSL_OUTPUT_UNIFORM_BLOCKS_H_

#include "compiler/translator/Common.h"
#include "compiler/translator/Compiler.h"
#include "compiler/translator/IntermNode.h"

namespace sh
{

const char kDefaultUniformBlockVarType[]           = "ANGLE_DefaultUniformBlock";
const char kDefaultUniformBlockVarName[]           = "ANGLE_defaultUniformBlock";

enum WgslBindGroupNumbers : uint32_t
{
    kDefaultUniformBlockBindGroup = 0,
    kTextureAndSamplerBindGroup   = 1,
    kDriverUniformBindGroup       = 2,
    kMaxBindGroup                 = 2
};

const uint32_t kDefaultVertexUniformBlockBinding   = 0;
const uint32_t kDefaultFragmentUniformBlockBinding = 1;

const uint32_t kDriverUniformBlockBinding = 0;

// The translator emits this dummy location, which needs to be replaced when linking the two
// separate shader stages.
const char kTextureSamplerBindingMarker[] = "@group(1) @binding(@@@@@@) var ";
// The translator emits split samplers/textures for GLSL's combined textures/samplers (combined
// textures/samplers are not supported by WGSL). The new variables are prefixed with these
// constants, respectively.
const char kAngleSamplerPrefix[] = "ANGLE_sampler_";
const char kAngleTexturePrefix[] = "ANGLE_texture_";

const char kWrappedStructFieldName[] = "elem";

struct UniformBlockMetadata
{
    // A list of structs used anywhere in the uniform address space. These will require special
    // handling (@align() attributes, wrapping of basic types, etc.) to ensure they fit WGSL's
    // uniform layout requirements.
    // The key is TSymbolUniqueId::get().
    TUnorderedSet<int> structsInUniformAddressSpace;
};

// Given a GLSL AST `root`, fills in `outMetadata`, to be used when outputting WGSL.
// If the AST is manipulated after calling this, it may be out of sync with the data recorded in
// `outMetadata`.
bool RecordUniformBlockMetadata(TIntermBlock *root, UniformBlockMetadata &outMetadata);

// Based on the GLSL, some extra WGSL will have to be generated so it can be referenced by
// the WGSL generated by the traverser, This tracks exactly which WGSL snippets will need to be
// generated,
struct WGSLGenerationMetadataForUniforms
{
    // Arrays must have a stride of at least 16 if used in the uniform address spaces. If the array
    // element type doesn't have an aligned size of a multiple of 16 (e.g. f32), the element type
    // must be wrapped in a struct which is then aligned to 16. Adding to
    // `arrayElementTypesInUniforms` will cause `OutputUniformWrapperStructsAndConversions` to
    // generate a WGSL wrapper struct of the form:
    //
    // struct ANGLE_wrapper_f32 {
    //   @align(16) elem : f32;
    // };
    TSet<TType> arrayElementTypesInUniforms;

    // If we need to convert arrays with wrapped element types into arrays with unwrapped element
    // types, the necessary conversions are listed here.
    TSet<TType> arrayElementTypesThatNeedUnwrappingConversions;

    // MatCx2 in a uniform will be represented as array<ANGLE_wrapped_vec2, C> to match std140 (WGSL
    // uniforms pack their matrices a bit tighter). These will have to be converted back to matCx2
    // for all other purposes (e.g. multiplication).
    // If this set isn't empty, ANGLE_wrapped_vec2 will be generated even if it
    // hasn't been yet.
    TSet<TType> outputMatCx2Conversion;
};
bool OutputUniformWrapperStructsAndConversions(
    TInfoSinkBase &output,
    const WGSLGenerationMetadataForUniforms &wgslGenerationMetadataForUniforms);

bool IsMatCx2(const TType *type);

ImmutableString MakeUnwrappingArrayConversionFunctionName(const TType *type);
ImmutableString MakeMatCx2ConversionFunctionName(const TType *type);

// Bools and bvecN are represented in WGSL's uniform address space with u32 and uvecN respectively.
// This outputs the beginning of a conversion necessary to convert the bool or bvecN represented by
// `type` into an actual WGSL bool or vecN<bool>. E.g. for regular bools this will output `bool(`.
// The caller is responsible for closing the parentheses, e.g. output <<
// "uniforms.bool_represented_by_u32)";
bool OutputUniformBoolOrBvecConversion(TInfoSinkBase &output, const TType &type);

// TODO(anglebug.com/42267100): for now does not output all uniform blocks,
// just the default block.
// `uniformBlock` is the variable declaring an interface block of default uniforms.
bool OutputUniformBlocksAndSamplers(TCompiler *compiler,
                                    TIntermBlock *root,
                                    const TVariable *defaultUniformBlock);

// GLSL sampler uniforms are extracted from structs. Given a GLSL sampler's associated name string,
// this function retrieves its new WGSL name and strips off array indices.
std::string WGSLGetMappedSamplerName(const std::string &originalName);

}  // namespace sh

#endif  // COMPILER_TRANSLATOR_WGSL_OUTPUT_UNIFORM_BLOCKS_H_