AST/Traversal.h

#ifndef ROSE_AST_Traversal_H
#define ROSE_AST_Traversal_H
 
#include <RoseFirst.h>
#include <vector>
 
// A version of AST traversal that doesn't require huge header files to be included into the compilation unit that's calling the
// traversal.
 
class SgNode;
 
namespace Rose {
namespace AST {
 
namespace Traversal {
 
enum class Order {                                      // bit flags
    PRE    = 0x01,                                      
    POST   = 0x02                                       
};
 
// Implementation
 
namespace detail {
 
// Base visitor class. The idea is that the base class is abstract so it can be implemented in a single .C file. It is implemented
// in terms of ROSE's `AstPrePostProcessing` which requires definitions for every Sg node. It's fine to #include all those
// definitions into that one .C file were we need them, but we want to avoid having to #include all those definitions into every
// .C file that needs to do an AST traversal.
//
// Classes derived from `VisitorBase` are defined by templates, but they've been designed so as to not require any class definitions
// except for `T` (and T's base classes).
class VisitorBase {
public:
    virtual ~VisitorBase() {}
    virtual void operator()(SgNode*, Order) const = 0;
};
 
// Calls functor before and after (pre and post) with the node pointer and an indication of when it is being called.
template<class T, class Functor>
class VisitorOrdered: public VisitorBase {
    Functor &f;
public:
    explicit VisitorOrdered(Functor &f)
        : f(f) {}
 
    void operator()(SgNode *node, const Order order) const override {
        if (T *typedNode = dynamic_cast<T*>(node))
            f(typedNode, order);
    }
};
 
// Calls functor only before visiting related nodes. Functor is called with only the node pointer argument.
template<class T, class Functor>
class VisitorPre: public VisitorBase {
    Functor &f;
public:
    explicit VisitorPre(Functor &f)
        : f(f) {}
 
    void operator()(SgNode *node, const Order order) const override {
        if (Order::PRE == order) {
            if (T *typedNode = dynamic_cast<T*>(node))
                f(typedNode);
        }
    }
};
 
// Calls functor only after visiting related nodes. Functor is called with only the node pointer argument.
template<class T, class Functor>
class VisitorPost: public VisitorBase {
    Functor &f;
public:
    explicit VisitorPost(Functor &f)
        : f(f) {}
 
    void operator()(SgNode *node, const Order order) const override {
        if (Order::POST == order) {
            if (T *typedNode = dynamic_cast<T*>(node))
                f(typedNode);
        }
    }
};
 
// Base class for finding a particular node. This works similarly to the `VisitorBase` in that this abstract base class is
// what's passed across compilation units so that only the one .C file that implements the traversal needs to include the whole
// world.
//
// The derived classes, which are class templates, only need the definition for their `T` type (and it's base classes), and do not
// need the definitions for all AST node classes.
class FinderBase {
public:
    SgNode *found = nullptr;
 
    virtual ~FinderBase() {}
    virtual bool operator()(SgNode *node) = 0;
};
 
// Calls the specified predicate for each visited node of the correct dynamic type.
template<class T, class Functor>
class Finder: public FinderBase {
    Functor &f;
public:
    explicit Finder(Functor &f)
        : f(f) {}
 
    bool operator()(SgNode *node) override {
        if (T *typedNode = dynamic_cast<T*>(node)) {
            return f(typedNode);
        } else {
            return false;
        }
    }
};
 
// Traverse an AST in a forward direction, calling the visitor for a node before and then after visiting the node's children.
void forwardHelper(SgNode *ast, const VisitorBase&);
 
// Traverse an AST in the reverse direction, calling the visitor for a node before and then after visiting the node's parents.
void reverseHelper(SgNode *ast, const VisitorBase&);
 
// Traverse an AST in reverse to find the first occurrence of a particular node, and saves that node in the `FinderBase`.
void findReverseHelper(SgNode *ast, FinderBase&);
 
} // namespace
 
// Traversals
 
template<class T, class Functor>
void forward(SgNode *ast, Functor functor) {
    detail::forwardHelper(ast, detail::VisitorOrdered<T, Functor>(functor));
}
 
template<class T, class Functor>
void forwardPre(SgNode *ast, Functor functor) {
    detail::forwardHelper(ast, detail::VisitorPre<T, Functor>(functor));
}
 
template<class T, class Functor>
void forwardPost(SgNode *ast, Functor functor) {
    detail::forwardHelper(ast, detail::VisitorPost<T, Functor>(functor));
}
 
template<class T, class Functor>
void reverse(SgNode *ast, Functor functor) {
    detail::reverseHelper(ast, detail::VisitorOrdered<T, Functor>(functor));
}
 
template<class T, class Functor>
void reversePre(SgNode *ast, Functor functor) {
    detail::reverseHelper(ast, detail::VisitorPre<T, Functor>(functor));
}
 
template<class T, class Functor>
void reversePost(SgNode *ast, Functor functor) {
    detail::reverseHelper(ast, detail::VisitorPost<T, Functor>(functor));
}
 
template<class T, class Functor>
T*
findParentSatisying(SgNode *ast, Functor functor) {
    detail::Finder<T, Functor> finder(functor);
    detail::findReverseHelper(ast, finder /*in,out*/);
    return dynamic_cast<T*>(finder.found);
}
 
template<class T>
T*
findParentTyped(SgNode *ast) {
    auto yes = [](SgNode*) { return true; };
    detail::Finder<T, decltype(yes)> finder(yes);
    detail::findReverseHelper(ast, finder /*in,out*/);
    return dynamic_cast<T*>(finder.found);
}
 
template<class T>
std::vector<T*>
findDescendantsTyped(SgNode *ast) {
    std::vector<T*> retval;
    forwardPre<T>(ast, [&retval](T *node) {
        retval.push_back(node);
    });
    return retval;
}
 
} // namespace
} // namespace
} // namespace
 
#endif /* ROSE_AST_Traversal_H */