{-# LANGUAGE TupleSections #-}
module Bio.Chain
    ( ChainLike (..)
    , Chain
    , chain, fromList
    , (!), (//)
    ) where

import           Control.Lens
import           Data.Array      (Array, Ix, array, listArray, (!), (//))
import qualified Data.Array      as A (assocs, bounds)
import           Data.Array.Base (unsafeAt)
import qualified Data.Vector     as V

type Chain i a = Array i a

-- | Construct new chain from list
--
chain :: Ix i => (i, i) -> [(i, a)] -> Chain i a
chain :: (i, i) -> [(i, a)] -> Chain i a
chain = (i, i) -> [(i, a)] -> Chain i a
forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array

-- | Construct new int-labeled chain from list
--
fromList :: [a] -> Chain Int a
fromList :: [a] -> Chain Int a
fromList [a]
lst = (Int, Int) -> [a] -> Chain Int a
forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (Int
0, [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
lst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [a]
lst

-- | Chain-like sequence, by default it is an array or a list
--
class (Ixed m, Enum (Index m)) => ChainLike m where
    bounds       :: m -> (Index m, Index m)
    assocs       :: m -> [(Index m, IxValue m)]
    modify       :: Index m -> (IxValue m -> IxValue m) -> m -> m
    modifyBefore :: Index m -> (IxValue m -> IxValue m) -> m -> m
    modifyAfter  :: Index m -> (IxValue m -> IxValue m) -> m -> m

    unsafeRead   :: m -> Index m -> IxValue m
    unsafeRead m
ch Index m
i = m
ch m -> Getting (Endo (IxValue m)) m (IxValue m) -> IxValue m
forall s a. HasCallStack => s -> Getting (Endo a) s a -> a
^?! Index m -> Traversal' m (IxValue m)
forall m. Ixed m => Index m -> Traversal' m (IxValue m)
ix Index m
i

instance ChainLike [a] where
    bounds :: [a] -> (Index [a], Index [a])
bounds = (Int
0,) (Int -> (Int, Int)) -> ([a] -> Int) -> [a] -> (Int, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int
forall a. Enum a => a -> a
pred (Int -> Int) -> ([a] -> Int) -> [a] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length

    assocs :: [a] -> [(Index [a], IxValue [a])]
assocs  = [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..]

    modify :: Index [a] -> (IxValue [a] -> IxValue [a]) -> [a] -> [a]
modify       Index [a]
_ IxValue [a] -> IxValue [a]
_ []     = []
    modify       Index [a]
0 IxValue [a] -> IxValue [a]
f (a
x:[a]
xs) = IxValue [a] -> IxValue [a]
f a
IxValue [a]
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs
    modify       Index [a]
i IxValue [a] -> IxValue [a]
f (a
x:[a]
xs) = a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:Index [a] -> (IxValue [a] -> IxValue [a]) -> [a] -> [a]
forall m.
ChainLike m =>
Index m -> (IxValue m -> IxValue m) -> m -> m
modify (Int
Index [a]
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) IxValue [a] -> IxValue [a]
f [a]
xs

    modifyBefore :: Index [a] -> (IxValue [a] -> IxValue [a]) -> [a] -> [a]
modifyBefore Index [a]
i IxValue [a] -> IxValue [a]
f [a]
lst = (a -> a
IxValue [a] -> IxValue [a]
f (a -> a) -> [a] -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
Index [a]
i [a]
lst) [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
Index [a]
i [a]
lst
    modifyAfter :: Index [a] -> (IxValue [a] -> IxValue [a]) -> [a] -> [a]
modifyAfter  Index [a]
i IxValue [a] -> IxValue [a]
f [a]
lst = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take (Int
Index [a]
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
lst [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ (a -> a
IxValue [a] -> IxValue [a]
f (a -> a) -> [a] -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
Index [a]
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
lst)

    unsafeRead :: [a] -> Index [a] -> IxValue [a]
unsafeRead = [a] -> Index [a] -> IxValue [a]
forall a. [a] -> Int -> a
(!!)

instance (Ix i, Enum i) => ChainLike (Array i a) where
    bounds :: Array i a -> (Index (Array i a), Index (Array i a))
bounds = Array i a -> (Index (Array i a), Index (Array i a))
forall i e. Array i e -> (i, i)
A.bounds

    assocs :: Array i a -> [(Index (Array i a), IxValue (Array i a))]
assocs = Array i a -> [(Index (Array i a), IxValue (Array i a))]
forall i e. Ix i => Array i e -> [(i, e)]
A.assocs

    modify :: Index (Array i a)
-> (IxValue (Array i a) -> IxValue (Array i a))
-> Array i a
-> Array i a
modify       Index (Array i a)
i IxValue (Array i a) -> IxValue (Array i a)
f Array i a
ar = Array i a
ar Array i a -> [(i, a)] -> Array i a
forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(i
Index (Array i a)
i, IxValue (Array i a) -> IxValue (Array i a)
f (Array i a
ar Array i a -> i -> a
forall i e. Ix i => Array i e -> i -> e
! i
Index (Array i a)
i))]

    modifyBefore :: Index (Array i a)
-> (IxValue (Array i a) -> IxValue (Array i a))
-> Array i a
-> Array i a
modifyBefore Index (Array i a)
i IxValue (Array i a) -> IxValue (Array i a)
f Array i a
ar = let (i
mi, i
_) = Array i a -> (Index (Array i a), Index (Array i a))
forall m. ChainLike m => m -> (Index m, Index m)
bounds Array i a
ar
                          in Array i a
ar Array i a -> [(i, a)] -> Array i a
forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(i
j, IxValue (Array i a) -> IxValue (Array i a)
f (Array i a
ar Array i a -> i -> a
forall i e. Ix i => Array i e -> i -> e
! i
j)) | i
j <- [i
mi .. i -> i
forall a. Enum a => a -> a
pred i
Index (Array i a)
i]]
    modifyAfter :: Index (Array i a)
-> (IxValue (Array i a) -> IxValue (Array i a))
-> Array i a
-> Array i a
modifyAfter  Index (Array i a)
i IxValue (Array i a) -> IxValue (Array i a)
f Array i a
ar = let (i
_, i
ma) = Array i a -> (Index (Array i a), Index (Array i a))
forall m. ChainLike m => m -> (Index m, Index m)
bounds Array i a
ar
                          in Array i a
ar Array i a -> [(i, a)] -> Array i a
forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
// [(i
j, IxValue (Array i a) -> IxValue (Array i a)
f (Array i a
ar Array i a -> i -> a
forall i e. Ix i => Array i e -> i -> e
! i
j)) | i
j <- [i -> i
forall a. Enum a => a -> a
succ i
Index (Array i a)
i .. i
ma]]

    {-# INLINE unsafeRead #-}
    unsafeRead :: Array i a -> Index (Array i a) -> IxValue (Array i a)
unsafeRead = Array i a -> Index (Array i a) -> IxValue (Array i a)
forall m. UnsafeReadArray m => m -> Index m -> IxValue m
unsafeReadArray

instance ChainLike (V.Vector a) where
    bounds :: Vector a -> (Index (Vector a), Index (Vector a))
bounds Vector a
v = (Index (Vector a)
0, Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

    assocs :: Vector a -> [(Index (Vector a), IxValue (Vector a))]
assocs = [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] ([a] -> [(Int, a)]) -> (Vector a -> [a]) -> Vector a -> [(Int, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> [a]
forall a. Vector a -> [a]
V.toList

    modify :: Index (Vector a)
-> (IxValue (Vector a) -> IxValue (Vector a))
-> Vector a
-> Vector a
modify Index (Vector a)
i IxValue (Vector a) -> IxValue (Vector a)
f Vector a
ar = Vector a
ar Vector a -> [(Int, a)] -> Vector a
forall a. Vector a -> [(Int, a)] -> Vector a
V.// [(Int
Index (Vector a)
i, IxValue (Vector a) -> IxValue (Vector a)
f (Vector a
ar Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.! Int
Index (Vector a)
i))]

    modifyBefore :: Index (Vector a)
-> (IxValue (Vector a) -> IxValue (Vector a))
-> Vector a
-> Vector a
modifyBefore Index (Vector a)
i IxValue (Vector a) -> IxValue (Vector a)
f Vector a
ar = (a -> a) -> Vector a -> Vector a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
IxValue (Vector a) -> IxValue (Vector a)
f Vector a
before Vector a -> Vector a -> Vector a
forall a. Semigroup a => a -> a -> a
<> Vector a
after
      where
        (Vector a
before, Vector a
after) = Int -> Vector a -> (Vector a, Vector a)
forall a. Int -> Vector a -> (Vector a, Vector a)
V.splitAt Int
Index (Vector a)
i Vector a
ar

    modifyAfter :: Index (Vector a)
-> (IxValue (Vector a) -> IxValue (Vector a))
-> Vector a
-> Vector a
modifyAfter Index (Vector a)
i IxValue (Vector a) -> IxValue (Vector a)
f Vector a
ar = Vector a
before Vector a -> Vector a -> Vector a
forall a. Semigroup a => a -> a -> a
<> (a -> a) -> Vector a -> Vector a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
IxValue (Vector a) -> IxValue (Vector a)
f Vector a
after
      where
        (Vector a
before, Vector a
after) = Int -> Vector a -> (Vector a, Vector a)
forall a. Int -> Vector a -> (Vector a, Vector a)
V.splitAt (Int
Index (Vector a)
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Vector a
ar

    {-# INLINE unsafeRead #-}
    unsafeRead :: Vector a -> Index (Vector a) -> IxValue (Vector a)
unsafeRead = Vector a -> Index (Vector a) -> IxValue (Vector a)
forall a. Vector a -> Int -> a
V.unsafeIndex 


class (Ixed m) => UnsafeReadArray m where
    unsafeReadArray :: m -> Index m -> IxValue m

instance (Ix i, Enum i) => UnsafeReadArray (Array i a) where
    {-# INLINE unsafeReadArray #-}
    unsafeReadArray :: Array i a -> Index (Array i a) -> IxValue (Array i a)
unsafeReadArray = (!)

instance {-# OVERLAPPING #-} UnsafeReadArray (Array Int a) where
    {-# INLINE unsafeReadArray #-}
    unsafeReadArray :: Array Int a -> Index (Array Int a) -> IxValue (Array Int a)
unsafeReadArray = Array Int a -> Index (Array Int a) -> IxValue (Array Int a)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
unsafeAt