1 files changed, 56 insertions, 61 deletions

diff --git a/embassy-rp/src/pio_programs/onewire.rs b/embassy-rp/src/pio_programs/onewire.rs
index 9f2ed5695..82fd98b96 100644
--- a/embassy-rp/src/pio_programs/onewire.rs
+++ b/embassy-rp/src/pio_programs/onewire.rs
@@ -12,7 +12,6 @@ pub struct PioOneWireProgram<'a, PIO: Instance> {
     prg: LoadedProgram<'a, PIO>,
     reset_addr: u8,
     next_bit_addr: u8,
-    search_addr: u8,
 }
 
 impl<'a, PIO: Instance> PioOneWireProgram<'a, PIO> {
@@ -53,35 +52,17 @@ impl<'a, PIO: Instance> PioOneWireProgram<'a, PIO> {
                     ; The low pulse was already done, we only need to delay and poll the bit in case we are reading
                     write_1:
                         nop                     side 0 [( 6 / CLK) - 1]     ; Delay before sampling the input pin
-                        in pins, 1              side 0 [(54 / CLK) - 1]     ; This writes the state of the pin into the ISR
+                        in pins, 1              side 0 [(48 / CLK) - 1]     ; This writes the state of the pin into the ISR
                         ; Fallthrough
 
                 ; This is the entry point when reading and writing data
                 public next_bit:
                     .wrap_target
-                        out x, 1                side 0 [(18 / CLK) - 1]     ; Stalls if no data available in TX FIFO and OSR
-                        jmp x--, write_1        side 1 [(12 / CLK) - 1]     ; Do the always low part of a bit, jump to write_1 if we want to write a 1 bit
-                        in null, 1              side 1 [(60 / CLK) - 1]     ; Do the remainder of the low part of a 0 bit
+                        out x, 1                side 0 [(12 / CLK) - 1]     ; Stalls if no data available in TX FIFO and OSR
+                        jmp x--, write_1        side 1 [( 6 / CLK) - 1]     ; Do the always low part of a bit, jump to write_1 if we want to write a 1 bit
+                        in null, 1              side 1 [(54 / CLK) - 1]     ; Do the remainder of the low part of a 0 bit
                                                                             ; This writes 0 into the ISR so that the shift count stays in sync
                     .wrap
-
-                public search:
-                        set x, 1                side 0 [(78 / CLK) - 1]     ; Set x to 1 for the inner loop
-                    search_inner:
-                    ; Read 2 bits
-                        nop                     side 1 [(12 / CLK) - 1]     ; Do the always low part of a bit
-                        nop                     side 0 [( 6 / CLK) - 1]     ; Delay before sampling the input pin
-                        in pins, 1              side 0 [(54 / CLK) - 1]     ; This writes the state of the pin into the ISR
-                        jmp x--, search_inner   side 0 [(18 / CLK) - 1]
-                        ; Fallthrough
-
-                    ; Write output
-                        out x, 1                side 0 [( 6 / CLK) - 1]     ; Stalls if no data available in TX FIFO and OSR
-                        jmp x--, search         side 1 [(12 / CLK) - 1]     ; Do the always low part of a bit, jump to search to write a 1 bit
-                        ; Fallthrough
-
-                        set x, 1                side 1 [(60 / CLK) - 1]     ; Set x to 1 for the inner loop, write the remainder of the low part of a 0 bit
-                        jmp search_inner        side 0 [(18 / CLK) - 1]
             "#
         );
 
@@ -89,7 +70,6 @@ impl<'a, PIO: Instance> PioOneWireProgram<'a, PIO> {
             prg: common.load_program(&prg.program),
             reset_addr: prg.public_defines.reset as u8,
             next_bit_addr: prg.public_defines.next_bit as u8,
-            search_addr: prg.public_defines.search as u8,
         }
     }
 }
