From 9cdab91b527b9c7741dbdf7abc0e8db031230193 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Mon, 16 Dec 2024 13:37:17 -0600
Subject: [PATCH] Re-work surface signatures to support passing particle
 geometry states

---
 include/openmc/surface.h | 45 +++++++++++++++++++++++-----------------
 src/dagmc.cpp            |  6 +++---
 src/surface.cpp          | 34 +++++++++++++++---------------
 3 files changed, 46 insertions(+), 39 deletions(-)

diff --git a/include/openmc/surface.h b/include/openmc/surface.h
index 12b59c607a7..b73701ae24b 100644
--- a/include/openmc/surface.h
+++ b/include/openmc/surface.h
@@ -57,12 +57,18 @@ class Surface {
   //! Determine the direction of a ray reflected from the surface.
   //! \param[in] r The point at which the ray is incident.
   //! \param[in] u Incident direction of the ray
-  //! \param[inout] p Pointer to the particle. Only DAGMC uses this.
+  //! \param[inout] p Pointer to a particle geometry state. Only DAGMC uses this.
   //! \return Outgoing direction of the ray
-  virtual Direction reflect(Position r, Direction u) const;
+  virtual Direction reflect(Position r, Direction u, GeometryState* p = nullptr) const;
 
+  //! Determine the direction of a ray reflected from the surface.
+  //! \param[in] r The point at which the ray is incident.
+  //! \param[in] u Incident direction of the ray
+  //! \param[inout] seed Random number seed
+  //! \param[inout] p Pointer to a particle geometry state. Only DAGMC uses this.
+  //! \return Outgoing direction of the ray
   virtual Direction diffuse_reflect(
-    Position r, Direction u, uint64_t* seed) const;
+    Position r, Direction u, uint64_t* seed, GeometryState* p = nullptr) const;
 
   //! Evaluate the equation describing the surface.
   //!
