{-# 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 => HFix f  :=> [E (HFix f)]
subterms :: HFix f :=> [E (HFix f)]
subterms t :: HFix f i
t = (forall b. (E (HFix f) -> b -> b) -> b -> b) -> [E (HFix f)]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (HFix f i -> (E (HFix f) -> b -> b) -> b -> b
forall i b. HFix f i -> (E (HFix f) -> b -> b) -> b -> b
f HFix f i
t)
    where f :: forall i b. HFix f i -> (E (HFix f) -> b -> b) -> b -> b
          f :: HFix f i -> (E (HFix f) -> b -> b) -> b -> b
f t :: HFix f i
t cons :: E (HFix f) -> b -> b
cons nil :: b
nil = HFix f i -> E (HFix f)
forall (f :: * -> *) i. f i -> E f
E HFix f i
t E (HFix f) -> b -> b
`cons` (b -> HFix f :=> b) -> b -> f (HFix f) i -> b
forall (h :: (* -> *) -> * -> *) b (a :: * -> *).
HFoldable h =>
(b -> a :=> b) -> b -> h a :=> b
hfoldl (\u :: b
u s :: HFix f i
s -> HFix f i -> (E (HFix f) -> b -> b) -> b -> b
forall i b. HFix f i -> (E (HFix f) -> b -> b) -> b -> b
f HFix f i
s E (HFix f) -> b -> b
cons b
u) b
nil (HFix f i -> f (HFix f) i
forall (f :: (* -> *) -> * -> *) t. HFix f t -> f (HFix f) t
unTerm HFix 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) => HFix f :=> [E (g (HFix f))]
subterms' :: HFix f :=> [E (g (HFix f))]
subterms' (Term t :: f (HFix f) i
t) = (forall b. (E (g (HFix f)) -> b -> b) -> b -> b)
-> [E (g (HFix f))]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (f (HFix f) i -> (E (g (HFix f)) -> b -> b) -> b -> b
forall i b. f (HFix f) i -> (E (g (HFix f)) -> b -> b) -> b -> b
f f (HFix f) i
t)
    where f :: forall i b. f (HFix f) i -> (E (g (HFix f)) -> b -> b) -> b -> b
          f :: f (HFix f) i -> (E (g (HFix f)) -> b -> b) -> b -> b
f t :: f (HFix f) i
t cons :: E (g (HFix f)) -> b -> b
cons nil :: b
nil = let rest :: b
rest = (b -> HFix f :=> b) -> b -> f (HFix f) i -> b
forall (h :: (* -> *) -> * -> *) b (a :: * -> *).
HFoldable h =>
(b -> a :=> b) -> b -> h a :=> b
hfoldl (\u :: b
u (Term s) -> f (HFix f) i -> (E (g (HFix f)) -> b -> b) -> b -> b
forall i b. f (HFix f) i -> (E (g (HFix f)) -> b -> b) -> b -> b
f f (HFix f) i
s E (g (HFix f)) -> b -> b
cons b
u) b
nil f (HFix f) i
t
                         in case f (HFix f) i -> Maybe (g (HFix f) i)
forall (f :: (* -> *) -> * -> *) (g :: (* -> *) -> * -> *)
       (a :: * -> *).
(f :<: g) =>
NatM Maybe (g a) (f a)
proj f (HFix f) i
t of
                              Just t' :: g (HFix f) i
t' -> g (HFix f) i -> E (g (HFix f))
forall (f :: * -> *) i. f i -> E f
E g (HFix f) i
t' E (g (HFix f)) -> b -> b
`cons` b
rest
                              Nothing -> b
rest

-- |
-- @
-- transform :: (forall i. Term fs i -> Term fs i) -> Term f l -> Term f l
-- @
--
-- If @f :: Term fs i -> Term fs i@ rewrites a single node, then @transform f t@
-- is the result of running @f@ on all nodes within @t@ in a bottom-up fashion.
transform :: forall f . (HFunctor f) => (HFix f :-> HFix f) -> HFix f :-> HFix f
transform :: (HFix f :-> HFix f) -> HFix f :-> HFix f
transform f :: HFix f :-> HFix f
f = HFix f i -> HFix f i
HFix f :-> HFix f
run
    where run :: HFix f :-> HFix f
          run :: HFix f i -> HFix f i
