From 6d94310182ae5cf9d1644c21ef4ebc1da7583195 Mon Sep 17 00:00:00 2001
From: max-dfinity <100170574+max-dfinity@users.noreply.github.com>
Date: Fri, 8 Nov 2024 15:50:05 -0800
Subject: [PATCH] fix(nns): Fix ballot generation inefficiency (for stable
 neurons) (#2526)

This fixes an inefficiency in how ballots are generated when neurons are
in stable memory.
---
 rs/nns/governance/src/neuron_store.rs | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/rs/nns/governance/src/neuron_store.rs b/rs/nns/governance/src/neuron_store.rs
index 5cb901df309..4083721bb74 100644
--- a/rs/nns/governance/src/neuron_store.rs
+++ b/rs/nns/governance/src/neuron_store.rs
@@ -734,6 +734,14 @@ impl NeuronStore {
     pub fn with_active_neurons_iter<R>(
         &self,
         callback: impl for<'b> FnOnce(Box<dyn Iterator<Item = Neuron> + 'b>) -> R,
+    ) -> R {
+        self.with_active_neurons_iter_sections(callback, NeuronSections::all())
+    }
+
+    fn with_active_neurons_iter_sections<R>(
+        &self,
+        callback: impl for<'b> FnOnce(Box<dyn Iterator<Item = Neuron> + 'b>) -> R,
+        sections: NeuronSections,
     ) -> R {
         if self.use_stable_memory_for_all_neurons {
             // Note, during migration, we still need heap_neurons, so we chain them onto the iterator
@@ -741,7 +749,7 @@ impl NeuronStore {
                 let now = self.now();
                 let iter = Box::new(
                     stable_store
-                        .range_neurons(..)
+                        .range_neurons_sections(.., sections)
                         .filter(|n| !n.is_inactive(now))
                         .chain(self.heap_neurons.values().cloned()),
                 ) as Box<dyn Iterator<Item = Neuron>>;
@@ -871,11 +879,14 @@ impl NeuronStore {
         };
 
         // Active neurons iterator already makes distinctions between stable and heap neurons.
-        self.with_active_neurons_iter(|iter| {
-            for neuron in iter {
-                process_neuron(neuron);
-            }
-        });
+        self.with_active_neurons_iter_sections(
+            |iter| {
+                for neuron in iter {
+                    process_neuron(neuron);
+                }
+            },
+            NeuronSections::default(),
+        );
 
         (ballots, deciding_voting_power, potential_voting_power)
     }