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

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

import GHC.Exts
import GHC.Prelude
import GHC.ST

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 #)


-- | 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


-- | 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 #)