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 
18 namespace ark::es2panda::compiler {
19 
20 static bool OnLeftSideOfAssignment(ir::Identifier const *const ident) noexcept
21 {
22     return ident->Parent()->IsAssignmentExpression() && ident->Parent()->AsAssignmentExpression()->Left() == ident;
23 }
24 
25 static 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 
39 static 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 
67 static 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 
103 static 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 
Perform(public_lib::Context *ctx, parser::Program *program)115 bool 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