{-# LANGUAGE ConstraintKinds           #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE Rank2Types                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TypeOperators             #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  Data.Comp.Multi.Generic
-- Copyright   :  (c) 2011 Patrick Bahr
-- License     :  BSD3
-- Maintainer  :  Patrick Bahr <paba@diku.dk>
-- Stability   :  experimental
-- Portability :  non-portable (GHC Extensions)
--
-- This module defines type generic functions and recursive schemes
-- along the lines of the Uniplate library. All definitions are
-- generalised versions of those in "Data.Comp.Generic".
--
--------------------------------------------------------------------------------

module Data.Comp.Multi.Generic where

import Control.Monad
import Data.Comp.Multi.HFoldable
import Data.Comp.Multi.HFunctor
import Data.Comp.Multi.HTraversable
import Data.Comp.Multi.Sum
import Data.Comp.Multi.Term
import GHC.Exts
import Prelude

import Data.Maybe

-- | This function returns a list of all subterms of the given
-- term. This function is similar to Uniplate's @universe@ function.
subterms :: forall f  . HFoldable f => Term f  :=> [E (Term f)]
subterms :: Term f :=> [E (Term f)]
subterms Term f i
t = (forall b. (E (Term f) -> b -> b) -> b -> b) -> [E (Term f)]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (Term f i -> (E (Term f) -> b -> b) -> b -> b
forall i b. Term f i -> (E (Term f) -> b -> b) -> b -> b
f Term f i
t)
    where f :: forall i b. Term f i -> (E (Term f) -> b -> b) -> b -> b
          f :: Term f i -> (E (Term f) -> b -> b) -> b -> b
f Term f i
t E (Term f) -> b -> b
cons b
nil = Term f i -> E (Term f)
forall (f :: * -> *) i. f i -> E f
E Term f i
t E (Term f) -> b -> b
`cons` (b -> Term f :=> b) -> b -> f (Term f) i -> b
forall (h :: (* -> *) -> * -> *) b (a :: * -> *).
HFoldable h =>
(b -> a :=> b) -> b -> h a :=> b
hfoldl (\b
u Term f i
s -> Term f i -> (E (Term f) -> b -> b) -> b -> b
forall i b. Term f i -> (E (Term f) -> b -> b) -> b -> b
f Term f i
s E (Term f) -> b -> b
cons b
u) b
nil (Term f i -> f (Term f) i
forall (f :: (* -> *) -> * -> *) t. Term f t -> f (Term f) t
unTerm Term f i
t)

-- | This function returns a list of all subterms of the given term
-- that are constructed from a particular functor.
subterms' :: forall f g . (HFoldable f, g :<: f) => Term f :=> [E (g (Term f))]
subterms' :: Term f :=> [E (g (Term f))]
subterms' (Term f (Term f) i
t) = (forall b. (E (g (Term f)) -> b -> b) -> b -> b)
-> [E (g (Term f))]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (f (Term f) i -> (E (g (Term f)) -> b -> b) -> b -> b
forall i b. f (Term f) i -> (E (g (Term f)) -> b -> b) -> b -> b
f f (Term f) i
t)
    where f :: forall i b. f (Term f) i -> (E (g (Term f)) -> b -> b) -> b -> b
          f :: f (Term f) i -> (E (g (Term f)) -> b -> b) -> b -> b
f f (Term f) i
t E (g (Term f)) -> b -> b
cons b
nil = let rest :: b
rest = (b -> Term f :=> b) -> b -> f (Term f) i -> b
forall (h :: (* -> *) -> * -> *) b (a :: * -> *).
HFoldable h =>
(b -> a :=> b) -> b -> h a :=> b
hfoldl (\b
u (Term s) -> f (Term f) i -> (E (g (Term f)) -> b -> b) -> b -> b
forall i b. f (Term f) i -> (E (g (Term f)) -> b -> b) -> b -> b
f f (Term f) i
s E (g (Term f)) -> b -> b
cons b
u) b
nil f (Term f) i
t
                         in case f (Term f) i -> Maybe (g (Term f) i)
forall (f :: (* -> *) -> * -> *) (g :: (* -> *) -> * -> *)
       (a :: * -> *).
(f :<: g) =>
NatM Maybe (g a) (f a)
proj f (Term f) i
t of
                              Just g (Term f) i
