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 "ambientLowering.h"
17#include <string_view>
18
19#include "ir/expressions/dummyNode.h"
20#include "ir/astNode.h"
21
22namespace ark::es2panda::compiler {
23std::string_view AmbientLowering::Name() const
24{
25    static std::string const NAME = "AmbientLowering";
26    return NAME;
27}
28
29bool AmbientLowering::Postcondition(public_lib::Context *ctx, const parser::Program *program)
30{
31    for (auto &[_, extPrograms] : program->ExternalSources()) {
32        (void)_;
33        for (auto *extProg : extPrograms) {
34            if (!Postcondition(ctx, extProg)) {
35                return false;
36            }
37        }
38    }
39
40    return !program->Ast()->IsAnyChild(
41        [](const ir::AstNode *node) -> bool { return node->IsDummyNode() && node->AsDummyNode()->IsAmbientIndexer(); });
42}
43
44bool AmbientLowering::Perform(public_lib::Context *ctx, parser::Program *program)
45{
46    for (auto &[_, extPrograms] : program->ExternalSources()) {
47        (void)_;
48        for (auto *extProg : extPrograms) {
49            Perform(ctx, extProg);
50        }
51    }
52
53    // Generate $_get and $_set for ambient Indexer
54    program->Ast()->TransformChildrenRecursively(
55        [this, ctx](ir::AstNode *ast) -> ir::AstNode * {
56            if (ast->IsClassDefinition()) {
57                return CreateIndexerMethodIfNeeded(ast->AsClassDefinition(), ctx);
58            }
59            return ast;
60        },
61        Name());
62
63    return true;
64}
65
66ir::MethodDefinition *CreateMethodFunctionDefinition(ir::DummyNode *node, public_lib::Context *ctx,
67                                                     ir::MethodDefinitionKind funcKind)
68{
69    auto parser = ctx->parser->AsETSParser();
70
71    auto const indexName = node->GetIndexName();
72    auto const returnType = node->GetReturnTypeLiteral()->AsETSTypeReferencePart()->Name()->AsIdentifier()->Name();
73    std::string sourceCode;
74    if (funcKind == ir::MethodDefinitionKind::GET) {
75        sourceCode = "$_get(" + std::string(indexName) + " : number) : " + std::string(returnType);
76    } else if (funcKind == ir::MethodDefinitionKind::SET) {
77        sourceCode =
78            "$_set(" + std::string(indexName) + " : number, " + "value : " + std::string(returnType) + " ) : void";
79    } else {
80        UNREACHABLE();
81    }
82
83    auto methodDefinition = parser->CreateFormattedClassMethodDefinition(sourceCode);
84
85    methodDefinition->SetRange(node->Range());
86    methodDefinition->SetParent(node->Parent());
87    methodDefinition->AddModifier(ir::ModifierFlags::DECLARE);
88    return methodDefinition->AsMethodDefinition();
89}
90
91ir::ClassDefinition *AmbientLowering::CreateIndexerMethodIfNeeded(ir::ClassDefinition *classDef,
92                                                                  public_lib::Context *ctx)
93{
94    auto &classBody = classDef->Body();
95    auto it = classBody.begin();
96    // Only one DummyNode is allowed in classBody for now
97    ASSERT(std::count_if(classBody.begin(), classBody.end(), [](ir::AstNode *node) { return node->IsDummyNode(); }) <=
98           1);
99    while (it != classBody.end()) {
100        if ((*it)->IsDummyNode() && (*it)->AsDummyNode()->IsAmbientIndexer()) {
101            auto setDefinition =
102                CreateMethodFunctionDefinition((*it)->AsDummyNode(), ctx, ir::MethodDefinitionKind::SET);
103            auto getDefinition =
104                CreateMethodFunctionDefinition((*it)->AsDummyNode(), ctx, ir::MethodDefinitionKind::GET);
105
106            classBody.erase(it);
107            classBody.emplace_back(setDefinition);
108            classBody.emplace_back(getDefinition);
109            break;
110        }
111        ++it;
112    }
113
114    return classDef;
115}
116}  // namespace ark::es2panda::compiler
117