// Copyright 2022 The Tint Authors.
//
// 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 "src/tint/transform/unwind_discard_functions.h"

#include <memory>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>

#include "src/tint/ast/discard_statement.h"
#include "src/tint/ast/return_statement.h"
#include "src/tint/ast/traverse_expressions.h"
#include "src/tint/sem/block_statement.h"
#include "src/tint/sem/call.h"
#include "src/tint/sem/for_loop_statement.h"
#include "src/tint/sem/function.h"
#include "src/tint/sem/if_statement.h"
#include "src/tint/transform/utils/get_insertion_point.h"

TINT_INSTANTIATE_TYPEINFO(tint::transform::UnwindDiscardFunctions);

namespace tint::transform {
namespace {

bool ShouldRun(const Program* program) {
    auto& sem = program->Sem();
    for (auto* f : program->AST().Functions()) {
        if (sem.Get(f)->Behaviors().Contains(sem::Behavior::kDiscard)) {
            return true;
        }
    }
    return false;
}

}  // namespace

/// PIMPL state for the transform
struct UnwindDiscardFunctions::State {
    /// Constructor
    /// @param ctx_in the context
    explicit State(CloneContext& ctx_in) : ctx(ctx_in), b(*ctx_in.dst), sem(ctx_in.src->Sem()) {}

    /// Runs the transform
    void Run() {
        ctx.ReplaceAll([&](const ast::BlockStatement* block) -> const ast::Statement* {
            // Iterate block statements and replace them as needed.
            for (auto* stmt : block->statements) {
                if (auto* new_stmt = Statement(stmt)) {
                    ctx.Replace(stmt, new_stmt);
                }

                // Handle for loops, as they are the only other AST node that
                // contains statements outside of BlockStatements.
                if (auto* fl = stmt->As<ast::ForLoopStatement>()) {
                    if (auto* new_stmt = Statement(fl->initializer)) {
                        ctx.Replace(fl->initializer, new_stmt);
                    }
                    if (auto* new_stmt = Statement(fl->continuing)) {
                        // NOTE: Should never reach here as we cannot discard in a
                        // continuing block.
                        ctx.Replace(fl->continuing, new_stmt);
                    }
                }
            }

            return nullptr;
        });
    }

  private:
    CloneContext& ctx;
    ProgramBuilder& b;
    const sem::Info& sem;
    Symbol module_discard_var_name;   // Use ModuleDiscardVarName() to read
    Symbol module_discard_func_name;  // Use ModuleDiscardFuncName() to read

    // Returns true if `sem_expr` contains a call expression that may
    // (transitively) execute a discard statement.
    bool MayDiscard(const sem::Expression* sem_expr) {
        return sem_expr && sem_expr->Behaviors().Contains(sem::Behavior::kDiscard);
    }

    // Lazily creates and returns the name of the module bool variable for whether
    // to discard: "tint_discard".
    Symbol ModuleDiscardVarName() {
        if (!module_discard_var_name.IsValid()) {
            module_discard_var_name = b.Symbols().New("tint_discard");
            ctx.dst->GlobalVar(module_discard_var_name, b.ty.bool_(), b.Expr(false),
                               ast::AddressSpace::kPrivate);
        }
        return module_discard_var_name;
    }

    // Lazily creates and returns the name of the function that contains a single
    // discard statement: "tint_discard_func".
    // We do this to avoid having multiple discard statements in a single program,
    // which causes problems in certain backends (see crbug.com/1118).
    Symbol ModuleDiscardFuncName() {
        if (!module_discard_func_name.IsValid()) {
            module_discard_func_name = b.Symbols().New("tint_discard_func");
            b.Func(module_discard_func_name, tint::utils::Empty, b.ty.void_(),
                   tint::utils::Vector{b.Discard()});
        }
        return module_discard_func_name;
    }

    // Creates "return <default return value>;" based on the return type of
    // `stmt`'s owning function.
    const ast::ReturnStatement* Return(const ast::Statement* stmt) {
        const ast::Expression* ret_val = nullptr;
        auto* ret_type = sem.Get(stmt)->Function()->Declaration()->return_type;
        if (!ret_type->Is<ast::Void>()) {
            ret_val = b.Construct(ctx.Clone(ret_type));
        }
        return b.Return(ret_val);
    }

