update storage extension implementation

ports/espressif/supervisor/internal_flash.c

@@ -24,6 +24,7 @@
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  * THE SOFTWARE.
  */
+
 #include "supervisor/internal_flash.h"
 
 #include <stdint.h>
@@ -32,31 +33,46 @@
 
 #include "extmod/vfs.h"
 #include "extmod/vfs_fat.h"
+
 #include "py/mphal.h"
 #include "py/obj.h"
 #include "py/runtime.h"
-#include "lib/oofatfs/ff.h"
 
 #include "esp_ota_ops.h"
 #include "esp_partition.h"
 
+#include "supervisor/filesystem.h"
 #include "supervisor/flash.h"
 #include "supervisor/usb.h"
 
-#if CIRCUITPY_STORAGE_EXTEND
+#define OP_READ     0
-#define PARTITION_NUM (2)
+#define OP_WRITE    1
-#else
-#define PARTITION_NUM (1)
-#endif
-
-STATIC const esp_partition_t *_partition[PARTITION_NUM];
 
 // TODO: Split the caching out of supervisor/shared/external_flash so we can use it.
 #define SECTOR_SIZE 4096
 STATIC uint8_t _cache[SECTOR_SIZE];
 STATIC uint32_t _cache_lba = 0xffffffff;
 
