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 "spreadLowering.h"
17#include "checker/ETSchecker.h"
18#include "compiler/lowering/scopesInit/scopesInitPhase.h"
19#include "compiler/lowering/util.h"
20#include "ir/expressions/literals/numberLiteral.h"
21
22namespace ark::es2panda::compiler {
23
24using AstNodePtr = ir::AstNode *;
25
26std::string SpreadConstructionPhase::CreateLengthString(ir::ArrayExpression *array)
27{
28    int spreadElementCount = 0;
29    std::stringstream lengthCalculationString;
30
31    for (const auto *element : array->Elements()) {
32        if (element->Type() == ir::AstNodeType::SPREAD_ELEMENT) {
33            spreadElementCount++;
34            lengthCalculationString << element->AsSpreadElement()->Argument()->AsIdentifier()->Name() << ".length + ";
35        }
36    }
37
38    lengthCalculationString << "0";
39    int newArrayLength = array->Elements().size() - spreadElementCount;
40    std::stringstream lengthString;
41    lengthString << "let length : int = " << newArrayLength << " + " << lengthCalculationString.str();
42
43    return lengthString.str();
44}
45
46std::string SpreadConstructionPhase::CreateETSCode(ir::ArrayExpression *array, std::vector<ir::AstNode *> &node,
47                                                   public_lib::Context *ctx)
48{
49    std::stringstream src;
50    std::string lengthString = CreateLengthString(array);
51    std::string arrayType = array->TsType()->AsETSArrayType()->ElementType()->ToString();
52
53    src.clear();
54    std::string newArrayName = "tempArrayVar";
55    src << lengthString << std::endl;
56    src << "type typeOfTempArray = " << arrayType << std::endl;
57    src << "let " << newArrayName << ": typeOfTempArray[] = new typeOfTempArray[length]" << std::endl;
58    src << "let newArrayIndex = 0" << std::endl;
59    size_t argumentCount = 1;
60
61    for (std::uint32_t i = 0; i < array->Elements().size(); ++i) {
62        if (array->Elements()[i]->Type() == ir::AstNodeType::SPREAD_ELEMENT) {
63            std::string spreadArrayName =
64                array->Elements()[i]->AsSpreadElement()->Argument()->AsIdentifier()->Name().Mutf8();
65            src << "let elementOfSpread" << i << ": " << array->Elements()[i]->TsType()->ToString() << std::endl;
66            src << "for (elementOfSpread" << i << " of " << spreadArrayName << ") {" << std::endl;
67            src << newArrayName << "[newArrayIndex] = "
68                << "elementOfSpread" << i << std::endl;
69            src << "newArrayIndex++" << std::endl;
70            src << "}" << std::endl;
71        } else {
72            src << newArrayName << "[newArrayIndex] = "
73                << "(@@E" << argumentCount << ")" << std::endl;
74            src << "newArrayIndex++" << std::endl;
75            argumentCount++;
76            node.emplace_back(array->Elements()[i]->Clone(ctx->allocator, nullptr));
77        }
78    }
79    src << newArrayName << ";" << std::endl;
80
81    return src.str();
82}
83
84bool SpreadConstructionPhase::Perform(public_lib::Context *ctx, parser::Program *program)
85{
86    for (auto &[_, ext_programs] : program->ExternalSources()) {
87        (void)_;
88        for (auto *extProg : ext_programs) {
89            Perform(ctx, extProg);
90        }
91    }
92
93    auto *const parser = ctx->parser->AsETSParser();
94    checker::ETSChecker *const checker = ctx->checker->AsETSChecker();
95
96    program->Ast()->TransformChildrenRecursively(
97        [&parser, &checker, &ctx, this](ir::AstNode *const node) -> AstNodePtr {
98            if (node->IsArrayExpression() &&
99                std::any_of(node->AsArrayExpression()->Elements().begin(), node->AsArrayExpression()->Elements().end(),
100                            [](const auto *param) { return param->Type() == ir::AstNodeType::SPREAD_ELEMENT; })) {
101                auto scopeCtx =
102                    varbinder::LexicalScope<varbinder::Scope>::Enter(checker->VarBinder(), NearestScope(node));
103                std::vector<ir::AstNode *> normalElements {};
104                std::string src = CreateETSCode(node->AsArrayExpression(), normalElements, ctx);
105
106                ir::BlockExpression *blockExpression =
107                    parser->CreateFormattedExpression(src, normalElements)->AsBlockExpression();
108                blockExpression->SetParent(node->Parent());
109                InitScopesPhaseETS::RunExternalNode(blockExpression, checker->VarBinder());
110                checker->VarBinder()->AsETSBinder()->ResolveReferencesForScope(blockExpression,
111                                                                               NearestScope(blockExpression));
112                blockExpression->Check(checker);
113
114                return blockExpression;
115            }
116
117            return node;
118        },
119        Name());
120    return true;
121}
122
123bool SpreadConstructionPhase::Postcondition(public_lib::Context *ctx, const parser::Program *program)
124{
125    for (auto &[_, ext_programs] : program->ExternalSources()) {
126        (void)_;
127        for (auto *extProg : ext_programs) {
128            if (!Postcondition(ctx, extProg)) {
129                return false;
130            }
131        }
132    }
133    return true;
134}
135
136}  // namespace ark::es2panda::compiler
137