From babfa674d2e249d38d548a38fa9bb5cc3219a881 Mon Sep 17 00:00:00 2001
From: Max Summe <maximilian.summe@dfinity.org>
Date: Wed, 20 Nov 2024 11:26:54 -0800
Subject: [PATCH] Fix types to avoid performance penalty before migration to
 stable memory

---
 rs/nns/governance/src/governance/tla/mod.rs |  4 +++-
 rs/nns/governance/src/neuron_store.rs       | 24 ++++++++++++---------
 rs/nns/governance/tests/governance.rs       | 18 ++++++++++------
 3 files changed, 28 insertions(+), 18 deletions(-)

diff --git a/rs/nns/governance/src/governance/tla/mod.rs b/rs/nns/governance/src/governance/tla/mod.rs
index 7b608e28608..32702e36edf 100644
--- a/rs/nns/governance/src/governance/tla/mod.rs
+++ b/rs/nns/governance/src/governance/tla/mod.rs
@@ -1,6 +1,7 @@
 use itertools::Itertools;
 use std::{
     collections::{BTreeMap, BTreeSet},
+    ops::Deref,
     thread,
 };
 
@@ -33,7 +34,8 @@ pub use claim_neuron::claim_neuron_desc;
 fn neuron_global(gov: &Governance) -> TlaValue {
     let neuron_map: BTreeMap<u64, TlaValue> = with_stable_neuron_store(|store| {
         gov.neuron_store.with_active_neurons_iter(|iter| {
-            iter.chain(store.range_neurons(std::ops::RangeFull))
+            iter.map(|n| n.deref().clone())
+                .chain(store.range_neurons(std::ops::RangeFull))
                 .map(|neuron| {
                     (
                         neuron.id().id,
diff --git a/rs/nns/governance/src/neuron_store.rs b/rs/nns/governance/src/neuron_store.rs
index f2793daf096..619c470119a 100644
--- a/rs/nns/governance/src/neuron_store.rs
+++ b/rs/nns/governance/src/neuron_store.rs
@@ -807,14 +807,14 @@ impl NeuronStore {
     }
     pub fn with_active_neurons_iter<R>(
         &self,
-        callback: impl for<'b> FnOnce(Box<dyn Iterator<Item = Neuron> + 'b>) -> R,
+        callback: impl for<'b> FnOnce(Box<dyn Iterator<Item = Cow<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,
+        callback: impl for<'b> FnOnce(Box<dyn Iterator<Item = Cow<Neuron>> + 'b>) -> R,
         sections: NeuronSections,
     ) -> R {
         if self.use_stable_memory_for_all_neurons {
@@ -825,13 +825,13 @@ impl NeuronStore {
                     stable_store
                         .range_neurons_sections(.., sections)
                         .filter(|n| !n.is_inactive(now))
-                        .chain(self.heap_neurons.values().cloned()),
-                ) as Box<dyn Iterator<Item = Neuron>>;
+                        .map(Cow::Owned)
+                        .chain(self.heap_neurons.values().map(Cow::Borrowed)),
+                );
                 callback(iter)
             })
         } else {
-            let iter =
-                Box::new(self.heap_neurons.values().cloned()) as Box<dyn Iterator<Item = Neuron>>;
+            let iter = Box::new(self.heap_neurons.values().map(Cow::Borrowed));
             callback(iter)
         }
     }
@@ -870,9 +870,13 @@ impl NeuronStore {
     fn filter_map_active_neurons<R>(
         &self,
         filter: impl Fn(&Neuron) -> bool,
-        f: impl FnMut(Neuron) -> R,
+        f: impl Fn(&Neuron) -> R,
     ) -> Vec<R> {
-        self.with_active_neurons_iter(|iter| iter.filter(|n| filter(n)).map(f).collect())
+        self.with_active_neurons_iter(|iter| {
+            iter.filter(|n| filter(n.as_ref()))
+                .map(|n| f(n.as_ref()))
+                .collect()
+        })
     }
 
     fn is_active_neurons_fund_neuron(neuron: &Neuron, now: u64) -> bool {
@@ -932,7 +936,7 @@ impl NeuronStore {
         let mut deciding_voting_power: u128 = 0;
         let mut potential_voting_power: u128 = 0;
 
-        let mut process_neuron = |neuron: Neuron| {
+        let mut process_neuron = |neuron: &Neuron| {
             if neuron.is_inactive(now_seconds)
                 || neuron.dissolve_delay_seconds(now_seconds)
                     < MIN_DISSOLVE_DELAY_FOR_VOTE_ELIGIBILITY_SECONDS
@@ -956,7 +960,7 @@ impl NeuronStore {
         self.with_active_neurons_iter_sections(
             |iter| {
                 for neuron in iter {
-                    process_neuron(neuron);
+                    process_neuron(neuron.as_ref());
                 }
             },
             NeuronSections::default(),
diff --git a/rs/nns/governance/tests/governance.rs b/rs/nns/governance/tests/governance.rs
index 349bf15f518..64fbc4b283b 100644
--- a/rs/nns/governance/tests/governance.rs
+++ b/rs/nns/governance/tests/governance.rs
@@ -126,7 +126,7 @@ use std::{
     collections::{BTreeMap, HashSet, VecDeque},
     convert::{TryFrom, TryInto},
     iter::{self, once},
-    ops::Div,
+    ops::{Deref, Div},
     path::PathBuf,
     sync::{Arc, Mutex},
 };
@@ -7200,9 +7200,11 @@ fn test_manage_and_reward_node_providers() {
         ProposalStatus::Executed
     );
     // Find the neuron...
-    let neuron = gov
-        .neuron_store
-        .with_active_neurons_iter(|mut iter| iter.find(|x| x.controller() == np_pid).unwrap());
+    let neuron = gov.neuron_store.with_active_neurons_iter(|mut iter| {
+        iter.find(|x| x.controller() == np_pid)
+            .map(|n| n.deref().clone())
+            .unwrap()
+    });
     assert_eq!(neuron.stake_e8s(), 99_999_999);
     // Find the transaction in the ledger...
     driver.assert_account_contains(
@@ -7508,9 +7510,11 @@ fn test_manage_and_reward_multiple_node_providers() {
 
     // Check third reward
     // Find the neuron...
-    let neuron = gov
-        .neuron_store
-        .with_active_neurons_iter(|mut iter| iter.find(|x| x.controller() == np_pid_2).unwrap());
+    let neuron = gov.neuron_store.with_active_neurons_iter(|mut iter| {
+        iter.find(|x| x.controller() == np_pid_2)
+            .map(|n| n.deref().clone())
+            .unwrap()
+    });
     assert_eq!(neuron.stake_e8s(), 99_999_999);
     // Find the transaction in the ledger...
     driver.assert_account_contains(