    // Returns true if the function `stmt` is in is an entry point
    bool IsInEntryPointFunc(const ast::Statement* stmt) {
        return sem.Get(stmt)->Function()->Declaration()->IsEntryPoint();
    }

    // Creates "tint_discard_func();"
    const ast::CallStatement* CallDiscardFunc() {
        auto func_name = ModuleDiscardFuncName();
        return b.CallStmt(b.Call(func_name));
    }

    // Creates and returns a new if-statement of the form:
    //
    //    if (tint_discard) {
    //      return <default value>;
    //    }
    //
    // or if `stmt` is in a entry point function:
    //
    //    if (tint_discard) {
    //      tint_discard_func();
    //      return <default value>;
    //    }
    //
    const ast::IfStatement* IfDiscardReturn(const ast::Statement* stmt) {
        tint::utils::Vector<const ast::Statement*, 2> stmts;

        // For entry point functions, also emit the discard statement
        if (IsInEntryPointFunc(stmt)) {
            stmts.Push(CallDiscardFunc());
        }

        stmts.Push(Return(stmt));

        auto var_name = ModuleDiscardVarName();
        return b.If(var_name, b.Block(stmts));
    }

    // Hoists `sem_expr` to a let followed by an `IfDiscardReturn` before `stmt`.
    // For example, if `stmt` is:
    //
    //    return f();
    //
    // This function will transform this to:
    //
    //    let t1 = f();
    //    if (tint_discard) {
    //      return;
    //    }
    //    return t1;
    //
    const ast::Statement* HoistAndInsertBefore(const ast::Statement* stmt,
                                               const sem::Expression* sem_expr) {
        auto* expr = sem_expr->Declaration();

        auto ip = utils::GetInsertionPoint(ctx, stmt);
        auto var_name = b.Sym();
        auto* decl = b.Decl(b.Var(var_name, ctx.Clone(expr)));
        ctx.InsertBefore(ip.first->Declaration()->statements, ip.second, decl);

        ctx.InsertBefore(ip.first->Declaration()->statements, ip.second, IfDiscardReturn(stmt));

        auto* var_expr = b.Expr(var_name);

        // Special handling for CallStatement as we can only replace its expression
        // with a CallExpression.
        if (stmt->Is<ast::CallStatement>()) {
            // We could replace the call statement with no statement, but we can't do
            // that with transforms (yet), so just return a phony assignment.
            return b.Assign(b.Phony(), var_expr);
        }

        ctx.Replace(expr, var_expr);
        return ctx.CloneWithoutTransform(stmt);
    }

    // Returns true if `stmt` is a for-loop initializer statement.
    bool IsForLoopInitStatement(const ast::Statement* stmt) {
        if (auto* sem_stmt = sem.Get(stmt)) {
            if (auto* sem_fl = tint::As<sem::ForLoopStatement>(sem_stmt->Parent())) {
                return sem_fl->Declaration()->initializer == stmt;
            }
        }
        return false;
    }

    // Inserts an `IfDiscardReturn` after `stmt` if possible (i.e. `stmt` is not
    // in a for-loop init), otherwise falls back to HoistAndInsertBefore, hoisting
    // `sem_expr` to a let followed by an `IfDiscardReturn` before `stmt`.
    //
    // For example, if `stmt` is:
    //
    //    let r = f();
    //
    // This function will transform this to:
    //
    //    let r = f();
    //    if (tint_discard) {
    //      return;
    //    }
    const ast::Statement* TryInsertAfter(const ast::Statement* stmt,
                                         const sem::Expression* sem_expr) {
        // If `stmt` is the init of a for-loop, hoist and insert before instead.
        if (IsForLoopInitStatement(stmt)) {
            return HoistAndInsertBefore(stmt, sem_expr);
        }

        auto ip = utils::GetInsertionPoint(ctx, stmt);
        ctx.InsertAfter(ip.first->Declaration()->statements, ip.second, IfDiscardReturn(stmt));
        return nullptr;  // Don't replace current statement
    }

