module Data.PrimitiveArray.Index.Int where

import qualified Data.Vector.Fusion.Stream.Monadic as SM

import           Data.PrimitiveArray.Index.Class



instance Index Int where
  newtype LimitType Int = LtInt Int
  linearIndex :: LimitType Int -> Int -> Int
linearIndex LimitType Int
_ Int
k = Int
k
  {-# Inline linearIndex #-}
  size :: LimitType Int -> Int
size (LtInt h) = Int
hInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1
  {-# Inline size #-}
  inBounds :: LimitType Int -> Int -> Bool
inBounds (LtInt h) Int
k = Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
k Bool -> Bool -> Bool
&& Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
h
  {-# Inline inBounds #-}
  zeroBound :: Int
zeroBound = Int
0
  {-# Inline [0] zeroBound #-}
  zeroBound' :: LimitType Int
zeroBound' = Int -> LimitType Int
LtInt Int
0
  {-# Inline [0] zeroBound' #-}
  totalSize :: LimitType Int -> [Integer]
totalSize (LtInt h) = [Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ Int
hInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1]
  {-# Inline [0] totalSize #-}
  fromLinearIndex :: LimitType Int -> Int -> Int
fromLinearIndex LimitType Int
_ = Int -> Int
forall a. a -> a
id
  {-# Inline [0] fromLinearIndex #-}
  showBound :: LimitType Int -> [String]
showBound (LtInt b) = [String
"LtInt " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
b]
  showIndex :: Int -> [String]
showIndex Int
i = [String
"Int " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
i]

deriving instance Show (LimitType Int)

instance IndexStream z => IndexStream (z:.Int) where
  streamUp :: LimitType (z :. Int) -> LimitType (z :. Int) -> Stream m (z :. Int)
streamUp (ls:.. LtInt l) (hs:.. LtInt h) = (z -> m (z, Int))
-> ((z, Int) -> m (Step (z, Int) (z :. Int)))
-> Stream m z
-> Stream m (z :. Int)
forall (m :: * -> *) a s b.
Monad m =>
(a -> m s) -> (s -> m (Step s b)) -> Stream m a -> Stream m b
SM.flatten z -> m (z, Int)
mk (z, Int) -> m (Step (z, Int) (z :. Int))
step (Stream m z -> Stream m (z :. Int))
-> Stream m z -> Stream m (z :. Int)
forall a b. (a -> b) -> a -> b
$ LimitType z -> LimitType z -> Stream m z
forall i (m :: * -> *).
(IndexStream i, Monad m) =>
LimitType i -> LimitType i -> Stream m i
streamUp LimitType z
ls LimitType z
hs
    where mk :: z -> m (z, Int)
mk z
z = (z, Int) -> m (z, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (z
z,Int
l)
          step :: (z, Int) -> m (Step (z, Int) (z :. Int))
step (z
z,Int
k)
            | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
h     = Step (z, Int) (z :. Int) -> m (Step (z, Int) (z :. Int))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (z, Int) (z :. Int) -> m (Step (z, Int) (z :. Int)))
-> Step (z, Int) (z :. Int) -> m (Step (z, Int) (z :. Int))
forall a b. (a -> b) -> a -> b
$ Step (z, Int) (z :. Int)
forall s a. Step s a
SM.Done
            | Bool
otherwise = Step (z, Int) (z :. Int) -> m (Step (z, Int) (z :. Int))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (z, Int) (z :. Int) -> m (Step (z, Int) (z :. Int)))
-> Step (z, Int) (z :. Int) -> m (Step (z, Int) (z :. Int))
forall a b. (a -> b) -> a -> b
$ (z :. Int) -> (z, Int) -> Step (z, Int) (z :. Int)
forall a s. a -> s -> Step s a
SM.Yield (z
zz -> Int -> z :. Int
forall a b. a -> b -> a :. b
:.Int
k) (z
z,Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
          {-# Inline [0] mk   #-}
          {-# Inline [0] step #-}
  {-# Inline streamUp #-}
  streamDown :: LimitType (z :. Int) -> LimitType (z :. Int) -> Stream m (z :. Int)
streamDown (ls:..LtInt l) (hs:..LtInt h) = (z -> m (z, Int))
-> ((z, Int) -> m (Step (z, Int) (z :. Int)))
-> Stream m z
-> Stream m (z :. Int)
forall (m :: * -> *) a s b.
Monad m =>
(a -> m s) -> (s -> m (Step s b)) -> Stream m a -> Stream m b
SM.flatten z -> m (z, Int)
mk (z, Int) -> m (Step (z, Int) (z :. Int))
step (Stream m z -> Stream m (z :. Int))
-> Stream m z -> Stream m (z :. Int)
forall a b. (a -> b) -> a -> b
$ LimitType z -> LimitType z -> Stream m z
forall i (m :: * -> *).
(IndexStream i, Monad m) =>
LimitType i -> LimitType i -> Stream m i
streamDown LimitType z
ls LimitType z
hs
    where mk :: z -> m (z, Int)
mk z
z = (z, Int) -> m (z, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (z
z,Int
h)
          step :: (z, Int) -> m (Step (z, Int) (z :. Int))
step (z
z,Int
k)
            | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l     = Step (z, Int) (z :. Int) -> m (Step (z, Int) (z :. Int))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (z, Int) (z :. Int) -> m (Step (z, Int) (z :. Int)))
-> Step (z, Int) (z :. Int) -> m (Step (z, Int) (z :. Int))
forall a b. (a -> b) -> a -> b
$ Step (z, Int) (z :. Int)
forall s a. Step s a
SM.Done
            | Bool
otherwise = Step (z, Int) (z :. Int) -> m (Step (z, Int) (z :. Int))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (z, Int) (z :. Int) -> m (Step (z, Int) (z :. Int)))
-> Step (z, Int) (z :. Int) -> m (Step (z, Int) (z :. Int))
forall a b. (a -> b) -> a -> b
$ (z :. Int) -> (z, Int) -> Step (z, Int) (z :. Int)
forall a s. a -> s -> Step s a
SM.Yield (z
zz -> Int -> z :. Int
forall a b. a -> b -> a :. b
:.Int
k) (z
z,Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
          {-# Inline [0] mk   #-}
          {-# Inline [0] step #-}
  {-# Inline streamDown #-}

instance IndexStream Int where
  streamUp :: LimitType Int -> LimitType Int -> Stream m Int
streamUp LimitType Int
l LimitType Int
h = ((Z :. Int) -> Int) -> Stream m (Z :. Int) -> Stream m Int
forall (m :: * -> *) a b.
Monad m =>
(a -> b) -> Stream m a -> Stream m b
SM.map (\(Z
Z:.Int
k) -> Int
k) (Stream m (Z :. Int) -> Stream m Int)
-> Stream m (Z :. Int) -> Stream m Int
forall a b. (a -> b) -> a -> b
$ LimitType (Z :. Int) -> LimitType (Z :. Int) -> Stream m (Z :. Int)
forall i (m :: * -> *).
(IndexStream i, Monad m) =>
LimitType i -> LimitType i -> Stream m i
streamUp (LimitType Z
ZZLimitType Z -> LimitType Int -> LimitType (Z :. Int)
forall zs z. LimitType zs -> LimitType z -> LimitType (zs :. z)
:..LimitType Int
l) (LimitType Z
ZZLimitType Z -> LimitType Int -> LimitType (Z :. Int)
forall zs z. LimitType zs -> LimitType z -> LimitType (zs :. z)
:..LimitType Int
h)
  {-# Inline streamUp #-}
  streamDown :: LimitType Int -> LimitType Int -> Stream m Int
streamDown LimitType Int
l LimitType Int
h = ((Z :. Int) -> Int) -> Stream m (Z :. Int) -> Stream m Int
forall (m :: * -> *) a b.
Monad m =>
(a -> b) -> Stream m a -> Stream m b
SM.map (\(Z
Z:.Int
k) -> Int
k) (Stream m (Z :. Int) -> Stream m Int)
-> Stream m (Z :. Int) -> Stream m Int
forall a b. (a -> b) -> a -> b
$ LimitType (Z :. Int) -> LimitType (Z :. Int) -> Stream m (Z :. Int)
forall i (m :: * -> *).
(IndexStream i, Monad m) =>
LimitType i -> LimitType i -> Stream m i
streamDown (LimitType Z
ZZLimitType Z -> LimitType Int -> LimitType (Z :. Int)
forall zs z. LimitType zs -> LimitType z -> LimitType (zs :. z)
:..LimitType Int
l) (LimitType Z
ZZLimitType Z -> LimitType Int -> LimitType (Z :. Int)
forall zs z. LimitType zs -> LimitType z -> LimitType (zs :. z)
:..LimitType Int
h)
  {-# Inline streamDown #-}