From 0817a792d92b89bb44b03b4e120be248fc3a4b23 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Mon, 16 Aug 2021 16:21:12 -0500
Subject: [PATCH] Applying a new stack to simplify parent cell traversal alg.

---
 include/openmc/cell.h | 28 ++++++++++++++++++++++++++++
 src/cell.cpp          | 23 ++++++++++-------------
 2 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/include/openmc/cell.h b/include/openmc/cell.h
index 0502c7517..856337e46 100644
--- a/include/openmc/cell.h
+++ b/include/openmc/cell.h
@@ -6,6 +6,7 @@
 #include <limits>
 #include <string>
 #include <unordered_map>
+#include <set>
 
 #include "hdf5.h"
 #include "pugixml.hpp"
@@ -319,6 +320,33 @@ struct ParentCell {
   gsl::index lattice_index;
 };
 
+struct ParentCellStack {
+
+  void push_back(int32_t search_universe, const ParentCell& pc) {
+    parent_cells_.push_back(pc);
+    // add parent cell to the set of cells we've visited for this search universe
+    visited_cells_[search_universe].insert(pc);
+  }
+
+  void pop_back() {
+    visited_cells_[this->current_univ()].clear();
+    parent_cells_.pop_back();
+  }
+
+  bool visited(int32_t univ_idx, const ParentCell& parent_cell) {
+    return visited_cells_[univ_idx].count(parent_cell) != 0;
+  }
+
+  int32_t current_univ() const { return model::cells[parent_cells_.back().cell_index]->universe_; }
+
+  bool empty() const { return parent_cells_.empty(); }
+
+  std::vector<ParentCell> parent_cells() { return parent_cells_; }
+
+  std::vector<ParentCell> parent_cells_;
+  std::unordered_map<int32_t, std::set<ParentCell>> visited_cells_;
+};
+
 struct ParentCellHash {
   std::size_t operator()(const ParentCell& p) const
   {
diff --git a/src/cell.cpp b/src/cell.cpp
index 01282e35b..568ec7901 100644
--- a/src/cell.cpp
+++ b/src/cell.cpp
@@ -5,7 +5,6 @@
 #include <cctype>
 #include <cmath>
 #include <iterator>
-#include <set>
 #include <sstream>
 #include <string>
 
@@ -1223,26 +1222,21 @@ extern "C" int openmc_cell_set_name(int32_t index, const char* name)
 void
 Cell::find_parent_cells(vector<ParentCell>& parent_cells, int32_t instance) const
 {
-  // clear out the list of parent cells
-  parent_cells.clear();
+  ParentCellStack stack;
   // start with this cell's universe
-  int32_t visited_univ = -1;
-  int32_t prev_univ_idx = -1;
+  int32_t prev_univ_idx = C_NONE;
   int32_t univ_idx = this->universe_;
-  int32_t search_instance = -1;
-
-  std::set<ParentCell> visited_cells;
 
   while (true) {
     const auto& univ = model::universes[univ_idx];
     prev_univ_idx = univ_idx;
     for (const auto& cell : model::cells) {
       // if this is in our current set of visited cells, move on
-      if (visited_cells.count({model::cell_map[cell->id_], -1})) continue;
+      if (stack.visited(univ_idx, {model::cell_map[cell->id_], C_NONE})) continue;
 
       if (cell->type_ == Fill::UNIVERSE) {
         if(cell->fill_ == univ_idx) {
-          parent_cells.push_back({model::cell_map[cell->id_], -1});
+          stack.push_back(univ_idx, {model::cell_map[cell->id_], C_NONE});
           univ_idx = cell->universe_;
         }
       } else if (cell->type_ == Fill::LATTICE) {
@@ -1258,9 +1252,9 @@ Cell::find_parent_cells(vector<ParentCell>& parent_cells, int32_t instance) cons
 
           lat_it++;
 
-          if (visited_cells.count({model::cell_map[cell->id_], lattice_idx})) continue;
+          if (stack.visited(univ_idx, {model::cell_map[cell->id_], lattice_idx})) continue;
 
-          parent_cells.push_back({model::cell_map[cell->id_], lattice_idx});
+          stack.push_back(univ_idx, {model::cell_map[cell->id_], lattice_idx});
           univ_idx = cell->universe_;
         }
       }
@@ -1279,10 +1273,13 @@ Cell::find_parent_cells(vector<ParentCell>& parent_cells, int32_t instance) cons
 
     // if we don't find a suitable update, adjust the stack and continue
     if (univ_idx == model::root_universe || univ_idx == prev_univ_idx) {
-     }
+      stack.pop_back();
+      univ_idx = stack.empty() ? this->universe_ : stack.current_univ();
+    }
 
   } // end while
 
+  parent_cells = stack.parent_cells();
   std::reverse(parent_cells.begin(), parent_cells.end());
 }