include/geocel/VolumeVisitor.hh

0001 //------------------------------- -*- C++ -*- -------------------------------//
0002 // Copyright Celeritas contributors: see top-level COPYRIGHT file for details
0003 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
0004 //---------------------------------------------------------------------------//
0005 //! \file geocel/VolumeVisitor.hh
0006 //---------------------------------------------------------------------------//
0007 #pragma once
0008
0009 #include <unordered_set>
0010 #include <vector>
0011
0012 namespace celeritas
0013 {
0014 //---------------------------------------------------------------------------//
0015 //! Traits class to access children and associated logical volume
0016 template<class PV>
0017 struct VolumeVisitorTraits
0018 {
0019     using LV = void;
0020
0021     static void get_children(PV const& parent, std::vector<PV const*>& dst);
0022     static LV get_lv(PV const& pv);
0023 };
0024
0025 //---------------------------------------------------------------------------//
0026 /*!
0027  * Recursively visit volumes.
0028  * \tparam T Volume type
0029  *
0030  * This class can be used for both Geant4 and VecGeom to give the same visiting
0031  * behavior across the two.
0032  *
0033  * The function must have the signature
0034  * <code>bool(*)(T const&, int)</code>
0035  * where the return value indicates whether the volume's children should be
0036  * visited, and the integer is the depth of the volume being visited.
0037  *
0038  * By default this will visit the entire "touchable" hierarchy: this may be
0039  * very expensive! If it's desired to only visit single physical volumes, mark
0040  * them as visited using a set (see unit test for example).
0041  */
0042 template<class T>
0043 class VolumeVisitor
0044 {
0045   public:
0046     //! Construct from top-level volume
0047     explicit VolumeVisitor(T const& world) : world_{world} {}
0048
0049     // Apply this visitor
0050     template<class F>
0051     inline void operator()(F&& visit);
0052
0053   private:
0054     T const& world_;
0055
0056     struct QueuedVolume
0057     {
0058         T const* pv{nullptr};
0059         int depth{0};
0060     };
0061 };
0062
0063 //---------------------------------------------------------------------------//
0064
0065 // Visit all logical volumes
0066 template<class F, class T>
0067 inline void visit_logical_volumes(F&& vis, T const& parent_vol);
0068
0069 //---------------------------------------------------------------------------//
0070 // INLINE DEFINITIONS
0071 //---------------------------------------------------------------------------//
0072 /*!
0073  * Apply this visitor.
0074  */
0075 template<class T>
0076 template<class F>
0077 void VolumeVisitor<T>::operator()(F&& visit)
0078 {
0079     using TraitsT = VolumeVisitorTraits<T>;
0080
0081     std::vector<QueuedVolume> queue;
0082     std::vector<T const*> temp_children;
0083     auto visit_impl = [&](T const& pv, int depth) {
0084         if (visit(pv, depth))
0085         {
0086             temp_children.clear();
0087             TraitsT::get_children(pv, temp_children);
0088
0089             // Append children in reverse order since we pop back
0090             for (auto iter = temp_children.rbegin();
0091                  iter != temp_children.rend();
0092                  ++iter)
0093             {
0094                 queue.push_back({*iter, depth + 1});
0095             }
0096         }
0097     };
0098
0099     // Visit the top-level physical volume
0100     visit_impl(world_, 0);
0101
0102     while (!queue.empty())
0103     {
0104         QueuedVolume qv = queue.back();
0105         queue.pop_back();
0106
0107         // Visit popped daughter
0108         visit_impl(*qv.pv, qv.depth);
0109     }
0110 }
0111
0112 //---------------------------------------------------------------------------//
0113 /*!
0114  * Visit all logical volumes, once, depth-first.
0115  */
0116 template<class F, class T>
0117 inline void visit_logical_volumes(F&& vis, T const& parent_vol)
0118 {
0119     using TraitsT = VolumeVisitorTraits<T>;
0120     using LV = typename TraitsT::LV;
0121
0122     VolumeVisitor visit_impl{parent_vol};
0123
0124     std::unordered_set<LV const*> visited;
0125     visit_impl([&vis, &visited](T const& pv, int) -> bool {
0126         auto const& lv = TraitsT::get_lv(pv);
0127         if (!visited.insert(&lv).second)
0128         {
0129             // Already visited
0130             return false;
0131         }
0132         vis(lv);
0133         return true;
0134     });
0135 }
0136
0137 //---------------------------------------------------------------------------//
0138 }  // namespace celeritas