{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Data.Stack.Circular
(
CStack (..),
empty,
toVector,
toVectorN,
fromVector,
get,
pop,
push,
unsafePush,
isFull,
foldl1',
sum,
mean,
product,
)
where
import Control.Monad.ST
import Data.Aeson
import Data.Aeson.Types
import qualified Data.Vector.Generic as V
import Data.Vector.Generic (Vector)
import qualified Data.Vector.Generic.Mutable as M
import Prelude hiding (product, sum)
data CStack v a = CStack
{ stack :: v a,
index :: !Int,
curSize :: !Int
}
instance (Eq (v a), Vector v a) => Eq (CStack v a) where
(CStack v1 i1 m1) == (CStack v2 i2 m2) = (v1 == v2) && (i1 == i2) && (m1 == m2)
instance (Show (v a), Vector v a) => Show (CStack v a) where
show c@(CStack _ i m) = "CStack {" ++ show (toVector c) ++ ", " ++ show i ++ ", " ++ show m ++ "}"
instance (ToJSON a, ToJSON (v a), Vector v a) => ToJSON (CStack v a) where
toJSON c = object ["stack" .= toVector c, "maxSize" .= n]
where
n = V.length $ stack c
toEncoding c = pairs ("stack" .= toVector c <> "maxSize" .= n)
where
n = V.length $ stack c
instance (FromJSON a, FromJSON (v a), Vector v a) => FromJSON (CStack v a) where
parseJSON = withObject "CStack" fromObject
fromObject :: forall v a. (FromJSON (v a), Vector v a) => Object -> Parser (CStack v a)
fromObject o = do
v <- o .: "stack" :: Parser (v a)
n <- o .: "maxSize" :: Parser Int
let c = empty n
pure $ V.foldr' unsafePush c v
startIndex :: Int -> Int -> Int -> Int
startIndex i m n
| m == 0 = error "startIndex: empty stack"
| m <= i + 1 = i + 1 - m
| otherwise = i + 1 - m + n
empty :: Vector v a => Int -> CStack v a
empty n
| n <= 0 = error "empty: maximum size must be 1 or larger"
| otherwise = CStack (V.create $ M.unsafeNew n) 0 0
toVector :: Vector v a => CStack v a -> v a
toVector (CStack v i m)
| m == 0 = V.empty
| i' + m <= n = V.unsafeSlice i' m v
| otherwise = V.unsafeDrop i' v V.++ V.unsafeTake (i + 1) v
where
n = V.length v
i' = startIndex i m n
toVectorN :: Vector v a => Int -> CStack v a -> v a
toVectorN k (CStack v i m)
| k < 0 = error "toVectorN: negative n"
| k > m = error "toVectorN: stack too small"
| k == 0 = V.empty
| i' + k <= n = V.unsafeSlice i' k v
| otherwise = V.unsafeDrop i' v V.++ V.unsafeTake (i + 1) v
where
n = V.length v
i' = startIndex i k n
fromVector :: Vector v a => v a -> CStack v a
fromVector v
| V.null v = error "fromVector: empty vector"
| otherwise = CStack v (n - 1) n
where
n = V.length v
get :: Vector v a => CStack v a -> a
get (CStack v i _) = V.unsafeIndex v i
previous :: Vector v a => CStack v a -> CStack v a
previous (CStack v i m)
| m == 0 = error "previous: empty stack"
| i == 0 = CStack v (n - 1) (m - 1)
| otherwise = CStack v (i - 1) (m - 1)
where
n = V.length v
pop :: Vector v a => CStack v a -> (a, CStack v a)
pop c = (get c, previous c)
set :: Vector v a => Int -> a -> v a -> v a
set i x = V.modify (\v -> M.write v i x)
{-# INLINE set #-}
put :: Vector v a => a -> CStack v a -> CStack v a
put x (CStack v i m) = CStack (set i x v) i m
next :: Vector v a => CStack v a -> CStack v a
next (CStack v i m)
| i == (n - 1) = CStack v 0 (min (m + 1) n)
| otherwise = CStack v (i + 1) (min (m + 1) n)
where
n = V.length v
push :: Vector v a => a -> CStack v a -> CStack v a
push x c = put x $ next c
unsafeSet :: Vector v a => Int -> a -> v a -> v a
unsafeSet i x v = runST $ do
mv <- V.unsafeThaw v
M.unsafeWrite mv i x
V.unsafeFreeze mv
unsafePut :: Vector v a => a -> CStack v a -> CStack v a
unsafePut x (CStack v i m) = CStack (unsafeSet i x v) i m
unsafePush :: Vector v a => a -> CStack v a -> CStack v a
unsafePush x c = unsafePut x $ next c
isFull :: Vector v a => CStack v a -> Bool
isFull (CStack v _ m) = V.length v == m
foldl1' :: Vector v a => (a -> a -> a) -> CStack v a -> a
foldl1' f (CStack v i m)
| m == n = V.foldl1' f v
| i' + m <= n = V.foldl1' f $ V.unsafeSlice i' m v
| otherwise = f (V.foldl1' f (V.unsafeDrop i' v)) (V.foldl1' f (V.unsafeTake (i + 1) v))
where
n = V.length v
i' = startIndex i m n
sum :: (Num a, Vector v a) => CStack v a -> a
sum (CStack v i m)
| m == n = V.sum v
| i' + m <= n = V.sum $ V.unsafeSlice i' m v
| otherwise = V.sum (V.unsafeDrop i' v) + V.sum (V.unsafeTake (i + 1) v)
where
n = V.length v
i' = startIndex i m n
mean :: (Real a, Vector v a, Fractional b) => CStack v a -> b
mean c = realToFrac (sum c) / fromIntegral (curSize c)
product :: (Num a, Vector v a) => CStack v a -> a
product (CStack v i m)
| m == n = V.product v
| i' + m <= n = V.product $ V.unsafeSlice i' m v
| otherwise = V.product (V.unsafeDrop i' v) * V.product (V.unsafeTake (i + 1) v)
where
n = V.length v
i' = startIndex i m n