+#if CIRCUITPY_STORAGE_EXTEND
+#if FF_MAX_SS == FF_MIN_SS
+#define SECSIZE(fs) (FF_MIN_SS)
+#else
+#define SECSIZE(fs) ((fs)->ssize)
+#endif // FF_MAX_SS == FF_MIN_SS
+STATIC DWORD fatfs_bytes(void) {
+    FATFS *fatfs = filesystem_circuitpy();
+    return (fatfs->csize * SECSIZE(fatfs)) * (fatfs->n_fatent - 2);
+}
+STATIC bool storage_extended = true;
+STATIC const esp_partition_t *_partition[2];
+#else
+STATIC const esp_partition_t *_partition[1];
+#endif // CIRCUITPY_STORAGE_EXTEND
+
 void supervisor_flash_init(void) {
+    if (_partition[0] != NULL) {
+        return;
+    }
     _partition[0] = esp_partition_find_first(ESP_PARTITION_TYPE_DATA,
         ESP_PARTITION_SUBTYPE_DATA_FAT,
         NULL);
@@ -71,38 +87,60 @@ uint32_t supervisor_flash_get_block_size(void) {
 
 uint32_t supervisor_flash_get_block_count(void) {
     #if CIRCUITPY_STORAGE_EXTEND
-    return (_partition[0]->size + _partition[1]->size) / FILESYSTEM_BLOCK_SIZE;
+    return ((storage_extended) ? (_partition[0]->size + _partition[1]->size) : _partition[0]->size) / FILESYSTEM_BLOCK_SIZE;
     #else
     return _partition[0]->size / FILESYSTEM_BLOCK_SIZE;
     #endif
 }
 
 void port_internal_flash_flush(void) {
-
 }
 
-mp_uint_t supervisor_flash_read_blocks(uint8_t *dest, uint32_t block, uint32_t num_blocks) {
+STATIC void single_partition_rw(const esp_partition_t *partition, uint8_t *data,
-    uint32_t offset = block * FILESYSTEM_BLOCK_SIZE;
+    const uint32_t offset, const uint32_t size_total, const bool op) {
-    uint32_t read_total = num_blocks * FILESYSTEM_BLOCK_SIZE;
+    if (op == OP_READ) {
-
+        esp_partition_read(partition, offset, data, size_total);
-    #if CIRCUITPY_STORAGE_EXTEND
+    } else {
-    if (offset > _partition[0]->size) {
+        esp_partition_erase_range(partition, offset, size_total);
-        // only read from partition 1
+        esp_partition_write(partition, offset, _cache, size_total);
-        esp_partition_read(_partition[1], (offset - _partition[0]->size), dest, read_total);
-    } else if ((offset + read_total) > _partition[0]->size) {
-        // first read from partition 0, then partition 1
-        uint32_t read_0 = _partition[0]->size - offset;
-        uint32_t read_1 = read_total - read_0;
-        esp_partition_read(_partition[0], offset, dest, read_0);
-        esp_partition_read(_partition[1], 0, (dest + read_0), read_1);
-    } else
-    #endif
-    {
-        // only read from partition 0
-        esp_partition_read(_partition[0], offset, dest, read_total);
     }
+}
 
-    return 0;
+#if CIRCUITPY_STORAGE_EXTEND
+STATIC void multi_partition_rw(uint8_t *data,
+    const uint32_t offset, const uint32_t size_total, const bool op) {
+    if (offset > _partition[0]->size) {
+        // only r/w partition 1
+        single_partition_rw(_partition[1], data, (offset - _partition[0]->size), size_total, op);
+    } else if ((offset + size_total) > _partition[0]->size) {
+        // first r/w partition 0, then partition 1
+        uint32_t size_0 = _partition[0]->size - offset;
+        uint32_t size_1 = size_total - size_0;
+        if (op == OP_READ) {
+            esp_partition_read(_partition[0], offset, data, size_0);
+            esp_partition_read(_partition[1], 0, (data + size_0), size_1);
+        } else {
+            esp_partition_erase_range(_partition[0], offset, size_0);
+            esp_partition_write(_partition[0], offset, _cache, size_0);
+            esp_partition_erase_range(_partition[1], 0, size_1);
+            esp_partition_write(_partition[1], 0, (_cache + size_0), size_1);
+        }
+    } else {
+        // only r/w partition 0
+        single_partition_rw(_partition[0], data, offset, size_total, op);
+    }
+}
+#endif
+
+mp_uint_t supervisor_flash_read_blocks(uint8_t *dest, uint32_t block, uint32_t num_blocks) {
+    const uint32_t offset = block * FILESYSTEM_BLOCK_SIZE;
+    const uint32_t read_total = num_blocks * FILESYSTEM_BLOCK_SIZE;
+    #if CIRCUITPY_STORAGE_EXTEND
+    multi_partition_rw(dest, offset, read_total, OP_READ);
+    #else
+    single_partition_rw(_partition[0], dest, offset, read_total, OP_READ);
+    #endif
+    return 0; // success
 }
 
 mp_uint_t supervisor_flash_write_blocks(const uint8_t *src, uint32_t lba, uint32_t num_blocks) {
@@ -112,12 +150,10 @@ mp_uint_t supervisor_flash_write_blocks(const uint8_t *src, uint32_t lba, uint32
         uint32_t block_address = lba + block;
         uint32_t sector_offset = block_address / blocks_per_sector * SECTOR_SIZE;
         uint8_t block_offset = block_address % blocks_per_sector;
-
         if (_cache_lba != block_address) {
             supervisor_flash_read_blocks(_cache, sector_offset / FILESYSTEM_BLOCK_SIZE, blocks_per_sector);
             _cache_lba = sector_offset;
         }
-
         for (uint8_t b = block_offset; b < blocks_per_sector; b++) {
             // Stop copying after the last block.
             if (block >= num_blocks) {
@@ -128,31 +164,34 @@ mp_uint_t supervisor_flash_write_blocks(const uint8_t *src, uint32_t lba, uint32
                 FILESYSTEM_BLOCK_SIZE);
             block++;
         }
-
         #if CIRCUITPY_STORAGE_EXTEND
-        if (sector_offset > _partition[0]->size) {
+        multi_partition_rw(_cache, sector_offset, SECTOR_SIZE, OP_WRITE);
-            // only write to partition 1
+        #else
-            esp_partition_erase_range(_partition[1], sector_offset - _partition[0]->size, SECTOR_SIZE);
+        single_partition_rw(_partition[0], _cache, sector_offset, SECTOR_SIZE, OP_READ);
-            esp_partition_write(_partition[1], sector_offset - _partition[0]->size, _cache, SECTOR_SIZE);
-        } else if ((sector_offset + SECTOR_SIZE) > _partition[0]->size) {
-            // first write to partition 0, then partition 1
-            uint32_t write_0 = _partition[0]->size - sector_offset;
-            uint32_t write_1 = SECTOR_SIZE - write_0;
-            esp_partition_erase_range(_partition[0], sector_offset, write_0);
-            esp_partition_write(_partition[0], sector_offset, _cache, write_0);
-            esp_partition_erase_range(_partition[1], 0, write_1);
-            esp_partition_write(_partition[1], 0, _cache + write_0, write_1);
-        } else
         #endif
-        {
-            // only write to partition 0
-            esp_partition_erase_range(_partition[0], sector_offset, SECTOR_SIZE);
-            esp_partition_write(_partition[0], sector_offset, _cache, SECTOR_SIZE);
-        }
     }
-
     return 0; // success
 }
 
 void supervisor_flash_release_cache(void) {
 }
+
+void supervisor_flash_set_extended(bool extended) {
+    #if CIRCUITPY_STORAGE_EXTEND
+    storage_extended = extended;
+    #endif
+}
+
+bool supervisor_flash_get_extended(void) {
+    #if CIRCUITPY_STORAGE_EXTEND
+    return storage_extended;
+    #else
+    return false;
+    #endif
+}
+
+void supervisor_flash_update_extended(void) {
+    #if CIRCUITPY_STORAGE_EXTEND
+    storage_extended = (_partition[0]->size < fatfs_bytes());
+    #endif
+}

supervisor/flash.h

@@ -50,4 +50,8 @@ void supervisor_flash_init_vfs(struct _fs_user_mount_t *vfs);
 void supervisor_flash_flush(void);
 void supervisor_flash_release_cache(void);
 
+void supervisor_flash_set_extended(bool extended);
+bool supervisor_flash_get_extended(void);
+void supervisor_flash_update_extended(void);
+
 #endif  // MICROPY_INCLUDED_SUPERVISOR_FLASH_H

supervisor/shared/filesystem.c

@@ -143,16 +143,22 @@ bool filesystem_init(bool create_allowed, bool force_create) {
     } else if (res != FR_OK) {
         return false;
     }
+
     vfs->str = "/";
     vfs->len = 1;
     vfs->obj = MP_OBJ_FROM_PTR(vfs_fat);
     vfs->next = NULL;
+
     MP_STATE_VM(vfs_mount_table) = vfs;
 
     // The current directory is used as the boot up directory.
     // It is set to the internal flash filesystem by default.
     MP_STATE_PORT(vfs_cur) = vfs;
 
+    #if CIRCUITPY_STORAGE_EXTEND
+    supervisor_flash_update_extended();
+    #endif
+
     return true;
 }