From 238fbfe5698fb1791d47e8772ba1c5a86f9300ca Mon Sep 17 00:00:00 2001
From: Jesse Plamondon-Willard <Pathoschild@users.noreply.github.com>
Date: Sun, 8 Dec 2019 12:20:59 -0500
Subject: [PATCH] let mods use Read/WriteSaveData while a save is being loaded

---
 docs/release-notes.md                        |  1 +
 src/SMAPI/Framework/ModHelpers/DataHelper.cs | 48 +++++++++++++++-----
 2 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/docs/release-notes.md b/docs/release-notes.md
index 8754e777..fc3cd4f7 100644
--- a/docs/release-notes.md
+++ b/docs/release-notes.md
@@ -9,6 +9,7 @@
 
 * For modders:
   * Added asset propagation for grass textures.
+  * `helper.Read/WriteSaveData` can now be used while a save is being loaded (e.g. within a `Specialized.LoadStageChanged` event).
 
 * For the web UI:
   * If a JSON validator upload can't be saved to Pastebin (e.g. due to rate limits), it's now uploaded to Amazon S3 instead. Files uploaded to S3 expire after one month.
diff --git a/src/SMAPI/Framework/ModHelpers/DataHelper.cs b/src/SMAPI/Framework/ModHelpers/DataHelper.cs
index cc08c42b..3d43c539 100644
--- a/src/SMAPI/Framework/ModHelpers/DataHelper.cs
+++ b/src/SMAPI/Framework/ModHelpers/DataHelper.cs
@@ -1,6 +1,8 @@
 using System;
+using System.Collections.Generic;
 using System.IO;
 using Newtonsoft.Json;
+using StardewModdingAPI.Enums;
 using StardewModdingAPI.Toolkit.Serialization;
 using StardewModdingAPI.Toolkit.Utilities;
 using StardewValley;
@@ -77,33 +79,45 @@ namespace StardewModdingAPI.Framework.ModHelpers
         /// <exception cref="InvalidOperationException">The player hasn't loaded a save file yet or isn't the main player.</exception>
         public TModel ReadSaveData<TModel>(string key) where TModel : class
         {
-            if (!Game1.hasLoadedGame)
+            if (Context.LoadStage == LoadStage.None)
                 throw new InvalidOperationException($"Can't use {nameof(IMod.Helper)}.{nameof(IModHelper.Data)}.{nameof(this.ReadSaveData)} when a save file isn't loaded.");
             if (!Game1.IsMasterGame)
                 throw new InvalidOperationException($"Can't use {nameof(IMod.Helper)}.{nameof(IModHelper.Data)}.{nameof(this.ReadSaveData)} because this isn't the main player. (Save files are stored on the main player's computer.)");
 
-            return Game1.CustomData.TryGetValue(this.GetSaveFileKey(key), out string value)
-                ? this.JsonHelper.Deserialize<TModel>(value)
-                : null;
+
+            string internalKey = this.GetSaveFileKey(key);
+            foreach (IDictionary<string, string> dataField in this.GetDataFields(Context.LoadStage))
+            {
+                if (dataField.TryGetValue(internalKey, out string value))
+                    return this.JsonHelper.Deserialize<TModel>(value);
+            }
+            return null;
         }
 
         /// <summary>Save arbitrary data to the current save slot. This is only possible if a save has been loaded, and the data will be lost if the player exits without saving the current day.</summary>
         /// <typeparam name="TModel">The model type. This should be a plain class that has public properties for the data you want. The properties can be complex types.</typeparam>
         /// <param name="key">The unique key identifying the data.</param>
-        /// <param name="data">The arbitrary data to save.</param>
+        /// <param name="model">The arbitrary data to save.</param>
         /// <exception cref="InvalidOperationException">The player hasn't loaded a save file yet or isn't the main player.</exception>
-        public void WriteSaveData<TModel>(string key, TModel data) where TModel : class
+        public void WriteSaveData<TModel>(string key, TModel model) where TModel : class
         {
-            if (!Game1.hasLoadedGame)
+            if (Context.LoadStage == LoadStage.None)
                 throw new InvalidOperationException($"Can't use {nameof(IMod.Helper)}.{nameof(IModHelper.Data)}.{nameof(this.WriteSaveData)} when a save file isn't loaded.");
             if (!Game1.IsMasterGame)
                 throw new InvalidOperationException($"Can't use {nameof(IMod.Helper)}.{nameof(IModHelper.Data)}.{nameof(this.WriteSaveData)} because this isn't the main player. (Save files are stored on the main player's computer.)");
 
             string internalKey = this.GetSaveFileKey(key);
-            if (data != null)
-                Game1.CustomData[internalKey] = this.JsonHelper.Serialize(data, Formatting.None);
-            else
-                Game1.CustomData.Remove(internalKey);
+            string data = model != null
+                ? this.JsonHelper.Serialize(model, Formatting.None)
+                : null;
+
+            foreach (IDictionary<string, string> dataField in this.GetDataFields(Context.LoadStage))
+            {
+                if (data != null)
+                    dataField[internalKey] = data;
+                else
+                    dataField.Remove(internalKey);
+            }
         }
 
         /****
@@ -146,6 +160,18 @@ namespace StardewModdingAPI.Framework.ModHelpers
             return $"smapi/mod-data/{this.ModID}/{key}".ToLower();
         }
 
+        /// <summary>Get the data fields to read/write for save data.</summary>
+        /// <param name="stage">The current load stage.</param>
+        private IEnumerable<IDictionary<string, string>> GetDataFields(LoadStage stage)
+        {
+            if (stage == LoadStage.None)
+                yield break;
+
+            yield return Game1.CustomData;
+            if (SaveGame.loaded != null)
+                yield return SaveGame.loaded.CustomData;
+        }
+
         /// <summary>Get the absolute path for a global data file.</summary>
         /// <param name="key">The unique key identifying the data.</param>
         private string GetGlobalDataPath(string key)