Branch data     Line data    Source code

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