spi25: Integrate 4BA support

Back port the upstream `commit f43c654a` reworked for our tree.

BUG=chromium:478356
BRANCH=none
TEST=still builds

Change-Id: I95c12d54fcc4acfbce00e59e11d00a324cbfaf7f
Signed-off-by: Edward O'Callaghan <quasisec@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/flashrom/+/1667513
Reviewed-by: Stefan Reinauer <reinauer@google.com>
diff --git a/flash.h b/flash.h
index b74b428..6fa0ead 100644
--- a/flash.h
+++ b/flash.h
@@ -22,6 +22,7 @@
 
 #include <stdint.h>
 #include <stddef.h>
+#include <stdbool.h>
 #include "hwaccess.h"
 #ifdef _WIN32
 #include <windows.h>
@@ -127,6 +128,8 @@
 #define FEATURE_UNBOUND_READ	(1 << 10)
 #define FEATURE_NO_ERASE	(1 << 11)
 #define FEATURE_4BA_SUPPORT	(1 << 12)
+#define FEATURE_4BA_EXT_ADDR	(1 << 13) /**< Regular 3-byte operations can be used by writing the most
+						significant address byte into an extended address register. */
 
 struct voltage_range {
 	uint16_t min, max;
@@ -248,6 +251,14 @@
 	/* Some flash devices have an additional register space. */
 	chipaddr virtual_registers;
 	struct registered_master *mst;
+
+	/* We cache the state of the extended address register (highest byte
+	 * of a 4BA for 3BA instructions) and the state of the 4BA mode here.
+	 * If possible, we enter 4BA mode early. If that fails, we make use
+	 * of the extended address register.
+	 */
+	int address_high_byte;
+	bool in_4ba_mode;
 };
 
 
diff --git a/flashrom.c b/flashrom.c
index 4b59f82..2ab9d63 100644
--- a/flashrom.c
+++ b/flashrom.c
@@ -2030,6 +2030,9 @@
 	if (flash->chip->unlock)
 		flash->chip->unlock(flash);
 
+	flash->address_high_byte = -1;
+	flash->in_4ba_mode = false;
+
 	/* Enable/disable 4-byte addressing mode if flash chip supports it */
 	if ((flash->chip->feature_bits & FEATURE_4BA_SUPPORT) &&
 	    flash->chip->four_bytes_addr_funcs.set_4ba) {
diff --git a/spi25.c b/spi25.c
index 6b9dce6..bc54c6d 100644
--- a/spi25.c
+++ b/spi25.c
@@ -340,14 +340,37 @@
 	return result ? result : status;
 }
 
+static int spi_set_extended_address(struct flashctx *const flash, const uint8_t addr_high)
+{
+	if (flash->address_high_byte != addr_high &&
+	    spi_write_extended_address_register(flash, addr_high))
+		return -1;
+	flash->address_high_byte = addr_high;
+	return 0;
+}
+
 static int spi_prepare_address(struct flashctx *const flash,
 			       uint8_t cmd_buf[], const unsigned int addr)
 {
-	/* TODO: extend for 4BA */
-	cmd_buf[1] = (addr >> 16) & 0xff;
-	cmd_buf[2] = (addr >>  8) & 0xff;
-	cmd_buf[3] = (addr >>  0) & 0xff;
-	return 3;
+	if (flash->in_4ba_mode) {
+		cmd_buf[1] = (addr >> 24) & 0xff;
+		cmd_buf[2] = (addr >> 16) & 0xff;
+		cmd_buf[3] = (addr >>  8) & 0xff;
+		cmd_buf[4] = (addr >>  0) & 0xff;
+		return 4;
+	} else {
+		if (flash->chip->feature_bits & FEATURE_4BA_EXT_ADDR) {
+			if (spi_set_extended_address(flash, addr >> 24))
+				return -1;
+		} else {
+			if (addr >> 24)
+				return -1;
+		}
+		cmd_buf[1] = (addr >> 16) & 0xff;
+		cmd_buf[2] = (addr >>  8) & 0xff;
+		cmd_buf[3] = (addr >>  0) & 0xff;
+		return 3;
+	}
 }
 
 /**
diff --git a/spi4ba.c b/spi4ba.c
index 74fe507..703f30b 100644
--- a/spi4ba.c
+++ b/spi4ba.c
@@ -11,7 +11,6 @@
  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  * GNU General Public License for more details.
- *
  */
 
 /*
@@ -36,12 +35,17 @@
 /* Enter 4-bytes addressing mode (without sending WREN before) */
 int spi_enter_4ba_b7(struct flashctx *flash)
 {
+	int result;
 	const unsigned char cmd[JEDEC_ENTER_4_BYTE_ADDR_MODE_OUTSIZE] = { JEDEC_ENTER_4_BYTE_ADDR_MODE };
 
 	msg_trace("-> %s\n", __func__);
 
 	/* Switch to 4-bytes addressing mode */
-	return spi_send_command(flash, sizeof(cmd), 0, cmd, NULL);
+	result = spi_send_command(flash, sizeof(cmd), 0, cmd, NULL);
+	if (!result)
+		flash->in_4ba_mode = true;
+
+	return result;
 }
 
 /* Enter 4-bytes addressing mode with sending WREN before */
@@ -72,18 +76,25 @@
 	result = spi_send_multicommand(flash, cmds);
 	if (result)
 		msg_cerr("%s failed during command execution\n", __func__);
+	else
+		flash->in_4ba_mode = true;
 	return result;
 }
 
 /* Exit 4-bytes addressing mode (without sending WREN before) */
 int spi_exit_4ba_e9(struct flashctx *flash)
 {
+	int result;
 	const unsigned char cmd[JEDEC_EXIT_4_BYTE_ADDR_MODE_OUTSIZE] = { JEDEC_EXIT_4_BYTE_ADDR_MODE };
 
 	msg_trace("-> %s\n", __func__);
 
 	/* Switch to 3-bytes addressing mode  */
-	return spi_send_command(flash, sizeof(cmd), 0, cmd, NULL);
+	result = spi_send_command(flash, sizeof(cmd), 0, cmd, NULL);
+	if (!result)
+		flash->in_4ba_mode = false;
+
+	return result;
 }
 
 /* Exit 4-bytes addressing mode with sending WREN before */
@@ -112,9 +123,10 @@
 
 	/* Switch to 3-bytes addressing mode  */
 	result = spi_send_multicommand(flash, cmds);
-	if (result) {
+	if (result)
 		msg_cerr("%s failed during command execution\n", __func__);
-	}
+	else
+		flash->in_4ba_mode = false;
 	return result;
 }
 
diff --git a/spi4ba.h b/spi4ba.h
index efe1021..6202953 100644
--- a/spi4ba.h
+++ b/spi4ba.h
@@ -111,5 +111,6 @@
 int spi_block_erase_5c_4ba_direct(struct flashctx *flash, unsigned int addr, unsigned int blocklen);
 int spi_block_erase_dc_4ba_direct(struct flashctx *flash, unsigned int addr, unsigned int blocklen);
 
+int spi_write_extended_address_register(struct flashctx *flash, uint8_t regdata);
 
 #endif /* __SPI_4BA_H__ */