#pragma once #include #include "sir/sir.h" #include "visitor.h" #define LAMBDA_VISIT(x) \ virtual void handle(codon::ir::x *v) {} \ void visit(codon::ir::x *v) override { \ if (childrenFirst) \ processChildren(v); \ preHook(v); \ handle(v); \ postHook(v); \ if (!childrenFirst) \ processChildren(v); \ } namespace codon { namespace ir { namespace util { /// Pass that visits all values in a module. class Operator : public Visitor { private: /// IDs of previously visited nodes std::unordered_set seen; /// stack of IR nodes being visited std::vector nodeStack; /// stack of iterators std::vector itStack; /// true if should visit children first bool childrenFirst; protected: void defaultVisit(Node *) override {} public: /// Constructs an operator. /// @param childrenFirst true if children should be visited first explicit Operator(bool childrenFirst = false) : childrenFirst(childrenFirst) {} virtual ~Operator() noexcept = default; /// This function is applied to all nodes before handling the node /// itself. It provides a way to write one function that gets /// applied to every visited node. /// @param node the node virtual void preHook(Node *node) {} /// This function is applied to all nodes after handling the node /// itself. It provides a way to write one function that gets /// applied to every visited node. /// @param node the node virtual void postHook(Node *node) {} void visit(Module *m) override { nodeStack.push_back(m); nodeStack.push_back(m->getMainFunc()); process(m->getMainFunc()); nodeStack.pop_back(); for (auto *s : *m) { nodeStack.push_back(s); process(s); nodeStack.pop_back(); } nodeStack.pop_back(); } void visit(BodiedFunc *f) override { seen.insert(f->getBody()->getId()); process(f->getBody()); } LAMBDA_VISIT(VarValue); LAMBDA_VISIT(PointerValue); void visit(codon::ir::SeriesFlow *v) override { if (childrenFirst) processSeriesFlowChildren(v); preHook(v); handle(v); postHook(v); if (!childrenFirst) processSeriesFlowChildren(v); } virtual void handle(codon::ir::SeriesFlow *v) {} LAMBDA_VISIT(IfFlow); LAMBDA_VISIT(WhileFlow); LAMBDA_VISIT(ForFlow); LAMBDA_VISIT(ImperativeForFlow); LAMBDA_VISIT(TryCatchFlow); LAMBDA_VISIT(PipelineFlow); LAMBDA_VISIT(dsl::CustomFlow); LAMBDA_VISIT(TemplatedConst); LAMBDA_VISIT(TemplatedConst); LAMBDA_VISIT(TemplatedConst); LAMBDA_VISIT(TemplatedConst); LAMBDA_VISIT(dsl::CustomConst); LAMBDA_VISIT(Instr); LAMBDA_VISIT(AssignInstr); LAMBDA_VISIT(ExtractInstr); LAMBDA_VISIT(InsertInstr); LAMBDA_VISIT(CallInstr); LAMBDA_VISIT(StackAllocInstr); LAMBDA_VISIT(TypePropertyInstr); LAMBDA_VISIT(YieldInInstr); LAMBDA_VISIT(TernaryInstr); LAMBDA_VISIT(BreakInstr); LAMBDA_VISIT(ContinueInstr); LAMBDA_VISIT(ReturnInstr); LAMBDA_VISIT(YieldInstr); LAMBDA_VISIT(ThrowInstr); LAMBDA_VISIT(FlowInstr); LAMBDA_VISIT(dsl::CustomInstr); template void process(Node *v) { v->accept(*this); } /// Return the parent of the current node. /// @param level the number of levels up from the current node template Desired *getParent(int level = 0) { return cast(nodeStack[nodeStack.size() - level - 1]); } /// @return current depth in the tree int depth() const { return nodeStack.size(); } /// @tparam Desired the desired type /// @return the last encountered example of the desired type template Desired *findLast() { for (auto it = nodeStack.rbegin(); it != nodeStack.rend(); ++it) { if (auto *v = cast(*it)) return v; } return nullptr; } /// @return the last encountered function Func *getParentFunc() { return findLast(); } /// @return an iterator to the first parent auto parent_begin() const { return nodeStack.begin(); } /// @return an iterator beyond the last parent auto parent_end() const { return nodeStack.end(); } /// @param v the value /// @return whether we have visited ("seen") the given value bool saw(const Value *v) const { return seen.find(v->getId()) != seen.end(); } /// Avoid visiting the given value in the future. /// @param v the value void see(const Value *v) { seen.insert(v->getId()); } /// Inserts the new value before the current position in the last seen SeriesFlow. /// @param v the new value auto insertBefore(Value *v) { return findLast()->insert(itStack.back(), v); } /// Inserts the new value after the current position in the last seen SeriesFlow. /// @param v the new value, which is marked seen auto insertAfter(Value *v) { auto newPos = itStack.back(); ++newPos; see(v); return findLast()->insert(newPos, v); } /// Resets the operator. void reset() { seen.clear(); nodeStack.clear(); itStack.clear(); } private: void processChildren(Value *v) { nodeStack.push_back(v); for (auto *c : v->getUsedValues()) { if (saw(c)) continue; see(c); process(c); } nodeStack.pop_back(); } void processSeriesFlowChildren(codon::ir::SeriesFlow *v) { nodeStack.push_back(v); for (auto it = v->begin(); it != v->end(); ++it) { itStack.push_back(it); process(*it); itStack.pop_back(); } nodeStack.pop_back(); } }; } // namespace util } // namespace ir } // namespace codon #undef LAMBDA_VISIT