From 284c0a95bfb8637515d3dc49f18db4002864d08c Mon Sep 17 00:00:00 2001
From: mtkennerly <mtkennerly@gmail.com>
Date: Thu, 22 Dec 2022 05:03:51 +0800
Subject: [PATCH] #166: Fix registry restoration

---
 CHANGELOG.md    |  1 +
 src/layout.rs   | 21 +++++++++++----------
 src/registry.rs | 27 ++++++++++-----------------
 3 files changed, 22 insertions(+), 27 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9e54de2d..aa4a7e7e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,7 @@
 ## Unreleased
 
 * Fixed:
+  * Native registry saves on Windows were not restored.
   * When switching between dropdowns, they would briefly flicker with incorrect content.
   * Game titles starting with a lowercase letter were listed after all titles starting with an uppercase letter.
   * When using the folder picker for roots and custom game files, glob special characters were not escaped.
diff --git a/src/layout.rs b/src/layout.rs
index e036cb73..2f28085c 100644
--- a/src/layout.rs
+++ b/src/layout.rs
@@ -894,7 +894,7 @@ impl GameLayout {
         #[cfg(target_os = "windows")]
         {
             use crate::registry::Hives;
-            let hives = Hives::from(&scan.found_registry_keys);
+            let hives = Hives::incorporated(&scan.found_registry_keys);
             if !hives.is_empty() {
                 registry.hash = Some(crate::prelude::sha1(hives.serialize()));
             }
@@ -932,7 +932,7 @@ impl GameLayout {
         #[cfg(target_os = "windows")]
         {
             use crate::registry::Hives;
-            let hives = Hives::from(&scan.found_registry_keys);
+            let hives = Hives::incorporated(&scan.found_registry_keys);
             if !hives.is_empty() {
                 registry = Some(IndividualMappingRegistry {
                     hash: Some(crate::prelude::sha1(hives.serialize())),
@@ -1005,7 +1005,7 @@ impl GameLayout {
             let target_registry_file = self.registry_file_in(backup.name());
 
             if backup.includes_registry() {
-                let hives = Hives::from(&scan.found_registry_keys);
+                let hives = Hives::incorporated(&scan.found_registry_keys);
                 hives.save(&target_registry_file);
             } else {
                 let _ = target_registry_file.remove();
@@ -1156,7 +1156,7 @@ impl GameLayout {
             use crate::registry::Hives;
 
             if backup.includes_registry() {
-                let hives = Hives::from(&scan.found_registry_keys);
+                let hives = Hives::incorporated(&scan.found_registry_keys);
                 if zip.start_file("registry.yaml", options).is_ok() {
                     let _ = zip.write_all(hives.serialize().as_bytes());
                 }
@@ -1299,12 +1299,13 @@ impl GameLayout {
         {
             use crate::registry::Hives;
 
-            let mut hives = Hives::default();
-            let (found, _) = hives.incorporate(&scan.found_registry_keys);
-
-            if found {
-                // TODO: Track failed keys.
-                let _ = hives.restore();
+            if let Some(backup) = scan.backup.as_ref() {
+                if let Some(registry_content) = self.registry_content(&backup.id()) {
+                    if let Some(hives) = Hives::deserialize(&registry_content) {
+                        // TODO: Track failed keys.
+                        let _ = hives.restore();
+                    }
+                }
             }
         }
 
diff --git a/src/registry.rs b/src/registry.rs
index 3a23fd4a..c8b3e7fd 100644
--- a/src/registry.rs
+++ b/src/registry.rs
@@ -2,7 +2,7 @@ use std::collections::HashSet;
 
 use crate::{
     config::{BackupFilter, ToggledRegistry},
-    prelude::{Error, RegistryItem, ScanChange, ScanInfo, ScannedRegistry, StrictPath},
+    prelude::{Error, RegistryItem, ScanChange, ScannedRegistry, StrictPath},
 };
 use winreg::types::{FromRegValue, ToRegValue};
 
@@ -154,6 +154,9 @@ impl Hives {
         serde_yaml::from_str(content).ok()
     }
 
+    /// This only incorporates the keys, not the values.
+    /// It can be used during backup since we know the keys exist, so we can look up the values when needed.
+    /// It should not be used during restore since the keys may not exist.
     pub fn incorporate(&mut self, scan: &HashSet<ScannedRegistry>) -> (bool, HashSet<RegistryItem>) {
         let mut failed = HashSet::new();
         let mut found = false;
@@ -176,6 +179,12 @@ impl Hives {
         (found, failed)
     }
 
+    pub fn incorporated(scan: &HashSet<ScannedRegistry>) -> Self {
+        let mut hives = Hives::default();
+        hives.incorporate(scan);
+        hives
+    }
+
     pub fn store_key_from_full_path(&mut self, path: &str) -> Result<(), Error> {
         let path = RegistryItem::new(path.to_string()).interpreted();
 
@@ -265,22 +274,6 @@ impl Hives {
     }
 }
 
-impl From<&HashSet<ScannedRegistry>> for Hives {
-    fn from(source: &HashSet<ScannedRegistry>) -> Self {
-        let mut hives = Self::default();
-        hives.incorporate(source);
-        hives
-    }
-}
-
-impl From<&ScanInfo> for Hives {
-    fn from(scan: &ScanInfo) -> Self {
-        let mut hives = Self::default();
-        hives.incorporate(&scan.found_registry_keys);
-        hives
-    }
-}
-
 impl Entry {
     pub fn is_set(&self) -> bool {
         self.sz.is_some()