spi25: Move 4BA preparations into spi_prepare_4ba() hook

These preparations are specific to 4BA SPI chips and don't have to
clutter `flashprog.c`.

Change-Id: I842244c57e575f93b9c505e16f1f20c7afd23733
Signed-off-by: Nico Huber <nico.h@gmx.de>
Reviewed-on: https://review.sourcearcade.org/c/flashprog/+/72517
diff --git a/flashchips.c b/flashchips.c
index dfce11d..fc646b3 100644
--- a/flashchips.c
+++ b/flashchips.c
@@ -6857,6 +6857,7 @@
 			.tb     = {STATUS1, 6, RW},
 		},
 		.decode_range	= decode_range_spi25,
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -7565,6 +7566,7 @@
 		.write		= spi_chip_write_256,
 		.read		= spi_chip_read,
 		.voltage	= {2300, 3600},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -7741,6 +7743,7 @@
 		.write		= spi_chip_write_256,
 		.read		= spi_chip_read,
 		.voltage	= {1650, 1950},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -8797,6 +8800,7 @@
 		.write		= spi_chip_write_256,
 		.read		= spi_chip_read, /* Fast read (0x0B) supported */
 		.voltage	= {2700, 3600},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -9535,6 +9539,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {1650, 2000},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -9626,6 +9631,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {1650, 2000},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -10206,6 +10212,7 @@
 		.write		= spi_chip_write_256,
 		.read		= spi_chip_read, /* Fast read (0x0B) supported */
 		.voltage	= {2700, 3600},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -10255,6 +10262,7 @@
 		.write		= spi_chip_write_256,
 		.read		= spi_chip_read, /* Fast read (0x0B) supported */
 		.voltage	= {2700, 3600},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	/* The ST M25P05 is a bit of a problem. It has the same ID as the
@@ -11109,6 +11117,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {1700, 2000},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -11149,6 +11158,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {2700, 3600},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -11474,6 +11484,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {1700, 2000},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -11514,6 +11525,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {2700, 3600},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -11554,6 +11566,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {1700, 2000},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -11594,6 +11607,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {2700, 3600},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -11640,6 +11654,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {2700, 3600},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -11686,6 +11701,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {1700, 2000},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -11732,6 +11748,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {2700, 3600},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -11778,6 +11795,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {1700, 2000},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -11907,6 +11925,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {2700, 3600},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -11956,6 +11975,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {1700, 2000},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -12012,6 +12032,7 @@
 			.tb     = {STATUS1, 5, RW},
 		},
 		.decode_range	= decode_range_spi25,
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -12061,6 +12082,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {1700, 2000},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -16854,6 +16876,7 @@
 			.wps	= {STATUS3, 2, RW},
 		},
 		.decode_range	= decode_range_spi25,
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -16900,6 +16923,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {2700, 3600},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -16937,6 +16961,7 @@
 		.write		= spi_chip_write_256, /* Multi I/O supported, IGNORE for now */
 		.read		= spi_chip_read, /* Fast read (0x0B) and multi I/O supported */
 		.voltage	= {2700, 3600},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -17671,6 +17696,7 @@
 			.cmp    = {STATUS2, 6, RW},
 		},
 		.decode_range	= decode_range_spi25,
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -17726,6 +17752,7 @@
 			.cmp    = {STATUS2, 6, RW},
 		},
 		.decode_range	= decode_range_spi25,
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -17781,6 +17808,7 @@
 			.cmp    = {STATUS2, 6, RW},
 		},
 		.decode_range	= decode_range_spi25,
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -17827,6 +17855,7 @@
 		.write		= spi_chip_write_256,
 		.read		= spi_chip_read,
 		.voltage	= {1650, 1950},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -17884,6 +17913,7 @@
 			.wps	= {STATUS3, 2, RW},
 		},
 		.decode_range	= decode_range_spi25,
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -18406,6 +18436,7 @@
 		.write		= spi_chip_write_256,
 		.read		= spi_chip_read,
 		.voltage	= {2700, 3600},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -18460,6 +18491,7 @@
 			.wps	= {STATUS3, 2, RW},
 		},
 		.decode_range	= decode_range_spi25,
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -20012,6 +20044,7 @@
 			.tb     = {STATUS1, 6, RW},
 		},
 		.decode_range	= decode_range_spi25,
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
@@ -20134,6 +20167,7 @@
 		.write		= spi_chip_write_256,
 		.read		= spi_chip_read,
 		.voltage	= {1650, 1950},