@@ -98,7 +78,6 @@ pub struct PioOneWire<'d, PIO: Instance, const SM: usize> {
     sm: StateMachine<'d, PIO, SM>,
     cfg: Config<'d, PIO>,
     reset_addr: u8,
-    search_addr: u8,
     next_bit_addr: u8,
 }
 
@@ -119,13 +98,16 @@ impl<'d, PIO: Instance, const SM: usize> PioOneWire<'d, PIO, SM> {
         cfg.use_program(&program.prg, &[&pin]);
         cfg.set_in_pins(&[&pin]);
 
-        let byte_shift = ShiftConfig {
+        cfg.shift_in = ShiftConfig {
+            auto_fill: true,
+            direction: ShiftDirection::Right,
+            threshold: 8,
+        };
+        cfg.shift_out = ShiftConfig {
             auto_fill: true,
             direction: ShiftDirection::Right,
             threshold: 8,
         };
-        cfg.shift_in = byte_shift;
-        cfg.shift_out = byte_shift;
 
         let divider = (clk_sys_freq() / 1000000) as u16 * 6;
         cfg.clock_divider = divider.into();
@@ -142,11 +124,11 @@ impl<'d, PIO: Instance, const SM: usize> PioOneWire<'d, PIO, SM> {
             sm,
             cfg,
             reset_addr: program.reset_addr,
-            search_addr: program.search_addr,
             next_bit_addr: program.next_bit_addr,
         }
     }
 
+    /// Perform an initialization sequence, will return true if a presence pulse was detected from a device
     pub async fn reset(&mut self) -> bool {
         // The state machine immediately starts running when jumping to this address
         unsafe {
@@ -164,18 +146,18 @@ impl<'d, PIO: Instance, const SM: usize> PioOneWire<'d, PIO, SM> {
         found
     }
 
-    /// Write bytes over the wire
+    /// Write bytes to the onewire bus
     pub async fn write_bytes(&mut self, data: &[u8]) {
         let (rx, tx) = self.sm.rx_tx();
         for b in data {
             tx.wait_push(*b as u32).await;
 
-            // Empty the buffer that is always filled
+            // Empty the buffer that is being filled with every write
             let _ = rx.wait_pull().await;
         }
     }
 
-    /// Read bytes from the wire
+    /// Read bytes from the onewire bus
     pub async fn read_bytes(&mut self, data: &mut [u8]) {
         let (rx, tx) = self.sm.rx_tx();
         for b in data {
@@ -187,19 +169,29 @@ impl<'d, PIO: Instance, const SM: usize> PioOneWire<'d, PIO, SM> {
     }
 
     async fn search(&mut self, state: &mut PioOneWireSearch) -> Option<u64> {
-        let _ = self.reset().await;
-        self.write_bytes(&[0xF0]).await;
+        if !self.reset().await {
+            // No device present, no use in searching
+            state.finished = true;
+            return None;
+        }
+        self.write_bytes(&[0xF0]).await; // 0xF0 is the search rom command
 
-        let shift_cfg = self.prepare_search();
+        self.prepare_search();
 
         let (rx, tx) = self.sm.rx_tx();
 
-        let mut value = 0u64;
+        let mut value = 0;
         let mut last_zero = 0;
 
         for bit in 0..64 {
-            let push = match rx.wait_pull().await {
-                0b00 => {
+            // Write 2 dummy bits to read a bit and its complement
+            tx.wait_push(0x1).await;
+            tx.wait_push(0x1).await;
+            let in1 = rx.wait_pull().await;
+            let in2 = rx.wait_pull().await;
+            let push = match (in1, in2) {
+                (0, 0) => {
+                    // If both are 0, it means we have devices with 0 and 1 bits in this position
                     let write_value = if bit < state.last_discrepancy {
                         (state.last_rom & (1 << bit)) != 0
                     } else {
@@ -213,10 +205,11 @@ impl<'d, PIO: Instance, const SM: usize> PioOneWire<'d, PIO, SM> {
                         0
                     }
                 }
-                0b01 => 0,
-                0b10 => 1,
+                (0, 1) => 0, // Only devices with a 0 bit in this position
+                (1, 0) => 1, // Only devices with a 1 bit in this position
                 _ => {
-                    self.restore_after_search(&shift_cfg);
+                    // If both are 1, it means there is no device active and there is no point in continuing
+                    self.restore_after_search();
                     state.finished = true;
                     return None;
                 }
@@ -226,9 +219,10 @@ impl<'d, PIO: Instance, const SM: usize> PioOneWire<'d, PIO, SM> {
                 value |= 1 << 63;
             }
             tx.wait_push(push).await;
+            let _ = rx.wait_pull().await; // Discard the result of the write action
         }
 
-        self.restore_after_search(&shift_cfg);
+        self.restore_after_search();
 
         state.last_discrepancy = last_zero;
         state.finished = last_zero == 0;
@@ -236,44 +230,42 @@ impl<'d, PIO: Instance, const SM: usize> PioOneWire<'d, PIO, SM> {
         Some(value)
     }
 
-    fn prepare_search(&mut self) -> ShiftConfig {
-        let shift_cfg = self.cfg.shift_in;
-        self.cfg.shift_in = ShiftConfig {
-            auto_fill: true,
-            direction: ShiftDirection::Left,
-            threshold: 2,
-        };
-        self.cfg.shift_out = ShiftConfig {
-            auto_fill: true,
-            direction: ShiftDirection::Right,
-            threshold: 1,
-        };
+    fn prepare_search(&mut self) {
+        self.cfg.shift_in.threshold = 1;
+        self.cfg.shift_in.direction = ShiftDirection::Left;
+        self.cfg.shift_out.threshold = 1;
 
         self.sm.set_enable(false);
         self.sm.set_config(&self.cfg);
 
+        // set_config jumps to the wrong address so jump to the right one here
         unsafe {
-            self.sm.exec_jmp(self.search_addr);
+            self.sm.exec_jmp(self.next_bit_addr);
         }
         self.sm.set_enable(true);
-        shift_cfg
     }
 
-    fn restore_after_search(&mut self, cfg: &ShiftConfig) {
-        self.cfg.shift_in = *cfg;
-        self.cfg.shift_out = *cfg;
+    fn restore_after_search(&mut self) {
+        self.cfg.shift_in.threshold = 8;
+        self.cfg.shift_in.direction = ShiftDirection::Right;
+        self.cfg.shift_out.threshold = 8;
 
         self.sm.set_enable(false);
         self.sm.set_config(&self.cfg);
+
+        // Clear the state in case we aborted prematurely with some bits still in the shift registers
+        self.sm.clear_fifos();
+        self.sm.restart();
+
+        // set_config jumps to the wrong address so jump to the right one here
         unsafe {
             self.sm.exec_jmp(self.next_bit_addr);
         }
-        self.sm.clear_fifos();
-        self.sm.restart();
         self.sm.set_enable(true);
     }
 }
 
+/// Onewire search state
 pub struct PioOneWireSearch {
     last_rom: u64,
     last_discrepancy: u8,
@@ -281,6 +273,7 @@ pub struct PioOneWireSearch {
 }
 
 impl PioOneWireSearch {
+    /// Create a new Onewire search state
     pub fn new() -> Self {
         Self {
             last_rom: 0,
@@ -289,6 +282,7 @@ impl PioOneWireSearch {
         }
     }
 
+    /// Search for the next address on the bus
     pub async fn next<PIO: Instance, const SM: usize>(&mut self, pio: &mut PioOneWire<'_, PIO, SM>) -> Option<u64> {
         if self.finished {
             None
@@ -297,6 +291,7 @@ impl PioOneWireSearch {
         }
     }
 
+    /// Is finished when all devices have been found
     pub fn is_finished(&self) -> bool {
         self.finished
     }