2 files changed, 10 insertions, 8 deletions

diff --git a/docs/pages/bootloader.adoc b/docs/pages/bootloader.adoc
index b0f0331aa..c010b0622 100644
--- a/docs/pages/bootloader.adoc
+++ b/docs/pages/bootloader.adoc
@@ -43,14 +43,14 @@ Partition Size~dfu~= Partition Size~active~+ Page Size~active~
 +
 All values are specified in bytes.
 
-* BOOTLOADER STATE - Where the bootloader stores the current state describing if the active and dfu partitions need to be swapped. When the new firmware has been written to the DFU partition, a magic field is written to instruct the bootloader that the partitions should be swapped. This partition must be able to store a magic field as well as the partition swap progress. The partition size given by:
+* BOOTLOADER STATE - Where the bootloader stores the current state describing if the active and dfu partitions need to be swapped. When the new firmware has been written to the DFU partition, a magic field is written to instruct the bootloader that the partitions should be swapped. This partition must be able to store a magic field as well as the partition swap progress. The partition size is given by:
 +
-Partition Size~state~ = Write Size~state~ + (2 × Partition Size~active~ / Page Size~active~)
+Partition Size~state~ = (2 × Write Size~state~) + (4 × Write Size~state~ × Partition Size~active~ / Page Size~active~)
 +
 All values are specified in bytes.
 
 The partitions for ACTIVE (+BOOTLOADER), DFU and BOOTLOADER_STATE may be placed in separate flash. The page size used by the bootloader is determined by the lowest common multiple of the ACTIVE and DFU page sizes.
-The BOOTLOADER_STATE partition must be big enough to store one word per page in the ACTIVE and DFU partitions combined.
+The BOOTLOADER_STATE partition must be big enough to store two words, plus four words per page in the ACTIVE partition.
 
 The bootloader has a platform-agnostic part, which implements the power fail safe swapping algorithm given the boundaries set by the partitions. The platform-specific part is a minimal shim that provides additional functionality such as watchdogs or supporting the nRF52 softdevice.
 
diff --git a/embassy-boot/src/boot_loader.rs b/embassy-boot/src/boot_loader.rs
index c38940d6e..a3a307051 100644
--- a/embassy-boot/src/boot_loader.rs
+++ b/embassy-boot/src/boot_loader.rs
@@ -135,10 +135,12 @@ pub struct BootLoader<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash> {
     dfu: DFU,
     /// The state partition has the following format:
     /// All ranges are in multiples of WRITE_SIZE bytes.
-    /// | Range    | Description                                                                      |
-    /// | 0..1     | Magic indicating bootloader state. BOOT_MAGIC means boot, SWAP_MAGIC means swap. |
-    /// | 1..2     | Progress validity. ERASE_VALUE means valid, !ERASE_VALUE means invalid.          |
-    /// | 2..2 + N | Progress index used while swapping or reverting      
+    /// N = Active partition size divided by WRITE_SIZE.
+    /// | Range              | Description                                                                      |
+    /// | 0..1               | Magic indicating bootloader state. BOOT_MAGIC means boot, SWAP_MAGIC means swap. |
+    /// | 1..2               | Progress validity. ERASE_VALUE means valid, !ERASE_VALUE means invalid.          |
+    /// | 2..(2 + 2N)        | Progress index used while swapping                                               |
+    /// | (2 + 2N)..(2 + 4N) | Progress index used while reverting
     state: STATE,
 }
 
@@ -429,7 +431,7 @@ fn assert_partitions<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash>(
     assert_eq!(dfu.capacity() as u32 % page_size, 0);
     // DFU partition has to be bigger than ACTIVE partition to handle swap algorithm
     assert!(dfu.capacity() as u32 - active.capacity() as u32 >= page_size);
-    assert!(2 + 2 * (active.capacity() as u32 / page_size) <= state.capacity() as u32 / STATE::WRITE_SIZE as u32);
+    assert!(2 + 4 * (active.capacity() as u32 / page_size) <= state.capacity() as u32 / STATE::WRITE_SIZE as u32);
 }
 
 #[cfg(test)]