{-# LANGUAGE CPP #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ParallelListComp #-}
{-# LANGUAGE MagicHash #-}
module Data.Discrimination.Internal
  ( runs
  , groupNum
  , bdiscNat
  , updateBag
  , updateSet
  , spanEither
  , integerCases
  , naturalCases
  ) where

import Data.Array as Array
import Data.Functor
import Data.Int
import qualified Data.List as List
import Prelude hiding (read, concat)

import GHC.Word
import GHC.Exts
import Data.Primitive.Types (Prim)
import Data.Primitive.PrimArray

#ifdef MIN_VERSION_ghc_bignum
import GHC.Num.Integer
import GHC.Num.Natural
#else
import GHC.Natural
import GHC.Integer.GMP.Internals
#endif

--------------------------------------------------------------------------------
-- * Utilities
--------------------------------------------------------------------------------

bdiscNat :: Int -> ([v] -> v -> [v]) -> [(Int,v)] -> [[v]]
bdiscNat :: Int -> ([v] -> v -> [v]) -> [(Int, v)] -> [[v]]
bdiscNat Int
n [v] -> v -> [v]
update [(Int, v)]
xs = [v] -> [v]
forall a. [a] -> [a]
reverse ([v] -> [v]) -> [[v]] -> [[v]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array Int [v] -> [[v]]
forall i e. Array i e -> [e]
Array.elems (([v] -> v -> [v])
-> [v] -> (Int, Int) -> [(Int, v)] -> Array Int [v]
forall i e a.
Ix i =>
(e -> a -> e) -> e -> (i, i) -> [(i, a)] -> Array i e
Array.accumArray [v] -> v -> [v]
update [] (Int
0,Int
n) [(Int, v)]
xs)
{-# INLINE bdiscNat #-}

runs :: Eq a => [(a,b)] -> [[b]]
runs :: [(a, b)] -> [[b]]
runs [] = []
runs ((a
a,b
b):[(a, b)]
xs0) = (b
bb -> [b] -> [b]
forall a. a -> [a] -> [a]
:[b]
ys0) [b] -> [[b]] -> [[b]]
forall a. a -> [a] -> [a]
: [(a, b)] -> [[b]]
forall a b. Eq a => [(a, b)] -> [[b]]
runs [(a, b)]
zs0
  where
    ([b]
ys0,[(a, b)]
zs0) = [(a, b)] -> ([b], [(a, b)])
forall a. [(a, a)] -> ([a], [(a, a)])
go [(a, b)]
xs0
    go :: [(a, a)] -> ([a], [(a, a)])
go [] = ([],[])
    go xs :: [(a, a)]
xs@((a
a', a
b'):[(a, a)]
xs')
      | a
a a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
a' = case [(a, a)] -> ([a], [(a, a)])
go [(a, a)]
xs' of
         ([a]
ys, [(a, a)]
zs) -> (a
b'a -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys,[(a, a)]
zs)
      | Bool
otherwise = ([], [(a, a)]
xs)

groupNum :: [[k]] -> [(k,Int)]
groupNum :: [[k]] -> [(k, Int)]
groupNum [[k]]
kss = [[(k, Int)]] -> [(k, Int)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
List.concat [ (,Int
n) (k -> (k, Int)) -> [k] -> [(k, Int)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [k]
ks | Int
n <- [Int
0..] | [k]
ks <- [[k]]
kss ]

updateBag :: [Int] -> Int -> [Int]
updateBag :: [Int] -> Int -> [Int]
updateBag [Int]
vs Int
v = Int
v Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
vs

updateSet :: [Int] -> Int -> [Int]
updateSet :: [Int] -> Int -> [Int]
updateSet [] Int
w = [Int
w]
updateSet vs :: [Int]
vs@(Int
v:[Int]
_) Int
w
  | Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
w    = [Int]
vs
  | Bool
otherwise = Int
w Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
vs

-- | Optimized and CPS'd version of 'Data.Either.partitionEithers', where all lefts are known to come before all rights
spanEither :: ([a] -> [b] -> c) -> [Either a b] -> c
spanEither :: ([a] -> [b] -> c) -> [Either a b] -> c
spanEither [a] -> [b] -> c
k [Either a b]
xs0 = [a] -> [Either a b] -> c
go [] [Either a b]
xs0 where
  go :: [a] -> [Either a b] -> c
go [a]
acc (Left a
x:[Either a b]
xs) = [a] -> [Either a b] -> c
go (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
acc) [Either a b]
xs
  go [a]
acc [Either a b]
rights = [a] -> [b] -> c
k ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
acc) (Either a b -> b
forall a b. Either a b -> b
fromRight (Either a b -> b) -> [Either a b] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Either a b]
rights)

fromRight :: Either a b -> b
fromRight :: Either a b -> b
fromRight (Right b
y) = b
y
fromRight Either a b
_ = [Char] -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"unstable discriminator"

-------------------------------------------------------------------------------
-- * Integer and Natural
-------------------------------------------------------------------------------

integerCases :: Integer -> Either (Int,[Word]) (Either Int (Int,[Word]))
#ifdef MIN_VERSION_ghc_bignum
integerCases (IN b) = Left          $ decomposeBigNat b
integerCases (IS i) = Right . Left  $ I# i
integerCases (IP b) = Right . Right $ decomposeBigNat b
#else
integerCases :: Integer -> Either (Int, [Word]) (Either Int (Int, [Word]))
integerCases (Jn# BigNat
b) = (Int, [Word]) -> Either (Int, [Word]) (Either Int (Int, [Word]))
forall a b. a -> Either a b
Left          ((Int, [Word]) -> Either (Int, [Word]) (Either Int (Int, [Word])))
-> (Int, [Word]) -> Either (Int, [Word]) (Either Int (Int, [Word]))
forall a b. (a -> b) -> a -> b
$ BigNat -> (Int, [Word])
decomposeBigNat BigNat
b
integerCases (S#  Int#
i) = Either Int (Int, [Word])
-> Either (Int, [Word]) (Either Int (Int, [Word]))
forall a b. b -> Either a b
Right (Either Int (Int, [Word])
 -> Either (Int, [Word]) (Either Int (Int, [Word])))
-> (Int -> Either Int (Int, [Word]))
-> Int
-> Either (Int, [Word]) (Either Int (Int, [Word]))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Either Int (Int, [Word])
forall a b. a -> Either a b
Left  (Int -> Either (Int, [Word]) (Either Int (Int, [Word])))
-> Int -> Either (Int, [Word]) (Either Int (Int, [Word]))
forall a b. (a -> b) -> a -> b
$ Int# -> Int
I# Int#
i
integerCases (Jp# BigNat
b) = Either Int (Int, [Word])
-> Either (Int, [Word]) (Either Int (Int, [Word]))
forall a b. b -> Either a b
Right (Either Int (Int, [Word])
 -> Either (Int, [Word]) (Either Int (Int, [Word])))
-> ((Int, [Word]) -> Either Int (Int, [Word]))
-> (Int, [Word])
-> Either (Int, [Word]) (Either Int (Int, [Word]))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, [Word]) -> Either Int (Int, [Word])
forall a b. b -> Either a b
Right ((Int, [Word]) -> Either (Int, [Word]) (Either Int (Int, [Word])))
-> (Int, [Word]) -> Either (Int, [Word]) (Either Int (Int, [Word]))
forall a b. (a -> b) -> a -> b
$ BigNat -> (Int, [Word])
decomposeBigNat BigNat
b
#endif
{-# INLINE integerCases #-}

naturalCases :: Natural -> Either Word (Int,[Word])
#ifdef MIN_VERSION_ghc_bignum
naturalCases (NS w) = Left $ W# w
naturalCases (NB b) = Right $ decomposeBigNat b
#else
naturalCases :: Natural -> Either Word (Int, [Word])
naturalCases (NatS# GmpLimb#
w) = Word -> Either Word (Int, [Word])
forall a b. a -> Either a b
Left (Word -> Either Word (Int, [Word]))
-> Word -> Either Word (Int, [Word])
forall a b. (a -> b) -> a -> b
$ GmpLimb# -> Word
W# GmpLimb#
w
naturalCases (NatJ# BigNat
b) = (Int, [Word]) -> Either Word (Int, [Word])
forall a b. b -> Either a b
Right ((Int, [Word]) -> Either Word (Int, [Word]))
-> (Int, [Word]) -> Either Word (Int, [Word])
forall a b. (a -> b) -> a -> b
$ BigNat -> (Int, [Word])
decomposeBigNat BigNat
b
#endif
{-# INLINE naturalCases #-}

-- We need to reverse the limb array. Its stored least-significant word first
-- but for comparasion to work right we need most-significant words first.
#ifdef MIN_VERSION_ghc_bignum
decomposeBigNat :: ByteArray# -> (Int, [Word])
decomposeBigNat ba = let pa = PrimArray ba :: PrimArray Word in (sizeofPrimArray pa, primArrayToReverseList pa)
#else
decomposeBigNat :: BigNat -> (Int, [Word])
decomposeBigNat :: BigNat -> (Int, [Word])
decomposeBigNat (BN# ByteArray#
ba) = let pa :: PrimArray Word
pa = ByteArray# -> PrimArray Word
forall a. ByteArray# -> PrimArray a
PrimArray ByteArray#
ba :: PrimArray Word in (PrimArray Word -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Word
pa, PrimArray Word -> [Word]
forall a. Prim a => PrimArray a -> [a]
primArrayToReverseList PrimArray Word
pa)
#endif
{-# INLINE decomposeBigNat #-}

primArrayToReverseList :: Prim a => PrimArray a -> [a]
primArrayToReverseList :: PrimArray a -> [a]
primArrayToReverseList PrimArray a
xs = (forall b. (a -> b -> b) -> b -> b) -> [a]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (\a -> b -> b
c b
n -> (b -> a -> b) -> b -> PrimArray a -> b
forall a b. Prim a => (b -> a -> b) -> b -> PrimArray a -> b
foldlPrimArray ((a -> b -> b) -> b -> a -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> b
c) b
n PrimArray a
xs)
{-# INLINE primArrayToReverseList #-}