From adf2047b6ee8cc0791b331196264bc252dc2403d Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Thu, 9 Jun 2022 17:42:32 -0500
Subject: [PATCH] Refactor of DAGMC next_cell

---
 include/openmc/dagmc.h |  3 +--
 src/dagmc.cpp          | 25 +++++++++++++++----------
 src/particle.cpp       |  8 +-------
 3 files changed, 17 insertions(+), 19 deletions(-)

diff --git a/include/openmc/dagmc.h b/include/openmc/dagmc.h
index 993076347..5451474d0 100644
--- a/include/openmc/dagmc.h
+++ b/include/openmc/dagmc.h
@@ -168,8 +168,7 @@ private:
 // Non-member functions
 //==============================================================================
 
-int32_t next_cell(
-  DAGUniverse* dag_univ, DAGCell* cur_cell, DAGSurface* surf_xed);
+int32_t next_cell(int32_t surf, int32_t curr_cell, int32_t univ);
 
 } // namespace openmc
 
diff --git a/src/dagmc.cpp b/src/dagmc.cpp
index 01ce8ef11..243c62e7f 100644
--- a/src/dagmc.cpp
+++ b/src/dagmc.cpp
@@ -728,19 +728,21 @@ void check_dagmc_root_univ()
   }
 }
 
-int32_t next_cell(
-  DAGUniverse* dag_univ, DAGCell* cur_cell, DAGSurface* surf_xed)
-{
-  moab::EntityHandle surf =
-    surf_xed->dagmc_ptr()->entity_by_index(2, surf_xed->dag_index());
-  moab::EntityHandle vol =
-    cur_cell->dagmc_ptr()->entity_by_index(3, cur_cell->dag_index());
+int32_t next_cell(int32_t surf, int32_t curr_cell, int32_t univ) {
+  auto surfp = dynamic_cast<DAGSurface*>(model::surfaces[surf - 1].get());
+  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 new_vol;
-  cur_cell->dagmc_ptr()->next_vol(surf, vol, new_vol);
+  cellp->dagmc_ptr()->next_vol(surf_handle, curr_vol, new_vol);
 
-  return cur_cell->dagmc_ptr()->index_by_handle(new_vol) +
-         dag_univ->cell_idx_offset_;
+  return cellp->dagmc_ptr()->index_by_handle(new_vol) +
+         univp->cell_idx_offset_;
 }
 
 } // namespace openmc
@@ -756,8 +758,11 @@ void read_dagmc_universes(pugi::xml_node node)
                 "with DAGMC");
   }
 };
+
 void check_dagmc_root_univ() {};
 
+int32_t next_cell(int32_t surf, int32_t curr_cell, int32_t univ);
+
 } // namespace openmc
 
 #endif // DAGMC
diff --git a/src/particle.cpp b/src/particle.cpp
index cb959e14d..1ec95a4ae 100644
--- a/src/particle.cpp
+++ b/src/particle.cpp
@@ -433,13 +433,7 @@ void Particle::cross_surface()
 #ifdef DAGMC
   // in DAGMC, we know what the next cell should be
   if (surf->geom_type_ == GeometryType::DAG) {
-    auto surfp = dynamic_cast<DAGSurface*>(surf);
-    auto cellp =
-      dynamic_cast<DAGCell*>(model::cells[cell_last(n_coord() - 1)].get());
-    auto univp = static_cast<DAGUniverse*>(
-      model::universes[coord(n_coord() - 1).universe].get());
-    // determine the next cell for this crossing
-    int32_t i_cell = next_cell(univp, cellp, surfp) - 1;
+    int32_t i_cell = next_cell(i_surface, cell_last(n_coord() - 1), lowest_coord().universe) - 1;
     // save material and temp
     material_last() = material();
     sqrtkT_last() = sqrtkT();