From 58ab0056914876b01fe24e258ac855634b8f985a Mon Sep 17 00:00:00 2001
From: David Dal Busco <david.dalbusco@outlook.com>
Date: Sat, 25 Nov 2023 10:33:59 +0100
Subject: [PATCH] fix: orbiter upgrade

Signed-off-by: David Dal Busco <david.dalbusco@outlook.com>
---
 src/orbiter/src/lib.rs           | 11 ++---------
 src/orbiter/src/upgrade/impls.rs | 12 -----------
 src/orbiter/src/upgrade/mod.rs   |  2 --
 src/orbiter/src/upgrade/types.rs | 34 --------------------------------
 4 files changed, 2 insertions(+), 57 deletions(-)
 delete mode 100644 src/orbiter/src/upgrade/impls.rs
 delete mode 100644 src/orbiter/src/upgrade/mod.rs
 delete mode 100644 src/orbiter/src/upgrade/types.rs

diff --git a/src/orbiter/src/lib.rs b/src/orbiter/src/lib.rs
index 3c5cebfb2..3d5417aa8 100644
--- a/src/orbiter/src/lib.rs
+++ b/src/orbiter/src/lib.rs
@@ -10,7 +10,6 @@ mod msg;
 mod serializers;
 mod store;
 mod types;
-mod upgrade;
 
 use crate::assert::assert_enabled;
 use crate::config::store::{
@@ -32,7 +31,6 @@ use crate::types::interface::{
 };
 use crate::types::memory::Memory;
 use crate::types::state::{AnalyticKey, HeapState, PageView, SatelliteConfigs, State, TrackEvent};
-use crate::upgrade::types::upgrade::UpgradeState;
 use ciborium::{from_reader, into_writer};
 use ic_cdk::api::call::arg_data;
 use ic_cdk::trap;
@@ -103,15 +101,10 @@ fn post_upgrade() {
     memory.read(u64::try_from(OFFSET).unwrap(), &mut state_bytes);
 
     // Deserialize and set the state.
-    let upgrade_state: UpgradeState = from_reader(&*state_bytes)
+    let state: State = from_reader(&*state_bytes)
         .expect("Failed to decode the state of the satellite in post_upgrade hook.");
 
-    STATE.with(|s| {
-        *s.borrow_mut() = State {
-            stable: upgrade_state.stable,
-            heap: HeapState::from(&upgrade_state.heap),
-        }
-    });
+    STATE.with(|s| *s.borrow_mut() = state);
 }
 
 /// Data
diff --git a/src/orbiter/src/upgrade/impls.rs b/src/orbiter/src/upgrade/impls.rs
deleted file mode 100644
index ea3d92fac..000000000
--- a/src/orbiter/src/upgrade/impls.rs
+++ /dev/null
@@ -1,12 +0,0 @@
-use crate::types::state::HeapState;
-use crate::upgrade::types::upgrade::UpgradeHeapState;
-use std::collections::HashMap;
-
-impl From<&UpgradeHeapState> for HeapState {
-    fn from(state: &UpgradeHeapState) -> Self {
-        HeapState {
-            controllers: state.controllers.clone(),
-            config: HashMap::new(),
-        }
-    }
-}
diff --git a/src/orbiter/src/upgrade/mod.rs b/src/orbiter/src/upgrade/mod.rs
deleted file mode 100644
index fb880c018..000000000
--- a/src/orbiter/src/upgrade/mod.rs
+++ /dev/null
@@ -1,2 +0,0 @@
-pub mod impls;
-pub mod types;
diff --git a/src/orbiter/src/upgrade/types.rs b/src/orbiter/src/upgrade/types.rs
deleted file mode 100644
index 41c942d55..000000000
--- a/src/orbiter/src/upgrade/types.rs
+++ /dev/null
@@ -1,34 +0,0 @@
-pub mod upgrade {
-    use crate::memory::init_stable_state;
-    use crate::types::state::StableState;
-    use candid::{CandidType, Principal};
-    use serde::{Deserialize, Serialize};
-    use shared::types::state::{Controllers, SatelliteId};
-    use std::collections::HashMap;
-
-    #[derive(Serialize, Deserialize)]
-    pub struct UpgradeState {
-        // Direct stable state: State that is uses stable memory directly as its store. No need for pre/post upgrade hooks.
-        #[serde(skip, default = "init_stable_state")]
-        pub stable: StableState,
-
-        // Indirect stable state: State that lives on the heap, but is saved into stable memory on upgrades.
-        pub heap: UpgradeHeapState,
-    }
-
-    #[derive(Default, CandidType, Serialize, Deserialize, Clone)]
-    pub struct UpgradeHeapState {
-        pub controllers: Controllers,
-        pub origins: UpgradeOriginConfigs,
-    }
-
-    pub type UpgradeOriginConfigs = HashMap<SatelliteId, UpgradeOriginConfig>;
-
-    #[derive(CandidType, Serialize, Deserialize, Clone)]
-    pub struct UpgradeOriginConfig {
-        pub key: Principal,
-        pub filter: String,
-        pub created_at: u64,
-        pub updated_at: u64,
-    }
-}