run = HFix f i -> HFix f i
HFix f :-> HFix f
f (HFix f i -> HFix f i)
-> (HFix f i -> HFix f i) -> HFix f i -> HFix f i
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (HFix f) i -> HFix f i
forall (f :: (* -> *) -> * -> *) h (a :: * -> *) i.
f (Cxt h f a) i -> Cxt h f a i
Term (f (HFix f) i -> HFix f i)
-> (HFix f i -> f (HFix f) i) -> HFix f i -> HFix f i
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (HFix f :-> HFix f) -> f (HFix f) :-> f (HFix f)
forall (h :: (* -> *) -> * -> *) (f :: * -> *) (g :: * -> *).
HFunctor h =>
(f :-> g) -> h f :-> h g
hfmap HFix f :-> HFix f
run (f (HFix f) i -> f (HFix f) i)
-> (HFix f i -> f (HFix f) i) -> HFix f i -> f (HFix f) i
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HFix f i -> f (HFix f) i
forall (f :: (* -> *) -> * -> *) t. HFix f t -> f (HFix f) t
unTerm


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

-- |
-- @
-- query :: (forall i. Term fs i -> r) -> (r -> r -> r) -> Term fs l -> r
-- @
--
-- Example usage: let @getConsts :: (IntConst :-<: fs) => Term fs l -> [Int]@ be
-- a function where @getConsts (iIntConst n) = [n]@, and @getConsts t = []@
-- for everything that is not an @IntConst@. Then
--
-- @
-- query getConsts (++) term
-- @
--
-- returns a list of all integer constants in @term@.
query :: HFoldable f => (HFix f :=>  r) -> (r -> r -> r) -> HFix f :=> r
-- query q c = run
--     where run i@(Term t) = foldl (\s x -> s `c` run x) (q i) t
query :: (HFix f :=> r) -> (r -> r -> r) -> HFix f :=> r
query q :: HFix f :=> r
q c :: r -> r -> r
c i :: HFix f i
i@(Term t :: f (HFix f) i
t) = (r -> HFix f :=> r) -> r -> f (HFix f) i -> r
forall (h :: (* -> *) -> * -> *) b (a :: * -> *).
HFoldable h =>
(b -> a :=> b) -> b -> h a :=> b
hfoldl (\s :: r
s x :: Cxt NoHole f (K ()) i
x -> r
s r -> r -> r
`c` (HFix f :=> r) -> (r -> r -> r) -> Cxt NoHole f (K ()) i -> r
forall (f :: (* -> *) -> * -> *) r.
HFoldable f =>
(HFix f :=> r) -> (r -> r -> r) -> HFix f :=> r
query HFix f :=> r
q r -> r -> r
c Cxt NoHole f (K ()) i
x) (HFix f i -> r
HFix f :=> r
q HFix f i
i) f (HFix f) i
t

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

subs' :: (HFoldable f, g :<: f) => HFix f :=> [E (g (HFix f))]
subs' :: HFix f :=> [E (g (HFix f))]
subs' = (E (HFix f) -> Maybe (E (g (HFix f))))
-> [E (HFix f)] -> [E (g (HFix f))]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe E (HFix f) -> Maybe (E (g (HFix f)))
forall (g :: (* -> *) -> * -> *) (f :: (* -> *) -> * -> *) h
       (a :: * -> *).
(g :<: f) =>
E (Cxt h f a) -> Maybe (E (g (Cxt h f a)))
pr ([E (HFix f)] -> [E (g (HFix f))])
-> (HFix f i -> [E (HFix f)]) -> HFix f i -> [E (g (HFix f))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HFix f i -> [E (HFix f)]
forall (f :: (* -> *) -> * -> *).
HFoldable f =>
HFix f :=> [E (HFix f)]
subs
        where pr :: E (Cxt h f a) -> Maybe (E (g (Cxt h f a)))
pr (E v :: 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 {}) = 0
size (Term t :: 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 (\s :: Int
s x :: 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) 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 {}) = 0
depth (Term t :: f (Cxt h f a) i
t) = 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 (\s :: Int
s x :: 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) 0 f (Cxt h f a) i
t