From 91ec2d5e3bd161c882c4652fede45a18000bd714 Mon Sep 17 00:00:00 2001
From: rowanG077 <goemansrowan@gmail.com>
Date: Wed, 13 Sep 2023 09:41:58 +0200
Subject: [PATCH] gen/genlib/misc: Add CE option to WaitTimer

---
 litex/gen/genlib/misc.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/litex/gen/genlib/misc.py b/litex/gen/genlib/misc.py
index b39077860..036c8870c 100644
--- a/litex/gen/genlib/misc.py
+++ b/litex/gen/genlib/misc.py
@@ -74,9 +74,10 @@ def timeline(trigger, events):
 
 
 class WaitTimer(Module):
-    def __init__(self, t):
-        self.wait = Signal()
-        self.done = Signal()
+    def __init__(self, t, has_ce=False):
+        self.wait = Signal() # i
+        self.ce   = Signal() # i
+        self.done = Signal() # o
 
         # # #
 
@@ -84,10 +85,16 @@ class WaitTimer(Module):
         t     = int(t)
         count = Signal(bits_for(t), reset=t)
 
+        tick = Signal()
+        if has_ce:
+            self.comb += tick.eq(~self.done & self.ce)
+        else:
+            self.comb += tick.eq(~self.done)
+
         self.comb += self.done.eq(count == 0)
         self.sync += [
             If(self.wait,
-                If(~self.done,
+                If(tick,
                     count.eq(count - 1)
                 )
             ).Else(