From 34a35d832a1cfe1fec304eace56676265d5b6f3d Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Mon, 7 Mar 2022 22:02:24 -0600
Subject: [PATCH] Fix cell instance counting bug with same lattice in multiple
 cells

---
 src/geometry.cpp     |  1 +
 src/geometry_aux.cpp |  3 ++-
 src/lattice.cpp      | 12 ++++++++++++
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/geometry.cpp b/src/geometry.cpp
index acce3fc78e..216650ff18 100644
--- a/src/geometry.cpp
+++ b/src/geometry.cpp
@@ -86,6 +86,7 @@ int cell_instance_at_level(const Particle& p, int level)
     if (c_i.type_ == Fill::UNIVERSE) {
       instance += c_i.offset_[c.distribcell_index_];
     } else if (c_i.type_ == Fill::LATTICE) {
+      instance += c_i.offset_[c.distribcell_index_];
       auto& lat {*model::lattices[p.coord(i + 1).lattice]};
       const auto& i_xyz {p.coord(i + 1).lattice_i};
       if (lat.are_valid_indices(i_xyz)) {
diff --git a/src/geometry_aux.cpp b/src/geometry_aux.cpp
index 842822d05a..ddce2469b2 100644
--- a/src/geometry_aux.cpp
+++ b/src/geometry_aux.cpp
@@ -412,8 +412,9 @@ void prepare_distribcell(const std::vector<int32_t>* user_distribcells)
             search_univ, target_univ_id, univ_count_memo);
 
         } else if (c.type_ == Fill::LATTICE) {
+          c.offset_[map] = offset;
           Lattice& lat = *model::lattices[c.fill_];
-          offset =
+          offset +=
             lat.fill_offset_table(offset, target_univ_id, map, univ_count_memo);
         }
       }
diff --git a/src/lattice.cpp b/src/lattice.cpp
index 3875b2fa72..fa2e2828ec 100644
--- a/src/lattice.cpp
+++ b/src/lattice.cpp
@@ -101,6 +101,18 @@ void Lattice::adjust_indices()
 int32_t Lattice::fill_offset_table(int32_t offset, int32_t target_univ_id,
   int map, std::unordered_map<int32_t, int32_t>& univ_count_memo)
 {
+  // If the offsets have already been determined for this "map", don't bother
+  // recalculating all of them and just return the total offset. Note that the
+  // offsets_ array doesn't actually include the offset accounting for the last
+  // universe, so we get the before-last offset for the given map and then
+  // explicitly add the count for the last universe.
+  if (offsets_[map * universes_.size()] != C_NONE) {
+    int last_offset = offsets_[(map + 1) * universes_.size() - 1];
+    int last_univ = universes_.back();
+    return last_offset +
+           count_universe_instances(last_univ, target_univ_id, univ_count_memo);
+  }
+
   for (LatticeIter it = begin(); it != end(); ++it) {
     offsets_[map * universes_.size() + it.indx_] = offset;
     offset += count_universe_instances(*it, target_univ_id, univ_count_memo);