From f69af1661923c05149b2a998d5d074d85c51026d Mon Sep 17 00:00:00 2001
From: David Lechner <david@pybricks.com>
Date: Fri, 29 Jul 2022 15:45:13 -0500
Subject: [PATCH] stm32/boards/LEGO_HUB_NO6/spiflash: Pick command type at
 runtime.

This changes spiflash.py to read the flash chip ID at runtime to select the
read/write/erase commands.  This will allow the code to be shared with
LEGO_HUB_NO7 which doesn't use the 32-bit commands.

Also remove an unused constant while we are touching this.

Signed-off-by: David Lechner <david@pybricks.com>
---
 ports/stm32/boards/LEGO_HUB_NO6/spiflash.py | 24 +++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/ports/stm32/boards/LEGO_HUB_NO6/spiflash.py b/ports/stm32/boards/LEGO_HUB_NO6/spiflash.py
index e052f7738b..e483ace950 100644
--- a/ports/stm32/boards/LEGO_HUB_NO6/spiflash.py
+++ b/ports/stm32/boards/LEGO_HUB_NO6/spiflash.py
@@ -4,18 +4,34 @@
 from micropython import const
 
 _PAGE_SIZE = const(256)  # maximum bytes writable in one SPI transfer
+_CMD_WRITE = const(0x02)
+_CMD_READ = const(0x03)
 _CMD_RDSR = const(0x05)
 _CMD_WREN = const(0x06)
 _CMD_WRITE_32 = const(0x12)
 _CMD_READ_32 = const(0x13)
+_CMD_SEC_ERASE = const(0x20)
 _CMD_SEC_ERASE_32 = const(0x21)
-_CMD_C4READ_32 = const(0xEC)
+_CMD_JEDEC_ID = const(0x9F)
 
 
 class SPIFlash:
     def __init__(self, spi, cs):
         self.spi = spi
         self.cs = cs
+        self.id = self._get_id()
+        # flash chip on Hub No. 6 uses 32-bit addressing
+        _32_bit = self.id == b"\xef\x40\x19"
+        self._READ = _CMD_READ_32 if _32_bit else _CMD_READ
+        self._WRITE = _CMD_WRITE_32 if _32_bit else _CMD_WRITE
+        self._ERASE = _CMD_SEC_ERASE_32 if _32_bit else _CMD_SEC_ERASE
+
+    def _get_id(self):
+        self.cs(0)
+        self.spi.write(bytearray([_CMD_JEDEC_ID]))
+        buf = self.spi.read(3)
+        self.cs(1)
+        return buf
 
     def _wait_wel1(self):
         # wait WEL=1
@@ -52,11 +68,11 @@ class SPIFlash:
         self._wait_wip0()
 
     def erase_block(self, addr):
-        self._flash_modify(_CMD_SEC_ERASE_32, addr, None)
+        self._flash_modify(self._ERASE, addr, None)
 
     def readinto(self, addr, buf):
         self.cs(0)
-        self.spi.write(bytearray([_CMD_READ_32, addr >> 16, addr >> 8, addr]))
+        self.spi.write(bytearray([self._READ, addr >> 16, addr >> 8, addr]))
         self.spi.readinto(buf)
         self.cs(1)
 
@@ -67,7 +83,7 @@ class SPIFlash:
         buf_offset = 0
         while remain:
             l = min(_PAGE_SIZE - offset, remain)
-            self._flash_modify(_CMD_WRITE_32, addr, buf[buf_offset : buf_offset + l])
+            self._flash_modify(self._WRITE, addr, buf[buf_offset : buf_offset + l])
             remain -= l
             addr += l
             buf_offset += l