t' -> g (Term f) i -> E (g (Term f))
forall (f :: * -> *) i. f i -> E f
E g (Term f) i
t' E (g (Term f)) -> b -> b
`cons` b
rest
                              Maybe (g (Term f) i)
Nothing -> b
rest

-- | This function transforms every subterm according to the given
-- function in a bottom-up manner. This function is similar to
-- Uniplate's @transform@ function.
transform :: forall f . (HFunctor f) => (Term f :-> Term f) -> Term f :-> Term f
transform :: (Term f :-> Term f) -> Term f :-> Term f
transform Term f :-> Term f
f = Term f i -> Term f i
Term f :-> Term f
run
    where run :: Term f :-> Term f
          run :: Term f i -> Term f i
run = Term f i -> Term f i
Term f :-> Term f
f (Term f i -> Term f i)
-> (Term f i -> Term f i) -> Term f i -> Term f i
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (Term f) i -> Term f i
forall (f :: (* -> *) -> * -> *) h (a :: * -> *) i.
f (Cxt h f a) i -> Cxt h f a i
Term (f (Term f) i -> Term f i)
-> (Term f i -> f (Term f) i) -> Term f i -> Term f i
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Term f :-> Term f) -> f (Term f) :-> f (Term f)
forall (h :: (* -> *) -> * -> *) (f :: * -> *) (g :: * -> *).
HFunctor h =>
(f :-> g) -> h f :-> h g
hfmap Term f :-> Term f
run (f (Term f) i -> f (Term f) i)
-> (Term f i -> f (Term f) i) -> Term f i -> f (Term f) i
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term f i -> f (Term f) i
forall (f :: (* -> *) -> * -> *) t. Term f t -> f (Term f) t
unTerm


-- | Monadic version of 'transform'.
transformM :: forall f m . (HTraversable f, Monad m) =>
             NatM m (Term f) (Term f) -> NatM m (Term f) (Term f)
transformM :: NatM m (Term f) (Term f) -> NatM m (Term f) (Term f)
transformM  NatM m (Term f) (Term f)
f = Term f i -> m (Term f i)
NatM m (Term f) (Term f)
run
    where run :: NatM m (Term f) (Term f)
          run :: Term f i -> m (Term f i)
run Term f i
t = Term f i -> m (Term f i)
NatM m (Term f) (Term f)
f (Term f i -> m (Term f i)) -> m (Term f i) -> m (Term f i)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (f (Term f) i -> Term f i) -> m (f (Term f) i) -> m (Term f i)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM f (Term f) i -> Term f i
forall (f :: (* -> *) -> * -> *) h (a :: * -> *) i.
f (Cxt h f a) i -> Cxt h f a i
Term (NatM m (Term f) (Term f) -> NatM m (f (Term f)) (f (Term f))
forall (t :: (* -> *) -> * -> *) (m :: * -> *) (a :: * -> *)
       (b :: * -> *).
(HTraversable t, Monad m) =>
NatM m a b -> NatM m (t a) (t b)
hmapM NatM m (Term f) (Term f)
run (f (Term f) i -> m (f (Term f) i))
-> f (Term f) i -> m (f (Term f) i)
forall a b. (a -> b) -> a -> b
$ Term f i -> f (Term f) i
forall (f :: (* -> *) -> * -> *) t. Term f t -> f (Term f) t
unTerm Term f i
t)

query :: HFoldable f => (Term f :=>  r) -> (r -> r -> r) -> Term f :=> r
-- query q c = run
--     where run i@(Term t) = foldl (\s x -> s `c` run x) (q i) t
query :: (Term f :=> r) -> (r -> r -> r) -> Term f :=> r
query Term f :=> r
q r -> r -> r
c i :: Term f i
i@(Term f (Term f) i
t) = (r -> Term f :=> r) -> r -> f (Term f) i -> r
forall (h :: (* -> *) -> * -> *) b (a :: * -> *).
HFoldable h =>
(b -> a :=> b) -> b -> h a :=> b
hfoldl (\r
s Cxt NoHole f (K ()) i
x -> r
s r -> r -> r
`c` (Term f :=> r) -> (r -> r -> r) -> Cxt NoHole f (K ()) i -> r
forall (f :: (* -> *) -> * -> *) r.
HFoldable f =>
(Term f :=> r) -> (r -> r -> r) -> Term f :=> r
query Term f :=> r
q r -> r -> r
c Cxt NoHole f (K ()) i
x) (Term f i -> r
Term f :=> r
q Term f i
i) f (Term f) i
t

