flash_spi: Generalize to support chips with 24-bit addresses

Currently flash_spi.py is hardcoded for a small set of devices, all of
which can use 4-byte addressing. Fix this by allowing the command set and
address generation to be switched dynamically.

We also remove the whitelist of tested devices since the interfaces used
is very common amoung different flash devices. However as a safety measure
the user can provide the expected flash size and this will be checked
against the density field.

flash/flash_spi.py

@@ -1,5 +1,4 @@
-# flash_spi.py MicroPython driver for Cypress S25FL128L 16MiB and S25FL256L 32MiB
-# flash devices.
+# flash_spi.py MicroPython driver for SPI NOR flash devices.
 
 # Released under the MIT License (MIT). See LICENSE.
 # Copyright (c) 2019 Peter Hinch
@@ -9,10 +8,12 @@ from micropython import const
 from bdevice import FlashDevice
 
 # Supported instruction set:
-# 4 byte address commands
-_READ = const(0x13)
-_PP = const(0x12)  # Page program
-_SE = const(0x21)  # Sector erase
+# 3 and 4 byte address commands
+_READ = const(0)
+_PP = const(1)
+_SE = const(2)
+_CMDS3BA = b'\x03\x02\x20'
+_CMDS4BA = b'\x13\x12\x21'
 # No address
 _WREN = const(6)  # Write enable
 _RDSR1 = const(5)  # Read status register 1
@@ -25,23 +26,31 @@ _SEC_SIZE = const(4096)  # Flash sector size 0x1000
 # Logical Flash device comprising one or more physical chips sharing an SPI bus.
 class FLASH(FlashDevice):
 
-    def __init__(self, spi, cspins, size=16384, verbose=True, sec_size=_SEC_SIZE, block_size=9):
+    def __init__(self, spi, cspins, size=None, verbose=True, sec_size=_SEC_SIZE, block_size=9):
         # args: virtual block size in bits, no. of chips, bytes in each chip
-        if size not in (16384, 32768):
-            raise ValueError('Valid sizes: 16384 or 32768KiB')
-        super().__init__(block_size, len(cspins), size * 1024, sec_size)
         self._spi = spi
         self._cspins = cspins
         self._ccs = None  # Chip select Pin object for current chip
         self._bufp = bytearray(6)  # instruction + 4 byte address + 1 byte value
         self._mvp = memoryview(self._bufp)  # cost-free slicing
         self._page_size = 256  # Write uses 256 byte pages.
-        self.scan(verbose)
+
+        size = self.scan(verbose, size)
+        super().__init__(block_size, len(cspins), size * 1024, sec_size)
+
+        # Select the correct command set
+        if size <= 4096:
+            self._cmds = _CMDS3BA
+            self._cmdlen = 4
+        else:
+            self._cmds = _CMDS4BA
+            self._cmdlen = 5
+
         self.initialise()  # Initially cache sector 0
 
     # **** API SPECIAL METHODS ****
     # Scan: read manf ID
-    def scan(self, verbose):
+    def scan(self, verbose, size):
         mvp = self._mvp
         for n, cs in enumerate(self._cspins):
             mvp[:] = b'\0\0\0\0\0\0'
@@ -49,12 +58,16 @@ class FLASH(FlashDevice):
             cs(0)
             self._spi.write_readinto(mvp[:4], mvp[:4])
             cs(1)
-            if mvp[1] != 1 or mvp[2] != 0x60 or not (mvp[3] == 0x18 or mvp[3] == 0x19):
-                raise RuntimeError('Flash not found at cs[{}].'.format(n))
+            scansize = 1 << (mvp[3] - 10)
+            if not size:
+                size = scansize
+            if size != scansize:
+                raise ValueError('Flash size mismatch: expected {}KiB, found {}KiB'.format(size, scansize))
         if verbose:
             s = '{} chips detected. Total flash size {}MiB.'
-            print(s.format(n + 1, self._a_bytes // (1024 * 1024)))
-        return n
+            n += 1
+            print(s.format(n, (n * size) // 1024))
+        return size
 
     # Chip erase. Can take minutes.
     def erase(self):
@@ -86,9 +99,9 @@ class FLASH(FlashDevice):
             cs(0)
             self._spi.write(mvp[:1])  # Enable write
             cs(1)
-            mvp[0] = _PP
+            mvp[0] = self._cmds[_PP]
             cs(0)
-            self._spi.write(mvp[:5])  # Start write
+            self._spi.write(mvp[:self._cmdlen])  # Start write
             self._spi.write(cache[start : start + ps])
             cs(1)
             self._wait_rdy()  # Wait for write to complete
@@ -104,9 +117,9 @@ class FLASH(FlashDevice):
         while nbytes > 0:
             npage = self._getaddr(addr, nbytes)  # No. of bytes in current chip
             cs = self._ccs
-            mvp[0] = _READ
+            mvp[0] = self._cmds[_READ]
             cs(0)
-            self._spi.write(mvp[:5])
+            self._spi.write(mvp[:self._cmdlen])
             self._spi.readinto(mvb[start : start + npage])
             cs(1)
 #            print('addr {} npage {} data {}'.format(addr, npage, mvb[start]))
@@ -141,11 +154,13 @@ class FLASH(FlashDevice):
             raise RuntimeError("Flash Address is out of range")
         ca, la = divmod(addr, self._c_bytes)  # ca == chip no, la == offset into chip
         self._ccs = self._cspins[ca]  # Current chip select
-        mvp = self._mvp
-        mvp[1] = la >> 24
-        mvp[2] = la >> 16 & 0xff
-        mvp[3] = (la >> 8) & 0xff
-        mvp[4] = la & 0xff
+        cmdlen = self._cmdlen
+        mvp = self._mvp[:cmdlen]
+        if cmdlen > 3:
+            mvp[-4] = la >> 24
+        mvp[-3] = la >> 16 & 0xff
+        mvp[-2] = (la >> 8) & 0xff
+        mvp[-1] = la & 0xff
         pe = (addr & -self._c_bytes) + self._c_bytes  # Byte 0 of next chip
         return min(nbytes, pe - la)
 
@@ -160,8 +175,8 @@ class FLASH(FlashDevice):
             cs(0)
             self._spi.write(mvp[:1])  # Enable write
             cs(1)
-            mvp[0] = _SE
+            mvp[0] = self._cmds[_SE]
             cs(0)
-            self._spi.write(mvp[:5])  # Start erase
+            self._spi.write(mvp[:self._cmdlen])  # Start erase
             cs(1)
             self._wait_rdy()  # Wait for erase to complete