From ccb76b3dbf7afcc3f3c808a9511a03705a0553c4 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Wed, 13 Sep 2023 21:28:11 -0500
Subject: [PATCH] Refactor of translation from mesh handle to index for DAGMC
 geometry. (#2687)

---
 include/openmc/dagmc.h | 19 +++++++++++++++++--
 src/dagmc.cpp          | 32 +++++++++++++++++++++++++++-----
 2 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/include/openmc/dagmc.h b/include/openmc/dagmc.h
index 5451474d04..5ed426e5af 100644
--- a/include/openmc/dagmc.h
+++ b/include/openmc/dagmc.h
@@ -37,6 +37,8 @@ class DAGSurface : public Surface {
 public:
   DAGSurface(std::shared_ptr<moab::DagMC> dag_ptr, int32_t dag_idx);
 
+  moab::EntityHandle mesh_handle() const;
+
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
   Direction normal(Position r) const override;
@@ -45,7 +47,7 @@ public:
   inline void to_hdf5_inner(hid_t group_id) const override {};
 
   // Accessor methods
-  const std::shared_ptr<moab::DagMC>& dagmc_ptr() const { return dagmc_ptr_; }
+  moab::DagMC* dagmc_ptr() const { return dagmc_ptr_.get(); }
   int32_t dag_index() const { return dag_index_; }
 
 private:
@@ -57,6 +59,8 @@ class DAGCell : public Cell {
 public:
   DAGCell(std::shared_ptr<moab::DagMC> dag_ptr, int32_t dag_idx);
 
+  moab::EntityHandle mesh_handle() const;
+
   bool contains(Position r, Direction u, int32_t on_surface) const override;
 
   std::pair<double, int32_t> distance(
@@ -67,7 +71,7 @@ public:
   void to_hdf5_inner(hid_t group_id) const override;
 
   // Accessor methods
-  const std::shared_ptr<moab::DagMC>& dagmc_ptr() const { return dagmc_ptr_; }
+  moab::DagMC* dagmc_ptr() const { return dagmc_ptr_.get(); }
   int32_t dag_index() const { return dag_index_; }
 
 private:
@@ -122,6 +126,16 @@ public:
   void legacy_assign_material(
     std::string mat_string, std::unique_ptr<DAGCell>& c) const;
 
+  //! Return the index into the model cells vector for a given DAGMC volume
+  //! handle in the universe
+  //! \param[in] vol MOAB handle to the DAGMC volume set
+  int32_t cell_index(moab::EntityHandle vol) const;
+
+  //! Return the index into the model surfaces vector for a given DAGMC surface
+  //! handle in the universe
+  //! \param[in] surf MOAB handle to the DAGMC surface set
+  int32_t surface_index(moab::EntityHandle surf) const;
+
   //! Generate a string representing the ranges of IDs present in the DAGMC
   //! model. Contiguous chunks of IDs are represented as a range (i.e. 1-10). If
   //! there is a single ID a chunk, it will be represented as a single number
@@ -142,6 +156,7 @@ public:
                             //!< universe in OpenMC's surface vector
 
   // Accessors
+  moab::DagMC* dagmc_ptr() const { return dagmc_instance_.get(); }
   bool has_graveyard() const { return has_graveyard_; }
 
 private:
diff --git a/src/dagmc.cpp b/src/dagmc.cpp
index 35522d3caf..e245d4475d 100644
--- a/src/dagmc.cpp
+++ b/src/dagmc.cpp
@@ -329,6 +329,20 @@ void DAGUniverse::init_geometry()
   } // end surface loop
 }
 
+int32_t DAGUniverse::cell_index(moab::EntityHandle vol) const
+{
+  // return the index of the volume in the DAGMC instance and then
+  // adjust by the offset into the model cells for this DAGMC universe
+  return dagmc_ptr()->index_by_handle(vol) + cell_idx_offset_;
+}
+
+int32_t DAGUniverse::surface_index(moab::EntityHandle surf) const
+{
+  // return the index of the surface in the DAGMC instance and then
+  // adjust by the offset into the model cells for this DAGMC universe
+  return dagmc_ptr()->index_by_handle(surf) + surf_idx_offset_;
+}
+
 std::string DAGUniverse::dagmc_ids_for_dim(int dim) const
 {
   // generate a vector of ids
@@ -643,6 +657,11 @@ bool DAGCell::contains(Position r, Direction u, int32_t on_surface) const
   return result;
 }
 
+moab::EntityHandle DAGCell::mesh_handle() const
+{
+  return dagmc_ptr()->entity_by_index(3, dag_index());
+}
+
 void DAGCell::to_hdf5_inner(hid_t group_id) const
 {
   write_string(group_id, "geom_type", "dagmc", false);
@@ -668,6 +687,11 @@ DAGSurface::DAGSurface(std::shared_ptr<moab::DagMC> dag_ptr, int32_t dag_idx)
   geom_type_ = GeometryType::DAG;
 } // empty constructor
 
+moab::EntityHandle DAGSurface::mesh_handle() const
+{
+  return dagmc_ptr()->entity_by_index(2, dag_index());
+}
+
 double DAGSurface::evaluate(Position r) const
 {
   return 0.0;
@@ -747,10 +771,8 @@ int32_t next_cell(int32_t surf, int32_t curr_cell, int32_t univ)
   auto cellp = dynamic_cast<DAGCell*>(model::cells[curr_cell].get());
   auto univp = static_cast<DAGUniverse*>(model::universes[univ].get());
 
-  moab::EntityHandle surf_handle =
-    surfp->dagmc_ptr()->entity_by_index(2, surfp->dag_index());
-  moab::EntityHandle curr_vol =
-    cellp->dagmc_ptr()->entity_by_index(3, cellp->dag_index());
+  moab::EntityHandle surf_handle = surfp->mesh_handle();
+  moab::EntityHandle curr_vol = cellp->mesh_handle();
 
   moab::EntityHandle new_vol;
   moab::ErrorCode rval =
@@ -758,7 +780,7 @@ int32_t next_cell(int32_t surf, int32_t curr_cell, int32_t univ)
   if (rval != moab::MB_SUCCESS)
     return -1;
 
-  return cellp->dagmc_ptr()->index_by_handle(new_vol) + univp->cell_idx_offset_;
+  return univp->cell_index(new_vol);
 }
 
 } // namespace openmc