#include <string>
#include <vector>
#include <functional>
#include <algorithm>
class FormatVisitor : public BaseVisitor {
public:
void Visit(const BaseNode *node) override {
node->Visit(this);
}
void Visit(const ClassDeclarationNode *node) override {
formated.push_back(tab + "class " + node->ClassName() + " {");
if (!node->PublicFields().empty()) {
tab += " ";
formated.push_back(tab + "public:");
tab += " ";
for (auto it : node->PublicFields()) {
Visit(it);
}
tab.erase(0, 4);
}
if (!node->ProtectedFields().empty()) {
if (formated.back().back() != '{') {
formated.push_back("");
}
tab += " ";
formated.push_back(tab + "protected:");
tab += " ";
for (auto it : node->ProtectedFields()) {
Visit(it);
}
tab.erase(0, 4);
}
if (!node->PrivateFields().empty()) {
if (formated.back().back() != '{') {
formated.push_back("");
}
tab += " ";
formated.push_back(tab + "private:");
tab += " ";
for (auto it : node->PrivateFields()) {
Visit(it);
}
tab.erase(0, 4);
}
formated.push_back(tab + "};");
}
void Visit(const VarDeclarationNode *node) override {
std::string tmp = node->TypeName() + " " + node->VarName();
if (!args.empty()) {
args += tmp;
} else {
formated.push_back(tab + tmp + ";");
}
}
void Visit(const MethodDeclarationNode *node) override {
formated.push_back(tab + node->ReturnTypeName() + " " + node->MethodName());
args = "(";
if (!node->Arguments().empty()) {
for (auto it : node->Arguments()) {
Visit(it);
args += ", ";
}
args.pop_back();
args.pop_back();
}
formated.back() += args + ");";
args = "";
}
[[nodiscard]] const std::vector<std::string> &GetFormattedCode() const {
return formated;
}
private:
std::vector<std::string> formated;
std::string args;
std::string tab = "";
};