From 7884babb9d1557549e5732d6ab42a4f4c28f25e7 Mon Sep 17 00:00:00 2001
From: Giona Imperatori <giona.imperatori@gmail.com>
Date: Mon, 1 Jul 2024 17:20:42 +0200
Subject: [PATCH] feat(pwm): allow specifying OutputDrive for PWM channels

---
 embassy-nrf/CHANGELOG.md |  1 +
 embassy-nrf/src/pwm.rs   | 58 ++++++++++++++++++++++++++++++++++++----
 2 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/embassy-nrf/CHANGELOG.md b/embassy-nrf/CHANGELOG.md
index 773a1a108..6f07a8c6d 100644
--- a/embassy-nrf/CHANGELOG.md
+++ b/embassy-nrf/CHANGELOG.md
@@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - spi: Add support for configuring bit order for bus
 - pwm: Expose `pwm::PWM_CLK_HZ` and add `is_enabled` method
 - gpio: Drop GPIO Pin generics (API break)
+- pwm: Allow specifying OutputDrive for PWM channels
 
 ## 0.1.0 - 2024-01-12
 
diff --git a/embassy-nrf/src/pwm.rs b/embassy-nrf/src/pwm.rs
index 12057f7dd..a2853de3f 100644
--- a/embassy-nrf/src/pwm.rs
+++ b/embassy-nrf/src/pwm.rs
@@ -6,7 +6,7 @@ use core::sync::atomic::{compiler_fence, Ordering};
 
 use embassy_hal_internal::{into_ref, PeripheralRef};
 
-use crate::gpio::{AnyPin, Pin as GpioPin, PselBits, SealedPin as _};
+use crate::gpio::{convert_drive, AnyPin, OutputDrive, Pin as GpioPin, PselBits, SealedPin as _};
 use crate::ppi::{Event, Task};
 use crate::util::slice_in_ram_or;
 use crate::{interrupt, pac, Peripheral};
@@ -128,19 +128,23 @@ impl<'d, T: Instance> SequencePwm<'d, T> {
 
         if let Some(pin) = &ch0 {
             pin.set_low();
-            pin.conf().write(|w| w.dir().output());
+            pin.conf()
+                .write(|w| w.dir().output().drive().variant(convert_drive(config.ch0_drive)));
         }
         if let Some(pin) = &ch1 {
             pin.set_low();
-            pin.conf().write(|w| w.dir().output());
+            pin.conf()
+                .write(|w| w.dir().output().drive().variant(convert_drive(config.ch1_drive)));
         }
         if let Some(pin) = &ch2 {
             pin.set_low();
-            pin.conf().write(|w| w.dir().output());
+            pin.conf()
+                .write(|w| w.dir().output().drive().variant(convert_drive(config.ch2_drive)));
         }
         if let Some(pin) = &ch3 {
             pin.set_low();
-            pin.conf().write(|w| w.dir().output());
+            pin.conf()
+                .write(|w| w.dir().output().drive().variant(convert_drive(config.ch3_drive)));
         }
 
         r.psel.out[0].write(|w| unsafe { w.bits(ch0.psel_bits()) });
@@ -319,6 +323,14 @@ pub struct Config {
     pub prescaler: Prescaler,
     /// How a sequence is read from RAM and is spread to the compare register
     pub sequence_load: SequenceLoad,
+    /// Drive strength for the channel 0 line.
+    pub ch0_drive: OutputDrive,
+    /// Drive strength for the channel 1 line.
+    pub ch1_drive: OutputDrive,
+    /// Drive strength for the channel 2 line.
+    pub ch2_drive: OutputDrive,
+    /// Drive strength for the channel 3 line.
+    pub ch3_drive: OutputDrive,
 }
 
 impl Default for Config {
@@ -328,6 +340,10 @@ impl Default for Config {
             max_duty: 1000,
             prescaler: Prescaler::Div16,
             sequence_load: SequenceLoad::Common,
+            ch0_drive: OutputDrive::Standard,
+            ch1_drive: OutputDrive::Standard,
+            ch2_drive: OutputDrive::Standard,
+            ch3_drive: OutputDrive::Standard,
         }
     }
 }
@@ -815,6 +831,38 @@ impl<'d, T: Instance> SimplePwm<'d, T> {
         let max_duty = self.max_duty() as u32;
         clk / max_duty
     }
+
+    /// Sets the PWM-Channel0 output drive strength
+    #[inline(always)]
+    pub fn set_ch0_drive(&self, drive: OutputDrive) {
+        if let Some(pin) = &self.ch0 {
+            pin.conf().write(|w| w.drive().variant(convert_drive(drive)));
+        }
+    }
+
+    /// Sets the PWM-Channel1 output drive strength
+    #[inline(always)]
+    pub fn set_ch1_drive(&self, drive: OutputDrive) {
+        if let Some(pin) = &self.ch1 {
+            pin.conf().write(|w| w.drive().variant(convert_drive(drive)));
+        }
+    }
+
+    /// Sets the PWM-Channel2 output drive strength
+    #[inline(always)]
+    pub fn set_ch2_drive(&self, drive: OutputDrive) {
+        if let Some(pin) = &self.ch2 {
+            pin.conf().write(|w| w.drive().variant(convert_drive(drive)));
+        }
+    }
+
+    /// Sets the PWM-Channel3 output drive strength
+    #[inline(always)]
+    pub fn set_ch3_drive(&self, drive: OutputDrive) {
+        if let Some(pin) = &self.ch3 {
+            pin.conf().write(|w| w.drive().variant(convert_drive(drive)));
+        }
+    }
 }
 
 impl<'a, T: Instance> Drop for SimplePwm<'a, T> {