{-# LANGUAGE TypeFamilies #-}
module Data.Array.Comfort.Shape.Tuple where

import qualified Data.Array.Comfort.Shape as Shape
import Data.Complex (Complex((:+)))

import qualified Control.Monad.Trans.State as MS
import qualified Control.Applicative.HT as App
import Control.Applicative ((<$>))


get :: MS.State [a] a
get :: forall a. State [a] a
get =
   ([a] -> (a, [a])) -> StateT [a] Identity a
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
MS.state (([a] -> (a, [a])) -> StateT [a] Identity a)
-> ([a] -> (a, [a])) -> StateT [a] Identity a
forall a b. (a -> b) -> a -> b
$ \[a]
at ->
      case [a]
at of
         a
a:[a]
as -> (a
a,[a]
as)
         [] -> [Char] -> (a, [a])
forall a. HasCallStack => [Char] -> a
error [Char]
"Shape.Tuple.get: no element left"

cons ::
   (Shape.ElementTuple shape) =>
   shape -> MS.State [a] (Shape.DataTuple shape a)
cons :: forall shape a.
ElementTuple shape =>
shape -> State [a] (DataTuple shape a)
cons = (Element -> StateT [a] Identity a)
-> shape -> StateT [a] Identity (DataTuple shape a)
forall tuple (f :: * -> *) a.
(ElementTuple tuple, Applicative f) =>
(Element -> f a) -> tuple -> f (DataTuple tuple a)
forall (f :: * -> *) a.
Applicative f =>
(Element -> f a) -> shape -> f (DataTuple shape a)
Shape.indexTupleA (StateT [a] Identity a -> Element -> StateT [a] Identity a
forall a b. a -> b -> a
const StateT [a] Identity a
forall a. State [a] a
get)


next :: MS.State Shape.Element Shape.Element
next :: State Element Element
next = do
   Element
ix <- State Element Element
forall (m :: * -> *) s. Monad m => StateT s m s
MS.get
   (Element -> Element) -> StateT Element Identity ()
forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
MS.modify (\(Shape.Element Int
k) -> Int -> Element
Shape.Element (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))
   Element -> State Element Element
forall a. a -> StateT Element Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return Element
ix


class (Shape.ElementTuple shape) => NestedTuple shape where
   decons :: Shape.DataTuple shape a -> MS.State Shape.Element (shape, [a])

instance NestedTuple () where
   decons :: forall a. DataTuple () a -> State Element ((), [a])
decons () = ((), [a]) -> StateT Element Identity ((), [a])
forall a. a -> StateT Element Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return ((),[])

instance NestedTuple Shape.Element where
   decons :: forall a. DataTuple Element a -> State Element (Element, [a])
decons DataTuple Element a
a = (Element -> [a] -> (Element, [a]))
-> [a] -> Element -> (Element, [a])
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) [a
DataTuple Element a
a] (Element -> (Element, [a]))
-> State Element Element -> StateT Element Identity (Element, [a])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> State Element Element
next

instance (NestedTuple a, NestedTuple b) => NestedTuple (a,b) where
   decons :: forall a. DataTuple (a, b) a -> State Element ((a, b), [a])
decons (DataTuple a a
a,DataTuple b a
b) =
      ((a, [a]) -> (b, [a]) -> ((a, b), [a]))
