Typed.hs

{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE NoFieldSelectors #-}

-- |
-- Module        : Control.NumerusClosus.Typed
-- Copyright     : Gautier DI FOLCO
-- License       : ISC
--
-- Maintainer    : Gautier DI FOLCO <gautier.difolco@gmail.com>
-- Stability     : Stable
-- Portability   : Portable
--
-- Simple, composable, pure rate-limiting primitives supporting finite buckets,
-- fixed windows, sliding windows, and sliding window counters. Includes AND/OR
-- combinators and scheduling helpers.
--
-- > import Control.NumerusClosus
-- > import Data.Time (getCurrentTime)
-- >
-- > main :: IO ()
-- > main = do
-- >   now <- getCurrentTime
-- >   let rl = fixedWindow 60 5 now  -- 5 requests per 60 seconds
-- >   result <- schedule rl (putStrLn "Request allowed")
-- >   print result
module Control.NumerusClosus.Typed
  ( RateLimiter (..),
    NextDebitable (..),

    -- * Base helpers
    AlwaysAllow (..),
    alwaysAllow,
    AlwaysDeny (..),
    alwaysDeny,

    -- * Strategies
    BucketSize (..),
    FiniteBucket (..),
    finiteBucket,
    WindowSize (..),
    FixedWindow (..),
    fixedWindow,
    SlidingWindow (..),
    slidingWindow,
    SlidingWindowCount (..),
    slidingWindowCount,
    BucketsCount (..),
    SlidingWindowBucketed (..),
    slidingWindowBucketed,
    RefillRate (..),
    TokenBucket (..),
    tokenBucket,
    DrainRate (..),
    LeakyBucket (..),
    leakyBucket,
    GCRA (..),
    gcra,

    -- * Combinators
    type (:&&) (..),
    type (:||) (..),

    -- * Scheduling
    PositionInTime (..),
    ioPositionInTime,
    schedule,
    scheduleWith,
    loopSchedule,
    loopScheduleWith,

    -- * Reexports
    First (..),
    Last (..),
    Max (..),
    Min (..),
  )
where

import Control.Concurrent.Thread.Delay (delay)
import Control.Monad (mfilter)
import qualified Data.Map.Strict as Map
import Data.Maybe (fromMaybe)
import Data.Semigroup (First (..), Last (..), Max (..), Min (..))
import qualified Data.Set as Set
import Data.Time (FormatTime, NominalDiffTime, ParseTime, UTCTime, addUTCTime, diffUTCTime, getCurrentTime)
import GHC.Generics (Generic)

-- * Main logic

-- | The core rate limiter type.
-- Given the current time, returns either the next debitable time or a new rate limiter state.
class RateLimiter a where
  -- | Function to determine if a request is allowed at a given time.
  debit :: a -> UTCTime -> Either NextDebitable a

-- | Represents when the next request can be debited.
data NextDebitable
  = -- | The request can be debited from this time onwards.
    DebitableFrom UTCTime
  | -- | The request will never be debitable again.
    Never
  deriving stock (Eq, Ord, Show)

-- * Base helpers

-- | A rate limiter that always allows requests.
data AlwaysAllow = AlwaysAllow
  deriving stock (Eq, Show, Generic)

instance RateLimiter AlwaysAllow where
  debit AlwaysAllow _ = Right AlwaysAllow

-- | Smart constructor for 'AlwaysAllow'.
alwaysAllow :: AlwaysAllow
alwaysAllow = AlwaysAllow

-- | A rate limiter that always denies requests.
data AlwaysDeny = AlwaysDeny
  deriving stock (Eq, Show, Generic)

instance RateLimiter AlwaysDeny where
  debit AlwaysDeny _ = Left Never

-- | Smart constructor for 'AlwaysDeny'.
alwaysDeny :: AlwaysDeny
alwaysDeny = AlwaysDeny

-- * Strategies

-- | The maximum number of requests a bucket can hold.
newtype BucketSize a = BucketSize
  { unBucketSize :: a
  }
  deriving stock (Eq, Ord, Show)
  deriving newtype (Num)

-- | A simple finite bucket that allows exactly @n@ requests.
newtype FiniteBucket
  = FiniteBucket {fullBucket :: BucketSize Integer}
  deriving stock (Eq, Show, Generic)

instance RateLimiter FiniteBucket where
  debit (FiniteBucket (BucketSize n)) _ =
    if n > 0
      then Right $ FiniteBucket $ BucketSize (n - 1)
      else Left Never

finiteBucket :: BucketSize Integer -> FiniteBucket
finiteBucket fullBucket = FiniteBucket {..}

-- | The size of a time window.
newtype WindowSize = WindowSize
  { unWindowSize :: NominalDiffTime
  }
  deriving stock (Eq, Ord, Show)
  deriving newtype (Num, Fractional, Real, RealFrac, FormatTime, ParseTime)

-- | A fixed window rate limiter that allows a given number of requests per window.
data FixedWindow = FixedWindow
  { windowSize :: WindowSize,
    maxBucket :: BucketSize Integer,
    -- \* Iteration fields
    bucket :: BucketSize Integer,
    windowStart :: UTCTime,
    windowEnd :: UTCTime
  }
  deriving stock (Eq, Show, Generic)

instance RateLimiter FixedWindow where
  debit (FixedWindow {..}) now =
    if refreshedBucket > 0
      then Right $ FixedWindow {bucket = refreshedBucket - 1, windowEnd = refreshedEndTime, ..}
      else Left $ DebitableFrom $ nextTimeUnit refreshedEndTime
    where
      (refreshedBucket, refreshedEndTime) =
        if now > windowEnd
          then (maxBucket, addUTCTime (fromIntegral (floor (diffUTCTime now windowEnd / windowSize.unWindowSize) + 1 :: Integer) * windowSize.unWindowSize) windowEnd)
          else (bucket, windowEnd)

fixedWindow :: WindowSize -> BucketSize Integer -> UTCTime -> FixedWindow
fixedWindow windowSize@(WindowSize window) maxBucket windowStart =
  FixedWindow {bucket = maxBucket, windowEnd = addUTCTime window windowStart, ..}

-- | A sliding window rate limiter that allows a given number of requests within the sliding window.
data SlidingWindow = SlidingWindow
  { windowSize :: WindowSize,
    maxBucket :: BucketSize Int,
    -- \* Iteration field
    bucket :: Set.Set UTCTime
  }
  deriving stock (Eq, Show, Generic)

instance RateLimiter SlidingWindow where
  debit SlidingWindow {..} now =
    if Set.size refreshedBucket < maxBucket.unBucketSize
      then Right $ SlidingWindow {bucket = Set.insert now refreshedBucket, ..}
      else Left $ DebitableFrom $ nextTimeUnit $ addUTCTime windowSize.unWindowSize $ fromMaybe now $ Set.lookupMin refreshedBucket
    where
      refreshedBucket = snd $ Set.split (addUTCTime ((-1) * windowSize.unWindowSize) now) bucket

slidingWindow :: WindowSize -> BucketSize Int -> SlidingWindow
slidingWindow windowSize maxBucket =
  SlidingWindow {bucket = Set.empty, ..}

-- | The number of sub-windows for bucketed sliding windows.
newtype BucketsCount a = BucketsCount
  { unBucketsCount :: a
  }
  deriving stock (Eq, Ord, Show)
  deriving newtype (Num)

-- | Sliding window rate limiter using sub-bucket counters.
-- Divides the window into sub-buckets keyed by request arrival times and sums
-- their counters. This is a bucketed variant rather than the canonical
-- two-window weighted interpolation.
data SlidingWindowBucketed = SlidingWindowBucketed
  { windowSize :: WindowSize,
    bucketsCount :: BucketsCount Int,
    maxBucket :: BucketSize Integer,
    -- \* Iteration field
    buckets :: Map.Map UTCTime Integer
  }
  deriving stock (Eq, Show, Generic)

instance RateLimiter SlidingWindowBucketed where
  debit SlidingWindowBucketed {..} now =
    if sum refreshedBucket < maxBucket.unBucketSize
      then Right $ SlidingWindowBucketed {buckets = Map.alter (Just . maybe 1 (+ 1)) lastBucket refreshedBucket, ..}
      else Left $ DebitableFrom $ nextTimeUnit $ addUTCTime windowSize.unWindowSize $ maybe now fst $ Map.lookupMin refreshedBucket
    where
      refreshedBucket =
        Map.restrictKeys buckets $
          snd $
            Set.split (addUTCTime ((-1) * windowSize.unWindowSize) now) $
              Map.keysSet buckets
      lastBucket =
        fromMaybe now $
          mfilter (> addUTCTime ((-1) * bucketWindow) now) $
            fst <$> Map.lookupMax refreshedBucket
      bucketWindow = windowSize.unWindowSize / fromIntegral bucketsCount.unBucketsCount

slidingWindowBucketed :: WindowSize -> BucketsCount Int -> BucketSize Integer -> SlidingWindowBucketed
slidingWindowBucketed windowSize bucketsCount maxBucket =
  SlidingWindowBucketed {buckets = Map.empty, ..}

-- | Sliding window counter using two-window weighted interpolation.
-- Tracks counters for the current and previous fixed windows, then estimates
-- the request rate as: @prev * (1 - elapsed\/window) + current@.
-- This is the canonical sliding window counter algorithm with O(1) memory.
data SlidingWindowCount = SlidingWindowCount
  { windowSize :: WindowSize,
    maxBucket :: BucketSize Int,
    -- \* Iteration field
    prevCount :: Integer,
    currentCount :: Integer,
    windowStart :: UTCTime,
    windowEnd :: UTCTime
  }
  deriving stock (Eq, Show, Generic)

instance RateLimiter SlidingWindowCount where
  debit SlidingWindowCount {..} now =
    if estimate < fromIntegral maxBucket.unBucketSize
      then Right $ SlidingWindowCount {prevCount = refreshedPrev, currentCount = refreshedCurrent + 1, windowEnd = refreshedEnd, ..}
      else Left $ DebitableFrom $ nextTimeUnit refreshedEnd
    where
      (refreshedPrev, refreshedCurrent, refreshedEnd) =
        if now > windowEnd
          then
            let windowsElapsed = floor (diffUTCTime now windowEnd / windowSize.unWindowSize) :: Integer
                newEnd = addUTCTime (fromIntegral (windowsElapsed + 1) * windowSize.unWindowSize) windowEnd
             in if windowsElapsed == 0
                  then (currentCount, 0, newEnd)
                  else (0, 0, newEnd)
          else (prevCount, currentCount, windowEnd)
      newWindowStart = addUTCTime ((-1) * windowSize.unWindowSize) refreshedEnd
      fraction = realToFrac (diffUTCTime now newWindowStart) / realToFrac windowSize.unWindowSize :: Double
      estimate = fromIntegral refreshedPrev * (1 - fraction) + fromIntegral refreshedCurrent :: Double

slidingWindowCount :: WindowSize -> BucketSize Int -> UTCTime -> SlidingWindowCount
slidingWindowCount windowSize@(WindowSize window) maxBucket windowStart =
  SlidingWindowCount {prevCount = 0, currentCount = 0, windowEnd = addUTCTime window windowStart, ..}

-- | The rate at which tokens are refilled (tokens per second).
newtype RefillRate = RefillRate
  { unRefillRate :: Double
  }
  deriving stock (Eq, Ord, Show)

-- | Token bucket with continuous refill.
-- Tokens are consumed on each request and refilled at a steady rate over time.
-- Allows bursts up to the bucket capacity while maintaining a long-term average rate.
data TokenBucket = TokenBucket
  { maxBucket :: BucketSize Double,
    refillRate :: RefillRate,
    -- \* Iteration fields
    tokens :: Double,
    lastRefill :: UTCTime
  }
  deriving stock (Eq, Show, Generic)

instance RateLimiter TokenBucket where
  debit TokenBucket {..} now =
    if refreshedTokens >= 1
      then Right $ TokenBucket {tokens = refreshedTokens - 1, lastRefill = now, ..}
      else
        let tokensNeeded = 1 - refreshedTokens
            secondsToWait = tokensNeeded / refillRate.unRefillRate
         in Left $ DebitableFrom $ nextTimeUnit $ addUTCTime (realToFrac secondsToWait) now
    where
      elapsed = realToFrac (diffUTCTime now lastRefill) :: Double
      refreshedTokens = min maxBucket.unBucketSize (tokens + elapsed * refillRate.unRefillRate)

-- | Create a token bucket rate limiter.
-- @maxBucket@ is the burst capacity, @refillRate@ is tokens restored per second.
tokenBucket :: BucketSize Double -> RefillRate -> UTCTime -> TokenBucket
tokenBucket maxBucket refillRate lastRefill =
  TokenBucket {tokens = maxBucket.unBucketSize, ..}

-- | The rate at which the bucket drains (requests per second).
newtype DrainRate = DrainRate
  { unDrainRate :: Double
  }
  deriving stock (Eq, Ord, Show)

-- | Leaky bucket rate limiter.
-- Requests fill a bucket that drains at a constant rate. If the bucket overflows
-- the request is denied. This produces a perfectly smooth output rate.
data LeakyBucket = LeakyBucket
  { maxBucket :: BucketSize Double,
    drainRate :: DrainRate,
    -- \* Iteration fields
    level :: Double,
    lastDrain :: UTCTime
  }
  deriving stock (Eq, Show, Generic)

instance RateLimiter LeakyBucket where
  debit LeakyBucket {..} now =
    let newLevel = refreshedLevel + 1
     in if newLevel <= maxBucket.unBucketSize
          then Right $ LeakyBucket {level = newLevel, lastDrain = now, ..}
          else
            let excess = newLevel - maxBucket.unBucketSize
                secondsToWait = excess / drainRate.unDrainRate
             in Left $ DebitableFrom $ nextTimeUnit $ addUTCTime (realToFrac secondsToWait) now
    where
      elapsed = realToFrac (diffUTCTime now lastDrain) :: Double
      refreshedLevel = max 0 (level - elapsed * drainRate.unDrainRate)

-- | Create a leaky bucket rate limiter.
-- @maxBucket@ is the queue capacity, @drainRate@ is how fast it empties.
leakyBucket :: BucketSize Double -> DrainRate -> UTCTime -> LeakyBucket
leakyBucket maxBucket drainRate lastDrain =
  LeakyBucket {level = 0, ..}

-- | Generic Cell Rate Algorithm (GCRA).
-- A memoryless variant of leaky bucket that tracks only a single \"theoretical
-- arrival time\" (TAT). Each request advances the TAT by the emission interval
-- (@1 / rate@). A request is allowed if the TAT is not too far in the future
-- (bounded by the burst tolerance, derived from @limit / rate@).
data GCRA = GCRA
  { emissionInterval :: NominalDiffTime,
    burstTolerance :: NominalDiffTime,
    -- \* Iteration field
    tat :: UTCTime
  }
  deriving stock (Eq, Show, Generic)

instance RateLimiter GCRA where
  debit GCRA {..} now =
    let newTat = addUTCTime emissionInterval (max now tat)
        allowAt = addUTCTime (negate burstTolerance) newTat
     in if allowAt <= now
          then Right $ GCRA {tat = newTat, ..}
          else Left $ DebitableFrom $ nextTimeUnit allowAt

-- | Create a GCRA rate limiter.
-- @limit@ is the number of requests allowed in the period, @period@ is the
-- time window. The emission interval is @period / limit@ and the burst
-- tolerance is @period@.
gcra :: Int -> NominalDiffTime -> UTCTime -> GCRA
gcra limit period now =
  GCRA
    { emissionInterval = period / fromIntegral limit,
      burstTolerance = period,
      tat = now
    }

-- | Helper to get the smallest next time unit for debiting.
nextTimeUnit :: UTCTime -> UTCTime
nextTimeUnit = addUTCTime 0.000001

-- * Combinators

-- | AND combinator type: allows a request if both rate limiters allow it.
data a :&& b = a :&& b
  deriving stock (Eq, Show, Generic)

infixr 3 :&&

instance (RateLimiter a, RateLimiter b) => RateLimiter (a :&& b) where
  debit (x :&& y) at =
    case (debit x at, debit y at) of
      (Right x', Right y') -> Right $ x' :&& y'
      (Left x', Left y') -> Left $ getLast $ Last x' <> Last y'
      (Left x', _) -> Left x'
      (_, Left y') -> Left y'

-- | OR combinator type: allows a request if either rate limiter allows it.
data a :|| b = a :|| b
  deriving stock (Eq, Show, Generic)

infixr 3 :||

instance (RateLimiter a, RateLimiter b) => RateLimiter (a :|| b) where
  debit (x :|| y) at =
    case (debit x at, debit y at) of
      (Right x', Right y') -> Right $ x' :|| y'
      (Left x', Left y') -> Left $ getFirst $ First x' <> First y'
      (Right x', _) -> Right $ x' :|| y
      (_, Right y') -> Right $ x :|| y'

-- * Scheduling

-- | Fetch the time and compute the delay time.
data PositionInTime m = PositionInTime
  { -- | How to fetch the current time.
    getTime :: m UTCTime,
    -- | How to delay until a specific time.
    delayUntil :: UTCTime -> m ()
  }

-- | Run an IO action if the rate limiter allows it, using the current time.
schedule :: (RateLimiter a) => a -> (a -> IO b) -> IO (Either NextDebitable (a, b))
schedule = scheduleWith ioPositionInTime

-- | Default time strategy for IO.
ioPositionInTime :: PositionInTime IO
ioPositionInTime = PositionInTime getCurrentTime diffDelay
  where
    diffDelay to = do
      now <- getCurrentTime
      delay $ 0 `max` round (diffUTCTime to now * 1000000)

-- | Run a monadic action if the rate limiter allows it, using a custom time strategy.
scheduleWith :: (Monad m, RateLimiter a) => PositionInTime m -> a -> (a -> m b) -> m (Either NextDebitable (a, b))
scheduleWith pit rl action = do
  now <- pit.getTime
  case debit rl now of
    Right rl' -> Right . (rl',) <$> action rl'
    Left nd -> return $ Left nd

-- | Repeatedly run an IO action, respecting the rate limiter. Will sleep until the next available slot if rate limited.
loopSchedule :: (RateLimiter a) => a -> (a -> IO b) -> IO ()
loopSchedule = loopScheduleWith ioPositionInTime

-- | Repeatedly run a monadic action, respecting the rate limiter and using a custom time strategy. Will sleep until the next available slot if rate limited.
loopScheduleWith :: (Monad m, RateLimiter a) => PositionInTime m -> a -> (a -> m b) -> m ()
loopScheduleWith pit rl action = go rl
  where
    go rl' = do
      result <- scheduleWith pit rl' action
      case result of
        Right (rl'', _) -> go rl''
        Left Never -> return ()
        Left (DebitableFrom at) -> pit.delayUntil at >> go rl'