From 0a950eb6236c1524e128a4a4f60114611367e085 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Mon, 1 May 2023 22:00:45 -0500
Subject: [PATCH] Reworking Tally::add_filter a bit (#2501)

---
 include/openmc/tallies/tally.h      |  2 +-
 src/tallies/tally.cpp               | 29 +++++++++------
 tests/cpp_unit_tests/CMakeLists.txt |  1 +
 tests/cpp_unit_tests/test_tally.cpp | 55 +++++++++++++++++++++++++++++
 4 files changed, 75 insertions(+), 12 deletions(-)
 create mode 100644 tests/cpp_unit_tests/test_tally.cpp

diff --git a/include/openmc/tallies/tally.h b/include/openmc/tallies/tally.h
index 56a51370a8..e840a0aa14 100644
--- a/include/openmc/tallies/tally.h
+++ b/include/openmc/tallies/tally.h
@@ -67,7 +67,7 @@ public:
   //----------------------------------------------------------------------------
   // Other methods.
 
-  void add_filter(Filter* filter) { set_filters({&filter, 1}); }
+  void add_filter(Filter* filter);
 
   void init_triggers(pugi::xml_node node);
 
diff --git a/src/tallies/tally.cpp b/src/tallies/tally.cpp
index 0610c8cf13..e29384772f 100644
--- a/src/tallies/tally.cpp
+++ b/src/tallies/tally.cpp
@@ -364,20 +364,27 @@ void Tally::set_filters(gsl::span<Filter*> filters)
   auto n = filters.size();
   filters_.reserve(n);
 
-  for (int i = 0; i < n; ++i) {
-    // Add index to vector of filters
-    auto& f {filters[i]};
-    filters_.push_back(model::filter_map.at(f->id()));
-
-    // Keep track of indices for special filters.
-    if (dynamic_cast<const EnergyoutFilter*>(f)) {
-      energyout_filter_ = i;
-    } else if (dynamic_cast<const DelayedGroupFilter*>(f)) {
-      delayedgroup_filter_ = i;
-    }
+  for (auto* filter : filters) {
+    add_filter(filter);
   }
 }
 
+void Tally::add_filter(Filter* filter)
+{
+  int32_t filter_idx = model::filter_map.at(filter->id());
+  // if this filter is already present, do nothing and return
+  if (std::find(filters_.begin(), filters_.end(), filter_idx) != filters_.end())
+    return;
+
+  // Keep track of indices for special filters
+  if (dynamic_cast<const EnergyoutFilter*>(filter)) {
+    energyout_filter_ = filters_.size();
+  } else if (dynamic_cast<const DelayedGroupFilter*>(filter)) {
+    delayedgroup_filter_ = filters_.size();
+  }
+  filters_.push_back(filter_idx);
+}
+
 void Tally::set_strides()
 {
   // Set the strides.  Filters are traversed in reverse so that the last filter
diff --git a/tests/cpp_unit_tests/CMakeLists.txt b/tests/cpp_unit_tests/CMakeLists.txt
index e017b87002..50ae9a2082 100644
--- a/tests/cpp_unit_tests/CMakeLists.txt
+++ b/tests/cpp_unit_tests/CMakeLists.txt
@@ -1,6 +1,7 @@
 set(TEST_NAMES
   test_distribution
   test_file_utils
+  test_tally
   # Add additional unit test files here
 )
 
diff --git a/tests/cpp_unit_tests/test_tally.cpp b/tests/cpp_unit_tests/test_tally.cpp
new file mode 100644
index 0000000000..964d30cc42
--- /dev/null
+++ b/tests/cpp_unit_tests/test_tally.cpp
@@ -0,0 +1,55 @@
+#include "openmc/tallies/tally.h"
+#include <catch2/catch_test_macros.hpp>
+
+using namespace openmc;
+
+TEST_CASE("Test add/set_filter")
+{
+  // create a new tally object
+  Tally* tally = Tally::create();
+
+  // create a new particle filter
+  Filter* particle_filter = Filter::create("particle");
+
+  // add the particle filter to the tally
+  tally->add_filter(particle_filter);
+
+  // the filter should be added to the tally
+  REQUIRE(tally->filters().size() == 1);
+  REQUIRE(model::filter_map[particle_filter->id()] == tally->filters(0));
+
+  // add the particle filter to the tally again
+  tally->add_filter(particle_filter);
+  // the tally should have the same number of filters
+  REQUIRE(tally->filters().size() == 1);
+
+  // create a cell filter
+  Filter* cell_filter = Filter::create("cell");
+  tally->add_filter(cell_filter);
+
+  // now the size of the filters should have increased
+  REQUIRE(tally->filters().size() == 2);
+  REQUIRE(model::filter_map[cell_filter->id()] == tally->filters(1));
+
+  // if we set the filters explicitly there shouldn't be extra filters hanging
+  // around
+  tally->set_filters({&cell_filter, 1});
+
+  REQUIRE(tally->filters().size() == 1);
+  REQUIRE(model::filter_map[cell_filter->id()] == tally->filters(0));
+
+  // set filters again using both filters
+  std::vector<Filter*> filters = {cell_filter, particle_filter};
+  tally->set_filters(filters);
+
+  REQUIRE(tally->filters().size() == 2);
+  REQUIRE(model::filter_map[cell_filter->id()] == tally->filters(0));
+  REQUIRE(model::filter_map[particle_filter->id()] == tally->filters(1));
+
+  // set filters with a duplicate filter, should only add the filter to the tally once
+  filters = {cell_filter, cell_filter};
+  tally->set_filters(filters);
+  REQUIRE(tally->filters().size() == 1);
+  REQUIRE(model::filter_map[cell_filter->id()] == tally->filters(0));
+
+}
\ No newline at end of file