14 files changed, 100 insertions, 39 deletions

diff --git a/embassy-nrf/src/chips/nrf51.rs b/embassy-nrf/src/chips/nrf51.rs
index a0365a6ee..fd13ae5c4 100644
--- a/embassy-nrf/src/chips/nrf51.rs
+++ b/embassy-nrf/src/chips/nrf51.rs
@@ -145,6 +145,8 @@ impl_pin!(P0_31, 0, 31);
 
 impl_radio!(RADIO, RADIO, RADIO);
 
+impl_wdt!(WDT, WDT, WDT, 0);
+
 embassy_hal_internal::interrupt_mod!(
     CLOCK_POWER,
     RADIO,
diff --git a/embassy-nrf/src/chips/nrf52805.rs b/embassy-nrf/src/chips/nrf52805.rs
index a9cbccec4..7e72df8fc 100644
--- a/embassy-nrf/src/chips/nrf52805.rs
+++ b/embassy-nrf/src/chips/nrf52805.rs
@@ -221,6 +221,8 @@ impl_radio!(RADIO, RADIO, RADIO);
 impl_egu!(EGU0, EGU0, EGU0_SWI0);
 impl_egu!(EGU1, EGU1, EGU1_SWI1);
 
+impl_wdt!(WDT, WDT, WDT, 0);
+
 embassy_hal_internal::interrupt_mod!(
     CLOCK_POWER,
     RADIO,
diff --git a/embassy-nrf/src/chips/nrf52810.rs b/embassy-nrf/src/chips/nrf52810.rs
index ca31c35e1..e388e44e8 100644
--- a/embassy-nrf/src/chips/nrf52810.rs
+++ b/embassy-nrf/src/chips/nrf52810.rs
@@ -247,6 +247,8 @@ impl_radio!(RADIO, RADIO, RADIO);
 impl_egu!(EGU0, EGU0, EGU0_SWI0);
 impl_egu!(EGU1, EGU1, EGU1_SWI1);
 
+impl_wdt!(WDT, WDT, WDT, 0);
+
 embassy_hal_internal::interrupt_mod!(
     CLOCK_POWER,
     RADIO,
diff --git a/embassy-nrf/src/chips/nrf52811.rs b/embassy-nrf/src/chips/nrf52811.rs
index caf3fdcf8..96b8df30b 100644
--- a/embassy-nrf/src/chips/nrf52811.rs
+++ b/embassy-nrf/src/chips/nrf52811.rs
@@ -249,6 +249,8 @@ impl_radio!(RADIO, RADIO, RADIO);
 impl_egu!(EGU0, EGU0, EGU0_SWI0);
 impl_egu!(EGU1, EGU1, EGU1_SWI1);
 
+impl_wdt!(WDT, WDT, WDT, 0);
+
 embassy_hal_internal::interrupt_mod!(
     CLOCK_POWER,
     RADIO,
diff --git a/embassy-nrf/src/chips/nrf52820.rs b/embassy-nrf/src/chips/nrf52820.rs
index 39573e4a5..ad461b153 100644
--- a/embassy-nrf/src/chips/nrf52820.rs
+++ b/embassy-nrf/src/chips/nrf52820.rs
@@ -244,6 +244,8 @@ impl_egu!(EGU3, EGU3, EGU3_SWI3);
 impl_egu!(EGU4, EGU4, EGU4_SWI4);
 impl_egu!(EGU5, EGU5, EGU5_SWI5);
 
+impl_wdt!(WDT, WDT, WDT, 0);
+
 embassy_hal_internal::interrupt_mod!(
     CLOCK_POWER,
     RADIO,
diff --git a/embassy-nrf/src/chips/nrf52832.rs b/embassy-nrf/src/chips/nrf52832.rs
index 2d9346229..cf2abf23a 100644
--- a/embassy-nrf/src/chips/nrf52832.rs
+++ b/embassy-nrf/src/chips/nrf52832.rs
@@ -287,6 +287,8 @@ impl_egu!(EGU3, EGU3, EGU3_SWI3);
 impl_egu!(EGU4, EGU4, EGU4_SWI4);
 impl_egu!(EGU5, EGU5, EGU5_SWI5);
 
+impl_wdt!(WDT, WDT, WDT, 0);
+
 embassy_hal_internal::interrupt_mod!(
     CLOCK_POWER,
     RADIO,
diff --git a/embassy-nrf/src/chips/nrf52833.rs b/embassy-nrf/src/chips/nrf52833.rs
index 4e4915c32..e46eb1d2b 100644
--- a/embassy-nrf/src/chips/nrf52833.rs
+++ b/embassy-nrf/src/chips/nrf52833.rs
@@ -329,6 +329,8 @@ impl_egu!(EGU3, EGU3, EGU3_SWI3);
 impl_egu!(EGU4, EGU4, EGU4_SWI4);
 impl_egu!(EGU5, EGU5, EGU5_SWI5);
 
+impl_wdt!(WDT, WDT, WDT, 0);
+
 embassy_hal_internal::interrupt_mod!(
     CLOCK_POWER,
     RADIO,
diff --git a/embassy-nrf/src/chips/nrf52840.rs b/embassy-nrf/src/chips/nrf52840.rs
index 70a56971b..88747843d 100644
--- a/embassy-nrf/src/chips/nrf52840.rs
+++ b/embassy-nrf/src/chips/nrf52840.rs
@@ -334,6 +334,8 @@ impl_egu!(EGU3, EGU3, EGU3_SWI3);
 impl_egu!(EGU4, EGU4, EGU4_SWI4);
 impl_egu!(EGU5, EGU5, EGU5_SWI5);
 
+impl_wdt!(WDT, WDT, WDT, 0);
+
 embassy_hal_internal::interrupt_mod!(
     CLOCK_POWER,
     RADIO,
diff --git a/embassy-nrf/src/chips/nrf5340_app.rs b/embassy-nrf/src/chips/nrf5340_app.rs
index 86274bf48..0103fa7ae 100644
--- a/embassy-nrf/src/chips/nrf5340_app.rs
+++ b/embassy-nrf/src/chips/nrf5340_app.rs
@@ -171,7 +171,8 @@ embassy_hal_internal::peripherals! {
     RTC1,
 
     // WDT
-    WDT,
+    WDT0,
+    WDT1,
 
     // NVMC
     NVMC,
@@ -473,6 +474,9 @@ impl_egu!(EGU3, EGU3, EGU3);
 impl_egu!(EGU4, EGU4, EGU4);
 impl_egu!(EGU5, EGU5, EGU5);
 
+impl_wdt!(WDT0, WDT0, WDT0, 0);
+impl_wdt!(WDT1, WDT1, WDT1, 1);
+
 embassy_hal_internal::interrupt_mod!(
     FPU,
     CACHE,
diff --git a/embassy-nrf/src/chips/nrf5340_net.rs b/embassy-nrf/src/chips/nrf5340_net.rs
index 6c6ac3fbb..22d33d080 100644
--- a/embassy-nrf/src/chips/nrf5340_net.rs
+++ b/embassy-nrf/src/chips/nrf5340_net.rs
@@ -299,6 +299,8 @@ impl_radio!(RADIO, RADIO, RADIO);
 
 impl_egu!(EGU0, EGU0, EGU0);
 
+impl_wdt!(WDT, WDT, WDT, 0);
+
 embassy_hal_internal::interrupt_mod!(
     CLOCK_POWER,
     RADIO,
diff --git a/embassy-nrf/src/chips/nrf9120.rs b/embassy-nrf/src/chips/nrf9120.rs
index b02b8c6d8..e8ddbf86f 100644
--- a/embassy-nrf/src/chips/nrf9120.rs
+++ b/embassy-nrf/src/chips/nrf9120.rs
@@ -342,6 +342,8 @@ impl_egu!(EGU3, EGU3, EGU3);
 impl_egu!(EGU4, EGU4, EGU4);
 impl_egu!(EGU5, EGU5, EGU5);
 
+impl_wdt!(WDT, WDT, WDT, 0);
+
 embassy_hal_internal::interrupt_mod!(
     SPU,
     CLOCK_POWER,
diff --git a/embassy-nrf/src/chips/nrf9160.rs b/embassy-nrf/src/chips/nrf9160.rs
index b0981e3b5..5d04a72e5 100644
--- a/embassy-nrf/src/chips/nrf9160.rs
+++ b/embassy-nrf/src/chips/nrf9160.rs
@@ -342,6 +342,8 @@ impl_egu!(EGU3, EGU3, EGU3);
 impl_egu!(EGU4, EGU4, EGU4);
 impl_egu!(EGU5, EGU5, EGU5);
 
+impl_wdt!(WDT, WDT, WDT, 0);
+
 embassy_hal_internal::interrupt_mod!(
     SPU,
     CLOCK_POWER,
diff --git a/embassy-nrf/src/lib.rs b/embassy-nrf/src/lib.rs
index d2ff054f4..07ba2f6d4 100644
--- a/embassy-nrf/src/lib.rs
+++ b/embassy-nrf/src/lib.rs
@@ -169,7 +169,6 @@ pub mod uarte;
 ))]
 pub mod usb;
 #[cfg(not(feature = "_nrf54l"))] // TODO
-#[cfg(not(feature = "_nrf5340"))]
 pub mod wdt;
 
 // This mod MUST go last, so that it sees all the `impl_foo!` macros
diff --git a/embassy-nrf/src/wdt.rs b/embassy-nrf/src/wdt.rs
index 11cfa398e..baaf1c801 100644
--- a/embassy-nrf/src/wdt.rs
+++ b/embassy-nrf/src/wdt.rs
@@ -3,11 +3,15 @@
 //! This HAL implements a basic watchdog timer with 1..=8 handles.
 //! Once the watchdog has been started, it cannot be stopped.
 
-use core::marker::PhantomData;
+#![macro_use]
+
+use core::hint::unreachable_unchecked;
+
+use embassy_hal_internal::PeripheralType;
 
 use crate::pac::wdt::vals;
 pub use crate::pac::wdt::vals::{Halt as HaltConfig, Sleep as SleepConfig};
-use crate::{peripherals, Peri};
+use crate::{interrupt, pac, peripherals, Peri};
 
 const MIN_TICKS: u32 = 15;
 
@@ -30,12 +34,12 @@ pub struct Config {
 impl Config {
     /// Create a config structure from the current configuration of the WDT
     /// peripheral.
-    pub fn try_new(_wdt: &peripherals::WDT) -> Option<Self> {
-        let r = crate::pac::WDT;
+    pub fn try_new<T: Instance>(_wdt: &Peri<'_, T>) -> Option<Self> {
+        let r = T::REGS;
 
-        #[cfg(not(feature = "_nrf91"))]
+        #[cfg(not(any(feature = "_nrf91", feature = "_nrf5340")))]
         let runstatus = r.runstatus().read().runstatus();
-        #[cfg(feature = "_nrf91")]
+        #[cfg(any(feature = "_nrf91", feature = "_nrf5340"))]
         let runstatus = r.runstatus().read().runstatuswdt();
 
         if runstatus {
@@ -62,11 +66,11 @@ impl Default for Config {
 }
 
 /// Watchdog driver.
-pub struct Watchdog {
-    _wdt: Peri<'static, peripherals::WDT>,
+pub struct Watchdog<T: Instance> {
+    _wdt: Peri<'static, T>,
 }
 
-impl Watchdog {
+impl<T: Instance> Watchdog<T> {
     /// Try to create a new watchdog driver.
     ///
     /// This function will return an error if the watchdog is already active
@@ -76,19 +80,19 @@ impl Watchdog {
     /// `N` must be between 1 and 8, inclusive.
     #[inline]
     pub fn try_new<const N: usize>(
-        wdt: Peri<'static, peripherals::WDT>,
+        wdt: Peri<'static, T>,
         config: Config,
-    ) -> Result<(Self, [WatchdogHandle; N]), Peri<'static, peripherals::WDT>> {
+    ) -> Result<(Self, [WatchdogHandle; N]), Peri<'static, T>> {
         assert!(N >= 1 && N <= 8);
 
-        let r = crate::pac::WDT;
+        let r = T::REGS;
 
         let crv = config.timeout_ticks.max(MIN_TICKS);
         let rren = crate::pac::wdt::regs::Rren((1u32 << N) - 1);
 
-        #[cfg(not(feature = "_nrf91"))]
+        #[cfg(not(any(feature = "_nrf91", feature = "_nrf5340")))]
         let runstatus = r.runstatus().read().runstatus();
-        #[cfg(feature = "_nrf91")]
+        #[cfg(any(feature = "_nrf91", feature = "_nrf5340"))]
         let runstatus = r.runstatus().read().runstatuswdt();
 
         if runstatus {
@@ -114,17 +118,9 @@ impl Watchdog {
 
         let this = Self { _wdt: wdt };
 
-        let mut handles = [const {
-            WatchdogHandle {
-                _wdt: PhantomData,
-                index: 0,
-            }
-        }; N];
+        let mut handles = [const { WatchdogHandle { index: 0 } }; N];
         for i in 0..N {
-            handles[i] = WatchdogHandle {
-                _wdt: PhantomData,
-                index: i as u8,
-            };
+            handles[i] = unsafe { WatchdogHandle::steal(T::INDEX, i as u8) };
             handles[i].pet();
         }
 
@@ -139,7 +135,7 @@ impl Watchdog {
     /// interrupt has been enabled.
     #[inline(always)]
     pub fn enable_interrupt(&mut self) {
-        crate::pac::WDT.intenset().write(|w| w.set_timeout(true));
+        T::REGS.intenset().write(|w| w.set_timeout(true));
     }
 
     /// Disable the watchdog interrupt.
@@ -147,7 +143,7 @@ impl Watchdog {
     /// NOTE: This has no effect on the reset caused by the Watchdog.
     #[inline(always)]
     pub fn disable_interrupt(&mut self) {
-        crate::pac::WDT.intenclr().write(|w| w.set_timeout(true));
+        T::REGS.intenclr().write(|w| w.set_timeout(true));
     }
 
     /// Is the watchdog still awaiting pets from any handle?
@@ -156,7 +152,7 @@ impl Watchdog {
     /// handles to prevent a reset this time period.
     #[inline(always)]
     pub fn awaiting_pets(&self) -> bool {
-        let r = crate::pac::WDT;
+        let r = T::REGS;
         let enabled = r.rren().read().0;
         let status = r.reqstatus().read().0;
         (status & enabled) == 0
@@ -165,11 +161,26 @@ impl Watchdog {
 
 /// Watchdog handle.
 pub struct WatchdogHandle {
-    _wdt: PhantomData<Peri<'static, peripherals::WDT>>,
     index: u8,
 }
 
 impl WatchdogHandle {
+    fn regs(&self) -> pac::wdt::Wdt {
+        match self.index / 8 {
+            #[cfg(not(feature = "_multi_wdt"))]
+            peripherals::WDT::INDEX => peripherals::WDT::REGS,
+            #[cfg(feature = "_multi_wdt")]
+            peripherals::WDT0::INDEX => peripherals::WDT0::REGS,
+            #[cfg(feature = "_multi_wdt")]
+            peripherals::WDT1::INDEX => peripherals::WDT1::REGS,
+            _ => unsafe { unreachable_unchecked() },
+        }
+    }
+
+    fn rr_index(&self) -> usize {
+        usize::from(self.index % 8)
+    }
+
     /// Pet the watchdog.
     ///
     /// This function pets the given watchdog handle.
@@ -178,25 +189,50 @@ impl WatchdogHandle {
     /// prevent a reset from occurring.
     #[inline]
     pub fn pet(&mut self) {
-        let r = crate::pac::WDT;
-        r.rr(self.index as usize).write(|w| w.set_rr(vals::Rr::RELOAD));
+        let r = self.regs();
+        r.rr(self.rr_index()).write(|w| w.set_rr(vals::Rr::RELOAD));
     }
 
     /// Has this handle been pet within the current window?
     pub fn is_pet(&self) -> bool {
-        let r = crate::pac::WDT;
-        !r.reqstatus().read().rr(self.index as usize)
+        let r = self.regs();
+        !r.reqstatus().read().rr(self.rr_index())
     }
 
     /// Steal a watchdog handle by index.
     ///
     /// # Safety
-    /// Watchdog must be initialized and `index` must be between `0` and `N-1`
-    /// where `N` is the handle count when initializing.
-    pub unsafe fn steal(index: u8) -> Self {
+    /// Watchdog must be initialized, `instance_index` must be the index
+    /// [`Instance::INDEX`], and `index` must be between `0` and `N-1` where `N`
+    /// is the handle count when initializing.
+    pub unsafe fn steal(instance_index: u8, index: u8) -> Self {
         Self {
-            _wdt: PhantomData,
-            index,
+            index: instance_index * 8 + index,
         }
     }
 }
+
+pub(crate) trait SealedInstance {
+    const REGS: pac::wdt::Wdt;
+}
+
+/// WDT instance.
+#[allow(private_bounds)]
+pub trait Instance: SealedInstance + PeripheralType + 'static + Send {
+    /// Interrupt for this peripheral.
+    type Interrupt: interrupt::typelevel::Interrupt;
+    /// Index of the watchdog instance.
+    const INDEX: u8;
+}
+
+macro_rules! impl_wdt {
+    ($type:ident, $pac_type:ident, $irq:ident, $index:literal) => {
+        impl crate::wdt::SealedInstance for peripherals::$type {
+            const REGS: pac::wdt::Wdt = pac::$pac_type;
+        }
+        impl crate::wdt::Instance for peripherals::$type {
+            type Interrupt = crate::interrupt::typelevel::$irq;
+            const INDEX: u8 = $index;
+        }
+    };
+}