From 50e640ea8435848ebc4777b72c4c9d3b2d2957ae Mon Sep 17 00:00:00 2001
From: Imbus <>
Date: Thu, 13 Feb 2025 00:35:28 +0100
Subject: [PATCH] Seed saving to flash with regular intervals

---
 rand.c | 71 ++++++++++++++++++++++++++++++++++++++++++++++++++++------
 rand.h | 25 ++++++++++++++++-----
 2 files changed, 84 insertions(+), 12 deletions(-)

diff --git a/rand.c b/rand.c
index 74cb7ad..062a64a 100644
--- a/rand.c
+++ b/rand.c
@@ -1,22 +1,79 @@
+#include "rand.h"
+#include "ch32fun.h"
 #include <stdint.h>
 
-#include <stdint.h>
+#define BUILD_SEED                                                             \
+    ((uint64_t)(__TIME__[0]) * (uint64_t)(__TIME__[1]) *                       \
+     (uint64_t)(__TIME__[3]) * (uint64_t)(__TIME__[4]) *                       \
+     (uint64_t)(__TIME__[6]) * (uint64_t)(__TIME__[7]))
 
-#define BUILD_SEED ((uint64_t)(__TIME__[0]) * (uint64_t)(__TIME__[1]) * \
-                    (uint64_t)(__TIME__[3]) * (uint64_t)(__TIME__[4]) * \
-                    (uint64_t)(__TIME__[6]) * (uint64_t)(__TIME__[7]))
+#define FLASH_SEED_ADDR ((uintptr_t *)0x08003700) // PRNG state storage
+#define PRNG_SAVE_INTERVAL 50 // Save every 1000 calls to prand()
 
 static uint64_t seed = BUILD_SEED;
 
-void sprand(uint64_t s) {
-    seed = s ? s : 1; // Ensure the seed is never 0
-}
+// Initialize this to something close to interval
+static int prand_counter = PRNG_SAVE_INTERVAL - 10;
 
 uint64_t prand() {
     seed = seed * 6364136223846793005ULL + 1;
+
+    if (++prand_counter >= PRNG_SAVE_INTERVAL) {
+        rand_save_to_flash();
+        prand_counter = 0;
+    }
+
     return seed;
 }
 
 uint64_t prand_range(uint64_t min, uint64_t max) {
     return min + (prand() % (max - min + 1));
 }
+
+void sprand(uint64_t s) {
+    if (s) {
+        seed = s;
+    } else {
+        rand_reseed();
+    }
+}
+
+void rand_reseed() {
+    uint64_t stored_seed = *(volatile uint64_t *)FLASH_SEED_ADDR;
+
+    if (stored_seed == 0 || stored_seed == 0xFFFFFFFFFFFFFFFFULL) {
+        seed = BUILD_SEED;
+    } else {
+        seed = stored_seed;
+    }
+}
+
+// See:
+// https://github.com/cnlohr/ch32v003fun/blob/2ac62072272f2ccd2122e688a9e0566de3976a94/examples/flashtest/flashtest.c
+void rand_save_to_flash() {
+    FLASH->KEYR = 0x45670123; // Unlock flash
+    FLASH->KEYR = 0xCDEF89AB;
+
+    FLASH->MODEKEYR = 0x45670123; // Unlock programming mode
+    FLASH->MODEKEYR = 0xCDEF89AB;
+
+    // Erase the flash page
+    FLASH->CTLR = CR_PAGE_ER;
+    FLASH->ADDR = (intptr_t)FLASH_SEED_ADDR;
+    FLASH->CTLR = CR_STRT_Set | CR_PAGE_ER;
+    while (FLASH->STATR & FLASH_STATR_BSY); // Wait for erase
+
+    // Write new seed
+    FLASH->CTLR = CR_PAGE_PG;
+    FLASH->CTLR = CR_BUF_RST | CR_PAGE_PG;
+    FLASH->ADDR = (intptr_t)FLASH_SEED_ADDR;
+
+    ((uint32_t *)FLASH_SEED_ADDR)[0] = (uint32_t)seed;
+    ((uint32_t *)FLASH_SEED_ADDR)[1] = (uint32_t)(seed >> 32);
+
+    FLASH->CTLR = CR_PAGE_PG | FLASH_CTLR_BUF_LOAD;
+    while (FLASH->STATR & FLASH_STATR_BSY); // Wait for completion
+
+    FLASH->CTLR = CR_PAGE_PG | CR_STRT_Set; // Commit write
+    while (FLASH->STATR & FLASH_STATR_BSY); // Wait for completion
+}
diff --git a/rand.h b/rand.h
index 3722734..26c4194 100644
--- a/rand.h
+++ b/rand.h
@@ -2,18 +2,18 @@
 #include <stdint.h>
 
 /**
- * @brief Sets the seed for the custom random number generator.
+ * @brief Sets the seed for the PRNG.
  *
- * This function initializes the seed value used by rand_custom().
- * Providing the same seed will produce the same sequence of random numbers.
- *
- * @param s The seed value (must be nonzero for best results).
+ * @param s The specific seed value or zero. If zero is passed, it will call
+ * rand_reseed().
  */
 void sprand(uint64_t s);
 
 /**
  * @brief Generates a pseudo-random 64-bit number.
  *
+ * Saves PRNG state to flash periodically.
+ *
  * Uses a simple Linear Congruential Generator (LCG) to produce
  * a sequence of pseudo-random numbers.
  *
@@ -32,3 +32,18 @@ uint64_t prand();
  * @return A random number between min and max.
  */
 uint64_t prand_range(uint64_t min, uint64_t max);
+
+/**
+ * @brief Saves the current PRNG seed to flash memory.
+ *
+ * This function erases the designated flash page and writes the current seed
+ * to ensure the PRNG state persists across resets.
+ */
+void rand_save_to_flash();
+
+/**
+ * @brief Re-seeds the PRNG seed state from either flash or BUILD_SEED.
+ *
+ * This function will not write to flash.
+ */
+void rand_reseed();