Range.hs

{-# OPTIONS_HADDOCK hide #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}

-- |
-- Module: Crypto.Range
-- Copyright: (c) 2026 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- Unbiased integers in a range, via Lemire's nearly-divisionless
-- method ("Fast Random Integer Generation in an Interval", 2019).

module Crypto.Range (
    upto
  , upto_pure
  , W2(..)
  , mul_wide
  , _mul_wide_limbs
  ) where

import Data.Bits ((.&.), (.>>.))
import qualified Data.Bits as B
import Data.Word (Word64)
import GHC.Exts (Word(..), timesWord2#)

-- | A 128-bit product, as its high and low words.
data W2 = W2 {-# UNPACK #-} !Word64 {-# UNPACK #-} !Word64
  deriving (Eq, Show)

-- | The full 128-bit product of two words.
mul_wide :: Word64 -> Word64 -> W2
mul_wide a b
  | B.finiteBitSize (0 :: Word) == 64 =
      let !(W# x) = fromIntegral a
          !(W# y) = fromIntegral b
          !(# h, l #) = timesWord2# x y
      in  W2 (fromIntegral (W# h)) (fromIntegral (W# l))
  | otherwise = _mul_wide_limbs a b
{-# INLINE mul_wide #-}

-- | The full 128-bit product of two words, via 32-bit limbs, for
--   platforms with 32-bit machine words.
_mul_wide_limbs :: Word64 -> Word64 -> W2
_mul_wide_limbs a b =
  let !a0  = a .&. 0xffffffff
      !a1  = a .>>. 32
      !b0  = b .&. 0xffffffff
      !b1  = b .>>. 32
      !p00 = a0 * b0
      !p01 = a0 * b1
      !p10 = a1 * b0
      !p11 = a1 * b1
      !mid = (p00 .>>. 32) + (p01 .&. 0xffffffff) + (p10 .&. 0xffffffff)
      !hi  = p11 + (p01 .>>. 32) + (p10 .>>. 32) + (mid .>>. 32)
  in  W2 hi (a * b)
{-# INLINE _mul_wide_limbs #-}

-- | A uniformly-distributed word in [0, n], given an action producing
--   uniformly-distributed words.
--
--   Each draw x is mapped to the high word of x * (n + 1); draws whose
--   low word falls below 2 ^ 64 mod (n + 1) are rejected, so that
--   every output is equally likely. Rejection is rare for small n, and
--   the remainder is only computed when it might occur.
upto :: Monad m => m Word64 -> Word64 -> m Word64
upto gen n
  | n == maxBound = gen
  | otherwise = do
      !x <- gen
      let !(W2 h l) = mul_wide x s
      if   l >= s
      then pure h
      else reject (negate s `rem` s) h l
  where
    !s = n + 1
    reject !t = loop where
      loop !h !l
        | l >= t    = pure h
        | otherwise = do
            !x <- gen
            let !(W2 nh nl) = mul_wide x s
            loop nh nl
{-# INLINE upto #-}

-- | 'upto' for a pure generator, given its step function.
upto_pure :: (g -> (Word64, g)) -> Word64 -> g -> (Word64, g)
upto_pure step n g0
  | n == maxBound = step g0
  | otherwise =
      let !(x, g1) = step g0
          !(W2 h l) = mul_wide x s
      in  if   l >= s
          then (h, g1)
          else reject (negate s `rem` s) h l g1
  where
    !s = n + 1
    reject !t = loop where
      loop !h !l !g
        | l >= t    = (h, g)
        | otherwise =
            let !(x, ng) = step g
                !(W2 nh nl) = mul_wide x s
            in  loop nh nl ng
{-# INLINE upto_pure #-}