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,

    -- * 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, ..}

-- | 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'