1/*
2 * Copyright (c) 2023-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 "localClassLowering.h"
17#include "checker/ETSchecker.h"
18#include "varbinder/ETSBinder.h"
19#include "../util.h"
20
21namespace ark::es2panda::compiler {
22
23std::string_view LocalClassConstructionPhase::Name() const
24{
25    return "LocalClassConstructionPhase";
26}
27
28static ir::ClassProperty *CreateCapturedField(checker::ETSChecker *checker, const varbinder::Variable *capturedVar,
29                                              varbinder::ClassScope *scope, size_t &idx,
30                                              const lexer::SourcePosition &pos)
31{
32    auto *allocator = checker->Allocator();
33    auto *varBinder = checker->VarBinder();
34
35    // Enter the lambda class instance field scope, every property will be bound to the lambda instance itself
36    auto fieldCtx = varbinder::LexicalScope<varbinder::LocalScope>::Enter(varBinder, scope->InstanceFieldScope());
37
38    // Create the name for the synthetic property node
39    util::UString fieldName(util::StringView("field#"), allocator);
40    fieldName.Append(std::to_string(idx));
41    // SUPPRESS_CSA_NEXTLINE(alpha.core.AllocatorETSCheckerHint)
42    auto *fieldIdent = allocator->New<ir::Identifier>(fieldName.View(), allocator);
43
44    // Create the synthetic class property node
45    auto *field =
46        allocator->New<ir::ClassProperty>(fieldIdent, nullptr, nullptr, ir::ModifierFlags::NONE, allocator, false);
47    fieldIdent->SetParent(field);
48
49    // Add the declaration to the scope, and set the type based on the captured variable's scope
50    auto [decl, var] = varBinder->NewVarDecl<varbinder::LetDecl>(pos, fieldIdent->Name());
51    var->SetScope(scope->InstanceFieldScope());
52    var->AddFlag(varbinder::VariableFlags::PROPERTY);
53    var->SetTsType(capturedVar->TsType());
54
55    fieldIdent->SetVariable(var);
56    field->SetTsType(capturedVar->TsType());
57    decl->BindNode(field);
58    return field;
59}
60
61static ir::Statement *CreateCtorFieldInit(checker::ETSChecker *checker, util::StringView name, varbinder::Variable *var)
62{
63    // Create synthetic field initializers for the local class fields
64    // The node structure is the following: this.field0 = field0, where the left hand side refers to the local
65    // classes field, and the right hand side is refers to the constructors parameter
66    // SUPPRESS_CSA_NEXTLINE(alpha.core.AllocatorETSCheckerHint)
67    auto *allocator = checker->Allocator();
68
69    auto *thisExpr = allocator->New<ir::ThisExpression>();
70    auto *fieldAccessExpr = allocator->New<ir::Identifier>(name, allocator);
71    fieldAccessExpr->SetReference();
72    auto *leftHandSide = util::NodeAllocator::ForceSetParent<ir::MemberExpression>(
73        allocator, thisExpr, fieldAccessExpr, ir::MemberExpressionKind::PROPERTY_ACCESS, false, false);
74    auto *rightHandSide = allocator->New<ir::Identifier>(name, allocator);
75    rightHandSide->SetVariable(var);
76    auto *initializer = util::NodeAllocator::ForceSetParent<ir::AssignmentExpression>(
77        allocator, leftHandSide, rightHandSide, lexer::TokenType::PUNCTUATOR_SUBSTITUTION);
78    initializer->SetTsType(var->TsType());
79    return util::NodeAllocator::ForceSetParent<ir::ExpressionStatement>(allocator, initializer);
80}
81
82void LocalClassConstructionPhase::CreateClassPropertiesForCapturedVariables(
83    public_lib::Context *ctx, ir::ClassDefinition *classDef, ArenaSet<varbinder::Variable *> const &capturedVars,
84    ArenaMap<varbinder::Variable *, varbinder::Variable *> &variableMap,
85    ArenaMap<varbinder::Variable *, ir::ClassProperty *> &propertyMap)
86{
87    checker::ETSChecker *const checker = ctx->checker->AsETSChecker();
88    size_t idx = 0;
89    ArenaVector<ir::AstNode *> properties(ctx->allocator->Adapter());
90    for (auto var : capturedVars) {
91        ASSERT(classDef->Scope()->Type() == varbinder::ScopeType::CLASS);
92        auto *property = CreateCapturedField(checker, var, reinterpret_cast<varbinder::ClassScope *>(classDef->Scope()),
93                                             idx, classDef->Start());
94        LOG(DEBUG, ES2PANDA) << "  - Creating property (" << property->Id()->Name()
95                             << ") for captured variable: " << var->Name();
96        properties.push_back(property);
97        variableMap[var] = property->Id()->Variable();
98        propertyMap[var] = property;
99        idx++;
100    }
101
102    classDef->AddProperties(std::move(properties));
103}
104
105ir::ETSParameterExpression *LocalClassConstructionPhase::CreateParam(checker::ETSChecker *const checker,
106                                                                     varbinder::FunctionParamScope *scope,
107                                                                     util::StringView name, checker::Type *type)
108{
109    auto newParam = checker->AddParam(name, nullptr);
110    newParam->SetTsType(type);
111    newParam->Ident()->SetTsType(type);
112    auto paramCtx = varbinder::LexicalScope<varbinder::FunctionParamScope>::Enter(checker->VarBinder(), scope, false);
113
114    auto *paramVar = std::get<1>(checker->VarBinder()->AddParamDecl(newParam));
115    paramVar->SetTsType(newParam->TsType());
116    newParam->Ident()->SetVariable(paramVar);
117    return newParam;
118}
119
120void LocalClassConstructionPhase::ModifyConstructorParameters(
121    public_lib::Context *ctx, ir::ClassDefinition *classDef, ArenaSet<varbinder::Variable *> const &capturedVars,
122    ArenaMap<varbinder::Variable *, varbinder::Variable *> &variableMap,
123    ArenaMap<varbinder::Variable *, varbinder::Variable *> &parameterMap)
124
125{
126    auto *classType = classDef->TsType()->AsETSObjectType();
127    checker::ETSChecker *const checker = ctx->checker->AsETSChecker();
128
129    for (auto *signature : classType->ConstructSignatures()) {
130        LOG(DEBUG, ES2PANDA) << "  - Modifying Constructor: " << signature->InternalName();
131        auto constructor = signature->Function();
132        auto &parameters = constructor->Params();
133        auto &sigParams = signature->Params();
134        signature->GetSignatureInfo()->minArgCount += capturedVars.size();
135
136        ASSERT(signature == constructor->Signature());
137        for (auto var : capturedVars) {
138            auto *newParam = CreateParam(checker, constructor->Scope()->ParamScope(), var->Name(), var->TsType());
139            newParam->SetParent(constructor);
140            // NOTE(psiket) : Moving the parameter after the 'this'. Should modify the AddParam
141            // to be able to insert after the this.
142            auto &paramScopeParams = constructor->Scope()->ParamScope()->Params();
143            auto thisParamIt = ++paramScopeParams.begin();
144            paramScopeParams.insert(thisParamIt, paramScopeParams.back());
145            paramScopeParams.pop_back();
146
147            parameters.insert(parameters.begin(), newParam);
148            ASSERT(newParam->Variable()->Type() == varbinder::VariableType::LOCAL);
149            sigParams.insert(sigParams.begin(), newParam->Ident()->Variable()->AsLocalVariable());
150            parameterMap[var] = newParam->Ident()->Variable()->AsLocalVariable();
151        }
152        reinterpret_cast<varbinder::ETSBinder *>(checker->VarBinder())->BuildFunctionName(constructor);
153        LOG(DEBUG, ES2PANDA) << "    Transformed Constructor: " << signature->InternalName();
154
155        auto *body = constructor->Body();
156        ArenaVector<ir::Statement *> initStatements(ctx->allocator->Adapter());
157        for (auto var : capturedVars) {
158            auto *propertyVar = variableMap[var];
159            auto *initStatement = CreateCtorFieldInit(checker, propertyVar->Name(), propertyVar);
160            auto *fieldInit = initStatement->AsExpressionStatement()->GetExpression()->AsAssignmentExpression();
161            auto *ctorParamVar = parameterMap[var];
162            auto *fieldVar = variableMap[var];
163            auto *leftHandSide = fieldInit->Left();
164            leftHandSide->AsMemberExpression()->SetObjectType(classType);
165            leftHandSide->AsMemberExpression()->SetPropVar(fieldVar->AsLocalVariable());
166            leftHandSide->AsMemberExpression()->SetIgnoreBox();
167            leftHandSide->AsMemberExpression()->SetTsType(fieldVar->TsType());
168            leftHandSide->AsMemberExpression()->Object()->SetTsType(classType);
169            fieldInit->Right()->AsIdentifier()->SetVariable(ctorParamVar);
170            fieldInit->Right()->SetTsType(ctorParamVar->TsType());
171            initStatement->SetParent(body);
172            initStatements.push_back(initStatement);
173        }
174        auto &statements = body->AsBlockStatement()->Statements();
175        statements.insert(statements.begin(), initStatements.begin(), initStatements.end());
176    }
177}
178
179void LocalClassConstructionPhase::RemapReferencesFromCapturedVariablesToClassProperties(
180    ir::ClassDefinition *classDef, ArenaMap<varbinder::Variable *, varbinder::Variable *> &variableMap)
181{
182    auto *classType = classDef->TsType()->AsETSObjectType();
183    auto remapCapturedVariables = [&variableMap](ir::AstNode *childNode) {
184        if (childNode->Type() == ir::AstNodeType::IDENTIFIER) {
185            LOG(DEBUG, ES2PANDA) << "    checking var:" << (void *)childNode;
186            const auto &mapIt = variableMap.find(childNode->AsIdentifier()->Variable());
187            if (mapIt != variableMap.end()) {
188                LOG(DEBUG, ES2PANDA) << "      Remap: " << childNode->AsIdentifier()->Name()
189                                     << " (identifier:" << (void *)childNode
190                                     << ") variable:" << (void *)childNode->AsIdentifier()->Variable()
191                                     << " -> property variable:" << (void *)mapIt->second;
192                childNode->AsIdentifier()->SetVariable(mapIt->second);
193            } else {
194            }
195        }
196    };
197
198    for (auto *it : classDef->Body()) {
199        if (it->IsMethodDefinition() && !it->AsMethodDefinition()->IsConstructor()) {
200            LOG(DEBUG, ES2PANDA) << "  - Rebinding variable rerferences in: "
201                                 << it->AsMethodDefinition()->Id()->Name().Mutf8().c_str();
202            it->AsMethodDefinition()->Function()->Body()->IterateRecursively(remapCapturedVariables);
203        }
204    }
205    // Since the constructor with zero parameter is not listed in the class_def body the constructors
206    // processed separately
207    for (auto *signature : classType->ConstructSignatures()) {
208        auto *constructor = signature->Function();
209        LOG(DEBUG, ES2PANDA) << "  - Rebinding variable rerferences in: " << constructor->Id()->Name();
210        constructor->Body()->IterateRecursively(remapCapturedVariables);
211    }
212}
213
214bool LocalClassConstructionPhase::Perform(public_lib::Context *ctx, parser::Program *program)
215{
216    auto *allocator = ctx->allocator;
217    checker::ETSChecker *const checker = ctx->checker->AsETSChecker();
218    ArenaUnorderedMap<ir::ClassDefinition *, ArenaSet<varbinder::Variable *>> capturedVarsMap {allocator->Adapter()};
219
220    auto handleLocalClass = [this, ctx, &capturedVarsMap](ir::ClassDefinition *classDef) {
221        LOG(DEBUG, ES2PANDA) << "Altering local class with the captured variables: " << classDef->InternalName();
222        auto capturedVars = FindCaptured(ctx->allocator, classDef);
223        // Map the captured variable to the variable of the class property
224        ArenaMap<varbinder::Variable *, varbinder::Variable *> variableMap(ctx->allocator->Adapter());
225        // Map the captured variable to the class property
226        ArenaMap<varbinder::Variable *, ir::ClassProperty *> propertyMap(ctx->allocator->Adapter());
227        // Map the captured variable to the constructor parameter
228        ArenaMap<varbinder::Variable *, varbinder::Variable *> parameterMap(ctx->allocator->Adapter());
229
230        CreateClassPropertiesForCapturedVariables(ctx, classDef, capturedVars, variableMap, propertyMap);
231        ModifyConstructorParameters(ctx, classDef, capturedVars, variableMap, parameterMap);
232        RemapReferencesFromCapturedVariablesToClassProperties(classDef, variableMap);
233        capturedVarsMap.emplace(classDef, std::move(capturedVars));
234    };
235
236    program->Ast()->IterateRecursivelyPostorder([&](ir::AstNode *ast) {
237        if (ast->IsClassDefinition() && ast->AsClassDefinition()->IsLocal()) {
238            handleLocalClass(ast->AsClassDefinition());
239        }
240    });
241
242    // Alter the instantiations
243    auto handleLocalClassInstantiation = [ctx, checker, &capturedVarsMap](ir::ClassDefinition *classDef,
244                                                                          ir::ETSNewClassInstanceExpression *newExpr) {
245        LOG(DEBUG, ES2PANDA) << "Instantiating local class: " << classDef->Ident()->Name();
246        auto capturedVarsIt = capturedVarsMap.find(classDef);
247        ASSERT(capturedVarsIt != capturedVarsMap.cend());
248        auto &capturedVars = capturedVarsIt->second;
249        for (auto *var : capturedVars) {
250            LOG(DEBUG, ES2PANDA) << "  - Extending constructor argument with captured variable: " << var->Name();
251
252            auto *param = checker->AllocNode<ir::Identifier>(var->Name(), ctx->allocator);
253            param->SetVariable(var);
254            param->SetIgnoreBox();
255            param->SetTsType(param->Variable()->TsType());
256            param->SetParent(newExpr);
257            newExpr->AddToArgumentsFront(param);
258        }
259    };
260
261    program->Ast()->IterateRecursivelyPostorder([&](ir::AstNode *ast) {
262        if (ast->IsETSNewClassInstanceExpression()) {
263            auto *newExpr = ast->AsETSNewClassInstanceExpression();
264            checker::Type *calleeType = newExpr->GetTypeRef()->Check(checker);
265            auto *calleeObj = calleeType->AsETSObjectType();
266            auto *classDef = calleeObj->GetDeclNode()->AsClassDefinition();
267            if (classDef->IsLocal()) {
268                handleLocalClassInstantiation(classDef, newExpr);
269            }
270        }
271    });
272
273    return true;
274}
275
276}  // namespace ark::es2panda::compiler
277