{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE BlockArguments #-}

-- | Small-array
module GHC.Data.SmallArray
  ( SmallMutableArray (..)
  , SmallArray (..)
  , newSmallArray
  , writeSmallArray
  , freezeSmallArray
  , unsafeFreezeSmallArray
  , indexSmallArray
  , sizeofSmallArray
  , listToArray
  , mapSmallArray
  , foldMapSmallArray
  , rnfSmallArray
  )
where

import GHC.Exts
import GHC.Prelude
import GHC.ST
import Control.DeepSeq

data SmallArray a = SmallArray (SmallArray# a)

data SmallMutableArray s a = SmallMutableArray (SmallMutableArray# s a)

newSmallArray
  :: Int  -- ^ size
  -> a    -- ^ initial contents
  -> State# s
  -> (# State# s, SmallMutableArray s a #)
{-# INLINE newSmallArray #-}
newSmallArray :: forall a s.
Int -> a -> State# s -> (# State# s, SmallMutableArray s a #)
newSmallArray (I# Int#
sz) a
x State# s
s = case Int# -> a -> State# s -> (# State# s, SmallMutableArray# s a #)
forall a d.
Int# -> a -> State# d -> (# State# d, SmallMutableArray# d a #)
newSmallArray# Int#
sz a
x State# s
s of
  (# State# s
s', SmallMutableArray# s a
a #) -> (# State# s
s', SmallMutableArray# s a -> SmallMutableArray s a
forall s a. SmallMutableArray# s a -> SmallMutableArray s a
SmallMutableArray SmallMutableArray# s a
a #)

writeSmallArray
  :: SmallMutableArray s a -- ^ array
  -> Int                   -- ^ index
  -> a                     -- ^ new element
  -> State# s
  -> State# s
{-# INLINE writeSmallArray #-}
writeSmallArray :: forall s a.
SmallMutableArray s a -> Int -> a -> State# s -> State# s
writeSmallArray (SmallMutableArray SmallMutableArray# s a
a) (I# Int#
i) a
x = SmallMutableArray# s a -> Int# -> a -> State# s -> State# s
forall d a.
SmallMutableArray# d a -> Int# -> a -> State# d -> State# d
writeSmallArray# SmallMutableArray# s a
a Int#
i a
x


-- | Copy and freeze a slice of a mutable array.
freezeSmallArray
  :: SmallMutableArray s a -- ^ source
  -> Int                   -- ^ offset
  -> Int                   -- ^ length
  -> State# s
  -> (# State# s, SmallArray a #)
{-# INLINE freezeSmallArray #-}
freezeSmallArray :: forall s a.
SmallMutableArray s a
-> Int -> Int -> State# s -> (# State# s, SmallArray a #)
freezeSmallArray (SmallMutableArray SmallMutableArray# s a
ma) (I# Int#
offset) (I# Int#
len) State# s
s =
  case SmallMutableArray# s a
-> Int# -> Int# -> State# s -> (# State# s, SmallArray# a #)
forall d a.
SmallMutableArray# d a
-> Int# -> Int# -> State# d -> (# State# d, SmallArray# a #)
freezeSmallArray# SmallMutableArray# s a
ma Int#
offset Int#
len State# s
s of
    (# State# s
s', SmallArray# a
a #) -> (# State# s
s', SmallArray# a -> SmallArray a
forall a. SmallArray# a -> SmallArray a
SmallArray SmallArray# a
a #)

-- | Freeze a mutable array (no copy!)
unsafeFreezeSmallArray
  :: SmallMutableArray s a
  -> State# s
  -> (# State# s, SmallArray a #)
{-# INLINE unsafeFreezeSmallArray #-}
unsafeFreezeSmallArray :: forall s a.
SmallMutableArray s a -> State# s -> (# State# s, SmallArray a #)
unsafeFreezeSmallArray (SmallMutableArray SmallMutableArray# s a
ma) State# s
s =
  case SmallMutableArray# s a -> State# s -> (# State# s, SmallArray# a #)
forall d a.
SmallMutableArray# d a -> State# d -> (# State# d, SmallArray# a #)
unsafeFreezeSmallArray# SmallMutableArray# s a
ma State# s
s of
    (# State# s
s', SmallArray# a
a #) -> (# State# s
s', SmallArray# a -> SmallArray a
forall a. SmallArray# a -> SmallArray a
SmallArray SmallArray# a
a #)

-- | Get the size of a 'SmallArray'
sizeofSmallArray
  :: SmallArray a
  -> Int
{-# INLINE sizeofSmallArray #-}
sizeofSmallArray :: forall a. SmallArray a -> Int
sizeofSmallArray (SmallArray SmallArray# a
sa#) =
  case SmallArray# a -> Int#
forall a. SmallArray# a -> Int#
sizeofSmallArray# SmallArray# a
sa# of
    Int#
s -> Int# -> Int
I# Int#
s

-- | Index a small-array (no bounds checking!)
indexSmallArray
  :: SmallArray a -- ^ array
  -> Int          -- ^ index
  -> a
{-# INLINE indexSmallArray #-}
indexSmallArray :: forall a. SmallArray a -> Int -> a
indexSmallArray (SmallArray SmallArray# a
sa#) (I# Int#
i) =
  case SmallArray# a -> Int# -> (# a #)
forall a. SmallArray# a -> Int# -> (# a #)
indexSmallArray# SmallArray# a
sa# Int#
i of
    (# a
v #) -> a
v

-- | Map a function over the elements of a 'SmallArray'
--
mapSmallArray :: (a -> b) -> SmallArray a -> SmallArray b
{-# INLINE mapSmallArray #-}
mapSmallArray :: forall a b. (a -> b) -> SmallArray a -> SmallArray b
mapSmallArray a -> b
f SmallArray a
sa = (forall s. ST s (SmallArray b)) -> SmallArray b
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (SmallArray b)) -> SmallArray b)
-> (forall s. ST s (SmallArray b)) -> SmallArray b
forall a b. (a -> b) -> a -> b
$ STRep s (SmallArray b) -> ST s (SmallArray b)
forall s a. STRep s a -> ST s a
ST (STRep s (SmallArray b) -> ST s (SmallArray b))
-> STRep s (SmallArray b) -> ST s (SmallArray b)
forall a b. (a -> b) -> a -> b
$ \State# s
s ->
  let
    n :: Int
n = SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
sa
    go :: Int -> SmallMutableArray s b -> State# s -> State# s
go !Int
i SmallMutableArray s b
saMut# State# s
state#
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n =
        let
          a :: a
a = SmallArray a -> Int -> a
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray a
sa Int
i
          newState# :: State# s
newState# = SmallMutableArray s b -> Int -> b -> State# s -> State# s
forall s a.
SmallMutableArray s a -> Int -> a -> State# s -> State# s
writeSmallArray SmallMutableArray s b
saMut# Int
i (a -> b
f a
a) State# s
state#
        in
          Int -> SmallMutableArray s b -> State# s -> State# s
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) SmallMutableArray s b
saMut# State# s
newState#
      | Bool
otherwise = State# s
state#
  in
  case Int -> b -> State# s -> (# State# s, SmallMutableArray s b #)
forall a s.
Int -> a -> State# s -> (# State# s, SmallMutableArray s a #)
newSmallArray Int
n ([Char] -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"SmallArray: internal error, uninitialised elements") State# s
s of
    (# State# s
s', SmallMutableArray s b
mutArr #) ->
      case Int -> SmallMutableArray s b -> State# s -> State# s
go Int
0 SmallMutableArray s b
mutArr State# s
s' of
        State# s
s'' -> SmallMutableArray s b -> STRep s (SmallArray b)
forall s a.
SmallMutableArray s a -> State# s -> (# State# s, SmallArray a #)
unsafeFreezeSmallArray SmallMutableArray s b
mutArr State# s
s''

-- | Fold the values of a 'SmallArray' into a 'Monoid m' of choice
foldMapSmallArray :: Monoid m => (a -> m) -> SmallArray a -> m
{-# INLINE foldMapSmallArray #-}
foldMapSmallArray :: forall m a. Monoid m => (a -> m) -> SmallArray a -> m
foldMapSmallArray a -> m
f SmallArray a
sa = Int -> m
go Int
0
  where
    n :: Int
n = SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
sa
    go :: Int -> m
go Int
i
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = a -> m
f (SmallArray a -> Int -> a
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray a
sa Int
i) m -> m -> m
forall a. Monoid a => a -> a -> a
`mappend` Int -> m
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      | Bool
otherwise = m
forall a. Monoid a => a
mempty

-- | Force the elements of the given 'SmallArray'
--
rnfSmallArray :: NFData a => SmallArray a -> ()
{-# INLINE rnfSmallArray #-}
rnfSmallArray :: forall a. NFData a => SmallArray a -> ()
rnfSmallArray SmallArray a
sa = Int -> ()
go Int
0
  where
    n :: Int
n = SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
sa
    go :: Int -> ()
go !Int
i
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = a -> ()
forall a. NFData a => a -> ()
rnf (SmallArray a -> Int -> a
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray a
sa Int
i) () -> () -> ()
forall a b. a -> b -> b
`seq` Int -> ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      | Bool
otherwise = ()

-- | Convert a list into an array.
listToArray :: Int -> (e -> Int) -> (e -> a) -> [e] -> SmallArray a
{-# INLINE listToArray #-}
listToArray :: forall e a. Int -> (e -> Int) -> (e -> a) -> [e] -> SmallArray a
listToArray (I# Int#
size) e -> Int
index_of e -> a
value_of [e]
xs = (forall s. ST s (SmallArray a)) -> SmallArray a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (SmallArray a)) -> SmallArray a)
-> (forall s. ST s (SmallArray a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ STRep s (SmallArray a) -> ST s (SmallArray a)
forall s a. STRep s a -> ST s a
ST \State# s
s ->
  let
    index_of' :: e -> Int#
index_of' e
e = case e -> Int
index_of e
e of I# Int#
i -> Int#
i
    write_elems :: SmallMutableArray# s a -> [e] -> State# s -> State# s
write_elems SmallMutableArray# s a
ma [e]
es State# s
s = case [e]
es of
      []    -> State# s
s
      e
e:[e]
es' -> case SmallMutableArray# s a -> Int# -> a -> State# s -> State# s
forall d a.
SmallMutableArray# d a -> Int# -> a -> State# d -> State# d
writeSmallArray# SmallMutableArray# s a
ma (e -> Int#
index_of' e
e) (e -> a
value_of e
e) State# s
s of
                 State# s
s' -> SmallMutableArray# s a -> [e] -> State# s -> State# s
write_elems SmallMutableArray# s a
ma [e]
es' State# s
s'
  in
  case Int# -> a -> State# s -> (# State# s, SmallMutableArray# s a #)
forall a d.
Int# -> a -> State# d -> (# State# d, SmallMutableArray# d a #)
newSmallArray# Int#
size a
forall a. HasCallStack => a
undefined State# s
s of
    (# State# s
s', SmallMutableArray# s a
ma #) -> case SmallMutableArray# s a -> [e] -> State# s -> State# s
write_elems SmallMutableArray# s a
ma [e]
xs State# s
s' of
      State# s
s'' -> case SmallMutableArray# s a -> State# s -> (# State# s, SmallArray# a #)
forall d a.
SmallMutableArray# d a -> State# d -> (# State# d, SmallArray# a #)
unsafeFreezeSmallArray# SmallMutableArray# s a
ma State# s
s'' of
        (# State# s
s''', SmallArray# a
a #) -> (# State# s
s''', SmallArray# a -> SmallArray a
forall a. SmallArray# a -> SmallArray a
SmallArray SmallArray# a
a #)