@@ -80,8 +86,9 @@ class Surface {
 
   //! Compute the local outward normal direction of the surface.
   //! \param r A 3D Cartesian coordinate.
+  //! \param[inout] p Pointer to a particle geometry state. Only DAGMC uses this.
   //! \return Normal direction
-  virtual Direction normal(Position r) const = 0;
+  virtual Direction normal(Position r, GeometryState* p = nullptr) const = 0;
 
   //! Write all information needed to reconstruct the surface to an HDF5 group.
   //! \param group_id An HDF5 group id.
@@ -111,7 +118,7 @@ class SurfaceXPlane : public CSGSurface {
   explicit SurfaceXPlane(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
   BoundingBox bounding_box(bool pos_side) const override;
 
@@ -129,7 +136,7 @@ class SurfaceYPlane : public CSGSurface {
   explicit SurfaceYPlane(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
   BoundingBox bounding_box(bool pos_side) const override;
 
@@ -147,7 +154,7 @@ class SurfaceZPlane : public CSGSurface {
   explicit SurfaceZPlane(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
   BoundingBox bounding_box(bool pos_side) const override;
 
@@ -165,7 +172,7 @@ class SurfacePlane : public CSGSurface {
   explicit SurfacePlane(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
 
   double A_, B_, C_, D_;
@@ -183,7 +190,7 @@ class SurfaceXCylinder : public CSGSurface {
   explicit SurfaceXCylinder(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
   BoundingBox bounding_box(bool pos_side) const override;
 
@@ -202,7 +209,7 @@ class SurfaceYCylinder : public CSGSurface {
   explicit SurfaceYCylinder(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
   BoundingBox bounding_box(bool pos_side) const override;
 
@@ -221,7 +228,7 @@ class SurfaceZCylinder : public CSGSurface {
   explicit SurfaceZCylinder(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
   BoundingBox bounding_box(bool pos_side) const override;
 
@@ -240,7 +247,7 @@ class SurfaceSphere : public CSGSurface {
   explicit SurfaceSphere(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
   BoundingBox bounding_box(bool pos_side) const override;
 
@@ -259,7 +266,7 @@ class SurfaceXCone : public CSGSurface {
   explicit SurfaceXCone(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
 
   double x0_, y0_, z0_, radius_sq_;
@@ -277,7 +284,7 @@ class SurfaceYCone : public CSGSurface {
   explicit SurfaceYCone(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
 
   double x0_, y0_, z0_, radius_sq_;
@@ -295,7 +302,7 @@ class SurfaceZCone : public CSGSurface {
   explicit SurfaceZCone(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
 
   double x0_, y0_, z0_, radius_sq_;
@@ -313,7 +320,7 @@ class SurfaceQuadric : public CSGSurface {
   explicit SurfaceQuadric(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
 
   // Ax^2 + By^2 + Cz^2 + Dxy + Eyz + Fxz + Gx + Hy + Jz + K = 0
@@ -331,7 +338,7 @@ class SurfaceXTorus : public CSGSurface {
   explicit SurfaceXTorus(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
 
   double x0_, y0_, z0_, A_, B_, C_;
@@ -348,7 +355,7 @@ class SurfaceYTorus : public CSGSurface {
   explicit SurfaceYTorus(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
 
   double x0_, y0_, z0_, A_, B_, C_;
@@ -365,7 +372,7 @@ class SurfaceZTorus : public CSGSurface {
   explicit SurfaceZTorus(pugi::xml_node surf_node);
   double evaluate(Position r) const override;
   double distance(Position r, Direction u, bool coincident) const override;
-  Direction normal(Position r) const override;
+  Direction normal(Position r, GeometryState* p = nullptr) const override;
   void to_hdf5_inner(hid_t group_id) const override;
 
   double x0_, y0_, z0_, A_, B_, C_;
diff --git a/src/dagmc.cpp b/src/dagmc.cpp
index 525acd44143..0836307847a 100644
--- a/src/dagmc.cpp
+++ b/src/dagmc.cpp
@@ -747,7 +747,7 @@ double DAGSurface::distance(Position r, Direction u, bool coincident) const
   return dist;
 }
 
-Direction DAGSurface::normal(Position r) const
+Direction DAGSurface::normal(Position r, GeometryState* p) const
 {
   moab::ErrorCode rval;
   moab::EntityHandle surf = dagmc_ptr_->entity_by_index(2, dag_index_);
@@ -758,13 +758,13 @@ Direction DAGSurface::normal(Position r) const
   return dir;
 }
 
-Direction DAGSurface::reflect(Position r, Direction u) const
+Direction DAGSurface::reflect(Position r, Direction u, GeometryState* p) const
 {
   moab::ErrorCode rval;
   moab::EntityHandle surf = dagmc_ptr_->entity_by_index(2, dag_index_);
   double pnt[3] = {r.x, r.y, r.z};
   double dir[3];
-  rval = dagmc_ptr_->get_angle(surf, pnt, dir);
+  rval = dagmc_ptr_->get_angle(surf, pnt, dir, p->history());
   MB_CHK_ERR_CONT(rval);
   return u.reflect(dir);
 }
diff --git a/src/surface.cpp b/src/surface.cpp
index 110fdb41083..21cb74052a2 100644
--- a/src/surface.cpp
+++ b/src/surface.cpp
@@ -130,7 +130,7 @@ bool Surface::sense(Position r, Direction u) const
   return f > 0.0;
 }
 
-Direction Surface::reflect(Position r, Direction u) const
+Direction Surface::reflect(Position r, Direction u, GeometryState* p) const
 {
   // Determine projection of direction onto normal and squared magnitude of
   // normal.
@@ -141,7 +141,7 @@ Direction Surface::reflect(Position r, Direction u) const
 }
 
 Direction Surface::diffuse_reflect(
-  Position r, Direction u, uint64_t* seed) const
+  Position r, Direction u, uint64_t* seed, GeometryState* p) const
 {
   // Diffuse reflect direction according to the normal.
   // cosine distribution
@@ -233,7 +233,7 @@ double SurfaceXPlane::distance(Position r, Direction u, bool coincident) const
   return axis_aligned_plane_distance<0>(r, u, coincident, x0_);
 }
 
-Direction SurfaceXPlane::normal(Position r) const
+Direction SurfaceXPlane::normal(Position r, GeometryState* p) const
 {
   return {1., 0., 0.};
 }
@@ -273,7 +273,7 @@ double SurfaceYPlane::distance(Position r, Direction u, bool coincident) const
   return axis_aligned_plane_distance<1>(r, u, coincident, y0_);
 }
 
-Direction SurfaceYPlane::normal(Position r) const
+Direction SurfaceYPlane::normal(Position r, GeometryState* p) const
 {
   return {0., 1., 0.};
 }
@@ -313,7 +313,7 @@ double SurfaceZPlane::distance(Position r, Direction u, bool coincident) const
   return axis_aligned_plane_distance<2>(r, u, coincident, z0_);
 }
 
-Direction SurfaceZPlane::normal(Position r) const
+Direction SurfaceZPlane::normal(Position r, GeometryState* p) const
 {
   return {0., 0., 1.};
 }
@@ -362,7 +362,7 @@ double SurfacePlane::distance(Position r, Direction u, bool coincident) const
   }
 }
 
-Direction SurfacePlane::normal(Position r) const
+Direction SurfacePlane::normal(Position r, GeometryState* p) const
 {
   return {A_, B_, C_};
 }
@@ -474,7 +474,7 @@ double SurfaceXCylinder::distance(
     r, u, coincident, y0_, z0_, radius_);
 }
 
-Direction SurfaceXCylinder::normal(Position r) const
+Direction SurfaceXCylinder::normal(Position r, GeometryState* p) const
 {
   return axis_aligned_cylinder_normal<0, 1, 2>(r, y0_, z0_);
 }
@@ -517,7 +517,7 @@ double SurfaceYCylinder::distance(
     r, u, coincident, x0_, z0_, radius_);
 }
 
-Direction SurfaceYCylinder::normal(Position r) const
+Direction SurfaceYCylinder::normal(Position r, GeometryState* p) const
 {
   return axis_aligned_cylinder_normal<1, 0, 2>(r, x0_, z0_);
 }
@@ -561,7 +561,7 @@ double SurfaceZCylinder::distance(
     r, u, coincident, x0_, y0_, radius_);
 }
 
-Direction SurfaceZCylinder::normal(Position r) const
+Direction SurfaceZCylinder::normal(Position r, GeometryState* p) const
 {
   return axis_aligned_cylinder_normal<2, 0, 1>(r, x0_, y0_);
 }
@@ -639,7 +639,7 @@ double SurfaceSphere::distance(Position r, Direction u, bool coincident) const
   }
 }
 
-Direction SurfaceSphere::normal(Position r) const
+Direction SurfaceSphere::normal(Position r, GeometryState* p) const
 {
   return {2.0 * (r.x - x0_), 2.0 * (r.y - y0_), 2.0 * (r.z - z0_)};
 }
@@ -769,7 +769,7 @@ double SurfaceXCone::distance(Position r, Direction u, bool coincident) const
     r, u, coincident, x0_, y0_, z0_, radius_sq_);
 }
 
-Direction SurfaceXCone::normal(Position r) const
+Direction SurfaceXCone::normal(Position r, GeometryState* p) const
 {
   return axis_aligned_cone_normal<0, 1, 2>(r, x0_, y0_, z0_, radius_sq_);
 }
@@ -801,7 +801,7 @@ double SurfaceYCone::distance(Position r, Direction u, bool coincident) const
     r, u, coincident, y0_, x0_, z0_, radius_sq_);
 }
 
-Direction SurfaceYCone::normal(Position r) const
+Direction SurfaceYCone::normal(Position r, GeometryState* p) const
 {
   return axis_aligned_cone_normal<1, 0, 2>(r, y0_, x0_, z0_, radius_sq_);
 }
@@ -833,7 +833,7 @@ double SurfaceZCone::distance(Position r, Direction u, bool coincident) const
     r, u, coincident, z0_, x0_, y0_, radius_sq_);
 }
 
-Direction SurfaceZCone::normal(Position r) const
+Direction SurfaceZCone::normal(Position r, GeometryState* p) const
 {
   return axis_aligned_cone_normal<2, 0, 1>(r, z0_, x0_, y0_, radius_sq_);
 }
@@ -935,7 +935,7 @@ double SurfaceQuadric::distance(
   return d;
 }
 
-Direction SurfaceQuadric::normal(Position r) const
+Direction SurfaceQuadric::normal(Position r, GeometryState* p) const
 {
   const double& x = r.x;
   const double& y = r.y;
@@ -1038,7 +1038,7 @@ double SurfaceXTorus::distance(Position r, Direction u, bool coincident) const
   return torus_distance(y, z, x, u.y, u.z, u.x, A_, B_, C_, coincident);
 }
 
-Direction SurfaceXTorus::normal(Position r) const
+Direction SurfaceXTorus::normal(Position r, GeometryState* p) const
 {
   // reduce the expansion of the full form for torus
   double x = r.x - x0_;
@@ -1091,7 +1091,7 @@ double SurfaceYTorus::distance(Position r, Direction u, bool coincident) const
   return torus_distance(x, z, y, u.x, u.z, u.y, A_, B_, C_, coincident);
 }
 
-Direction SurfaceYTorus::normal(Position r) const
+Direction SurfaceYTorus::normal(Position r, GeometryState* p) const
 {
   // reduce the expansion of the full form for torus
   double x = r.x - x0_;
@@ -1144,7 +1144,7 @@ double SurfaceZTorus::distance(Position r, Direction u, bool coincident) const
   return torus_distance(x, y, z, u.x, u.y, u.z, A_, B_, C_, coincident);
 }
 
-Direction SurfaceZTorus::normal(Position r) const
+Direction SurfaceZTorus::normal(Position r, GeometryState* p) const
 {
   // reduce the expansion of the full form for torus
   double x = r.x - x0_;