subs :: HFoldable f => Term f  :=> [E (Term f)]
subs :: Term f :=> [E (Term f)]
subs = (Term f :=> [E (Term f)])
-> ([E (Term f)] -> [E (Term f)] -> [E (Term f)])
-> Term f :=> [E (Term f)]
forall (f :: (* -> *) -> * -> *) r.
HFoldable f =>
(Term f :=> r) -> (r -> r -> r) -> Term f :=> r
query (\Term f i
x-> [Term f i -> E (Term f)
forall (f :: * -> *) i. f i -> E f
E Term f i
x]) [E (Term f)] -> [E (Term f)] -> [E (Term f)]
forall a. [a] -> [a] -> [a]
(++)

subs' :: (HFoldable f, g :<: f) => Term f :=> [E (g (Term f))]
subs' :: Term f :=> [E (g (Term f))]
subs' = (E (Term f) -> Maybe (E (g (Term f))))
-> [E (Term f)] -> [E (g (Term f))]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe E (Term f) -> Maybe (E (g (Term f)))
forall (g :: (* -> *) -> * -> *) (f :: (* -> *) -> * -> *) h
       (a :: * -> *).
Subsume (ComprEmb (Elem g f)) g f =>
E (Cxt h f a) -> Maybe (E (g (Cxt h f a)))
pr ([E (Term f)] -> [E (g (Term f))])
-> (Term f i -> [E (Term f)]) -> Term f i -> [E (g (Term f))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term f i -> [E (Term f)]
forall (f :: (* -> *) -> * -> *).
HFoldable f =>
Term f :=> [E (Term f)]
subs
        where pr :: E (Cxt h f a) -> Maybe (E (g (Cxt h f a)))
pr (E Cxt h f a i
v) = (g (Cxt h f a) i -> E (g (Cxt h f a)))
-> Maybe (g (Cxt h f a) i) -> Maybe (E (g (Cxt h f a)))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap g (Cxt h f a) i -> E (g (Cxt h f a))
forall (f :: * -> *) i. f i -> E f
E (Cxt h f a i -> Maybe (g (Cxt h f a) i)
forall (g :: (* -> *) -> * -> *) (f :: (* -> *) -> * -> *) h
       (a :: * -> *).
(g :<: f) =>
NatM Maybe (Cxt h f a) (g (Cxt h f a))
project Cxt h f a i
v)

-- | This function computes the generic size of the given term,
-- i.e. the its number of subterm occurrences.
size :: HFoldable f => Cxt h f a :=> Int
size :: Cxt h f a :=> Int
size (Hole {}) = Int
0
size (Term f (Cxt h f a) i
t) = (Int -> Cxt h f a :=> Int) -> Int -> f (Cxt h f a) i -> Int
forall (h :: (* -> *) -> * -> *) b (a :: * -> *).
HFoldable h =>
(b -> a :=> b) -> b -> h a :=> b
hfoldl (\Int
s Cxt h f a i
x -> Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Cxt h f a i -> Int
forall (f :: (* -> *) -> * -> *) h (a :: * -> *).
HFoldable f =>
Cxt h f a :=> Int
size Cxt h f a i
x) Int
1 f (Cxt h f a) i
t

-- | This function computes the generic depth of the given term.
depth :: HFoldable f => Cxt h f a :=> Int
depth :: Cxt h f a :=> Int
depth (Hole {}) = Int
0
depth (Term f (Cxt h f a) i
t) = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int -> Cxt h f a :=> Int) -> Int -> f (Cxt h f a) i -> Int
forall (h :: (* -> *) -> * -> *) b (a :: * -> *).
HFoldable h =>
(b -> a :=> b) -> b -> h a :=> b
hfoldl (\Int
s Cxt h f a i
x -> Int
s Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` Cxt h f a i -> Int
forall (f :: (* -> *) -> * -> *) h (a :: * -> *).
HFoldable f =>
Cxt h f a :=> Int
depth Cxt h f a i
x) Int
0 f (Cxt h f a) i
t