/*

 * Copyright 2024 Google LLC

 *

 * Use of this source code is governed by a BSD-style license that can be

 * found in the LICENSE file.

 */

#include "src/sksl/SkSLModule.h"

#include "src/sksl/ir/SkSLConstructorSplat.h"

#include "src/sksl/ir/SkSLExpression.h"

#include "src/sksl/ir/SkSLFunctionDefinition.h"

#include "src/sksl/ir/SkSLLiteral.h"

#include "src/sksl/ir/SkSLProgramElement.h"

#include "src/sksl/ir/SkSLSwizzle.h"

#include "src/sksl/ir/SkSLType.h"

#include "src/sksl/transform/SkSLProgramWriter.h"

#include "src/sksl/transform/SkSLTransform.h"

#include <memory>

#include <utility>

#include <vector>

namespace SkSL {

class Context;

void Transform::ReplaceSplatCastsWithSwizzles(const Context& context, Module& module) {

    class SwizzleWriter : public ProgramWriter {

    public:

        SwizzleWriter(const Context& ctx) : fContext(ctx) {}

        bool visitExpressionPtr(std::unique_ptr<Expression>& expr) override {

            if (INHERITED::visitExpressionPtr(expr)) {

                return true;

            }

            if (expr->is<ConstructorSplat>()) {

                // If the argument is a literal, only allow floats. The swizzled-literal syntax only

                // works properly for floats.

                std::unique_ptr<Expression>& arg = expr->as<ConstructorSplat>().argument();

                if (!arg->is<Literal>() || (arg->type().isFloat() && arg->type().highPrecision())) {

                    // Synthesize a splat like `.xxxx`, matching the column count of the splat.

                    ComponentArray components;

                    components.push_back_n(expr->type().columns(), SwizzleComponent::X);

                    // Replace the splat expression with the swizzle.

                    expr = Swizzle::MakeExact(fContext, expr->position(), std::move(arg),

                                              std::move(components));

                }

            }

            return false;

        }

        const Context& fContext;

        using INHERITED = ProgramWriter;

    };

    for (std::unique_ptr<ProgramElement>& pe : module.fElements) {

        if (pe->is<FunctionDefinition>()) {

            SwizzleWriter writer{context};

            writer.visitStatementPtr(pe->as<FunctionDefinition>().body());

        }

    }

}

}  // namespace SkSL