packages feed

phino-0.0.95: src/Random.hs

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

module Random (randomString, shuffle) where

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)

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' -> show <$> randomRIO (0 :: Int, 9999)
    _ -> pure ['%', ch]
  next <- generate rest
  pure (rep ++ next)
generate (ch : rest) = do
  rest' <- generate rest
  pure (ch : rest')

regenerate :: String -> Set String -> IO String
regenerate pat set = do
  next <- generate pat
  if next `Set.member` set
    then regenerate pat set
    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