-- | Miscellaneous utility functions.

{-# LANGUAGE CPP, MagicHash #-}
module Twee.Utils where

import Control.Arrow((&&&))
import Control.Exception
import Data.List(groupBy, sortBy)
import Data.Ord(comparing)
import System.IO
import GHC.Prim
import GHC.Types
import Data.Bits
import System.Random
--import Test.QuickCheck hiding ((.&.))

repeatM :: Monad m => m a -> m [a]
repeatM :: m a -> m [a]
repeatM = [m a] -> m [a]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence ([m a] -> m [a]) -> (m a -> [m a]) -> m a -> m [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> [m a]
forall a. a -> [a]
repeat

partitionBy :: Ord b => (a -> b) -> [a] -> [[a]]
partitionBy :: (a -> b) -> [a] -> [[a]]
partitionBy a -> b
value =
  ([(a, b)] -> [a]) -> [[(a, b)]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (((a, b) -> a) -> [(a, b)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> a
forall a b. (a, b) -> a
fst) ([[(a, b)]] -> [[a]]) -> ([a] -> [[(a, b)]]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
  ((a, b) -> (a, b) -> Bool) -> [(a, b)] -> [[(a, b)]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (\(a, b)
x (a, b)
y -> (a, b) -> b
forall a b. (a, b) -> b
snd (a, b)
x b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== (a, b) -> b
forall a b. (a, b) -> b
snd (a, b)
y) ([(a, b)] -> [[(a, b)]]) -> ([a] -> [(a, b)]) -> [a] -> [[(a, b)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
  ((a, b) -> (a, b) -> Ordering) -> [(a, b)] -> [(a, b)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (((a, b) -> b) -> (a, b) -> (a, b) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (a, b) -> b
forall a b. (a, b) -> b
snd) ([(a, b)] -> [(a, b)]) -> ([a] -> [(a, b)]) -> [a] -> [(a, b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
  (a -> (a, b)) -> [a] -> [(a, b)]
forall a b. (a -> b) -> [a] -> [b]
map (a -> a
forall a. a -> a
id (a -> a) -> (a -> b) -> a -> (a, b)
forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& a -> b
value)

collate :: Ord a => ([b] -> c) -> [(a, b)] -> [(a, c)]
collate :: ([b] -> c) -> [(a, b)] -> [(a, c)]
collate [b] -> c
f = ([(a, b)] -> (a, c)) -> [[(a, b)]] -> [(a, c)]
forall a b. (a -> b) -> [a] -> [b]
map [(a, b)] -> (a, c)
forall a. [(a, b)] -> (a, c)
g ([[(a, b)]] -> [(a, c)])
-> ([(a, b)] -> [[(a, b)]]) -> [(a, b)] -> [(a, c)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, b) -> a) -> [(a, b)] -> [[(a, b)]]
forall b a. Ord b => (a -> b) -> [a] -> [[a]]
partitionBy (a, b) -> a
forall a b. (a, b) -> a
fst
  where
    g :: [(a, b)] -> (a, c)
g [(a, b)]
xs = ((a, b) -> a
forall a b. (a, b) -> a
fst ([(a, b)] -> (a, b)
forall a. [a] -> a
head [(a, b)]
xs), [b] -> c
f (((a, b) -> b) -> [(a, b)] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> b
forall a b. (a, b) -> b
snd [(a, b)]
xs))

isSorted :: Ord a => [a] -> Bool
isSorted :: [a] -> Bool
isSorted [a]
xs = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((a -> a -> Bool) -> [a] -> [a] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<=) [a]
xs ([a] -> [a]
forall a. [a] -> [a]
tail [a]
xs))

isSortedBy :: Ord b => (a -> b) -> [a] -> Bool
isSortedBy :: (a -> b) -> [a] -> Bool
isSortedBy a -> b
f [a]
xs = [b] -> Bool
forall a. Ord a => [a] -> Bool
isSorted ((a -> b) -> [a] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map a -> b
f [a]
xs)

usort :: Ord a => [a] -> [a]
usort :: [a] -> [a]
usort = (a -> a -> Ordering) -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a]
usortBy a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare

usortBy :: (a -> a -> Ordering) -> [a] -> [a]
usortBy :: (a -> a -> Ordering) -> [a] -> [a]
usortBy a -> a -> Ordering
f = ([a] -> a) -> [[a]] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map [a] -> a
forall a. [a] -> a
head ([[a]] -> [a]) -> ([a] -> [[a]]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Bool) -> [a] -> [[a]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (\a
x a
y -> a -> a -> Ordering
f a
x a
y Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ) ([a] -> [[a]]) -> ([a] -> [a]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Ordering) -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy a -> a -> Ordering
f

sortBy' :: Ord b => (a -> b) -> [a] -> [a]
sortBy' :: (a -> b) -> [a] -> [a]
sortBy' a -> b
f = ((b, a) -> a) -> [(b, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (b, a) -> a
forall a b. (a, b) -> b
snd ([(b, a)] -> [a]) -> ([a] -> [(b, a)]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((b, a) -> (b, a) -> Ordering) -> [(b, a)] -> [(b, a)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (((b, a) -> b) -> (b, a) -> (b, a) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (b, a) -> b
forall a b. (a, b) -> a
fst) ([(b, a)] -> [(b, a)]) -> ([a] -> [(b, a)]) -> [a] -> [(b, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> (b, a)) -> [a] -> [(b, a)]
forall a b. (a -> b) -> [a] -> [b]
map (\a
x -> (a -> b
f a
x, a
x))

usortBy' :: Ord b => (a -> b) -> [a] -> [a]
usortBy' :: (a -> b) -> [a] -> [a]
usortBy' a -> b
f = ((b, a) -> a) -> [(b, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (b, a) -> a
forall a b. (a, b) -> b
snd ([(b, a)] -> [a]) -> ([a] -> [(b, a)]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((b, a) -> (b, a) -> Ordering) -> [(b, a)] -> [(b, a)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
usortBy (((b, a) -> b) -> (b, a) -> (b, a) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (b, a) -> b
forall a b. (a, b) -> a
fst) ([(b, a)] -> [(b, a)]) -> ([a] -> [(b, a)]) -> [a] -> [(b, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> (b, a)) -> [a] -> [(b, a)]
forall a b. (a -> b) -> [a] -> [b]
map (\a
x -> (a -> b
f a
x, a
x))

orElse :: Ordering -> Ordering -> Ordering
Ordering
EQ orElse :: Ordering -> Ordering -> Ordering
`orElse` Ordering
x = Ordering
x
Ordering
x  `orElse` Ordering
_ = Ordering
x

unbuffered :: IO a -> IO a
unbuffered :: IO a -> IO a
unbuffered IO a
x = do
  BufferMode
buf <- Handle -> IO BufferMode
hGetBuffering Handle
stdout
  IO () -> IO () -> IO a -> IO a
forall a b c. IO a -> IO b -> IO c -> IO c
bracket_
    (Handle -> BufferMode -> IO ()
hSetBuffering Handle
stdout BufferMode
NoBuffering)
    (Handle -> BufferMode -> IO ()
hSetBuffering Handle
stdout BufferMode
buf)
    IO a
x

labelM :: Monad m => (a -> m b) -> [a] -> m [(a, b)]
labelM :: (a -> m b) -> [a] -> m [(a, b)]
labelM a -> m b
f = (a -> m (a, b)) -> [a] -> m [(a, b)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\a
x -> do { b
y <- a -> m b
f a
x; (a, b) -> m (a, b)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
x, b
y) })

#if __GLASGOW_HASKELL__ < 710
isSubsequenceOf :: Ord a => [a] -> [a] -> Bool
[] `isSubsequenceOf` ys = True
(x:xs) `isSubsequenceOf` [] = False
(x:xs) `isSubsequenceOf` (y:ys)
  | x == y = xs `isSubsequenceOf` ys
  | otherwise = (x:xs) `isSubsequenceOf` ys
#endif

fixpoint :: Eq a => (a -> a) -> a -> a
fixpoint :: (a -> a) -> a -> a
fixpoint = (a -> a) -> (a -> a) -> a -> a
forall b a. Eq b => (a -> b) -> (a -> a) -> a -> a
fixpointOn a -> a
forall a. a -> a
id

{-# INLINE fixpoint #-}
fixpointOn :: Eq b => (a -> b) -> (a -> a) -> a -> a
fixpointOn :: (a -> b) -> (a -> a) -> a -> a
fixpointOn a -> b
key a -> a
f a
x = a -> a
fxp a
x
  where
    fxp :: a -> a
fxp a
x
      | a -> b
key a
x b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== a -> b
key a
y = a
x
      | Bool
otherwise = a -> a
fxp a
y
      where
        y :: a
y = a -> a
f a
x

-- From "Bit twiddling hacks": branchless min and max
{-# INLINE intMin #-}
intMin :: Int -> Int -> Int
intMin :: Int -> Int -> Int
intMin Int
x Int
y =
  Int
y Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` ((Int
x Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` Int
y) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
forall a. Num a => a -> a
negate (Int
x Int -> Int -> Int
.<. Int
y))
  where
    I# Int#
x .<. :: Int -> Int -> Int
.<. I# Int#
y = Int# -> Int
I# (Int#
x Int# -> Int# -> Int#
<# Int#
y)

{-# INLINE intMax #-}
intMax :: Int -> Int -> Int
intMax :: Int -> Int -> Int
intMax Int
x Int
y =
  Int
x Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` ((Int
x Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` Int
y) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
forall a. Num a => a -> a
negate (Int
x Int -> Int -> Int
.<. Int
y))
  where
    I# Int#
x .<. :: Int -> Int -> Int
.<. I# Int#
y = Int# -> Int
I# (Int#
x Int# -> Int# -> Int#
<# Int#
y)

-- Split an interval (inclusive bounds) into a particular number of blocks
splitInterval :: Integral a => a -> (a, a) -> [(a, a)]
splitInterval :: a -> (a, a) -> [(a, a)]
splitInterval a
k (a
lo, a
hi) =
  [ (a
loa -> a -> a
forall a. Num a => a -> a -> a
+a
ia -> a -> a
forall a. Num a => a -> a -> a
*a
blockSize, (a
loa -> a -> a
forall a. Num a => a -> a -> a
+(a
ia -> a -> a
forall a. Num a => a -> a -> a
+a
1)a -> a -> a
forall a. Num a => a -> a -> a
*a
blockSizea -> a -> a
forall a. Num a => a -> a -> a
-a
1) a -> a -> a
forall a. Ord a => a -> a -> a
`min` a
hi)
  | a
i <- [a
0..a
ka -> a -> a
forall a. Num a => a -> a -> a
-a
1] ]
  where
    size :: a
size = (a
hia -> a -> a
forall a. Num a => a -> a -> a
-a
loa -> a -> a
forall a. Num a => a -> a -> a
+a
1)
    blockSize :: a
blockSize = (a
size a -> a -> a
forall a. Num a => a -> a -> a
+ a
k a -> a -> a
forall a. Num a => a -> a -> a
- a
1) a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
k -- division rounding up
{-
prop_split_1 (Positive k) (lo, hi) =
  -- Check that all elements occur exactly once
  concat [[x..y] | (x, y) <- splitInterval k (lo, hi)] === [lo..hi]

-- Check that we have the correct number and distribution of blocks
prop_split_2 (Positive k) (lo, hi) =
  counterexample (show splits) $ conjoin
    [counterexample "Reason: too many splits" $
       length splits <= k,
     counterexample "Reason: too few splits" $
       length [lo..hi] >= k ==> length splits == k,
     counterexample "Reason: uneven distribution" $
      not (null splits) ==>
       minimum (map length splits) + 1 >= maximum (map length splits)]
  where
    splits = splitInterval k (lo, hi)
-}

reservoir :: Int -> [(Integer, Int)]
reservoir :: Int -> [(Integer, Int)]
reservoir Int
k =
  [Integer] -> [Int] -> [(Integer, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Int]
prefix) [Int]
prefix [(Integer, Int)] -> [(Integer, Int)] -> [(Integer, Int)]
forall a. [a] -> [a] -> [a]
++
  [Integer] -> [Int] -> [(Integer, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) ((Integer -> Integer -> Integer) -> [Integer] -> [Integer]
forall a. (a -> a -> a) -> [a] -> [a]
scanl1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+) [Integer]
is)) [Int]
ks
  where
    xs, ys :: [Double]
    xs :: [Double]
xs = (Double, Double) -> StdGen -> [Double]
forall a g. (Random a, RandomGen g) => (a, a) -> g -> [a]
randomRs (Double
0, Double
1) (Int -> StdGen
mkStdGen Int
314159265)
    ys :: [Double]
ys = (Double, Double) -> StdGen -> [Double]
forall a g. (Random a, RandomGen g) => (a, a) -> g -> [a]
randomRs (Double
0, Double
1) (Int -> StdGen
mkStdGen Int
358979323)
    ks :: [Int]
ks = (Int, Int) -> StdGen -> [Int]
forall a g. (Random a, RandomGen g) => (a, a) -> g -> [a]
randomRs (Int
0, Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (Int -> StdGen
mkStdGen Int
846264338)

    ws :: [Double]
ws = (Double -> Double -> Double) -> [Double] -> [Double]
forall a. (a -> a -> a) -> [a] -> [a]
scanl1 Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [ Double
x Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) | Double
x <- [Double]
xs ]
    is :: [Integer]
is = (Double -> Double -> Integer) -> [Double] -> [Double] -> [Integer]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Integer
forall a a. (RealFrac a, Integral a, Floating a) => a -> a -> a
gen [Double]
ws [Double]
ys
    gen :: a -> a -> a
gen a
w a
y = a -> a
forall a b. (RealFrac a, Integral b) => a -> b
floor (a -> a
forall a. Floating a => a -> a
log a
y a -> a -> a
forall a. Fractional a => a -> a -> a
/ a -> a
forall a. Floating a => a -> a
log (a
1a -> a -> a
forall a. Num a => a -> a -> a
-a
w)) a -> a -> a
forall a. Num a => a -> a -> a
+ a
1
    prefix :: [Int]
prefix = [Int
0..Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]

data Sample a = Sample Integer [(Integer, Int)] [a]

emptySample :: Int -> Sample a
emptySample :: Int -> Sample a
emptySample Int
k = Integer -> [(Integer, Int)] -> [a] -> Sample a
forall a. Integer -> [(Integer, Int)] -> [a] -> Sample a
Sample Integer
0 (Int -> [(Integer, Int)]
reservoir Int
k) []

addSample :: (Int, [a]) -> Sample a -> Sample a
addSample :: (Int, [a]) -> Sample a -> Sample a
addSample (Int
m, [a]
xs) (Sample Integer
total ((Integer
n, Int
pos):[(Integer, Int)]
ps) [a]
sample)
  | Integer
idx Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m =
    (Int, [a]) -> Sample a -> Sample a
forall a. (Int, [a]) -> Sample a -> Sample a
addSample (Int
m, [a]
xs) (Sample a -> Sample a) -> Sample a -> Sample a
forall a b. (a -> b) -> a -> b
$
      Integer -> [(Integer, Int)] -> [a] -> Sample a
forall a. Integer -> [(Integer, Int)] -> [a] -> Sample a
Sample Integer
total [(Integer, Int)]
ps ([a] -> Sample a) -> [a] -> Sample a
forall a b. (a -> b) -> a -> b
$
        Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
pos [a]
sample [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++
        [[a]
xs [a] -> Int -> a
forall a. [a] -> Int -> a
!! Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
idx] [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++
        Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
posInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [a]
sample
  where
    idx :: Integer
idx = Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
total
addSample (Int
m, [a]
_) (Sample Integer
total [(Integer, Int)]
ps [a]
sample) =
  Integer -> [(Integer, Int)] -> [a] -> Sample a
forall a. Integer -> [(Integer, Int)] -> [a] -> Sample a
Sample (Integer
totalInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m) [(Integer, Int)]
ps [a]
sample

sampleValue :: Sample a -> [a]
sampleValue :: Sample a -> [a]
sampleValue (Sample Integer
_ [(Integer, Int)]
_ [a]
sample) = [a]
sample

mapSample :: (a -> b) -> Sample a -> Sample b
mapSample :: (a -> b) -> Sample a -> Sample b
mapSample a -> b
f (Sample Integer
total [(Integer, Int)]
ps [a]
sample) =
  Integer -> [(Integer, Int)] -> [b] -> Sample b
forall a. Integer -> [(Integer, Int)] -> [a] -> Sample a
Sample Integer
total [(Integer, Int)]
ps ((a -> b) -> [a] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map a -> b
f [a]
sample)

-- A combined inits/tails.
splits :: [a] -> [([a], [a])]
splits :: [a] -> [([a], [a])]
splits [] = [([], [])]
splits (a
x:[a]
xs) =
  [([], a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs)] [([a], [a])] -> [([a], [a])] -> [([a], [a])]
forall a. [a] -> [a] -> [a]
++
  [(a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys, [a]
zs) | ([a]
ys, [a]
zs) <- [a] -> [([a], [a])]
forall a. [a] -> [([a], [a])]
splits [a]
xs]

-- Fold over the natural numbers.
foldn :: (a -> a) -> a -> Int -> a
foldn :: (a -> a) -> a -> Int -> a
foldn a -> a
_ a
e Int
0 = a
e
foldn a -> a
op a
e Int
n | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = a -> a
op ((a -> a) -> a -> Int -> a
forall a. (a -> a) -> a -> Int -> a
foldn a -> a
op a
e (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))