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

--------------------------------------------------------------------------------
-- |
-- Module      :  Data.Comp.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.
--
--------------------------------------------------------------------------------

module Data.Comp.Generic where

import Control.Monad hiding (mapM)
import Data.Comp.Algebra
import Data.Comp.Sum
import Data.Comp.Term
import Data.Foldable
import Data.Maybe
import Data.Traversable
import GHC.Exts (build)
import Prelude hiding (foldl, mapM)


-- | This function returns the subterm of a given term at the position
-- specified by the given path or @Nothing@ if the input term has no
-- such subterm

getSubterm :: (Functor g, Foldable g) => [Int] -> Term g -> Maybe (Term g)
getSubterm :: [Int] -> Term g -> Maybe (Term g)
getSubterm [Int]
path Term g
t = Alg g ([Int] -> Maybe (Term g))
-> Term g -> [Int] -> Maybe (Term g)
forall (f :: * -> *) a. Functor f => Alg f a -> Term f -> a
cata Alg g ([Int] -> Maybe (Term g))
forall (g :: * -> *) h a.
(Functor g, Foldable g) =>
Alg g ([Int] -> Maybe (Cxt h g a))
alg Term g
t [Int]
path where
    alg :: (Functor g, Foldable g) => Alg g ([Int] -> Maybe (Cxt h g a))
    alg :: Alg g ([Int] -> Maybe (Cxt h g a))
alg g ([Int] -> Maybe (Cxt h g a))
t [] = Cxt h g a -> Maybe (Cxt h g a)
forall a. a -> Maybe a
Just (Cxt h g a -> Maybe (Cxt h g a)) -> Cxt h g a -> Maybe (Cxt h g a)
forall a b. (a -> b) -> a -> b
$ g (Cxt h g a) -> Cxt h g a
forall (f :: * -> *) h a. f (Cxt h f a) -> Cxt h f a
Term (g (Cxt h g a) -> Cxt h g a) -> g (Cxt h g a) -> Cxt h g a
forall a b. (a -> b) -> a -> b
$ (([Int] -> Maybe (Cxt h g a)) -> Cxt h g a)
-> g ([Int] -> Maybe (Cxt h g a)) -> g (Cxt h g a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Maybe (Cxt h g a) -> Cxt h g a
forall a. HasCallStack => Maybe a -> a
fromJust) (Maybe (Cxt h g a) -> Cxt h g a)
-> (([Int] -> Maybe (Cxt h g a)) -> Maybe (Cxt h g a))
-> ([Int] -> Maybe (Cxt h g a))
-> Cxt h g a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (([Int] -> Maybe (Cxt h g a)) -> [Int] -> Maybe (Cxt h g a)
forall a b. (a -> b) -> a -> b
$[])) g ([Int] -> Maybe (Cxt h g a))
t
    alg g ([Int] -> Maybe (Cxt h g a))
t (Int
i:[Int]
is) = case Int -> [[Int] -> Maybe (Cxt h g a)] -> [[Int] -> Maybe (Cxt h g a)]
forall a. Int -> [a] -> [a]
drop Int
i (g ([Int] -> Maybe (Cxt h g a)) -> [[Int] -> Maybe (Cxt h g a)]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList g ([Int] -> Maybe (Cxt h g a))
t) of
                     [] -> Maybe (Cxt h g a)
forall a. Maybe a
Nothing
                     [Int] -> Maybe (Cxt h g a)
x : [[Int] -> Maybe (Cxt h g a)]
_ -> [Int] -> Maybe (Cxt h g a)
x [Int]
is

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


-- | This function returns a list of all subterms of the given term
-- that are constructed from a particular functor.
subterms' :: forall f g . (Foldable f, g :<: f) => Term f -> [g (Term f)]
subterms' :: Term f -> [g (Term f)]
subterms' (Term f (Term f)
t) = (forall b. (g (Term f) -> b -> b) -> b -> b) -> [g (Term f)]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (f (Term f) -> (g (Term f) -> b -> b) -> b -> b
forall b. f (Term f) -> (g (Term f) -> b -> b) -> b -> b
f f (Term f)
t)
    where f :: f (Term f) -> (g (Term f) -> b -> b) -> b -> b
          f :: f (Term f) -> (g (Term f) -> b -> b) -> b -> b
