1/*
2 * Copyright (c) 2024 Huawei Device Co., Ltd.
3 * Licensed under the Apache License, Version 2.0 (the "License");
4 * you may not use this file except in compliance with the License.
5 * You may obtain a copy of the License at
6 *
7 * http://www.apache.org/licenses/LICENSE-2.0
8 *
9 * Unless required by applicable law or agreed to in writing, software
10 * distributed under the License is distributed on an "AS IS" BASIS,
11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 * See the License for the specific language governing permissions and
13 * limitations under the License.
14 */
15
16#include "capturedVariables.h"
17
18namespace ark::es2panda::compiler {
19
20static bool OnLeftSideOfAssignment(ir::Identifier const *const ident) noexcept
21{
22    return ident->Parent()->IsAssignmentExpression() && ident->Parent()->AsAssignmentExpression()->Left() == ident;
23}
24
25static void AddScopes(ir::AstNode *node, std::set<varbinder::Scope *> &scopes) noexcept
26{
27    if (node->Scope()->IsFunctionScope()) {
28        scopes.emplace(node->Scope()->AsFunctionScope()->ParamScope());
29    }
30    if (node->Scope()->IsCatchScope()) {
31        scopes.emplace(node->Scope()->AsCatchScope()->ParamScope());
32    }
33    if (node->Scope()->IsLoopScope()) {
34        scopes.emplace(node->Scope()->AsLoopScope()->DeclScope());
35    }
36    scopes.emplace(node->Scope());
37}
38
39static varbinder::Variable *FindVariable(ir::Identifier *ident, std::set<varbinder::Scope *> const &scopes) noexcept
40{
41    auto *var = ident->Variable();
42    //  NOTE! For some unknown reasons :) variables exist in scope collections but are not set to identifiers after
43    //        'varbinder->IdentifierAnalysis()' pass. Probably need to be investigated and fixed sometimes...
44    if (var == nullptr) {
45        //  We start from the innermost scope!
46        for (auto it = scopes.crbegin(); it != scopes.crend(); ++it) {
47            auto res = (*it)->Find(ident->Name(), varbinder::ResolveBindingOptions::VARIABLES);
48            if (res.variable != nullptr) {
49                var = res.variable;
50                break;
51            }
52        }
53    }
54
55    if (var != nullptr) {
56        auto *scope = var->GetScope();
57        ASSERT(scope != nullptr);
58        //  We are not interested in variables defined inside arrow function!
59        if (scopes.find(scope) != scopes.cend()) {
60            return nullptr;
61        }
62    }
63
64    return var;
65}
66
67static void FindModifiedCaptured(ir::ScriptFunction const *const scriptFunction,
68                                 std::set<varbinder::Variable *> &variables) noexcept
69{
70    auto scopes = std::set<varbinder::Scope *> {};
71    bool inLambda = false;
72
73    std::function<void(ir::AstNode *)> walker = [&](ir::AstNode *node) -> void {
74        if (node->IsArrowFunctionExpression() || node->IsClassDeclaration()) {
75            auto savedWL = inLambda;
76            auto savedScopes = std::set<varbinder::Scope *> {};
77            std::swap(scopes, savedScopes);
78
79            inLambda = true;
80            node->Iterate(walker);
81
82            inLambda = savedWL;
83            std::swap(scopes, savedScopes);
84            savedScopes.clear();
85
86            return;
87        }
88
89        if (inLambda && node->IsScopeBearer()) {
90            AddScopes(node, scopes);
91        } else if (inLambda && node->IsIdentifier() && OnLeftSideOfAssignment(node->AsIdentifier())) {
92            if (auto *var = FindVariable(node->AsIdentifier(), scopes); var != nullptr) {
93                variables.insert(var);
94            }
95        }
96
97        node->Iterate(walker);
98    };
99
100    scriptFunction->Iterate(walker);
101}
102
103static void HandleScriptFunction(ir::ScriptFunction const *const scriptFunction) noexcept
104{
105    auto variables = std::set<varbinder::Variable *> {};
106    FindModifiedCaptured(scriptFunction, variables);
107
108    for (auto *variable : variables) {
109        variable->AddFlag(varbinder::VariableFlags::CAPTURED_MODIFIED);
110    }
111
112    variables.clear();
113}
114
115bool CapturedVariables::Perform(public_lib::Context *ctx, parser::Program *program)
116{
117    if (ctx->config->options->CompilerOptions().compilationMode == CompilationMode::GEN_STD_LIB) {
118        for (auto &[_, ext_programs] : program->ExternalSources()) {
119            (void)_;
120            for (auto *extProg : ext_programs) {
121                Perform(ctx, extProg);
122            }
123        }
124    }
125
126    std::function<void(ir::AstNode *)> searchForFunctions = [&](ir::AstNode *ast) {
127        if (ast->IsScriptFunction()) {
128            HandleScriptFunction(ast->AsScriptFunction());  // no recursion
129        } else {
130            ast->Iterate(searchForFunctions);
131        }
132    };
133
134    program->Ast()->Iterate(searchForFunctions);
135    return true;
136}
137}  // namespace ark::es2panda::compiler
138