{-# 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 =
   forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
MS.state forall a b. (a -> b) -> a -> b
$ \[a]
at ->
      case [a]
at of
         a
a:[a]
as -> (a
a,[a]
as)
         [] -> 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 = forall tuple (f :: * -> *) a.
(ElementTuple tuple, Applicative f) =>
(Element -> f a) -> tuple -> f (DataTuple tuple a)
Shape.indexTupleA (forall a b. a -> b -> a
const forall a. State [a] a
get)


next :: MS.State Shape.Element Shape.Element
next :: State Element Element
next = do
   Element
ix <- forall (m :: * -> *) s. Monad m => StateT s m s
MS.get
   forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
MS.modify (\(Shape.Element Int
k) -> Int -> Element
Shape.Element (Int
kforall a. Num a => a -> a -> a
+Int
1))
   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 () = 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 = forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) [DataTuple Element a
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) =
      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]
asforall a. [a] -> [a] -> [a]
++[a]
bs)) (forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple a a
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) =
      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]
asforall a. [a] -> [a] -> [a]
++[a]
bsforall a. [a] -> [a] -> [a]
++[a]
cs))
         (forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple a a
a) (forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple b a
b) (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) =
      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]
asforall a. [a] -> [a] -> [a]
++[a]
bsforall a. [a] -> [a] -> [a]
++[a]
csforall a. [a] -> [a] -> [a]
++[a]
ds))
         (forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple a a
a) (forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple b a
b) (forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple c a
c) (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) =
      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
taforall a. a -> a -> Complex a
:+a
tb), [a]
asforall a. [a] -> [a] -> [a]
++[a]
bs)) (forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple a a
a) (forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
decons DataTuple a a
b)