Branch data     Line data    Source code

       1                 :            : /******************************************************************************
       2                 :            :  * This file is part of the cvc5 project.
       3                 :            :  *
       4                 :            :  * Copyright (c) 2009-2026 by the authors listed in the file AUTHORS
       5                 :            :  * in the top-level source directory and their institutional affiliations.
       6                 :            :  * All rights reserved.  See the file COPYING in the top-level source
       7                 :            :  * directory for licensing information.
       8                 :            :  * ****************************************************************************
       9                 :            :  *
      10                 :            :  * Iterators for traversing nodes.
      11                 :            :  */
      12                 :            : 
      13                 :            : #include "cvc5_private.h"
      14                 :            : 
      15                 :            : #ifndef CVC5__EXPR__NODE_TRAVERSAL_H
      16                 :            : #define CVC5__EXPR__NODE_TRAVERSAL_H
      17                 :            : 
      18                 :            : #include <iterator>
      19                 :            : #include <unordered_map>
      20                 :            : #include <vector>
      21                 :            : 
      22                 :            : #include "expr/node.h"
      23                 :            : 
      24                 :            : namespace cvc5::internal {
      25                 :            : 
      26                 :            : /**
      27                 :            :  * Enum that represents an order in which nodes are visited.
      28                 :            :  */
      29                 :            : enum class VisitOrder
      30                 :            : {
      31                 :            :   PREORDER,
      32                 :            :   POSTORDER
      33                 :            : };
      34                 :            : 
      35                 :            : // Iterator for traversing a node in pre-/post-order
      36                 :            : // It does DAG-traversal, so indentical sub-nodes will be visited once only.
      37                 :            : class NodeDfsIterator
      38                 :            : {
      39                 :            :  public:
      40                 :            :   // STL type definitions for an iterator
      41                 :            :   using value_type = TNode;
      42                 :            :   using pointer = TNode*;
      43                 :            :   using reference = TNode&;
      44                 :            :   using iterator_category = std::forward_iterator_tag;
      45                 :            :   using difference_type = std::ptrdiff_t;
      46                 :            : 
      47                 :            :   // Construct a traversal iterator beginning at `n`
      48                 :            :   NodeDfsIterator(TNode n, VisitOrder order, std::function<bool(TNode)> skipIf);
      49                 :            :   // Construct an end-of-traversal iterator
      50                 :            :   NodeDfsIterator(VisitOrder order);
      51                 :            : 
      52                 :            :   // Move/copy construction and assignment. Destructor.
      53                 :         20 :   NodeDfsIterator(NodeDfsIterator&&) = default;
      54                 :            :   NodeDfsIterator& operator=(NodeDfsIterator&&) = default;
      55                 :         58 :   NodeDfsIterator(const NodeDfsIterator&) = default;
      56                 :            :   NodeDfsIterator& operator=(const NodeDfsIterator&) = default;
      57                 :     339232 :   ~NodeDfsIterator() = default;
      58                 :            : 
      59                 :            :   // Preincrement
      60                 :            :   NodeDfsIterator& operator++();
      61                 :            :   // Postincrement
      62                 :            :   NodeDfsIterator operator++(int);
      63                 :            :   // Dereference
      64                 :            :   reference operator*();
      65                 :            :   // Equals
      66                 :            :   // It is not constant, because an unitilized node must be initialized before
      67                 :            :   // comparison
      68                 :            :   bool operator==(NodeDfsIterator&);
      69                 :            :   // Not equals
      70                 :            :   // It is not constant, because an unitilized node must be initialized before
      71                 :            :   // comparison
      72                 :            :   bool operator!=(NodeDfsIterator&);
      73                 :            : 
      74                 :            :  private:
      75                 :            :   // While we're not at an appropriate visit (see d_postorder), advance.
      76                 :            :   // In each step:
      77                 :            :   //   * enqueue children of a not-yet-pre-visited node (and mark it
      78                 :            :   //     previsited)
      79                 :            :   //   * pop a not-yet-post-visited node (and mark it post-visited)
      80                 :            :   //   * pop an already post-visited node.
      81                 :            :   // After calling this, `d_current` will be changed to the next node, if there
      82                 :            :   // is another node in the traversal.
      83                 :            :   void advanceToNextVisit();
      84                 :            : 
      85                 :            :   // If this iterator hasn't been dereferenced or incremented yet, advance to
      86                 :            :   // first visit.
      87                 :            :   // Necessary because we are lazy and don't find our first visit node at
      88                 :            :   // construction time.
      89                 :            :   void initializeIfUninitialized();
      90                 :            : 
      91                 :            :   // Stack of nodes to visit.
      92                 :            :   std::vector<TNode> d_stack;
      93                 :            : 
      94                 :            :   // Whether (and how) we've visited a node.
      95                 :            :   // Absent if we haven't visited it.
      96                 :            :   // Set to `false` if we've already pre-visited it (enqueued its children).
      97                 :            :   // Set to `true` if we've also already post-visited it.
      98                 :            :   std::unordered_map<TNode, bool> d_visited;
      99                 :            : 
     100                 :            :   // The visit order that this iterator is using
     101                 :            :   VisitOrder d_order;
     102                 :            : 
     103                 :            :   // Current referent node. A valid node to visit if non-null.
     104                 :            :   // Null after construction (but before first access) and at the end.
     105                 :            :   TNode d_current;
     106                 :            : 
     107                 :            :   // When to omit a node and its descendants from the traversal
     108                 :            :   std::function<bool(TNode)> d_skipIf;
     109                 :            : };
     110                 :            : 
     111                 :            : // Node wrapper that is iterable in DAG pre-/post-order
     112                 :            : class NodeDfsIterable
     113                 :            : {
     114                 :            :  public:
     115                 :            :   /**
     116                 :            :    * Creates a new node wrapper that can be used to iterate over the children
     117                 :            :    * of a node in pre-/post-order.
     118                 :            :    *
     119                 :            :    * @param n The node the iterate
     120                 :            :    * @param order The order in which the children are visited.
     121                 :            :    * @param skipIf Function that determines whether a given node and its
     122                 :            :    *               descendants should be skipped or not.
     123                 :            :    */
     124                 :            :   NodeDfsIterable(
     125                 :            :       TNode n,
     126                 :            :       VisitOrder order = VisitOrder::POSTORDER,
     127                 :         45 :       std::function<bool(TNode)> skipIf = [](TNode) { return false; });
     128                 :            : 
     129                 :            :   // Move/copy construction and assignment. Destructor.
     130                 :            :   NodeDfsIterable(NodeDfsIterable&&) = default;
     131                 :            :   NodeDfsIterable& operator=(NodeDfsIterable&&) = default;
     132                 :            :   NodeDfsIterable(const NodeDfsIterable&) = default;
     133                 :            :   NodeDfsIterable& operator=(const NodeDfsIterable&) = default;
     134                 :     169577 :   ~NodeDfsIterable() = default;
     135                 :            : 
     136                 :            :   NodeDfsIterator begin() const;
     137                 :            :   NodeDfsIterator end() const;
     138                 :            : 
     139                 :            :  private:
     140                 :            :   TNode d_node;
     141                 :            :   VisitOrder d_order;
     142                 :            :   std::function<bool(TNode)> d_skipIf;
     143                 :            : };
     144                 :            : 
     145                 :            : }  // namespace cvc5::internal
     146                 :            : 
     147                 :            : #endif  // CVC5__EXPR__NODE_TRAVERSAL_H