From c5cef78ade2197817bc36d674c9f6f6ff32db9f5 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Thu, 7 Nov 2019 14:37:23 -0600
Subject: [PATCH] Adding memo for get_all_cells/get_all_materials.

---
 openmc/cell.py     | 16 +++++++++++-----
 openmc/geometry.py |  9 +++++++--
 openmc/lattice.py  | 17 ++++++++++++-----
 openmc/universe.py | 16 +++++++++++-----
 4 files changed, 41 insertions(+), 17 deletions(-)

diff --git a/openmc/cell.py b/openmc/cell.py
index 3de20f259f..8c6b5a9f55 100644
--- a/openmc/cell.py
+++ b/openmc/cell.py
@@ -350,7 +350,7 @@ class Cell(IDManagerMixin):
 
         return nuclides
 
-    def get_all_cells(self):
+    def get_all_cells(self, memo=None):
         """Return all cells that are contained within this one if it is filled with a
         universe or lattice
 
@@ -364,12 +364,18 @@ class Cell(IDManagerMixin):
 
         cells = OrderedDict()
 
+        if memo and id(self) in memo:
+            return cells
+
+        if memo is not None:
+            memo.add(id(self))
+
         if self.fill_type in ('universe', 'lattice'):
-            cells.update(self.fill.get_all_cells())
+            cells.update(self.fill.get_all_cells(memo))
 
         return cells
 
-    def get_all_materials(self):
+    def get_all_materials(self, memo=None):
         """Return all materials that are contained within the cell
 
         Returns
@@ -388,9 +394,9 @@ class Cell(IDManagerMixin):
                     materials[m.id] = m
         else:
             # Append all Cells in each Cell in the Universe to the dictionary
-            cells = self.get_all_cells()
+            cells = self.get_all_cells(memo)
             for cell in cells.values():
-                materials.update(cell.get_all_materials())
+                materials.update(cell.get_all_materials(memo))
 
         return materials
 
diff --git a/openmc/geometry.py b/openmc/geometry.py
index 89016b722d..196c388452 100644
--- a/openmc/geometry.py
+++ b/openmc/geometry.py
@@ -273,8 +273,11 @@ class Geometry(object):
             Dictionary mapping cell IDs to :class:`openmc.Cell` instances
 
         """
+
+        memo = set()
+
         if self.root_universe is not None:
-            return self.root_universe.get_all_cells()
+            return self.root_universe.get_all_cells(memo)
         else:
             return []
 
@@ -303,7 +306,9 @@ class Geometry(object):
             instances
 
         """
-        return self.root_universe.get_all_materials()
+        memo = set()
+
+        return self.root_universe.get_all_materials(memo)
 
     def get_all_material_cells(self):
         """Return all cells filled by a material
diff --git a/openmc/lattice.py b/openmc/lattice.py
index 077fb26694..f4c4802302 100644
--- a/openmc/lattice.py
+++ b/openmc/lattice.py
@@ -332,7 +332,7 @@ class Lattice(IDManagerMixin, metaclass=ABCMeta):
 
         return nuclides
 
-    def get_all_cells(self):
+    def get_all_cells(self, memo=None):
         """Return all cells that are contained within the lattice
 
         Returns
@@ -344,14 +344,21 @@ class Lattice(IDManagerMixin, metaclass=ABCMeta):
         """
 
         cells = OrderedDict()
+
+        if memo and id(self) in memo:
+            return cells
+
+        if memo is not None:
+            memo.add(id(self))
+
         unique_universes = self.get_unique_universes()
 
         for universe_id, universe in unique_universes.items():
-            cells.update(universe.get_all_cells())
+            cells.update(universe.get_all_cells(memo))
 
         return cells
 
-    def get_all_materials(self):
+    def get_all_materials(self, memo=None):
         """Return all materials that are contained within the lattice
 
         Returns
@@ -365,9 +372,9 @@ class Lattice(IDManagerMixin, metaclass=ABCMeta):
         materials = OrderedDict()
 
         # Append all Cells in each Cell in the Universe to the dictionary
-        cells = self.get_all_cells()
+        cells = self.get_all_cells(memo)
         for cell_id, cell in cells.items():
-            materials.update(cell.get_all_materials())
+            materials.update(cell.get_all_materials(memo))
 
         return materials
 
diff --git a/openmc/universe.py b/openmc/universe.py
index 418fe1236d..24d1413a94 100644
--- a/openmc/universe.py
+++ b/openmc/universe.py
@@ -418,7 +418,7 @@ class Universe(IDManagerMixin):
 
         return nuclides
 
-    def get_all_cells(self):
+    def get_all_cells(self, memo=None):
         """Return all cells that are contained within the universe
 
         Returns
@@ -431,16 +431,22 @@ class Universe(IDManagerMixin):
 
         cells = OrderedDict()
 
+        if memo and id(self) in memo:
+            return cells
+
+        if memo is not None:
+            memo.add(id(self))
+
         # Add this Universe's cells to the dictionary
         cells.update(self._cells)
 
         # Append all Cells in each Cell in the Universe to the dictionary
         for cell in self._cells.values():
-            cells.update(cell.get_all_cells())
+            cells.update(cell.get_all_cells(memo))
 
         return cells
 
-    def get_all_materials(self):
+    def get_all_materials(self, memo=None):
         """Return all materials that are contained within the universe
 
         Returns
@@ -454,9 +460,9 @@ class Universe(IDManagerMixin):
         materials = OrderedDict()
 
         # Append all Cells in each Cell in the Universe to the dictionary
-        cells = self.get_all_cells()
+        cells = self.get_all_cells(memo)
         for cell in cells.values():
-            materials.update(cell.get_all_materials())
+            materials.update(cell.get_all_materials(memo))
 
         return materials