f f (Term f)
t g (Term f) -> b -> b
cons b
nil = let rest :: b
rest = (b -> Term f -> b) -> b -> f (Term f) -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\b
u (Term f (Term f)
s) -> f (Term f) -> (g (Term f) -> b -> b) -> b -> b
forall b. f (Term f) -> (g (Term f) -> b -> b) -> b -> b
f f (Term f)
s g (Term f) -> b -> b
cons b
u) b
nil f (Term f)
t
                         in case f (Term f) -> Maybe (g (Term f))
forall (f :: * -> *) (g :: * -> *) a.
(f :<: g) =>
g a -> Maybe (f a)
proj f (Term f)
t of
                              Just g (Term f)
t' -> g (Term f)
t'g (Term f) -> b -> b
`cons` b
rest
                              Maybe (g (Term f))
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 :: (Functor 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 -> Term f
run
    where run :: Term f -> Term f
run = Term f -> Term f
f (Term f -> Term f) -> (Term f -> Term f) -> Term f -> Term f
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (Term f) -> Term f
forall (f :: * -> *) h a. f (Cxt h f a) -> Cxt h f a
Term (f (Term f) -> Term f)
-> (Term f -> f (Term f)) -> Term f -> Term f
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Term f -> Term f) -> f (Term f) -> f (Term f)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Term f -> Term f
run (f (Term f) -> f (Term f))
-> (Term f -> f (Term f)) -> Term f -> f (Term f)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term f -> f (Term f)
forall (f :: * -> *) a. Cxt NoHole f a -> f (Cxt NoHole f a)
unTerm
-- transform f  = f . Term . fmap (transform f) . unTerm

transform' :: (Functor f) => (Term f -> Maybe (Term f)) -> Term f -> Term f
transform' :: (Term f -> Maybe (Term f)) -> Term f -> Term f
transform' Term f -> Maybe (Term f)
f = (Term f -> Term f) -> Term f -> Term f
forall (f :: * -> *).
Functor f =>
(Term f -> Term f) -> Term f -> Term f
transform Term f -> Term f
f' where
    f' :: Term f -> Term f
f' Term f
t = Term f -> Maybe (Term f) -> Term f
forall a. a -> Maybe a -> a
fromMaybe Term f
t (Term f -> Maybe (Term f)
f Term f
t)


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

query :: Foldable f => (Term f -> r) -> (r -> r -> r) -> Term f -> r
query :: (Term f -> r) -> (r -> r -> r) -> Term f -> r
query Term f -> r
q r -> r -> r
c = Term f -> r
run
    where run :: Term f -> r
run i :: Term f
i@(Term f (Term f)
t) = (r -> Term f -> r) -> r -> f (Term f) -> r
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\r
s Term f
x -> r
s r -> r -> r
`c` Term f -> r
run Term f
x) (Term f -> r
q Term f
i) f (Term f)
t
-- query q c i@(Term t) = foldl (\s x -> s `c` query q c x) (q i) t

gsize :: Foldable f => Term f -> Int
gsize :: Term f -> Int
gsize = (Term f -> Int) -> (Int -> Int -> Int) -> Term f -> Int
forall (f :: * -> *) r.
Foldable f =>
(Term f -> r) -> (r -> r -> r) -> Term f -> r
query (Int -> Term f -> Int
forall a b. a -> b -> a
const Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
(+)

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

-- | This function computes the generic height of the given term.
height :: Foldable f => Cxt h f a -> Int
height :: Cxt h f a -> Int
height (Hole {}) = Int
0
height (Term f (Cxt h f a)
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) -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Int
s Cxt h f a
x -> Int
s Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` Cxt h f a -> Int
forall (f :: * -> *) h a. Foldable f => Cxt h f a -> Int
height Cxt h f a
x) Int
0 f (Cxt h f a)
t