-> StateT Element Identity (a, [a])
-> StateT Element Identity (b, [a])
-> StateT Element Identity ((a, b), [a])
forall (m :: * -> *) a b r.
Applicative m =>
(a -> b -> r) -> m a -> m b -> m r
App.lift2 (\(a
ta,[a]
as) (b
tb,[a]
bs) -> ((a
ta,b
tb), [a]
as[a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++[a]
bs)) (DataTuple a a -> StateT Element Identity (a, [a])
forall a. DataTuple a a -> State Element (a, [a])
forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple a a
a) (DataTuple b a -> StateT Element Identity (b, [a])
forall a. DataTuple b a -> State Element (b, [a])
forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple b a
b)

instance
   (NestedTuple a, NestedTuple b, NestedTuple c) =>
      NestedTuple (a,b,c) where
   decons :: forall a. DataTuple (a, b, c) a -> State Element ((a, b, c), [a])
decons (DataTuple a a
a,DataTuple b a
b,DataTuple c a
c) =
      ((a, [a]) -> (b, [a]) -> (c, [a]) -> ((a, b, c), [a]))
-> StateT Element Identity (a, [a])
-> StateT Element Identity (b, [a])
-> StateT Element Identity (c, [a])
-> StateT Element Identity ((a, b, c), [a])
forall (m :: * -> *) a b c r.
Applicative m =>
(a -> b -> c -> r) -> m a -> m b -> m c -> m r
App.lift3
         (\(a
ta,[a]
as) (b
tb,[a]
bs) (c
tc,[a]
cs) -> ((a
ta,b
tb,c
tc), [a]
as[a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++[a]
bs[a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++[a]
cs))
         (DataTuple a a -> StateT Element Identity (a, [a])
forall a. DataTuple a a -> State Element (a, [a])
forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple a a
a) (DataTuple b a -> StateT Element Identity (b, [a])
forall a. DataTuple b a -> State Element (b, [a])
forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple b a
b) (DataTuple c a -> StateT Element Identity (c, [a])
forall a. DataTuple c a -> State Element (c, [a])
forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple c a
c)

instance
   (NestedTuple a, NestedTuple b, NestedTuple c, NestedTuple d) =>
      NestedTuple (a,b,c,d) where
   decons :: forall a.
DataTuple (a, b, c, d) a -> State Element ((a, b, c, d), [a])
decons (DataTuple a a
a,DataTuple b a
b,DataTuple c a
c,DataTuple d a
d) =
      ((a, [a])
 -> (b, [a]) -> (c, [a]) -> (d, [a]) -> ((a, b, c, d), [a]))
-> StateT Element Identity (a, [a])
-> StateT Element Identity (b, [a])
-> StateT Element Identity (c, [a])
-> StateT Element Identity (d, [a])
-> StateT Element Identity ((a, b, c, d), [a])
forall (m :: * -> *) a b c d r.
Applicative m =>
(a -> b -> c -> d -> r) -> m a -> m b -> m c -> m d -> m r
App.lift4
         (\(a
ta,[a]
as) (b
tb,[a]
bs) (c
tc,[a]
cs) (d
td,[a]
ds) -> ((a
ta,b
tb,c
tc,d
td), [a]
as[a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++[a]
bs[a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++[a]
cs[a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++[a]
ds))
         (DataTuple a a -> StateT Element Identity (a, [a])
forall a. DataTuple a a -> State Element (a, [a])
forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple a a
a) (DataTuple b a -> StateT Element Identity (b, [a])
forall a. DataTuple b a -> State Element (b, [a])
forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple b a
b) (DataTuple c a -> StateT Element Identity (c, [a])
forall a. DataTuple c a -> State Element (c, [a])
forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple c a
c) (DataTuple d a -> StateT Element Identity (d, [a])
forall a. DataTuple d a -> State Element (d, [a])
forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple d a
d)

instance (NestedTuple a) => NestedTuple (Complex a) where
   decons :: forall a. DataTuple (Complex a) a -> State Element (Complex a, [a])
decons (DataTuple a a
a:+DataTuple a a
b) =
      ((a, [a]) -> (a, [a]) -> (Complex a, [a]))
-> StateT Element Identity (a, [a])
-> StateT Element Identity (a, [a])
-> StateT Element Identity (Complex a, [a])
forall (m :: * -> *) a b r.
Applicative m =>
(a -> b -> r) -> m a -> m b -> m r
App.lift2 (\(a
ta,[a]
as) (a
tb,[a]
bs) -> ((a
taa -> a -> Complex a
forall a. a -> a -> Complex a
:+a
tb), [a]
as[a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++[a]
bs)) (DataTuple a a -> StateT Element Identity (a, [a])
forall a. DataTuple a a -> State Element (a, [a])
forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple a a
a) (DataTuple a a -> StateT Element Identity (a, [a])
forall a. DataTuple a a -> State Element (a, [a])
forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple a a
b)