{-# language DataKinds #-}
{-# language TypeOperators #-}
{-# language KindSignatures #-}
{-# language BangPatterns #-}
{-# language RoleAnnotations #-}
{-# language MagicHash #-}
{-# language UnboxedTuples #-}
{-# language NoStarIsType #-}
{-# language RankNTypes #-}
{-# language DeriveTraversable #-}
{-# language Unsafe #-}

module Data.CompactSequence.Internal.Array where
import Data.Primitive.SmallArray
import Control.Monad.ST.Strict

-- fixed-vector
-- unpacked-containers
-- contiguous

data Mult = Twice Mult | Mul1

newtype Array (n :: Mult) a = Array (SmallArray a)
  deriving (Functor, Foldable, Traversable)
type role Array nominal representational

newtype Size (n :: Mult) = Size Int
type role Size nominal

getSize :: Size n -> Int
getSize (Size n) = n

--halve :: Size (Twice m) -> Size m
--halve (Size n) = Size (n `quot` 2)

one :: Size Mul1
one = Size 1

twice :: Size n -> Size (Twice n)
twice (Size n) = Size (2*n)

singleton :: a -> Array Mul1 a
singleton x = Array (pure x)

-- | Unsafely convert a 'SmallArray' of size @n@
-- to an @'Array' n@. This is genuinely unsafe: if
-- @n@ is greater than the true array size, then
-- some operation will eventually violate memory safety.
unsafeSmallArrayToArray :: SmallArray a -> Array n a
unsafeSmallArrayToArray = Array

arrayToSmallArray :: Array n a -> SmallArray a
arrayToSmallArray (Array sa) = sa

getSingleton# :: Array Mul1 a -> (# a #)
getSingleton# (Array sa) = indexSmallArray## sa 0

getSingletonA :: Applicative f => Array Mul1 a -> f a
getSingletonA (Array sa)
  | (# a #) <- indexSmallArray## sa 0
  = pure a

splitArray :: Size n -> Array (Twice n) a -> (Array n a, Array n a)
splitArray (Size len) (Array sa1) = (Array sa2, Array sa3)
  where
    !sa2 = cloneSmallArray sa1 0 len
    !sa3 = cloneSmallArray sa1 len len

-- | Append two arrays of the same size. We take the size
-- of the argument arrays so we can build the result array
-- before loading the first argument array into cache. Is
-- this the right approach? Not sure. We *certainly* don't
-- want to just use `<>`, because 
append :: Size n -> Array n a -> Array n a -> Array (Twice n) a
append (Size n) (Array xs) (Array ys) = Array $
    createSmallArray (2*n)
      (error "Data.CompactSequence.Internal.Array.append: Internal error")
      $ \sma -> copySmallArray sma 0 xs 0 n
        *> copySmallArray sma n ys 0 n

-- Shamelessly stolen from primitive.
createSmallArray
  :: Int
  -> a
  -> (forall s. SmallMutableArray s a -> ST s ())
  -> SmallArray a
createSmallArray n x f = runSmallArray $ do
  mary <- newSmallArray n x
  f mary
  pure mary

arraySplitListN :: Size n -> [a] -> (Array n a, [a])
arraySplitListN (Size n) xs
  | (sa, xs') <- smallArraySplitListN n xs
  = (Array sa, xs')

smallArraySplitListN :: Int -> [a] -> (SmallArray a, [a])
smallArraySplitListN n l = runST $ do
  sma <- newSmallArray n (error "smallArraySplitListN: uninitialized")
  let go !ix [] = if ix == n
        then do
          sa <- unsafeFreezeSmallArray sma
          pure (sa, [])
        else error "smallArraySplitListN: list length less than specified size"
      go !ix xss@(x : xs) = if ix < n
        then do
          writeSmallArray sma ix x
          go (ix+1) xs
        else do
          sa <- unsafeFreezeSmallArray sma
          pure (sa, xss)
  go 0 l