Add range-based for and ND to BoxIterator

Author: AlexanderSinn

Src/Base/AMReX_Box.H

@@ -5,6 +5,7 @@
 
 #include <AMReX_Algorithm.H>
 #include <AMReX_ArrayLim.H>
+#include <AMReX_BoxIterator.H>
 #include <AMReX_ccse-mpi.H>
 #include <AMReX_IntVect.H>
 #include <AMReX_IndexType.H>
@@ -25,6 +26,8 @@ template<int dim>
 class BoxND;
 using Box = BoxND<AMREX_SPACEDIM>;
 class BoxCommHelper;
+template<int dim>
+class BoxIteratorND;
 
 /**
 * \brief A Rectangular Domain on an Integer Lattice
@@ -783,6 +786,21 @@ public:
         return *this;
     }
 
+    /**
+    * \brief Returns a BoxIteratorND that can be used to loop over the IntVects
+    * contained by the BoxND.
+    * \code{.cpp}
+    *   Box b;
+    *   ...
+    *   for (IntVect iv : b.iterator()) {
+    *       // do operations involving iv
+    *   }
+    * \endcode
+    */
+    [[nodiscard]] BoxIteratorND<dim> iterator () const noexcept {
+        return BoxIteratorND<dim>{*this};
+    }
+
     [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     static constexpr std::size_t ndims () noexcept {
         return static_cast<std::size_t>(dim);

Src/Base/AMReX_BoxIterator.H

@@ -3,7 +3,6 @@
 #include <AMReX_Config.H>
 
 #include <AMReX_BLassert.H>
-#include <AMReX_Box.H>
 #include <AMReX_REAL.H>
 #include <AMReX_SPACE.H>
 #include <AMReX_IntVect.H>
@@ -12,152 +11,179 @@
 
 namespace amrex
 {
-  /**
-     \brief iterates through the IntVects of a Box
-
-
-     BoxIterator iterates through the IntVects of a box.  The actual
-     sequence of IntVects is implementation-specific.
-     Typical usage:
-
-     Box b;
-     ...
-     BoxIterator bit (b);
-     for (bit.begin(); bit.ok(); ++bit)
-     {
-     IntVect iv = bit();
-     (do operations involving iv)
-     }
-  */
-  class BoxIterator
-  {
-  public:
-    ///
+
+template<int dim>
+class BoxND;
+
+/**
+* \brief iterates through the IntVects of a Box
+*
+* BoxIteratorND iterates through the IntVects of a box
+* in the same order as a series of nested for loops
+* with the innermost loop iterating over the first dimension.
+*
+* Typical usage using a range-based for loop:
+* \code{.cpp}
+*   Box b;
+*   ...
+*   for (IntVect iv : b.iterator()) {
+*       // do operations involving iv
+*   }
+* \endcode
+*
+* Using a traditional for loop:
+* \code{.cpp}
+*   Box b;
+*   ...
+*   BoxIterator bit (b);
+*   for (bit.begin(); bit.ok(); ++bit)
+*   {
+*       IntVect iv = bit();
+*       // do operations involving iv
+*   }
+* \endcode
+*
+* Note that while regular iteration of the range-based version is reasonably performant,
+* neither version is compatible with omp parallel for or simd autovectorization.
+*/
+template<int dim>
+class BoxIteratorND
+{
+public:
+
     /**
-       Default constructor.  This constructs an invalid iterator.
-       The user must call define before using.
+    * Default constructor. This constructs an invalid iterator.
+    * The user must call define before using.
     */
-    BoxIterator () noexcept = default;
+    BoxIteratorND () noexcept = default;
 
-    ///
     /**
-       Constructs a BoxIterator and associates it with a Box.
-       Arguments:
-       a_bx (not modified) the Box to iterate over.
+    * Constructs a BoxIteratorND and associates it with a Box.
+    * Arguments:
+    * a_bx (not modified) the Box to iterate over.
     */
-    explicit BoxIterator (const Box& a_bx) noexcept;
+    explicit BoxIteratorND (const BoxND<dim>& a_bx) noexcept {
+        define(a_bx);
+    }
 
-    void setBox (const Box& a_bx) noexcept;
+    void setBox (const BoxND<dim>& a_bx) noexcept {
+        define(a_bx);
+    }
+
+    /**
+    * Associates a Box with this BoxIteratorND.
+    * Arguments:
+    * a_bx (not modified) the Box to iterate over.
+    */
+    void define (const BoxND<dim>& a_bx) noexcept {
+        if (a_bx.ok() && a_bx.smallEnd().allLE(a_bx.bigEnd())) {
+            m_current = a_bx.smallEnd();
+            m_boxLo   = a_bx.smallEnd();
+            m_boxHi   = a_bx.bigEnd();
+        } else {
+            m_current = IntVectND<dim>::TheUnitVector();
+            m_boxLo   = IntVectND<dim>::TheUnitVector();
+            m_boxHi   = IntVectND<dim>::TheZeroVector();
+        }
+    }
 
-    ///
     /**
-       Associates a Box with this BoxIterator.
-       Arguments:
-       a_bx (not modified) the Box to iterate over.
+    * Sets this BoxIteratorND to the first IntVect in its Box.
+    * This is the smallEnd vector of the Box.
     */
-    void define (const Box& a_bx) noexcept;
+    BoxIteratorND begin () noexcept {
+        m_current = m_boxLo;
+        return *this;
+    }
 
-    ///
     /**
-       Sets this BoxIterator to the first IntVect in its Box.  The
-       definition of the "first" IntVect is
-       implementation-dependent.
+    * Get a BoxIteratorND with its vector set to the end of the Box.
+    * Should only be used internally by a range-based for loop.
     */
-    void begin () noexcept;
+    [[nodiscard]] BoxIteratorND end () const noexcept {
+        BoxIteratorND ret = *this;
+        // set current of end to the past-the-end IntVect
+        ret.m_current = m_boxLo;
+        if (m_boxLo.allLE(m_boxHi)) {
+            ret.m_current[dim-1] = m_boxHi[dim-1] + 1;
+        }
+        return ret;
+    }
 
-    ///
     /**
-       Sets this BoxIterator to the first IntVect in its Box.  The
-       definition of the "first" IntVect is
-       implementation-dependent.
+    * Sets this BoxIteratorND to the first IntVect in its Box.
+    * This is the smallEnd vector of the Box.
     */
-    void reset () noexcept;
+    void reset () noexcept {
+        begin();
+    }
 
-    ///
     /**
-       Modifies this BoxIterator to set it to the next location in its
-       Box.  The definition of the "next location" of a Box is
-       implementation-dependent.
+    * Modifies this BoxIteratorND to set it to the next location in its Box.
     */
-    void operator ++ () noexcept;
+    BoxIteratorND& operator++ () noexcept {
+        next();
+        return *this;
+    }
 
-    void next () noexcept;
+    /**
+    * Modifies this BoxIteratorND to set it to the next location in its Box.
+    */
+    void next () noexcept {
+        for (int i=0; i<dim; ++i) {
+            ++m_current[i];
+            if ((i+1<dim) && m_current[i] > m_boxHi[i]) {
+                m_current[i] = m_boxLo[i];
+                continue;
+            }
+            break;
+        }
+    }
 
-    ///
     /**
-       Returns the value of the InVect for the current location of this
-       BoxIterator.
+    * Returns the value of the InVect for the current location of this BoxIteratorND.
     */
-    [[nodiscard]] const IntVect& operator () () const noexcept;
+    [[nodiscard]] const IntVectND<dim>& operator () () const noexcept {
+        AMREX_ASSERT(m_current.allLE(m_boxHi));
+        AMREX_ASSERT(m_current.allGE(m_boxLo));
+        return m_current;
+    }
 
-    ///
     /**
-       Returns true if this BoxIterator's location is within its Box.
+    * Returns the value of the InVect for the current location of this BoxIteratorND.
     */
-    [[nodiscard]] bool ok () noexcept;
-
-  protected:
-    IntVect m_current = IntVect::TheUnitVector();
-    IntVect m_boxLo   = IntVect::TheUnitVector();
-    IntVect m_boxHi   = IntVect::TheZeroVector();
-  };
-
-  inline
-  BoxIterator::BoxIterator (const Box& a_bx) noexcept
-  {
-    define(a_bx);
-  }
-
-  inline
-  void BoxIterator::begin () noexcept
-  {
-    if (m_boxLo <= m_boxHi) { m_current = m_boxLo; }
-  }
-
-  inline
-  void BoxIterator::reset () noexcept
-  {
-    begin();
-  }
-
-  inline
-  void BoxIterator::operator ++ () noexcept
-  {
-    next();
-  }
-
-  inline
-  void BoxIterator::next () noexcept
-  {
-    m_current[0]++;
-#if AMREX_SPACEDIM >= 2
-    if (m_current[0] > m_boxHi[0])
-      {
-        m_current[0] = m_boxLo[0];
-        m_current[1]++;
-#if AMREX_SPACEDIM >= 3
-        if (m_current[1] > m_boxHi[1])
-          {
-            m_current[1] = m_boxLo[1];
-            m_current[2]++;
-          }
-#endif
-      }
-#endif
-  }
-
-  inline
-  const IntVect& BoxIterator::operator () () const noexcept
-  {
-    BL_ASSERT(m_current <= m_boxHi);
-    BL_ASSERT(m_current >= m_boxLo);
-    return m_current;
-  }
-
-  inline
-  bool BoxIterator::ok () noexcept
-  {
-    return (m_current  <= m_boxHi);
-  }
+    [[nodiscard]] IntVectND<dim> operator* () const noexcept {
+        AMREX_ASSERT(m_current.allLE(m_boxHi));
+        AMREX_ASSERT(m_current.allGE(m_boxLo));
+        // Returning by value gives cleaner assembly output.
+        return m_current;
+    }
+
+    /**
+    * Returns true if this BoxIteratorND's location is within its Box.
+    */
+    [[nodiscard]] bool ok () const noexcept {
+        return m_current.allLE(m_boxHi);
+    }
+
+    /**
+    * Special operator for range-based for loops.
+    * Only checks inequality of the last dimension of the current position for better performance.
+    * Should only be used internally by a range-based for loop.
+    */
+    [[nodiscard]] friend
+    bool operator != (const BoxIteratorND& b1, const BoxIteratorND& b2) noexcept {
+        // Only check one value for better performance.
+        return b1.m_current[dim-1] != b2.m_current[dim-1];
+    }
+
+private:
+    IntVectND<dim> m_current = IntVectND<dim>::TheUnitVector();
+    IntVectND<dim> m_boxLo   = IntVectND<dim>::TheUnitVector();
+    IntVectND<dim> m_boxHi   = IntVectND<dim>::TheZeroVector();
+};
+
+using BoxIterator = BoxIteratorND<AMREX_SPACEDIM>;
+
 }
 #endif