/*
 * Copyright 2023 Google LLC
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#include "include/private/base/SkAssert.h"
#include "include/private/base/SkTArray.h"
#include "src/sksl/SkSLAnalysis.h"
#include "src/sksl/SkSLDefines.h"
#include "src/sksl/SkSLPosition.h"
#include "src/sksl/ir/SkSLBlock.h"
#include "src/sksl/ir/SkSLExpression.h"
#include "src/sksl/ir/SkSLIRHelpers.h"
#include "src/sksl/ir/SkSLIRNode.h"
#include "src/sksl/ir/SkSLModifierFlags.h"
#include "src/sksl/ir/SkSLNop.h"
#include "src/sksl/ir/SkSLStatement.h"
#include "src/sksl/ir/SkSLSwitchStatement.h"
#include "src/sksl/ir/SkSLSymbolTable.h"
#include "src/sksl/ir/SkSLVarDeclarations.h"
#include "src/sksl/ir/SkSLVariable.h"
#include "src/sksl/transform/SkSLProgramWriter.h"
#include "src/sksl/transform/SkSLTransform.h"

#include <memory>
#include <utility>

using namespace skia_private;

namespace SkSL {

class Context;

std::unique_ptr<Statement> Transform::HoistSwitchVarDeclarationsAtTopLevel(
        const Context& context,
        std::unique_ptr<SwitchStatement> stmt) {
    struct HoistSwitchVarDeclsVisitor : public ProgramWriter {
        HoistSwitchVarDeclsVisitor(const Context& c) : fContext(c) {}

        bool visitExpressionPtr(std::unique_ptr<Expression>& expr) override {
            // We don't need to recurse into expressions.
            return false;
        }

        bool visitStatementPtr(std::unique_ptr<Statement>& stmt) override {
            switch (stmt->kind()) {
                case StatementKind::kSwitchCase:
                    // Recurse inward from the switch and its inner switch-cases.
                    return INHERITED::visitStatementPtr(stmt);

                case StatementKind::kBlock:
                    if (!stmt->as<Block>().isScope()) {
                        // Recurse inward from unscoped blocks.
                        return INHERITED::visitStatementPtr(stmt);
                    }
                    break;

                case StatementKind::kVarDeclaration:
                    // Keep track of variable declarations.
                    fVarDeclarations.push_back(&stmt);
                    break;

                default:
                    break;
            }

            // We don't need to recurse into other statement types; we're only interested in the top
            // level of the switch statement.
            return false;
        }

        const Context& fContext;
        TArray<std::unique_ptr<Statement>*> fVarDeclarations;

        using INHERITED = ProgramWriter;
    };

    // Visit every switch-case in the switch, looking for hoistable var-declarations.
    HoistSwitchVarDeclsVisitor visitor(context);
    for (std::unique_ptr<Statement>& sc : stmt->as<SwitchStatement>().cases()) {
        visitor.visitStatementPtr(sc);
    }

    // If no declarations were found, return the switch as-is.
    if (visitor.fVarDeclarations.empty()) {
        return stmt;
    }

    // Move all of the var-declaration statements into a separate block.
    SymbolTable* switchSymbols = stmt->caseBlock()->as<Block>().symbolTable();
    std::unique_ptr<SymbolTable> blockSymbols = switchSymbols->insertNewParent();

    StatementArray blockStmts;
    blockStmts.reserve_exact(visitor.fVarDeclarations.size() + 1);
    for (std::unique_ptr<Statement>* innerDeclaration : visitor.fVarDeclarations) {
        VarDeclaration& decl = (*innerDeclaration)->as<VarDeclaration>();
        Variable* var = decl.var();
        bool isConst = var->modifierFlags().isConst();

        std::unique_ptr<Statement> replacementStmt;
        if (decl.value() && !isConst) {
            // The inner variable-declaration has an initial-value; we must replace the declaration
            // with an assignment to the variable. This also has the helpful effect of stripping off
            // the initial-value from the declaration.
            struct AssignmentHelper : public IRHelpers {
                using IRHelpers::IRHelpers;

                std::unique_ptr<Statement> makeAssignmentStmt(VarDeclaration& decl) const {
                    return Assign(Ref(decl.var()), std::move(decl.value()));
                }
            };

            AssignmentHelper helper(context);
            replacementStmt = helper.makeAssignmentStmt(decl);
        } else {
            // The inner variable-declaration has no initial-value, or it's const and has a constant
            // value; we can move it upwards as-is and replace its statement with a no-op.
            SkASSERT(!isConst || Analysis::IsConstantExpression(*decl.value()));

            replacementStmt = Nop::Make();
        }

        // Move the var-declaration above the switch, and replace the existing statement with either
        // an assignment (if there was an initial-value) or a no-op (if there wasn't one).
        blockStmts.push_back(std::move(*innerDeclaration));
        *innerDeclaration = std::move(replacementStmt);

        // Hoist the variable's symbol outside of the switch's symbol table, and into the enclosing
        // block's symbol table.
        switchSymbols->moveSymbolTo(blockSymbols.get(), var, context);
    }

    // Return a scoped Block holding the switch.
    Position pos = stmt->fPosition;
    blockStmts.push_back(std::move(stmt));
    return Block::MakeBlock(pos, std::move(blockStmts), Block::Kind::kBracedScope,
                            std::move(blockSymbols));
}

}  // namespace SkSL