mirror of
https://github.com/ClickHouse/ClickHouse.git
synced 2024-12-05 22:12:21 +00:00
158 lines
4.1 KiB
C++
158 lines
4.1 KiB
C++
#pragma once
|
|
|
|
#include <typeinfo>
|
|
#include <vector>
|
|
#include <Common/typeid_cast.h>
|
|
#include <Common/checkStackSize.h>
|
|
#include <Parsers/DumpASTNode.h>
|
|
|
|
|
|
namespace DB
|
|
{
|
|
|
|
/// Visits AST tree in depth, call functions for nodes according to Matcher type data.
|
|
/// You need to define Data, visit() and needChildVisit() in Matcher class.
|
|
template <typename Matcher, bool _top_to_bottom, bool need_child_accept_data = false, typename T = ASTPtr>
|
|
class InDepthNodeVisitor
|
|
{
|
|
public:
|
|
using Data = typename Matcher::Data;
|
|
|
|
explicit InDepthNodeVisitor(Data & data_, WriteBuffer * ostr_ = nullptr)
|
|
: data(data_),
|
|
visit_depth(0),
|
|
ostr(ostr_)
|
|
{}
|
|
|
|
void visit(T & ast)
|
|
{
|
|
if (ostr)
|
|
visitImpl</* with_dump= */ true>(ast);
|
|
else
|
|
visitImpl</* with_dump= */ false>(ast);
|
|
}
|
|
|
|
private:
|
|
Data & data;
|
|
size_t visit_depth;
|
|
WriteBuffer * ostr;
|
|
|
|
template <bool with_dump>
|
|
void visitImpl(T & ast)
|
|
{
|
|
checkStackSize();
|
|
if constexpr (with_dump)
|
|
{
|
|
DumpASTNode dump(*ast, ostr, visit_depth, typeid(Matcher).name());
|
|
visitImplMain</* with_dump= */ true>(ast);
|
|
}
|
|
else
|
|
{
|
|
visitImplMain</* with_dump= */ false>(ast);
|
|
}
|
|
}
|
|
|
|
template <bool with_dump>
|
|
void visitImplMain(T & ast)
|
|
{
|
|
if constexpr (!_top_to_bottom)
|
|
visitChildren<with_dump>(ast);
|
|
|
|
doVisit(ast);
|
|
|
|
if constexpr (_top_to_bottom)
|
|
visitChildren<with_dump>(ast);
|
|
}
|
|
|
|
void doVisit(T & ast)
|
|
{
|
|
try
|
|
{
|
|
Matcher::visit(ast, data);
|
|
}
|
|
catch (Exception & e)
|
|
{
|
|
e.addMessage("While processing {}", ast->formatForErrorMessage());
|
|
throw;
|
|
}
|
|
}
|
|
|
|
template <bool with_dump>
|
|
void visitChildren(T & ast)
|
|
{
|
|
for (auto & child : ast->children)
|
|
{
|
|
bool need_visit_child = false;
|
|
if constexpr (need_child_accept_data)
|
|
need_visit_child = Matcher::needChildVisit(ast, child, data);
|
|
else
|
|
need_visit_child = Matcher::needChildVisit(ast, child);
|
|
|
|
if (need_visit_child)
|
|
visitImpl<with_dump>(child);
|
|
}
|
|
}
|
|
};
|
|
|
|
template <typename Matcher, bool top_to_bottom, bool need_child_accept_data = false>
|
|
using ConstInDepthNodeVisitor = InDepthNodeVisitor<Matcher, top_to_bottom, need_child_accept_data, const ASTPtr>;
|
|
|
|
struct NeedChild
|
|
{
|
|
using Condition = bool (*)(const ASTPtr & node, const ASTPtr & child);
|
|
|
|
static bool all(const ASTPtr &, const ASTPtr &) { return true; }
|
|
static bool none(const ASTPtr &, const ASTPtr &) { return false; }
|
|
};
|
|
|
|
/// Simple matcher for one node type. Use need_child function for complex traversal logic.
|
|
template <typename DataImpl, NeedChild::Condition need_child = NeedChild::all, typename T = ASTPtr>
|
|
class OneTypeMatcher
|
|
{
|
|
public:
|
|
using Data = DataImpl;
|
|
using TypeToVisit = typename Data::TypeToVisit;
|
|
|
|
static bool needChildVisit(const ASTPtr & node, const ASTPtr & child) { return need_child(node, child); }
|
|
|
|
static void visit(T & ast, Data & data)
|
|
{
|
|
if (auto * t = typeid_cast<TypeToVisit *>(ast.get()))
|
|
data.visit(*t, ast);
|
|
}
|
|
};
|
|
|
|
template <typename Data, NeedChild::Condition need_child = NeedChild::all>
|
|
using ConstOneTypeMatcher = OneTypeMatcher<Data, need_child, const ASTPtr>;
|
|
|
|
template <typename Visitor, typename T = ASTPtr>
|
|
struct InDepthNodeVisitorWithChildInfo : Visitor
|
|
{
|
|
using ChildInfo = typename Visitor::ChildInfo;
|
|
|
|
ChildInfo visit(T & ast, const T & parent = {})
|
|
{
|
|
ChildInfo all_children_info;
|
|
for (auto & child : ast->children)
|
|
{
|
|
if (Visitor::needVisitChild(ast, child))
|
|
{
|
|
ChildInfo child_info = visit(child, ast);
|
|
all_children_info.update(child_info);
|
|
}
|
|
}
|
|
|
|
try
|
|
{
|
|
return Visitor::visitNode(ast, parent, all_children_info);
|
|
}
|
|
catch (Exception & e)
|
|
{
|
|
e.addMessage("While processing {}", ast->formatForErrorMessage());
|
|
throw;
|
|
}
|
|
}
|
|
};
|
|
|
|
}
|