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 "boxingForLocals.h"
17
18#include "varbinder/ETSBinder.h"
19#include "checker/ETSchecker.h"
20
21namespace ark::es2panda::compiler {
22
23static constexpr std::string_view LOWERING_NAME = "boxing-for-locals";
24
25static ArenaSet<varbinder::Variable *> FindCaptured(public_lib::Context *ctx, ir::ScriptFunction *func)
26{
27    auto *allocator = ctx->allocator;
28    auto captured = ArenaSet<varbinder::Variable *>(allocator->Adapter());
29    bool withinLambda = false;
30
31    auto innermostArrowScopes = ArenaSet<varbinder::Scope *>(allocator->Adapter());
32    innermostArrowScopes.insert(func->Scope());
33    innermostArrowScopes.insert(func->Scope()->ParamScope());
34
35    std::function<void(ir::AstNode *)> walker = [&](ir::AstNode *ast) {
36        if (ast->IsArrowFunctionExpression() || ast->IsClassDeclaration()) {
37            auto savedWL = withinLambda;
38            auto savedScopes = ArenaSet<varbinder::Scope *>(allocator->Adapter());
39            std::swap(innermostArrowScopes, savedScopes);
40            withinLambda = true;
41
42            ast->Iterate(walker);
43
44            withinLambda = savedWL;
45            std::swap(innermostArrowScopes, savedScopes);
46
47            return;
48        }
49
50        if (withinLambda && ast->IsScopeBearer()) {
51            innermostArrowScopes.insert(ast->Scope());
52            if (ast->Scope()->IsFunctionScope()) {
53                innermostArrowScopes.insert(ast->Scope()->AsFunctionScope()->ParamScope());
54            }
55            if (ast->Scope()->IsCatchScope()) {
56                innermostArrowScopes.insert(ast->Scope()->AsCatchScope()->ParamScope());
57            }
58            if (ast->Scope()->IsLoopScope()) {
59                innermostArrowScopes.insert(ast->Scope()->AsLoopScope()->DeclScope());
60            }
61        } else if (withinLambda && ast->IsIdentifier()) {
62            auto *var = ast->AsIdentifier()->Variable();
63            if (var == nullptr) {
64                return;
65            }
66            auto *scope = var->GetScope();
67            if (scope != nullptr && !scope->IsClassScope() && !scope->IsGlobalScope() &&
68                innermostArrowScopes.count(scope) == 0) {
69                captured.insert(var);
70            }
71        }
72
73        ast->Iterate(walker);
74    };
75
76    func->Iterate(walker);
77
78    auto varsToBox = ArenaSet<varbinder::Variable *>(allocator->Adapter());
79    return captured;
80}
81
82static ArenaSet<varbinder::Variable *> FindModified(public_lib::Context *ctx, ir::ScriptFunction *func)
83{
84    auto *allocator = ctx->allocator;
85    auto modified = ArenaSet<varbinder::Variable *>(allocator->Adapter());
86
87    std::function<void(ir::AstNode *)> walker = [&](ir::AstNode *ast) {
88        if (ast->IsAssignmentExpression()) {
89            auto assignment = ast->AsAssignmentExpression();
90            if (assignment->Left()->IsIdentifier()) {
91                ASSERT(assignment->Left()->Variable() != nullptr);
92                auto *var = assignment->Left()->Variable();
93                var->AddFlag(varbinder::VariableFlags::INITIALIZED);
94                modified.insert(var);
95            }
96        }
97    };
98
99    func->IterateRecursively(walker);
100    return modified;
101}
102
103static ArenaSet<varbinder::Variable *> FindVariablesToBox(public_lib::Context *ctx, ir::ScriptFunction *func)
104{
105    auto *allocator = ctx->allocator;
106    auto captured = FindCaptured(ctx, func);
107    auto modified = FindModified(ctx, func);
108
109    auto varsToBox = ArenaSet<varbinder::Variable *>(allocator->Adapter());
110    std::set_intersection(captured.cbegin(), captured.cend(), modified.cbegin(), modified.cend(),
111                          std::inserter(varsToBox, varsToBox.begin()));
112
113    return varsToBox;
114}
115
116static void HandleFunctionParam(public_lib::Context *ctx, ir::ETSParameterExpression *param,
117                                ArenaMap<varbinder::Variable *, varbinder::Variable *> *varsMap)
118{
119    auto *allocator = ctx->allocator;
120    auto *checker = ctx->checker->AsETSChecker();
121    auto *varBinder = checker->VarBinder();
122
123    auto *id = param->Ident()->AsIdentifier();
124    auto *oldVar = id->Variable();
125    auto *oldType = oldVar->TsType();
126    auto *func = param->Parent()->AsScriptFunction();
127    ASSERT(func->Body()->IsBlockStatement());  // guaranteed after expressionLambdaLowering
128    auto *body = func->Body()->AsBlockStatement();
129    auto &bodyStmts = body->Statements();
130    auto *scope = body->Scope();
131
132    auto *initId = allocator->New<ir::Identifier>(id->Name(), allocator);
133    initId->SetVariable(id->Variable());
134    initId->SetTsType(oldType);
135
136    // The new variable will have the same name as the parameter. This is not representable in source code.
137    auto *boxedType = checker->GlobalBuiltinBoxType(oldType);
138    ArenaVector<ir::Expression *> newInitArgs {allocator->Adapter()};
139    newInitArgs.push_back(initId);
140    auto *newInit = util::NodeAllocator::ForceSetParent<ir::ETSNewClassInstanceExpression>(
141        allocator, allocator->New<ir::OpaqueTypeNode>(boxedType), std::move(newInitArgs), nullptr);
142    auto *newDeclarator = util::NodeAllocator::ForceSetParent<ir::VariableDeclarator>(
143        allocator, ir::VariableDeclaratorFlag::CONST, allocator->New<ir::Identifier>(id->Name(), allocator), newInit);
144    ArenaVector<ir::VariableDeclarator *> declVec {allocator->Adapter()};
145    declVec.push_back(newDeclarator);
146
147    auto *newDecl = allocator->New<varbinder::ConstDecl>(id->Name(), newDeclarator);
148    auto *newVar = allocator->New<varbinder::LocalVariable>(newDecl, oldVar->Flags());
149    newVar->SetTsType(boxedType);
150
151    newDeclarator->Id()->AsIdentifier()->SetVariable(newVar);
152    newVar->AddFlag(varbinder::VariableFlags::INITIALIZED);
153    newVar->SetScope(scope);
154    scope->EraseBinding(newVar->Name());
155    scope->InsertBinding(newVar->Name(), newVar);
156
157    auto *newDeclaration = util::NodeAllocator::ForceSetParent<ir::VariableDeclaration>(
158        allocator, ir::VariableDeclaration::VariableDeclarationKind::CONST, allocator, std::move(declVec), false);
159    newDeclaration->SetParent(body);
160    bodyStmts.insert(bodyStmts.begin(), newDeclaration);
161
162    auto lexScope = varbinder::LexicalScope<varbinder::Scope>::Enter(varBinder, scope);
163    auto savedContext = checker::SavedCheckerContext(checker, checker::CheckerStatus::NO_OPTS);
164    auto scopeContext = checker::ScopeContext(checker, scope);
165
166    newDeclaration->Check(checker);
167
168    varsMap->emplace(oldVar, newVar);
169}
170
171static ir::AstNode *HandleVariableDeclarator(public_lib::Context *ctx, ir::VariableDeclarator *declarator,
172                                             ArenaMap<varbinder::Variable *, varbinder::Variable *> *varsMap)
173{
174    auto *allocator = ctx->allocator;
175    auto *checker = ctx->checker->AsETSChecker();
176    auto *varBinder = checker->VarBinder();
177
178    auto *id = declarator->Id()->AsIdentifier();
179    auto *oldVar = id->Variable();
180    auto *scope = oldVar->GetScope();
181    auto *type = oldVar->TsType();
182    auto *boxedType = checker->GlobalBuiltinBoxType(type);
183
184    auto initArgs = ArenaVector<ir::Expression *>(allocator->Adapter());
185    if (declarator->Init() != nullptr) {
186        auto *arg = declarator->Init();
187        if (arg->TsType() != type) {
188            arg = util::NodeAllocator::ForceSetParent<ir::TSAsExpression>(
189                allocator, arg, allocator->New<ir::OpaqueTypeNode>(type), false);
190        }
191        initArgs.push_back(arg);
192    }
193    auto *newInit = util::NodeAllocator::ForceSetParent<ir::ETSNewClassInstanceExpression>(
194        allocator, allocator->New<ir::OpaqueTypeNode>(boxedType), std::move(initArgs), nullptr);
195    auto *newDeclarator = util::NodeAllocator::ForceSetParent<ir::VariableDeclarator>(
196        allocator, declarator->Flag(), allocator->New<ir::Identifier>(id->Name(), allocator), newInit);
197    newDeclarator->SetParent(declarator->Parent());
198
199    auto *newDecl = allocator->New<varbinder::ConstDecl>(oldVar->Name(), newDeclarator);
200    auto *newVar = allocator->New<varbinder::LocalVariable>(newDecl, oldVar->Flags());
201    newDeclarator->Id()->AsIdentifier()->SetVariable(newVar);
202    newVar->AddFlag(varbinder::VariableFlags::INITIALIZED);
203    newVar->SetScope(scope);
204
205    scope->EraseBinding(oldVar->Name());
206    scope->InsertBinding(newVar->Name(), newVar);
207
208    auto lexScope = varbinder::LexicalScope<varbinder::Scope>::Enter(varBinder, scope);
209    auto savedContext = checker::SavedCheckerContext(checker, checker::CheckerStatus::NO_OPTS);
210    auto scopeContext = checker::ScopeContext(checker, scope);
211
212    newDeclarator->Check(checker);
213
214    varsMap->emplace(oldVar, newVar);
215
216    return newDeclarator;
217}
218
219static bool IsBeingDeclared(ir::AstNode *ast)
220{
221    ASSERT(ast->IsIdentifier());
222    return (ast->Parent()->IsVariableDeclarator() && ast == ast->Parent()->AsVariableDeclarator()->Id()) ||
223           (ast->Parent()->IsETSParameterExpression() && ast == ast->Parent()->AsETSParameterExpression()->Ident());
224}
225
226static bool IsPartOfBoxInitializer(public_lib::Context *ctx, ir::AstNode *ast)
227{
228    ASSERT(ast->IsIdentifier());
229    auto *checker = ctx->checker->AsETSChecker();
230    auto *id = ast->AsIdentifier();
231
232    // NOTE(gogabr): rely on caching for type instantiations, so we can use pointer comparison.
233    return id->Parent()->IsETSNewClassInstanceExpression() &&
234           id->Parent()->AsETSNewClassInstanceExpression()->GetTypeRef()->TsType() ==
235               checker->GlobalBuiltinBoxType(id->TsType());
236}
237
238static bool OnLeftSideOfAssignment(ir::AstNode *ast)
239{
240    return ast->Parent()->IsAssignmentExpression() && ast->Parent()->AsAssignmentExpression()->Left() == ast;
241}
242
243static ir::AstNode *HandleReference(public_lib::Context *ctx, ir::Identifier *id, varbinder::Variable *var)
244{
245    auto *parser = ctx->parser->AsETSParser();
246    auto *checker = ctx->checker->AsETSChecker();
247
248    // `as` is needed to account for smart types
249    auto *res = parser->CreateFormattedExpression("@@I1.get() as @@T2", var->Name(), id->TsType());
250    res->SetParent(id->Parent());
251    res->AsTSAsExpression()
252        ->Expr()
253        ->AsCallExpression()
254        ->Callee()
255        ->AsMemberExpression()
256        ->Object()
257        ->AsIdentifier()
258        ->SetVariable(var);
259
260    // NOTE(gogabr) -- The `get` property remains without variable; this is OK for the current checker, but may need
261    // adjustment later.
262    res->Check(checker);
263
264    ASSERT(res->TsType() == id->TsType());
265    res->SetBoxingUnboxingFlags(id->GetBoxingUnboxingFlags());
266
267    return res;
268}
269
270static ir::AstNode *HandleAssignment(public_lib::Context *ctx, ir::AssignmentExpression *ass,
271                                     ArenaMap<varbinder::Variable *, varbinder::Variable *> const &varsMap)
272{
273    // Should be true after opAssignment lowering
274    ASSERT(ass->OperatorType() == lexer::TokenType::PUNCTUATOR_SUBSTITUTION);
275
276    auto *parser = ctx->parser->AsETSParser();
277    auto *varBinder = ctx->checker->VarBinder()->AsETSBinder();
278    auto *checker = ctx->checker->AsETSChecker();
279
280    auto *oldVar = ass->Left()->AsIdentifier()->Variable();
281    auto *newVar = varsMap.find(oldVar)->second;
282    auto *scope = newVar->GetScope();
283    newVar->AddFlag(varbinder::VariableFlags::INITIALIZED);
284
285    auto *res = parser->CreateFormattedExpression("@@I1.set(@@E2 as @@T3) as @@T4", newVar->Name(), ass->Right(),
286                                                  oldVar->TsType(), ass->TsType());
287    res->SetParent(ass->Parent());
288
289    // NOTE(gogabr) -- The `get` and `set` properties remain without variable; this is OK for the current checker, but
290    // may need adjustment later.
291    auto lexScope = varbinder::LexicalScope<varbinder::Scope>::Enter(varBinder, scope);
292    auto savedContext = checker::SavedCheckerContext(checker, checker::CheckerStatus::NO_OPTS);
293    auto scopeContext = checker::ScopeContext(checker, scope);
294
295    varBinder->ResolveReferencesForScopeWithContext(res, scope);
296    res->Check(checker);
297
298    ASSERT(res->TsType() == ass->TsType());
299    res->SetBoxingUnboxingFlags(ass->GetBoxingUnboxingFlags());
300
301    return res;
302}
303
304static void HandleScriptFunction(public_lib::Context *ctx, ir::ScriptFunction *func)
305{
306    auto *allocator = ctx->allocator;
307    auto varsToBox = FindVariablesToBox(ctx, func);
308    if (varsToBox.empty()) {
309        return;
310    }
311    auto varsMap = ArenaMap<varbinder::Variable *, varbinder::Variable *>(allocator->Adapter());
312
313    /*
314      The function relies on the following facts:
315      - TransformChildrenRecursively handles children in order
316      - local variables are never used before declaration.
317      This ensures that varsToMap has the appropriate record by the time the variable reference is processed.
318    */
319    auto handleNode = [ctx, &varsToBox, &varsMap](ir::AstNode *ast) {
320        if (ast->IsETSParameterExpression() && varsToBox.count(ast->AsETSParameterExpression()->Variable()) > 0) {
321            HandleFunctionParam(ctx, ast->AsETSParameterExpression(), &varsMap);
322            return ast;  // modifications occur in the containing function
323        }
324        if (ast->IsVariableDeclarator() && ast->AsVariableDeclarator()->Id()->IsIdentifier() &&
325            varsToBox.count(ast->AsVariableDeclarator()->Id()->AsIdentifier()->Variable()) > 0) {
326            return HandleVariableDeclarator(ctx, ast->AsVariableDeclarator(), &varsMap);
327        }
328        if (ast->IsAssignmentExpression() && ast->AsAssignmentExpression()->Left()->IsIdentifier() &&
329            varsToBox.count(ast->AsAssignmentExpression()->Left()->AsIdentifier()->Variable()) > 0) {
330            return HandleAssignment(ctx, ast->AsAssignmentExpression(), varsMap);
331        }
332        if (ast->IsIdentifier() && !IsBeingDeclared(ast) && !IsPartOfBoxInitializer(ctx, ast) &&
333            !OnLeftSideOfAssignment(ast) && varsToBox.count(ast->AsIdentifier()->Variable()) > 0) {
334            return HandleReference(ctx, ast->AsIdentifier(), varsMap.find(ast->AsIdentifier()->Variable())->second);
335        }
336
337        return ast;
338    };
339
340    func->TransformChildrenRecursivelyPostorder(handleNode, LOWERING_NAME);
341}
342
343bool BoxingForLocals::Perform(public_lib::Context *ctx, parser::Program *program)
344{
345    parser::SavedFormattingFileName savedFormattingName(ctx->parser->AsETSParser(), "boxing-for-lambdas");
346
347    if (ctx->config->options->CompilerOptions().compilationMode == CompilationMode::GEN_STD_LIB) {
348        for (auto &[_, ext_programs] : program->ExternalSources()) {
349            (void)_;
350            for (auto *extProg : ext_programs) {
351                Perform(ctx, extProg);
352            }
353        }
354    }
355
356    std::function<void(ir::AstNode *)> searchForFunctions = [&](ir::AstNode *ast) {
357        if (ast->IsScriptFunction()) {
358            HandleScriptFunction(ctx, ast->AsScriptFunction());  // no recursion
359        } else {
360            ast->Iterate(searchForFunctions);
361        }
362    };
363    program->Ast()->Iterate(searchForFunctions);
364    return true;
365}
366
367bool BoxingForLocals::Postcondition([[maybe_unused]] public_lib::Context *ctx, parser::Program const *program)
368{
369    for (auto &[_, ext_programs] : program->ExternalSources()) {
370        (void)_;
371        for (auto *extProg : ext_programs) {
372            if (!Postcondition(ctx, extProg)) {
373                return false;
374            }
375        }
376    }
377
378    return !program->Ast()->IsAnyChild([](const ir::AstNode *node) {
379        if (node->IsAssignmentExpression() && node->AsAssignmentExpression()->Left()->IsIdentifier()) {
380            auto asExpr = node->AsAssignmentExpression();
381            auto var = asExpr->Left()->AsIdentifier()->Variable();
382            if (var != nullptr && var->IsLocalVariable() && !var->HasFlag(varbinder::VariableFlags::INITIALIZED)) {
383                return true;
384            }
385        }
386        return false;
387    });
388}
389
390}  // namespace ark::es2panda::compiler
391