From 5f8d6a6f85e01c060eb636492170747d92f0d181 Mon Sep 17 00:00:00 2001
From: BlackMark <me@blackmark.me>
Date: Sun, 12 Apr 2020 15:12:02 +0200
Subject: [PATCH] Fix chip erase followed immediately by flash programming

---
 stk500v2/main.cpp | 42 ++++++++++++++++++++++++++----------------
 1 file changed, 26 insertions(+), 16 deletions(-)

diff --git a/stk500v2/main.cpp b/stk500v2/main.cpp
index 57d1a6c..ce48f13 100644
--- a/stk500v2/main.cpp
+++ b/stk500v2/main.cpp
@@ -405,17 +405,17 @@ static inline uint32_t getFlashSize()
 	return (FLASHEND - bootloaderSize + 1);
 }
 
-static inline void performChipErase()
+static inline void performChipErase(uint16_t flashStartAddress = 0x0000)
 {
 	constexpr auto getEepromEraseFuseBit = []() -> bool {
 		constexpr auto EESAVE = 3;
 		return boot_lock_fuse_bits_get(GET_HIGH_FUSE_BITS) & (1 << EESAVE);
 	};
 
-	constexpr auto eraseFlash = []() {
+	constexpr auto eraseFlash = [](const uint16_t &flashStartAddress) {
 		const auto flashSize = getFlashSize();
-
-		for (uint16_t i = 0; i < flashSize; i += SPM_PAGESIZE) {
+		const auto byteAddress = 2 * flashStartAddress;
+		for (uint16_t i = byteAddress; i < flashSize; i += SPM_PAGESIZE) {
 			boot_page_erase(i);
 			boot_spm_busy_wait();
 		}
@@ -433,7 +433,7 @@ static inline void performChipErase()
 		}
 	};
 
-	eraseFlash();
+	eraseFlash(flashStartAddress);
 	eraseEeprom();
 }
 
@@ -510,6 +510,8 @@ enum class ChipEraseState {
 	REQUEST = (1 << 1),
 	RESPONSE = (1 << 2),
 	PERFORM = REQUEST | RESPONSE,
+	PROGRAM = (1 << 3),
+	FINISH = REQUEST | RESPONSE | PROGRAM,
 };
 
 constexpr ChipEraseState operator|(const ChipEraseState &self, const ChipEraseState &other)
@@ -523,7 +525,8 @@ constexpr ChipEraseState &operator|=(ChipEraseState &self, const ChipEraseState
 	return self;
 }
 
-static inline void handleMessage(Message &msg, uint32_t &addr, ChipEraseState &chipEraseFlag)
+static inline void handleMessage(Message &msg, uint32_t &addr, uint16_t &finishEraseAddress,
+                                 ChipEraseState &chipEraseFlag)
 {
 	if (isSignOn(msg))
 		formatSignOnAnswer(msg);
@@ -550,11 +553,13 @@ static inline void handleMessage(Message &msg, uint32_t &addr, ChipEraseState &c
 	else if (isReadEepromIsp(msg))
 		formatReadEepromIspAnswer(msg, addr);
 	else if (isChipEraseIsp(msg)) {
-		chipEraseFlag = ChipEraseState::REQUEST;
+		chipEraseFlag |= ChipEraseState::REQUEST;
 		formatChipEraseIspAnswer(msg);
-	} else if (isProgramFlashIsp(msg))
+	} else if (isProgramFlashIsp(msg)) {
+		chipEraseFlag |= ChipEraseState::PROGRAM;
 		formatProgramFlashIspAnswer(msg, addr);
-	else if (isProgramEepromIsp(msg))
+		finishEraseAddress = addr;
+	} else if (isProgramEepromIsp(msg))
 		formatProgramEepromIspAnswer(msg, addr);
 	else if (isLeaveProgmodeIsp(msg)) {
 		chipEraseFlag |= ChipEraseState::RESPONSE;
@@ -563,11 +568,6 @@ static inline void handleMessage(Message &msg, uint32_t &addr, ChipEraseState &c
 		formatErrorAnswer(msg);
 
 	transmitMessage(msg);
-
-	if (chipEraseFlag == ChipEraseState::PERFORM) {
-		performChipErase();
-		chipEraseFlag = ChipEraseState::NONE;
-	}
 }
 
 int main()
@@ -576,16 +576,26 @@ int main()
 
 	Message msg;
 	uint32_t addr = 0x0000;
+	uint16_t finishEraseAddress = 0x0000;
 	ChipEraseState chipEraseFlag = ChipEraseState::NONE;
 	uint16_t timeout = TIMEOUT;
 
 	while (true) {
 		if (receiveMessage(msg, timeout)) {
-			handleMessage(msg, addr, chipEraseFlag);
+			handleMessage(msg, addr, finishEraseAddress, chipEraseFlag);
 		}
 
-		if (timeout == 0)
+		if (timeout == 0) {
 			timeout = TIMEOUT;
+
+			if (chipEraseFlag == ChipEraseState::PERFORM) {
+				performChipErase();
+				chipEraseFlag = ChipEraseState::NONE;
+			} else if (chipEraseFlag == ChipEraseState::FINISH) {
+				performChipErase(finishEraseAddress);
+				chipEraseFlag = ChipEraseState::NONE;
+			}
+		}
 	}
 
 	return 0;