From 70ec49c7acfc3b33a9ad49dbff552b9c4510e6f4 Mon Sep 17 00:00:00 2001
From: Nav <nav@navnav.me>
Date: Sat, 1 Feb 2025 23:04:02 +0000
Subject: [PATCH] Added support for flash page writes on UPDI targets, in the
 EDBG driver

---
 .../Microchip/Protocols/Edbg/Avr/Avr8Generic.hpp  |  2 ++
 .../Protocols/Edbg/Avr/EdbgAvr8Interface.cpp      | 15 ++++++++++++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/DebugToolDrivers/Microchip/Protocols/Edbg/Avr/Avr8Generic.hpp b/src/DebugToolDrivers/Microchip/Protocols/Edbg/Avr/Avr8Generic.hpp
index 02d32af7..ec293e70 100644
--- a/src/DebugToolDrivers/Microchip/Protocols/Edbg/Avr/Avr8Generic.hpp
+++ b/src/DebugToolDrivers/Microchip/Protocols/Edbg/Avr/Avr8Generic.hpp
@@ -165,6 +165,8 @@ namespace DebugToolDrivers::Microchip::Protocols::Edbg::Avr
         APPL_FLASH = 0xC0,
         BOOT_FLASH = 0xC1,
 
+        APPL_FLASH_ATOMIC = 0xC2,
+
         /**
          * The SPM memory type can be used to read memory from the target whilst in debugging mode.
          *
diff --git a/src/DebugToolDrivers/Microchip/Protocols/Edbg/Avr/EdbgAvr8Interface.cpp b/src/DebugToolDrivers/Microchip/Protocols/Edbg/Avr/EdbgAvr8Interface.cpp
index 60b52605..a9fa052e 100644
--- a/src/DebugToolDrivers/Microchip/Protocols/Edbg/Avr/EdbgAvr8Interface.cpp
+++ b/src/DebugToolDrivers/Microchip/Protocols/Edbg/Avr/EdbgAvr8Interface.cpp
@@ -759,6 +759,10 @@ namespace DebugToolDrivers::Microchip::Protocols::Edbg::Avr
                 }
             }
 
+            if (this->session.configVariant == Avr8ConfigVariant::UPDI) {
+                return this->writeMemory(Avr8MemoryType::APPL_FLASH_ATOMIC, startAddress, buffer);
+            }
+
             return this->writeMemory(Avr8MemoryType::FLASH_PAGE, startAddress, buffer);
         }
 
@@ -1516,6 +1520,7 @@ namespace DebugToolDrivers::Microchip::Protocols::Edbg::Avr
             memoryType == Avr8MemoryType::FLASH_PAGE
             || memoryType == Avr8MemoryType::SPM
             || memoryType == Avr8MemoryType::APPL_FLASH
+            || memoryType == Avr8MemoryType::APPL_FLASH_ATOMIC
             || memoryType == Avr8MemoryType::BOOT_FLASH
             || memoryType == Avr8MemoryType::EEPROM_ATOMIC
             || memoryType == Avr8MemoryType::EEPROM_PAGE
@@ -1529,6 +1534,7 @@ namespace DebugToolDrivers::Microchip::Protocols::Edbg::Avr
             case Avr8MemoryType::FLASH_PAGE:
             case Avr8MemoryType::SPM:
             case Avr8MemoryType::APPL_FLASH:
+            case Avr8MemoryType::APPL_FLASH_ATOMIC:
             case Avr8MemoryType::BOOT_FLASH: {
                 /*
                  * Although the EDBG documentation claims any number of bytes can be accessed via the FLASH_PAGE mem
@@ -1565,6 +1571,7 @@ namespace DebugToolDrivers::Microchip::Protocols::Edbg::Avr
             case Avr8MemoryType::FLASH_PAGE:
             case Avr8MemoryType::SPM:
             case Avr8MemoryType::APPL_FLASH:
+            case Avr8MemoryType::APPL_FLASH_ATOMIC:
             case Avr8MemoryType::BOOT_FLASH: {
                 // See comment in EdbgAvr8Interface::alignMemoryAddress()
                 alignTo = static_cast<std::uint16_t>(this->session.programMemorySegment.pageSize.value());
@@ -1593,6 +1600,7 @@ namespace DebugToolDrivers::Microchip::Protocols::Edbg::Avr
         if (
             memoryType == Avr8MemoryType::FLASH_PAGE
             || memoryType == Avr8MemoryType::APPL_FLASH
+            || memoryType == Avr8MemoryType::APPL_FLASH_ATOMIC
             || memoryType == Avr8MemoryType::BOOT_FLASH
             || (memoryType == Avr8MemoryType::SPM && this->session.configVariant == Avr8ConfigVariant::MEGAJTAG)
         ) {
@@ -1766,7 +1774,12 @@ namespace DebugToolDrivers::Microchip::Protocols::Edbg::Avr
                  * We can't just forward the memory type to readMemory(), because some memory types (such as
                  * EEPROM_ATOMIC) can only be used for writing.
                  */
-                const auto readMemType = type == Avr8MemoryType::EEPROM_ATOMIC ? Avr8MemoryType::EEPROM : type;
+                const auto readMemType = type == Avr8MemoryType::EEPROM_ATOMIC
+                    ? Avr8MemoryType::EEPROM
+                    : type == Avr8MemoryType::APPL_FLASH_ATOMIC
+                        ? Avr8MemoryType::FLASH_PAGE
+                        : type;
+
                 auto alignedBuffer = (alignedStartAddress < startAddress)
                     ? this->readMemory(readMemType, alignedStartAddress, startAddress - alignedStartAddress)
                     : TargetMemoryBuffer{};