edi: Split preparation/finalization out of edi_probe_kb9012()

This turns edi_probe_kb9012() into a pure probing function. To avoid
turning EDI off after probing, register edi_finish() only after full
preparation.

Change-Id: Icc342b8ab109d5621d8b65c79cecf71a44bea4bd
Signed-off-by: Nico Huber <nico.h@gmx.de>
Reviewed-on: https://review.sourcearcade.org/c/flashprog/+/423
diff --git a/edi.c b/edi.c
index 1f561f7..8adba16 100644
--- a/edi.c
+++ b/edi.c
@@ -466,34 +466,20 @@
 	return 0;
 }
 
-static int edi_shutdown(void *data)
+static void edi_finish(struct flashctx *flash)
 {
-	const struct spi_master *const spi = data;
-	int rc;
+	const struct spi_master *const spi = flash->mst.spi;
 
-	if (data == NULL)
-		return -1;
-
-	rc = edi_8051_execute(spi);
-	if (rc < 0) {
+	if (edi_8051_execute(spi) < 0)
 		msg_perr("%s: Unable to execute 8051!\n", __func__);
-		return -1;
-	}
 
-	rc = edi_disable(spi);
-	if (rc < 0) {
+	if (edi_disable(spi) < 0)
 		msg_perr("%s: Unable to disable EDI!\n", __func__);
-		return -1;
-	}
-
-	return 0;
 }
 
 int edi_probe_kb9012(struct flashctx *flash)
 {
 	const struct spi_master *const spi = flash->mst.spi;
-	int probe;
-	int rc;
 	unsigned char hwversion;
 
 	/*
@@ -507,17 +493,22 @@
 	 */
 	edi_read(spi, ENE_EC_HWVERSION, &hwversion);
 
-	probe = edi_chip_probe(spi, &ene_kb9012);
-	if (!probe)
+	return edi_chip_probe(spi, &ene_kb9012);
+}
+
+int edi_prepare(struct flashctx *flash, enum preparation_steps step)
+{
+	int rc;
+
+	if (step < PREPARE_FULL)
 		return 0;
 
-	rc = edi_8051_reset(spi);
+	rc = edi_8051_reset(flash->mst.spi);
 	if (rc < 0) {
 		msg_perr("%s: Unable to reset 8051!\n", __func__);
-		return 0;
+		return rc;
 	}
 
-	register_shutdown(edi_shutdown, (void *)spi);
-
-	return 1;
+	flash->chip->finish_access = edi_finish;
+	return 0;
 }
diff --git a/flashchips.c b/flashchips.c
index 37f85b6..02aaf88 100644
--- a/flashchips.c
+++ b/flashchips.c
@@ -4205,6 +4205,7 @@
 		.read		= edi_chip_read,
 		.voltage	= {2700, 3600},
 		.gran		= write_gran_128bytes,
+		.prepare_access	= edi_prepare,
 	},
 
 	{
diff --git a/include/chipdrivers/edi.h b/include/chipdrivers/edi.h
index 9cf2a28..48f24be 100644
--- a/include/chipdrivers/edi.h
+++ b/include/chipdrivers/edi.h
@@ -25,5 +25,6 @@
 int edi_chip_write(struct flashprog_flashctx *, const uint8_t *buf, unsigned int start, unsigned int len);
 int edi_chip_read(struct flashprog_flashctx *, uint8_t *buf, unsigned int start, unsigned int len);
 int edi_probe_kb9012(struct flashprog_flashctx *);
+int edi_prepare(struct flashprog_flashctx *, enum preparation_steps);
 
 #endif /* !__CHIPDRIVERS_EDI_H__ */