From b25606b95920bd06df87aff9202c7a15377d4a30 Mon Sep 17 00:00:00 2001
From: Roland Dobai <roland@espressif.com>
Date: Thu, 24 Nov 2022 14:09:34 +0100
Subject: [PATCH] feat(flash_id): Print the flash type if available for the
 chip

Closes https://github.com/espressif/esptool/issues/730
---
 esptool/cmds.py            |  5 +++++
 esptool/loader.py          |  4 ++++
 esptool/targets/esp32s2.py | 11 +++++++++++
 esptool/targets/esp32s3.py | 11 +++++++++++
 4 files changed, 31 insertions(+)

diff --git a/esptool/cmds.py b/esptool/cmds.py
index d1d92fa7f..61f2fd495 100644
--- a/esptool/cmds.py
+++ b/esptool/cmds.py
@@ -1012,6 +1012,11 @@ def flash_id(esp, args):
     print(
         "Detected flash size: %s" % (DETECTED_FLASH_SIZES.get(flid_lowbyte, "Unknown"))
     )
+    flash_type = esp.flash_type()
+    flash_type_dict = {0: "quad (4 data lines)", 1: "octal (8 data lines)"}
+    flash_type_str = flash_type_dict.get(flash_type)
+    if flash_type_str:
+        print(f"Flash type: {flash_type_str}")
 
 
 def read_flash(esp, args):
diff --git a/esptool/loader.py b/esptool/loader.py
index d32312386..6530182b3 100644
--- a/esptool/loader.py
+++ b/esptool/loader.py
@@ -867,6 +867,10 @@ def flash_id(self):
         SPIFLASH_RDID = 0x9F
         return self.run_spiflash_command(SPIFLASH_RDID, b"", 24)
 
+    def flash_type(self):
+        """Read flash type bit field from eFuse. Returns 0, 1, None (not present)"""
+        return None  # not implemented for all chip targets
+
     def get_security_info(self):
         res = self.check_command("get security info", self.ESP_GET_SECURITY_INFO, b"")
         esp32s2 = True if len(res) == 12 else False
diff --git a/esptool/targets/esp32s2.py b/esptool/targets/esp32s2.py
index 5bfefa6cf..810133ec9 100644
--- a/esptool/targets/esp32s2.py
+++ b/esptool/targets/esp32s2.py
@@ -68,6 +68,9 @@ class ESP32S2ROM(ESP32ROM):
     EFUSE_SECURE_BOOT_EN_REG = EFUSE_BASE + 0x038
     EFUSE_SECURE_BOOT_EN_MASK = 1 << 20
 
+    EFUSE_RD_REPEAT_DATA3_REG = EFUSE_BASE + 0x3C
+    EFUSE_RD_REPEAT_DATA3_REG_FLASH_TYPE_MASK = 1 << 9
+
     PURPOSE_VAL_XTS_AES256_KEY_1 = 2
     PURPOSE_VAL_XTS_AES256_KEY_2 = 3
     PURPOSE_VAL_XTS_AES128_KEY = 4
@@ -184,6 +187,14 @@ def read_mac(self):
         bitstring = struct.pack(">II", mac1, mac0)[2:]
         return tuple(bitstring)
 
+    def flash_type(self):
+        return (
+            1
+            if self.read_reg(self.EFUSE_RD_REPEAT_DATA3_REG)
+            & self.EFUSE_RD_REPEAT_DATA3_REG_FLASH_TYPE_MASK
+            else 0
+        )
+
     def get_flash_crypt_config(self):
         return None  # doesn't exist on ESP32-S2
 
diff --git a/esptool/targets/esp32s3.py b/esptool/targets/esp32s3.py
index 80469918e..4b9ef8e5d 100644
--- a/esptool/targets/esp32s3.py
+++ b/esptool/targets/esp32s3.py
@@ -71,6 +71,9 @@ class ESP32S3ROM(ESP32ROM):
     EFUSE_SECURE_BOOT_EN_REG = EFUSE_BASE + 0x038
     EFUSE_SECURE_BOOT_EN_MASK = 1 << 20
 
+    EFUSE_RD_REPEAT_DATA3_REG = EFUSE_BASE + 0x3C
+    EFUSE_RD_REPEAT_DATA3_REG_FLASH_TYPE_MASK = 1 << 9
+
     PURPOSE_VAL_XTS_AES256_KEY_1 = 2
     PURPOSE_VAL_XTS_AES256_KEY_2 = 3
     PURPOSE_VAL_XTS_AES128_KEY = 4
@@ -193,6 +196,14 @@ def read_mac(self):
         bitstring = struct.pack(">II", mac1, mac0)[2:]
         return tuple(bitstring)
 
+    def flash_type(self):
+        return (
+            1
+            if self.read_reg(self.EFUSE_RD_REPEAT_DATA3_REG)
+            & self.EFUSE_RD_REPEAT_DATA3_REG_FLASH_TYPE_MASK
+            else 0
+        )
+
     def uses_usb_otg(self, _cache=[]):
         """
         Check the UARTDEV_BUF_NO register to see if USB-OTG console is being used