    // Replaces the input discard statement with either setting the module level
    // discard bool ("tint_discard = true"), or calling the discard function
    // ("tint_discard_func()"), followed by a default return statement.
    //
    // Replaces "discard;" with:
    //
    //    tint_discard = true;
    //    return;
    //
    // Or if `stmt` is a entry point function, replaces with:
    //
    //    tint_discard_func();
    //    return;
    //
    const ast::Statement* ReplaceDiscardStatement(const ast::DiscardStatement* stmt) {
        const ast::Statement* to_insert = nullptr;
        if (IsInEntryPointFunc(stmt)) {
            to_insert = CallDiscardFunc();
        } else {
            auto var_name = ModuleDiscardVarName();
            to_insert = b.Assign(var_name, true);
        }

        auto ip = utils::GetInsertionPoint(ctx, stmt);
        ctx.InsertBefore(ip.first->Declaration()->statements, ip.second, to_insert);
        return Return(stmt);
    }

    // Handle statement
    const ast::Statement* Statement(const ast::Statement* stmt) {
        return Switch(
            stmt,
            [&](const ast::DiscardStatement* s) -> const ast::Statement* {
                return ReplaceDiscardStatement(s);
            },
            [&](const ast::AssignmentStatement* s) -> const ast::Statement* {
                auto* sem_lhs = sem.Get(s->lhs);
                auto* sem_rhs = sem.Get(s->rhs);
                if (MayDiscard(sem_lhs)) {
                    if (MayDiscard(sem_rhs)) {
                        TINT_ICE(Transform, b.Diagnostics())
                            << "Unexpected: both sides of assignment statement may "
                               "discard. Make sure transform::PromoteSideEffectsToDecl "
                               "was run first.";
                    }
                    return TryInsertAfter(s, sem_lhs);
                } else if (MayDiscard(sem_rhs)) {
                    return TryInsertAfter(s, sem_rhs);
                }
                return nullptr;
            },
            [&](const ast::CallStatement* s) -> const ast::Statement* {
                auto* sem_expr = sem.Get(s->expr);
                if (!MayDiscard(sem_expr)) {
                    return nullptr;
                }
                return TryInsertAfter(s, sem_expr);
            },
            [&](const ast::ForLoopStatement* s) -> const ast::Statement* {
                if (MayDiscard(sem.Get(s->condition))) {
                    TINT_ICE(Transform, b.Diagnostics())
                        << "Unexpected ForLoopStatement condition that may discard. "
                           "Make sure transform::PromoteSideEffectsToDecl was run "
                           "first.";
                }
                return nullptr;
            },
            [&](const ast::WhileStatement* s) -> const ast::Statement* {
                if (MayDiscard(sem.Get(s->condition))) {
                    TINT_ICE(Transform, b.Diagnostics())
                        << "Unexpected WhileStatement condition that may discard. "
                           "Make sure transform::PromoteSideEffectsToDecl was run "
                           "first.";
                }
                return nullptr;
            },
            [&](const ast::IfStatement* s) -> const ast::Statement* {
                auto* sem_expr = sem.Get(s->condition);
                if (!MayDiscard(sem_expr)) {
                    return nullptr;
                }
                return HoistAndInsertBefore(s, sem_expr);
            },
            [&](const ast::ReturnStatement* s) -> const ast::Statement* {
                auto* sem_expr = sem.Get(s->value);
                if (!MayDiscard(sem_expr)) {
                    return nullptr;
                }
                return HoistAndInsertBefore(s, sem_expr);
            },
            [&](const ast::SwitchStatement* s) -> const ast::Statement* {
                auto* sem_expr = sem.Get(s->condition);
                if (!MayDiscard(sem_expr)) {
                    return nullptr;
                }
                return HoistAndInsertBefore(s, sem_expr);
            },
            [&](const ast::VariableDeclStatement* s) -> const ast::Statement* {
                auto* var = s->variable;
                if (!var->initializer) {
                    return nullptr;
                }
                auto* sem_expr = sem.Get(var->initializer);
                if (!MayDiscard(sem_expr)) {
                    return nullptr;
                }
                return TryInsertAfter(s, sem_expr);
            });
    }
};

UnwindDiscardFunctions::UnwindDiscardFunctions() = default;
UnwindDiscardFunctions::~UnwindDiscardFunctions() = default;

Transform::ApplyResult UnwindDiscardFunctions::Apply(const Program* src,
                                                     const DataMap&,
                                                     DataMap&) const {
    if (!ShouldRun(src)) {
        return SkipTransform;
    }

    ProgramBuilder b;
    CloneContext ctx{&b, src, /* auto_clone_symbols */ true};

    State state(ctx);
    state.Run();

    ctx.Clone();
    return Program(std::move(b));
}

}  // namespace tint::transform