From b94afa1581a6cb035f02b0f98df5d9864aaa2d4e Mon Sep 17 00:00:00 2001
From: Astro <astro@spaceboyz.net>
Date: Sun, 15 Dec 2019 23:52:47 +0100
Subject: [PATCH] zynq::flash: add read_reg_until()

---
 src/zynq/flash/mod.rs                | 42 +++++++++++++++++++++++++---
 src/zynq/flash/spi_flash_register.rs | 17 +++++------
 2 files changed, 47 insertions(+), 12 deletions(-)

diff --git a/src/zynq/flash/mod.rs b/src/zynq/flash/mod.rs
index 70b62a3..64f3877 100644
--- a/src/zynq/flash/mod.rs
+++ b/src/zynq/flash/mod.rs
@@ -359,9 +359,40 @@ impl Flash<Manual> {
 
     pub fn read_reg<R: SpiFlashRegister>(&mut self) -> R {
         let args = Some(R::inst_code());
-        let transfer = self.transfer(args.into_iter(), R::transfer_len())
-            .bytes_transfer().skip(1);
-        R::new(transfer)
+        let transfer = self.transfer(args.into_iter(), 2)
+            .bytes_transfer();
+        R::new(transfer.skip(1).next().unwrap())
+    }
+
+    pub fn read_reg_until<R, F, A>(&mut self, f: F) -> A
+    where
+        R: SpiFlashRegister,
+        F: Fn(R) -> Option<A>,
+    {
+        let mut result = None;
+        while result.is_none() {
+            let args = Some(R::inst_code());
+            for b in self.transfer(args.into_iter(), 32)
+                .bytes_transfer().skip(1) {
+                    result = f(R::new(b));
+                    
+                    if result.is_none() {
+                        break;
+                    }
+                }
+        }
+        result.unwrap()
+    }
+
+    /// Status Register-1 remains `0x00` immediately after invoking a command.
+    fn wait_while_sr1_zeroed(&mut self) -> SR1 {
+        self.read_reg_until::<SR1, _, SR1>(|sr1|
+            if sr1.is_zeroed() {
+                None
+            } else {
+                Some(sr1)
+            }
+        )
     }
 
     /// Read Identification
@@ -385,7 +416,10 @@ impl Flash<Manual> {
         let args = Some(INST_WREN);
         self.transfer(args.into_iter(), 1);
         self.regs.gpio.modify(|_, w| w.wp_n(true));
-        while !self.read_reg::<SR1>().wel() {}
+        let sr1 = self.wait_while_sr1_zeroed();
+        if !sr1.wel() {
+            panic!("Cannot write-enable flash");
+        }
 
         let result = f(self);
 
diff --git a/src/zynq/flash/spi_flash_register.rs b/src/zynq/flash/spi_flash_register.rs
index ff03cb4..1f6fc5a 100644
--- a/src/zynq/flash/spi_flash_register.rs
+++ b/src/zynq/flash/spi_flash_register.rs
@@ -2,8 +2,7 @@ use bit_field::BitField;
 
 pub trait SpiFlashRegister {
     fn inst_code() -> u8;
-    fn transfer_len() -> usize;
-    fn new<I: Iterator<Item=u8>>(src: I) -> Self;
+    fn new(src: u8) -> Self;
 }
 
 macro_rules! u8_register {
@@ -18,16 +17,18 @@ macro_rules! u8_register {
                 $inst_code
             }
 
-            fn transfer_len() -> usize {
-                2
-            }
-
-            fn new<I: Iterator<Item=u8>>(mut src: I) -> Self {
+            fn new(src: u8) -> Self {
                 $name {
-                    inner: src.next().unwrap(),
+                    inner: src,
                 }
             }
         }
+
+        impl $name {
+            pub fn is_zeroed(&self) -> bool {
+                self.inner == 0
+            }
+        }
     };
 }