phino-0.0.116: src/Random.hs

-- SPDX-FileCopyrightText: Copyright (c) 2025 Objectionary.com
-- SPDX-License-Identifier: MIT

module Random (randomString, shuffle) where

import Control.Exception (throwIO)
import Control.Monad (forM_, replicateM)
import Data.Char (intToDigit)
import Data.IORef (IORef, modifyIORef', newIORef, readIORef)
import Data.Set (Set)
import qualified Data.Set as Set
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as M
import GHC.IO (unsafePerformIO)
import System.Random (newStdGen, randomRIO)
import System.Random.Stateful (newIOGenM, uniformRM)
import Text.Printf (printf)

strings :: IORef (Set String)
{-# NOINLINE strings #-}
strings = unsafePerformIO (newIORef Set.empty)

generate :: String -> IO String
generate [] = pure []
generate ('%' : ch : rest) = do
  rep <- case ch of
    'x' -> replicateM 8 $ do
      v <- randomRIO (0, 15)
      pure (intToDigit v)
    'd' -> printf "%04d" <$> randomRIO (0 :: Int, 9999)
    _ -> pure ['%', ch]
  next <- generate rest
  pure (rep ++ next)
generate (ch : rest) = do
  rest' <- generate rest
  pure (ch : rest')

-- The 'strings' set grows monotonically over a process, so a pattern with a
-- bounded space (e.g. '%d', which has exactly 10,000 values) eventually gets
-- exhausted. Trying again forever would hang, so the search gives up after a
-- bounded number of attempts and reports the collision space instead. The
-- limit is well above the largest realistic space (10,000) so that finding the
-- last free value of a nearly-full space still succeeds with overwhelming
-- probability: (9999/10000)^100000 ≈ 4.5e-5.
maxAttempts :: Int
maxAttempts = 100000

regenerate :: String -> Set String -> IO String
regenerate pat set = go maxAttempts
  where
    go :: Int -> IO String
    go 0 = throwIO (userError (printf "randomString() cannot produce a unique value for pattern '%s': the value space is exhausted" pat))
    go attempts = do
      next <- generate pat
      if next `Set.member` set
        then go (attempts - 1)
        else do
          modifyIORef' strings (Set.insert next)
          pure next

randomString :: String -> IO String
randomString pat
  | randomized pat = readIORef strings >>= regenerate pat
  | otherwise = generate pat
  where
    randomized :: String -> Bool
    randomized [] = False
    randomized ('%' : ch : rest) = ch == 'd' || ch == 'x' || randomized rest
    randomized (_ : rest) = randomized rest

-- Fast Fisher-Yates with mutable vectors.
-- The function is generated by ChatGPT and claimed as
-- fastest approach comparing to usage IOArray.
-- >>> shuffle [1..20]
-- [7,15,5,18,13,19,3,11,20,2,1,8,14,16,17,12,9,10,6,4]
shuffle :: [a] -> IO [a]
shuffle xs = do
  gen <- newIOGenM =<< newStdGen
  let n = length xs
  v <- V.thaw (V.fromList xs) -- Mutable copy
  forM_ [n - 1, n - 2 .. 1] $ \i -> do
    j <- uniformRM (0, i) gen
    M.swap v i j
  V.toList <$> V.freeze v