From c879272815d64bcf88ce3a817d11cab42b646667 Mon Sep 17 00:00:00 2001
From: Dragoljub Duric <dragoljub.duric@dfinity.org>
Date: Tue, 7 Mar 2023 16:24:56 +0000
Subject: [PATCH] fix: Fix implementation of PartialEq for ExecutionState

---
 .../src/canister_state/execution_state.rs     | 29 +++++-
 .../src/canister_state/tests.rs               | 98 +++++++++++++++++++
 2 files changed, 123 insertions(+), 4 deletions(-)

diff --git a/rs/replicated_state/src/canister_state/execution_state.rs b/rs/replicated_state/src/canister_state/execution_state.rs
index 1b1a7e64ad5..53b0569c1d3 100644
--- a/rs/replicated_state/src/canister_state/execution_state.rs
+++ b/rs/replicated_state/src/canister_state/execution_state.rs
@@ -410,16 +410,37 @@ pub struct ExecutionState {
 // equality (and doesn't need to be).
 impl PartialEq for ExecutionState {
     fn eq(&self, rhs: &Self) -> bool {
+        // Destruction is done on purpose, to ensure if the new
+        // field is added to 'ExecutionState' compiler will throw
+        // an error. Hence pointing to appropriate change here.
+        let &ExecutionState {
+            canister_root: _,
+            session_nonce: _,
+            ref wasm_binary,
+            ref wasm_memory,
+            ref stable_memory,
+            ref exported_globals,
+            ref exports,
+            ref metadata,
+            ref last_executed_round,
+        } = rhs;
+
         (
             &self.wasm_binary.binary,
             &self.wasm_memory,
+            &self.stable_memory,
             &self.exported_globals,
             &self.exports,
+            &self.metadata,
+            &self.last_executed_round,
         ) == (
-            &rhs.wasm_binary.binary,
-            &rhs.wasm_memory,
-            &rhs.exported_globals,
-            &rhs.exports,
+            &wasm_binary.binary,
+            wasm_memory,
+            stable_memory,
+            exported_globals,
+            exports,
+            metadata,
+            last_executed_round,
         )
     }
 }
diff --git a/rs/replicated_state/src/canister_state/tests.rs b/rs/replicated_state/src/canister_state/tests.rs
index 33845e95917..f102c042f93 100644
--- a/rs/replicated_state/src/canister_state/tests.rs
+++ b/rs/replicated_state/src/canister_state/tests.rs
@@ -1,4 +1,9 @@
+use std::collections::BTreeMap;
+use std::path::PathBuf;
+
 use super::*;
+use crate::canister_state::execution_state::CustomSection;
+use crate::canister_state::execution_state::CustomSectionType;
 use crate::canister_state::execution_state::WasmMetadata;
 use crate::canister_state::system_state::CyclesUseCase;
 use crate::CallOrigin;
@@ -608,3 +613,96 @@ fn canister_state_callback_round_trip() {
 
     assert_eq!(callback, round_trip);
 }
+
+#[test]
+fn execution_state_test_partial_eq() {
+    let state_1 = ExecutionState::new(
+        Default::default(),
+        execution_state::WasmBinary::new(CanisterModule::new(vec![1, 2, 3])),
+        ExportedFunctions::new(Default::default()),
+        Memory::new_for_testing(),
+        Memory::new_for_testing(),
+        vec![Global::I64(14)],
+        WasmMetadata::default(),
+    );
+
+    assert_eq!(state_1, state_1.clone());
+
+    assert_eq!(
+        ExecutionState {
+            canister_root: PathBuf::new(),
+            ..state_1.clone()
+        },
+        state_1
+    );
+
+    assert_eq!(
+        ExecutionState {
+            session_nonce: Some(([11; 32], 10)),
+            ..state_1.clone()
+        },
+        state_1
+    );
+
+    assert_ne!(
+        ExecutionState {
+            wasm_binary: execution_state::WasmBinary::new(CanisterModule::new(vec![1, 2, 4])),
+            ..state_1.clone()
+        },
+        state_1
+    );
+
+    assert_ne!(
+        ExecutionState {
+            exports: ExportedFunctions::new(BTreeSet::from([WasmMethod::System(
+                SystemMethod::CanisterGlobalTimer
+            )])),
+            ..state_1.clone()
+        },
+        state_1
+    );
+    let mut memory = Memory::new_for_testing();
+    memory.size = NumWasmPages::from(1);
+    assert_ne!(
+        ExecutionState {
+            wasm_memory: memory.clone(),
+            ..state_1.clone()
+        },
+        state_1
+    );
+    assert_ne!(
+        ExecutionState {
+            stable_memory: memory,
+            ..state_1.clone()
+        },
+        state_1
+    );
+
+    assert_ne!(
+        ExecutionState {
+            exported_globals: vec![Global::I64(13)],
+            ..state_1.clone()
+        },
+        state_1
+    );
+    let mut custom_sections: BTreeMap<String, CustomSection> = BTreeMap::new();
+    custom_sections.insert(
+        String::from("candid"),
+        CustomSection::new(CustomSectionType::Private, vec![0; 10 * 1024]),
+    );
+    assert_ne!(
+        ExecutionState {
+            metadata: WasmMetadata::new(custom_sections),
+            ..state_1.clone()
+        },
+        state_1
+    );
+
+    assert_ne!(
+        ExecutionState {
+            last_executed_round: ExecutionRound::from(12345),
+            ..state_1.clone()
+        },
+        state_1
+    );
+}