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/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;
+}