{-# LANGUAGE UndecidableInstances #-}

-- | Typed binary tree structures

module Data.TypedStruct where



import Control.Monad.Identity



--------------------------------------------------------------------------------
-- * Representation
--------------------------------------------------------------------------------

-- | Typed binary tree structure
--
-- The predicate @pred@ is assumed to rule out pairs. Functions like
-- 'extractSingle' and 'zipStruct' rely on this assumption.
data Struct pred con a
  where
    Single :: pred a => con a -> Struct pred con a
    Two    :: Struct pred con a -> Struct pred con b -> Struct pred con (a,b)

-- It would have been nice to add a constraint `IsPair a ~ False` to `Single`,
-- so that one wouldn't have to rely on @pred@ to rule out pairs. However,
-- attempting to do so lead to very strange problems in the rest of the Feldspar
-- implementation, so in the end I abandoned this extra safety.
--
-- The problems were strange enough that it seems likely they may be due to a
-- bug in GHC (7.10.2). So it might be worthwhile to try this again in a later
-- version.
--
-- Note however, that `IsPair a ~ False` on `Single` is not enough to please the
-- completeness checker for functions like `extractSingle` in GHC 7.10. Maybe
-- the new completeness checker in GHC 8 will be satisfied?

-- | Create a 'Struct' from a 'Struct' of any container @c@ and a structured
-- value @a@
--
-- For example:
--
-- @
-- `toStruct` (`Two` (`Single` `Proxy`) (`Single` `Proxy`)) (False,'a')
--   ==
-- Two (Single (Identity False)) (Single (Identity 'a'))
-- @
toStruct :: Struct p c a -> a -> Struct p Identity a
toStruct :: Struct p c a -> a -> Struct p Identity a
toStruct Struct p c a
rep = Struct p c a -> Identity a -> Struct p Identity a
forall (p :: * -> Constraint) (c :: * -> *) a.
Struct p c a -> Identity a -> Struct p Identity a
go Struct p c a
rep (Identity a -> Struct p Identity a)
-> (a -> Identity a) -> a -> Struct p Identity a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Identity a
forall a. a -> Identity a
Identity
  where
    go :: Struct p c a -> Identity a -> Struct p Identity a
    go :: Struct p c a -> Identity a -> Struct p Identity a
go (Single c a
_) Identity a
i = Identity a -> Struct p Identity a
forall (pred :: * -> Constraint) a (con :: * -> *).
pred a =>
con a -> Struct pred con a
Single Identity a
i
    go (Two Struct p c a
ra Struct p c b
rb) (Identity (a,b)) =
        Struct p Identity a
-> Struct p Identity b -> Struct p Identity (a, b)
forall (pred :: * -> Constraint) (con :: * -> *) a b.
Struct pred con a -> Struct pred con b -> Struct pred con (a, b)
Two (Struct p c a -> Identity a -> Struct p Identity a
forall (p :: * -> Constraint) (c :: * -> *) a.
Struct p c a -> Identity a -> Struct p Identity a
go Struct p c a
ra (a -> Identity a
forall a. a -> Identity a
Identity a
a)) (Struct p c b -> Identity b -> Struct p Identity b
forall (p :: * -> Constraint) (c :: * -> *) a.
Struct p c a -> Identity a -> Struct p Identity a
go Struct p c b
rb (b -> Identity b
forall a. a -> Identity a
Identity b
b))



--------------------------------------------------------------------------------
-- * Operations
--------------------------------------------------------------------------------

-- | Extract the value of a 'Single'
extractSingle :: pred a => Struct pred c a -> c a
extractSingle :: Struct pred c a -> c a
extractSingle (Single c a
a) = c a
a

-- | Map over a 'Struct'
mapStruct :: forall pred c1 c2 b
    .  (forall a . pred a => c1 a -> c2 a)
    -> Struct pred c1 b
    -> Struct pred c2 b
mapStruct :: (forall a. pred a => c1 a -> c2 a)
-> Struct pred c1 b -> Struct pred c2 b
mapStruct forall a. pred a => c1 a -> c2 a
f = Struct pred c1 b -> Struct pred c2 b
forall a. Struct pred c1 a -> Struct pred c2 a
go
  where
    go :: Struct pred c1 a -> Struct pred c2 a
    go :: Struct pred c1 a -> Struct pred c2 a
go (Single c1 a
a) = c2 a -> Struct pred c2 a
forall (pred :: * -> Constraint) a (con :: * -> *).
pred a =>
con a -> Struct pred con a
Single (c1 a -> c2 a
forall a. pred a => c1 a -> c2 a
f c1 a
a)
    go (Two Struct pred c1 a
a Struct pred c1 b
b)  = Struct pred c2 a -> Struct pred c2 b -> Struct pred c2 (a, b)
forall (pred :: * -> Constraint) (con :: * -> *) a b.
Struct pred con a -> Struct pred con b -> Struct pred con (a, b)
Two (Struct pred c1 a -> Struct pred c2 a
forall a. Struct pred c1 a -> Struct pred c2 a
go Struct pred c1 a
a) (Struct pred c1 b -> Struct pred c2 b
forall a. Struct pred c1 a -> Struct pred c2 a
go Struct pred c1 b
b)

-- | Monadic map over a 'Struct'
mapStructA :: forall m pred c1 c2 b . Applicative m
    => (forall a . pred a => c1 a -> m (c2 a))
    -> Struct pred c1 b -> m (Struct pred c2 b)
mapStructA :: (forall a. pred a => c1 a -> m (c2 a))
-> Struct pred c1 b -> m (Struct pred c2 b)
mapStructA forall a. pred a => c1 a -> m (c2 a)
f = Struct pred c1 b -> m (Struct pred c2 b)
forall a. Struct pred c1 a -> m (Struct pred c2 a)
go
  where
    go :: Struct pred c1 a -> m (Struct pred c2 a)
    go :: Struct pred c1 a -> m (Struct pred c2 a)
go (Single c1 a
a) = c2 a -> Struct pred c2 a
forall (pred :: * -> Constraint) a (con :: * -> *).
pred a =>
con a -> Struct pred con a
Single (c2 a -> Struct pred c2 a) -> m (c2 a) -> m (Struct pred c2 a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (c1 a -> m (c2 a)
forall a. pred a => c1 a -> m (c2 a)
f c1 a
a)
    go (Two Struct pred c1 a
a Struct pred c1 b
b)  = Struct pred c2 a -> Struct pred c2 b -> Struct pred c2 (a, b)
forall (pred :: * -> Constraint) (con :: * -> *) a b.
Struct pred con a -> Struct pred con b -> Struct pred con (a, b)
Two (Struct pred c2 a -> Struct pred c2 b -> Struct pred c2 (a, b))
-> m (Struct pred c2 a)
-> m (Struct pred c2 b -> Struct pred c2 (a, b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Struct pred c1 a -> m (Struct pred c2 a)
forall a. Struct pred c1 a -> m (Struct pred c2 a)
go Struct pred c1 a
a m (Struct pred c2 b -> Struct pred c2 (a, b))
-> m (Struct pred c2 b) -> m (Struct pred c2 (a, b))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Struct pred c1 b -> m (Struct pred c2 b)
forall a. Struct pred c1 a -> m (Struct pred c2 a)
go Struct pred c1 b
b

-- | Map over a 'Struct'
mapStructA_ :: forall m pred cont b . Applicative m =>
    (forall a . pred a => cont a -> m ()) -> Struct pred cont b -> m ()
mapStructA_ :: (forall a. pred a => cont a -> m ()) -> Struct pred cont b -> m ()
mapStructA_ forall a. pred a => cont a -> m ()
f = Struct pred cont b -> m ()
forall a. Struct pred cont a -> m ()
go
  where
    go :: Struct pred cont a -> m ()
    go :: Struct pred cont a -> m ()
go (Single cont a
a) = cont a -> m ()
forall a. pred a => cont a -> m ()
f cont a
a
    go (Two Struct pred cont a
a Struct pred cont b
b)  = Struct pred cont a -> m ()
forall a. Struct pred cont a -> m ()
go Struct pred cont a
a m () -> m () -> m ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Struct pred cont b -> m ()
forall a. Struct pred cont a -> m ()
go Struct pred cont b
b

-- mapStructM_ :: forall m pred cont b . Monad m =>
--     (forall a . pred a => cont a -> m ()) -> Struct pred cont b -> m ()
-- mapStructM_ f = sequence_ . listStruct f
  -- This doesn't work for some reason, only if `pred` is constrained to a
  -- concrete type. (On the other hand, using `listStruct` is probably less
  -- efficient due to the use of `++`.)

-- | Fold a 'Struct' to a list
listStruct :: forall pred cont b c .
    (forall y . pred y => cont y -> c) -> Struct pred cont b -> [c]
listStruct :: (forall y. pred y => cont y -> c) -> Struct pred cont b -> [c]
listStruct forall y. pred y => cont y -> c
f = Struct pred cont b -> [c]
forall a. Struct pred cont a -> [c]
go
  where
    go :: Struct pred cont a -> [c]
    go :: Struct pred cont a -> [c]
go (Single cont a
a) = [cont a -> c
forall y. pred y => cont y -> c
f cont a
a]
    go (Two Struct pred cont a
a Struct pred cont b
b)  = Struct pred cont a -> [c]
forall a. Struct pred cont a -> [c]
go Struct pred cont a
a [c] -> [c] -> [c]
forall a. [a] -> [a] -> [a]
++ Struct pred cont b -> [c]
forall a. Struct pred cont a -> [c]
go Struct pred cont b
b

-- | Zip two 'Struct's
zipStruct :: forall pred c1 c2 c3 b
    . (forall a . pred a => c1 a -> c2 a -> c3 a)
    -> Struct pred c1 b
    -> Struct pred c2 b
    -> Struct pred c3 b
zipStruct :: (forall a. pred a => c1 a -> c2 a -> c3 a)
-> Struct pred c1 b -> Struct pred c2 b -> Struct pred c3 b
zipStruct forall a. pred a => c1 a -> c2 a -> c3 a
f = Struct pred c1 b -> Struct pred c2 b -> Struct pred c3 b
forall a. Struct pred c1 a -> Struct pred c2 a -> Struct pred c3 a
go
  where
    go :: Struct pred c1 a -> Struct pred c2 a -> Struct pred c3 a
    go :: Struct pred c1 a -> Struct pred c2 a -> Struct pred c3 a
go (Single c1 a
a) (Single c2 a
b) = c3 a -> Struct pred c3 a
forall (pred :: * -> Constraint) a (con :: * -> *).
pred a =>
con a -> Struct pred con a
Single (c1 a -> c2 a -> c3 a
forall a. pred a => c1 a -> c2 a -> c3 a
f c1 a
a c2 a
b)
    go (Two Struct pred c1 a
a Struct pred c1 b
b) (Two Struct pred c2 a
c Struct pred c2 b
d)   = Struct pred c3 a -> Struct pred c3 b -> Struct pred c3 (a, b)
forall (pred :: * -> Constraint) (con :: * -> *) a b.
Struct pred con a -> Struct pred con b -> Struct pred con (a, b)
Two (Struct pred c1 a -> Struct pred c2 a -> Struct pred c3 a
forall a. Struct pred c1 a -> Struct pred c2 a -> Struct pred c3 a
go Struct pred c1 a
a Struct pred c2 a
Struct pred c2 a
c) (Struct pred c1 b -> Struct pred c2 b -> Struct pred c3 b
forall a. Struct pred c1 a -> Struct pred c2 a -> Struct pred c3 a
go Struct pred c1 b
b Struct pred c2 b
Struct pred c2 b
d)

-- | Zip two 'Struct's to a list
zipListStruct :: forall pred c1 c2 b r
    . (forall a . pred a => c1 a -> c2 a -> r)
    -> Struct pred c1 b
    -> Struct pred c2 b
    -> [r]
zipListStruct :: (forall a. pred a => c1 a -> c2 a -> r)
-> Struct pred c1 b -> Struct pred c2 b -> [r]
zipListStruct forall a. pred a => c1 a -> c2 a -> r
f = Struct pred c1 b -> Struct pred c2 b -> [r]
forall a. Struct pred c1 a -> Struct pred c2 a -> [r]
go
  where
    go :: Struct pred c1 a -> Struct pred c2 a -> [r]
    go :: Struct pred c1 a -> Struct pred c2 a -> [r]
go (Single c1 a
a) (Single c2 a
b) = [c1 a -> c2 a -> r
forall a. pred a => c1 a -> c2 a -> r
f c1 a
a c2 a
b]
    go (Two Struct pred c1 a
a Struct pred c1 b
b) (Two Struct pred c2 a
c Struct pred c2 b
d)   = Struct pred c1 a -> Struct pred c2 a -> [r]
forall a. Struct pred c1 a -> Struct pred c2 a -> [r]
go Struct pred c1 a
a Struct pred c2 a
Struct pred c2 a
c [r] -> [r] -> [r]
forall a. [a] -> [a] -> [a]
++ Struct pred c1 b -> Struct pred c2 b -> [r]
forall a. Struct pred c1 a -> Struct pred c2 a -> [r]
go Struct pred c1 b
b Struct pred c2 b
Struct pred c2 b
d

-- | Compare two 'Struct's using a function that compares the 'Single' elements.
-- If the structures don't match, 'False' is returned.
compareStruct :: forall pred c1 c2 c d
    . (forall a b . (pred a, pred b) => c1 a -> c2 b -> Bool)
    -> Struct pred c1 c
    -> Struct pred c2 d
    -> Bool
compareStruct :: (forall a b. (pred a, pred b) => c1 a -> c2 b -> Bool)
-> Struct pred c1 c -> Struct pred c2 d -> Bool
compareStruct forall a b. (pred a, pred b) => c1 a -> c2 b -> Bool
f = Struct pred c1 c -> Struct pred c2 d -> Bool
forall a b. Struct pred c1 a -> Struct pred c2 b -> Bool
go
  where
    go :: Struct pred c1 a -> Struct pred c2 b -> Bool
    go :: Struct pred c1 a -> Struct pred c2 b -> Bool
go (Single c1 a
a) (Single c2 b
b) = c1 a -> c2 b -> Bool
forall a b. (pred a, pred b) => c1 a -> c2 b -> Bool
f c1 a
a c2 b
b
    go (Two Struct pred c1 a
a Struct pred c1 b
b) (Two Struct pred c2 a
c Struct pred c2 b
d)   = Struct pred c1 a -> Struct pred c2 a -> Bool
forall a b. Struct pred c1 a -> Struct pred c2 b -> Bool
go Struct pred c1 a
a Struct pred c2 a
c Bool -> Bool -> Bool
&& Struct pred c1 b -> Struct pred c2 b -> Bool
forall a b. Struct pred c1 a -> Struct pred c2 b -> Bool
go Struct pred c1 b
b Struct pred c2 b
d

-- | Lift a function operating on containers @con@ to a function operating on
-- 'Struct's.
liftStruct :: (pred a, pred b) =>
    (con a -> con b) -> Struct pred con a -> Struct pred con b
liftStruct :: (con a -> con b) -> Struct pred con a -> Struct pred con b
liftStruct con a -> con b
f (Single con a
a) = con b -> Struct pred con b
forall (pred :: * -> Constraint) a (con :: * -> *).
pred a =>
con a -> Struct pred con a
Single (con a -> con b
f con a
a)

-- | Lift a function operating on containers @con@ to a function operating on
-- 'Struct's.
liftStruct2 :: (pred a, pred b, pred c)
    => (con a -> con b -> con c)
    -> Struct pred con a -> Struct pred con b -> Struct pred con c
liftStruct2 :: (con a -> con b -> con c)
-> Struct pred con a -> Struct pred con b -> Struct pred con c
liftStruct2 con a -> con b -> con c
f (Single con a
a) (Single con b
b) = con c -> Struct pred con c
forall (pred :: * -> Constraint) a (con :: * -> *).
pred a =>
con a -> Struct pred con a
Single (con a -> con b -> con c
f con a
a con b
b)