Rework scan result reading, fix C2

esp-wifi/src/wifi/mod.rs

@@ -866,60 +866,138 @@ impl<'d> WifiController<'d> {
 
         Ok(mode == wifi_mode_t_WIFI_MODE_AP || mode == wifi_mode_t_WIFI_MODE_APSTA)
     }
+}
+
+struct ScanResults {
+    count: usize,
+}
+
+impl ScanResults {
+    pub fn new() -> Self {
+        let mut bss_total = 0;
+
+        // `esp_wifi_scan_get_ap_num` shouldn't return an error. It's documentation lists state
+        // errors that would trip previous checks, and ESP_ERR_INVALID_ARG.
+        let count = esp_wifi_result!(unsafe {
+            crate::binary::include::esp_wifi_scan_get_ap_num(&mut bss_total)
+        })
+        .map(|_| bss_total as usize)
+        .unwrap_or(0);
 
-    fn scan_result_count(&mut self) -> Result<usize, WifiError> {
-        let mut bss_total: u16 = 0;
+        Self { count }
+    }
+
+    pub fn remaining(&self) -> usize {
+        self.count
+    }
+
+    #[cfg(not(esp32c2))]
+    fn read_into<const N: usize>(
+        &mut self,
+        scanned: &mut heapless::Vec<AccessPointInfo, N>,
+    ) -> Result<(), WifiError> {
+        while !scanned.is_full() {
+            if let Some(record) = self.next() {
+                // It is okay to ignore the result here, because we checked previously whether the
+                // vector is full.
+                _ = scanned.push(record?);
+            } else {
+                break;
+            }
+        }
+
+        Ok(())
+    }
+
+    #[cfg(esp32c2)]
+    fn read_into<const N: usize>(
+        &mut self,
+        scanned: &mut heapless::Vec<AccessPointInfo, N>,
+    ) -> Result<(), WifiError> {
+        let mut bss_total: u16 = N as u16;
+        let mut records: [MaybeUninit<crate::binary::include::wifi_ap_record_t>; N] =
+            [MaybeUninit::uninit(); N];
 
         unsafe {
-            esp_wifi_result!(crate::binary::include::esp_wifi_scan_get_ap_num(
-                &mut bss_total
-            ))
-            .map_err(|e| {
-                crate::binary::include::esp_wifi_clear_ap_list();
-                e
-            })?;
+            esp_wifi_result!(crate::binary::include::esp_wifi_scan_get_ap_records(
+                &mut bss_total,
+                records[0].as_mut_ptr(),
+            ))?;
+
+            for i in 0..(bss_total as usize).min(N) {
+                if !scanned.is_full() {
+                    let record = MaybeUninit::assume_init_ref(&records[i]);
+                    let ap_info = convert_ap_info(record);
+                    scanned.push(ap_info).ok();
+                } else {
+                    break;
+                }
+            }
         }
 
-        Ok(bss_total as usize)
+        // `esp_wifi_scan_get_ap_records` frees the whole list, regardless of how many records were
+        // actually returned.
+        self.count = 0;
+        Ok(())
     }
 
-    fn scan_results<const N: usize>(
+    pub fn take_results<const N: usize>(
         &mut self,
     ) -> Result<heapless::Vec<AccessPointInfo, N>, WifiError> {
         let mut scanned = heapless::Vec::<AccessPointInfo, N>::new();
 
-        unsafe {
-            let mut record: MaybeUninit<crate::binary::include::wifi_ap_record_t> =
-                MaybeUninit::uninit();
+        self.read_into(&mut scanned)?;
 
-            let count = self.scan_result_count()?;
+        Ok(scanned)
+    }
+}
 
-            extern "C" {
-                fn esp_mesh_scan_get_ap_record(
-                    ap_record: *mut crate::binary::include::wifi_ap_record_t,
-                    buffer: *mut crate::binary::c_types::c_void,
-                ) -> crate::binary::include::esp_err_t;
-            }
+#[cfg(not(esp32c2))]
+impl Iterator for ScanResults {
+    type Item = Result<AccessPointInfo, WifiError>;
 
-            for _ in 0..count {
-                esp_wifi_result!(esp_mesh_scan_get_ap_record(
-                    record.as_mut_ptr(),
-                    core::ptr::null_mut(),
-                ))
-                .map_err(|e| {
-                    // upon scan failure, list should be cleared to avoid memory leakage
-                    crate::binary::include::esp_wifi_clear_ap_list();
-                    e
-                })?;
+    fn next(&mut self) -> Option<Self::Item> {
+        extern "C" {
+            fn esp_mesh_scan_get_ap_record(
+                ap_record: *mut crate::binary::include::wifi_ap_record_t,
+                buffer: *mut crate::binary::c_types::c_void,
+            ) -> crate::binary::include::esp_err_t;
+        }
+
+        if self.count == 0 {
+            return None;
+        }
+
+        self.count -= 1;
 
-                let record = MaybeUninit::assume_init_ref(&record);
-                let ap_info = convert_ap_info(record);
+        let mut record: MaybeUninit<crate::binary::include::wifi_ap_record_t> =
+            MaybeUninit::uninit();
 
-                scanned.push(ap_info).ok();
+        unsafe {
+            let result = esp_wifi_result!(esp_mesh_scan_get_ap_record(
+                record.as_mut_ptr(),
+                core::ptr::null_mut(),
+            ));
+
+            match result {
+                Ok(_) => {
+                    let record = MaybeUninit::assume_init_ref(&record);
+                    let ap_info = convert_ap_info(record);
+
+                    Some(Ok(ap_info))
+                }
+
+                Err(e) => Some(Err(e)),
             }
         }
+    }
+}
 
-        Ok(scanned)
+impl Drop for ScanResults {
+    fn drop(&mut self) {
+        unsafe {
+            crate::binary::include::esp_wifi_clear_ap_list();
+        }
     }
 }
 
@@ -1067,10 +1145,12 @@ impl Wifi for WifiController<'_> {
     ) -> Result<(heapless::Vec<AccessPointInfo, N>, usize), Self::Error> {
         esp_wifi_result!(crate::wifi::wifi_start_scan(true))?;
 
-        let count = self.scan_result_count()?;
-        let result = self.scan_results()?;
+        let mut scan_result = ScanResults::new();
 
-        Ok((result, count))
+        let count = scan_result.remaining();
+        let ap_list = scan_result.take_results()?;
+
+        Ok((ap_list, count))
     }
 
     /// Get the currently used configuration.
@@ -1420,10 +1500,12 @@ mod asynch {
             esp_wifi_result!(wifi_start_scan(false))?;
             WifiEventFuture::new(WifiEvent::ScanDone).await;
 
-            let count = self.scan_result_count()?;
-            let result = self.scan_results()?;
+            let mut scan_result = ScanResults::new();
+
+            let count = scan_result.remaining();
+            let ap_list = scan_result.take_results()?;
 
-            Ok((result, count))
+            Ok((ap_list, count))
         }
 
         /// Async version of [`embedded_svc::wifi::Wifi`]'s `start` method