content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE CPP #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# OPTIONS_GHC -Wall #-}
  {-# OPTIONS_HADDOCK show-extensions #-}
  -----------------------------------------------------------------------------
  -- |
  -- Module      :  ToySolver.Internal.Data.Vec
  -- Copyright   :  (c) Masahiro Sakai 2014
  -- License     :  BSD-style
  --
  -- Maintainer  :  masahiro.sakai@gmail.com
  -- Stability   :  provisional
  -- Portability :  non-portable
  --
  -- Simple 1-dimentional resizable array
  --
  -----------------------------------------------------------------------------
  module ToySolver.Internal.Data.Vec
    (
    -- * Vec type
      GenericVec
    , Vec
    , UVec
    , Index
  
    -- * Constructors
    , new
    , clone
  
    -- * Operators
    , getSize
    , read
    , write
    , modify
    , modify'
    , unsafeRead
    , unsafeWrite
    , unsafeModify
    , unsafeModify'
    , resize
    , growTo
    , push
    , pop
    , popMaybe
    , unsafePop
    , peek
    , unsafePeek
    , clear
    , getElems
  
    -- * Low-level operators
    , getArray
    , getCapacity
    , resizeCapacity
    ) where
  
  import Prelude hiding (read)
  
  import Control.Loop
  import Control.Monad
  import Data.Ix
  import qualified Data.Array.Base as A
  import qualified Data.Array.IO as A
  import Data.IORef
  import ToySolver.Internal.Data.IOURef
  
  data GenericVec a e = GenericVec {-# UNPACK #-} !(IOURef Int) {-# UNPACK #-} !(IORef (a Index e))
    deriving Eq
  
  type Vec e = GenericVec A.IOArray e
  type UVec e = GenericVec A.IOUArray e
  
  type Index = Int
  
  {- INLINE readArray #-}
  readArray :: A.MArray a e m => a Index e -> Index -> m e
  #ifdef EXTRA_BOUNDS_CHECKING
  readArray = A.readArray
  #else
  readArray = A.unsafeRead
  #endif
  
  {- INLINE writeArray #-}
  writeArray :: A.MArray a e m => a Index e -> Index -> e -> m ()
  #ifdef EXTRA_BOUNDS_CHECKING
  writeArray = A.writeArray
  #else
  writeArray = A.unsafeWrite
  #endif
  
  new :: A.MArray a e IO => IO (GenericVec a e)
  new = do
    sizeRef <- newIOURef 0
    arrayRef <- newIORef =<< A.newArray_ (0,-1)
    return $ GenericVec sizeRef arrayRef
  
  {- INLINE getSize #-}
  -- | Get the internal representation array
  getSize :: GenericVec a e -> IO Int
  getSize (GenericVec sizeRef _) = readIOURef sizeRef
  
  {-# SPECIALIZE read :: Vec e -> Int -> IO e #-}
  {-# SPECIALIZE read :: UVec Int -> Int -> IO Int #-}
  {-# SPECIALIZE read :: UVec Double -> Int -> IO Double #-}
  {-# SPECIALIZE read :: UVec Bool -> Int -> IO Bool #-}
  read :: A.MArray a e IO => GenericVec a e -> Int -> IO e
  read !v !i = do
    a <- getArray v
    s <- getSize v
    if 0 <= i && i < s then
      readArray a i
    else
      error $ "ToySolver.Internal.Data.Vec.read: index " ++ show i ++ " out of bounds"
  
  {-# SPECIALIZE write :: Vec e -> Int -> e -> IO () #-}
  {-# SPECIALIZE write :: UVec Int -> Int -> Int -> IO () #-}
  {-# SPECIALIZE write :: UVec Double -> Int -> Double -> IO () #-}
  {-# SPECIALIZE write :: UVec Bool -> Int -> Bool -> IO () #-}
  write :: A.MArray a e IO => GenericVec a e -> Int -> e -> IO ()
  write !v !i e = do
    a <- getArray v
    s <- getSize v
    if 0 <= i && i < s then
      writeArray a i e
    else
      error $ "ToySolver.Internal.Data.Vec.write: index " ++ show i ++ " out of bounds"
  
  {-# INLINE modify #-}
  modify :: A.MArray a e IO => GenericVec a e -> Int -> (e -> e) -> IO ()
  modify !v !i f = do
    a <- getArray v
    s <- getSize v
    if 0 <= i && i < s then do
      x <- readArray a i
      writeArray a i (f x)
    else
      error $ "ToySolver.Internal.Data.Vec.modify: index " ++ show i ++ " out of bounds"
  
  {-# INLINE modify' #-}
  modify' :: A.MArray a e IO => GenericVec a e -> Int -> (e -> e) -> IO ()
  modify' !v !i f = do
    a <- getArray v
    s <- getSize v
    if 0 <= i && i < s then do
      x <- readArray a i
      writeArray a i $! f x
    else
      error $ "ToySolver.Internal.Data.Vec.modify': index " ++ show i ++ " out of bounds"
  
  {-# INLINE unsafeModify #-}
  unsafeModify :: A.MArray a e IO => GenericVec a e -> Int -> (e -> e) -> IO ()
  unsafeModify !v !i f = do
    a <- getArray v
    x <- readArray a i
    writeArray a i (f x)
  
  {-# INLINE unsafeModify' #-}
  unsafeModify' :: A.MArray a e IO => GenericVec a e -> Int -> (e -> e) -> IO ()
  unsafeModify' !v !i f = do
    a <- getArray v
    x <- readArray a i
    writeArray a i $! f x
  
  {-# INLINE unsafeRead #-}
  unsafeRead :: A.MArray a e IO => GenericVec a e -> Int -> IO e
  unsafeRead !v !i = do
    a <- getArray v
    readArray a i
  
  {-# INLINE unsafeWrite #-}
  unsafeWrite :: A.MArray a e IO => GenericVec a e -> Int -> e -> IO ()
  unsafeWrite !v !i e = do
    a <- getArray v
    writeArray a i e
  
  {-# SPECIALIZE resize :: Vec e -> Int -> IO () #-}
  {-# SPECIALIZE resize :: UVec Int -> Int -> IO () #-}
  {-# SPECIALIZE resize :: UVec Double -> Int -> IO () #-}
  {-# SPECIALIZE resize :: UVec Bool -> Int -> IO () #-}
  resize :: A.MArray a e IO => GenericVec a e -> Int -> IO ()
  resize v@(GenericVec sizeRef arrayRef) !n = do
    a <- getArray v
    capa <- getCapacity v
    unless (n <= capa) $ do
      let capa' = max 2 (capa * 3 `div` 2)
      a' <- A.newArray_ (0, capa'-1)
      copyTo a a' (0,capa-1)
      writeIORef arrayRef a'
    writeIOURef sizeRef n
  
  {-# SPECIALIZE growTo :: Vec e -> Int -> IO () #-}
  {-# SPECIALIZE growTo :: UVec Int -> Int -> IO () #-}
  {-# SPECIALIZE growTo :: UVec Double -> Int -> IO () #-}
  {-# SPECIALIZE growTo :: UVec Bool -> Int -> IO () #-}
  growTo :: A.MArray a e IO => GenericVec a e -> Int -> IO ()
  growTo v !n = do
    m <- getSize v
    when (m < n) $ resize v n
  
  {-# SPECIALIZE push :: Vec e -> e -> IO () #-}
  {-# SPECIALIZE push :: UVec Int -> Int -> IO () #-}
  {-# SPECIALIZE push :: UVec Double -> Double -> IO () #-}
  {-# SPECIALIZE push :: UVec Bool -> Bool -> IO () #-}
  push :: A.MArray a e IO => GenericVec a e -> e -> IO ()
  push v e = do
    s <- getSize v
    resize v (s+1)
    unsafeWrite v s e
  
  popMaybe :: A.MArray a e IO => GenericVec a e -> IO (Maybe e)
  popMaybe v = do
    s <- getSize v
    if s == 0 then
      return Nothing
    else do
      e <- unsafeRead v (s-1)
      resize v (s-1)
      return (Just e)
  
  {-# SPECIALIZE unsafePop :: Vec e -> IO e #-}
  {-# SPECIALIZE unsafePop :: UVec Int -> IO Int #-}
  {-# SPECIALIZE unsafePop :: UVec Double -> IO Double #-}
  {-# SPECIALIZE unsafePop :: UVec Bool -> IO Bool #-}
  pop :: A.MArray a e IO => GenericVec a e -> IO e
  pop v = do
    s <- getSize v
    if s == 0 then
      error $ "ToySolver.Internal.Data.Vec.pop: empty Vec"
    else do
      e <- unsafeRead v (s-1)
      resize v (s-1)
      return e
  
  {-# SPECIALIZE unsafePop :: Vec e -> IO e #-}
  {-# SPECIALIZE unsafePop :: UVec Int -> IO Int #-}
  {-# SPECIALIZE unsafePop :: UVec Double -> IO Double #-}
  {-# SPECIALIZE unsafePop :: UVec Bool -> IO Bool #-}
  unsafePop :: A.MArray a e IO => GenericVec a e -> IO e
  unsafePop v = do
    s <- getSize v
    e <- unsafeRead v (s-1)
    resize v (s-1)
    return e
  
  {-# INLINE peek #-}
  peek :: A.MArray a e IO => GenericVec a e -> IO e
  peek v = do
    s <- getSize v
    if s == 0 then
      error $ "ToySolver.Internal.Data.Vec.peek: empty Vec"
    else do
      unsafeRead v (s-1)
  
  {-# INLINE unsafePeek #-}
  unsafePeek :: A.MArray a e IO => GenericVec a e -> IO e
  unsafePeek v = do
    s <- getSize v
    unsafeRead v (s-1)
  
  clear :: A.MArray a e IO => GenericVec a e -> IO ()
  clear v = resize v 0
  
  getElems :: A.MArray a e IO => GenericVec a e -> IO [e]
  getElems v = do
    s <- getSize v
    forM [0..s-1] $ \i -> unsafeRead v i
  
  {-# SPECIALIZE clone :: Vec e -> IO (Vec e) #-}
  {-# SPECIALIZE clone :: UVec Int -> IO (UVec Int) #-}
  {-# SPECIALIZE clone :: UVec Double -> IO (UVec Double) #-}
  {-# SPECIALIZE clone :: UVec Bool -> IO (UVec Bool) #-}
  clone :: A.MArray a e IO => GenericVec a e -> IO (GenericVec a e)
  clone (GenericVec sizeRef arrayRef) = do
    a <- readIORef arrayRef
    arrayRef' <- newIORef =<< cloneArray a
    sizeRef'  <- newIOURef =<< readIOURef sizeRef
    return $ GenericVec sizeRef' arrayRef'
  
  {--------------------------------------------------------------------
  
  --------------------------------------------------------------------}
  
  {-# INLINE getArray #-}
  -- | Get the internal representation array
  getArray :: GenericVec a e -> IO (a Index e)
  getArray (GenericVec _ arrayRef) = readIORef arrayRef
  
  {-# INLINE getCapacity #-}
  -- | Get the internal representation array
  getCapacity :: A.MArray a e IO => GenericVec a e -> IO Int
  getCapacity vec = liftM rangeSize $ A.getBounds =<< getArray vec
  
  {-# SPECIALIZE resizeCapacity :: Vec e -> Int -> IO () #-}
  {-# SPECIALIZE resizeCapacity :: UVec Int -> Int -> IO () #-}
  {-# SPECIALIZE resizeCapacity :: UVec Double -> Int -> IO () #-}
  {-# SPECIALIZE resizeCapacity :: UVec Bool -> Int -> IO () #-}
  -- | Pre-allocate internal buffer for @n@ elements.
  resizeCapacity :: A.MArray a e IO => GenericVec a e -> Int -> IO ()
  resizeCapacity (GenericVec sizeRef arrayRef) capa = do
    n <- readIOURef sizeRef
    arr <- readIORef arrayRef
    capa0 <- liftM rangeSize $ A.getBounds arr
    when (capa0 < capa) $ do
      let capa' = max capa (n * 3 `div` 2)
      arr' <- A.newArray_ (0, capa'-1)
      copyTo arr arr' (0, n-1)
      writeIORef arrayRef arr'
  
  {--------------------------------------------------------------------
    utility
  --------------------------------------------------------------------}
  
  {-# INLINE cloneArray #-}
  cloneArray :: (A.MArray a e m) => a Index e -> m (a Index e)
  cloneArray arr = do
    b <- A.getBounds arr
    arr' <- A.newArray_ b
    copyTo arr arr' b
    return arr'
  
  {-# INLINE copyTo #-}
  copyTo :: (A.MArray a e m) => a Index e -> a Index e -> (Index,Index) -> m ()
  copyTo fromArr toArr (!lb,!ub) = do
    forLoop lb (<=ub) (+1) $ \i -> do
      val_i <- readArray fromArr i
      writeArray toArr i val_i
  