+		.prepare_access	= spi_prepare_4ba,
 	},
 
 	{
diff --git a/flashprog.c b/flashprog.c
index 6a4ce12..937a0a7 100644
--- a/flashprog.c
+++ b/flashprog.c
@@ -1640,32 +1640,6 @@
 	if (flash->chip->unlock)
 		flash->chip->unlock(flash);
 
-	flash->address_high_byte = -1;
-	flash->in_4ba_mode = false;
-
-	/* Be careful about 4BA chips and broken masters */
-	if (flash->chip->total_size > 16 * 1024 && spi_master_no_4ba_modes(flash)) {
-		/* If we can't use native instructions, bail out */
-		if ((flash->chip->feature_bits & FEATURE_4BA_NATIVE) != FEATURE_4BA_NATIVE
-		    || !spi_master_4ba(flash)) {
-			msg_cerr("Programmer doesn't support this chip. Aborting.\n");
-			return 1;
-		}
-	}
-
-	/* Enable/disable 4-byte addressing mode if flash chip supports it */
-	if (flash->chip->feature_bits & (FEATURE_4BA_ENTER | FEATURE_4BA_ENTER_WREN | FEATURE_4BA_ENTER_EAR7)) {
-		int ret;
-		if (spi_master_4ba(flash))
-			ret = spi_enter_4ba(flash);
-		else
-			ret = spi_exit_4ba(flash);
-		if (ret) {
-			msg_cerr("Failed to set correct 4BA mode! Aborting.\n");
-			return 1;
-		}
-	}
-
 	return 0;
 }
 
diff --git a/include/chipdrivers.h b/include/chipdrivers.h
index 31afc84..8e2c520 100644
--- a/include/chipdrivers.h
+++ b/include/chipdrivers.h
@@ -61,6 +61,7 @@
 int spi_enter_4ba(struct flashctx *flash);
 int spi_exit_4ba(struct flashctx *flash);
 int spi_set_extended_address(struct flashctx *, uint8_t addr_high);
+int spi_prepare_4ba(struct flashctx *, enum preparation_steps);
 
 
 /* spi25_statusreg.c */
diff --git a/spi25.c b/spi25.c
index e4efd1c..d6c00c2 100644
--- a/spi25.c
+++ b/spi25.c
@@ -835,3 +835,37 @@
 {
 	return spi_enter_exit_4ba(flash, false);
 }
+
+int spi_prepare_4ba(struct flashctx *const flash, const enum preparation_steps prep)
+{
+	if (prep != PREPARE_FULL)
+		return 0;
+
+	flash->address_high_byte = -1;
+	flash->in_4ba_mode = false;
+
+	/* Be careful about 4BA chips and broken masters */
+	if (flash->chip->total_size > 16 * 1024 && spi_master_no_4ba_modes(flash)) {
+		/* If we can't use native instructions, bail out */
+		if ((flash->chip->feature_bits & FEATURE_4BA_NATIVE) != FEATURE_4BA_NATIVE
+		    || !spi_master_4ba(flash)) {
+			msg_cerr("Programmer doesn't support this chip. Aborting.\n");
+			return 1;
+		}
+	}
+
+	/* Enable/disable 4-byte addressing mode if flash chip supports it */
+	if (flash->chip->feature_bits & (FEATURE_4BA_ENTER | FEATURE_4BA_ENTER_WREN | FEATURE_4BA_ENTER_EAR7)) {
+		int ret;
+		if (spi_master_4ba(flash))
+			ret = spi_enter_4ba(flash);
+		else
+			ret = spi_exit_4ba(flash);
+		if (ret) {
+			msg_cerr("Failed to set correct 4BA mode! Aborting.\n");
+			return 1;
+		}
+	}
+
+	return 0;
+}