{-# LANGUAGE DataKinds
           , GADTs
           , Rank2Types
           , FlexibleContexts
           #-}

----------------------------------------------------------------
--                                                    2016.07.19
-- |
-- Module      :  Language.Hakaru.Evaluation.Coalesce
-- Copyright   :  Copyright (c) 2016 the Hakaru team
-- License     :  BSD3
-- Maintainer  :  zsulliva@indiana.edu
-- Stability   :  experimental
-- Portability :  GHC-only
--
----------------------------------------------------------------

module Language.Hakaru.Evaluation.Coalesce
  ( coalesce )
  where

import qualified Language.Hakaru.Parser.AST as U
import           Language.Hakaru.Syntax.ABT
import qualified Data.Foldable              as F

import Language.Hakaru.Syntax.IClasses

coalesce
    :: U.AST
    -> U.AST
coalesce :: AST -> AST
coalesce =
    (forall (a :: Untyped). Term M a -> M '[] a)
-> forall (xs :: [Untyped]) (a :: Untyped). M xs a -> M xs a
cataABT_ forall (a :: Untyped). Term M a -> M '[] a
forall (abt :: [Untyped] -> Untyped -> *) (a :: Untyped).
ABT Term abt =>
Term abt a -> abt '[] a
alg
    where
    alg :: forall abt a. (ABT U.Term abt) => U.Term abt a -> abt '[] a
    alg :: Term abt a -> abt '[] a
alg (U.NaryOp_ NaryOp
op [abt '[] 'U]
args) = Term abt 'U -> abt '[] 'U
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn (Term abt 'U -> abt '[] 'U) -> Term abt 'U -> abt '[] 'U
forall a b. (a -> b) -> a -> b
$ NaryOp -> [abt '[] 'U] -> Term abt 'U
forall (abt :: [Untyped] -> Untyped -> *).
NaryOp -> [abt '[] 'U] -> Term abt 'U
U.NaryOp_ NaryOp
op (NaryOp -> [abt '[] 'U] -> [abt '[] 'U]
forall (abt :: [Untyped] -> Untyped -> *) (a :: Untyped).
ABT Term abt =>
NaryOp -> [abt '[] a] -> [abt '[] a]
coalesceNaryOp NaryOp
op [abt '[] 'U]
args)
    alg Term abt a
t                   = Term abt a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn Term abt a
t

coalesceNaryOp
    :: (ABT U.Term abt)
    => U.NaryOp
    -> [abt '[] a]
    -> [abt '[] a]
coalesceNaryOp :: NaryOp -> [abt '[] a] -> [abt '[] a]
coalesceNaryOp NaryOp
op = (abt '[] a -> [abt '[] a]) -> [abt '[] a] -> [abt '[] a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
F.concatMap ((abt '[] a -> [abt '[] a]) -> [abt '[] a] -> [abt '[] a])
-> (abt '[] a -> [abt '[] a]) -> [abt '[] a] -> [abt '[] a]
forall a b. (a -> b) -> a -> b
$ \abt '[] a
ast' ->
     abt '[] a
-> (Variable a -> [abt '[] a])
-> (Term abt a -> [abt '[] a])
-> [abt '[] a]
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) r.
ABT syn abt =>
abt '[] a -> (Variable a -> r) -> (syn abt a -> r) -> r
caseVarSyn abt '[] a
ast' (abt '[] a -> [abt '[] a]
forall (m :: * -> *) a. Monad m => a -> m a
return (abt '[] a -> [abt '[] a])
-> (Variable a -> abt '[] a) -> Variable a -> [abt '[] a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var) ((Term abt a -> [abt '[] a]) -> [abt '[] a])
-> (Term abt a -> [abt '[] a]) -> [abt '[] a]
forall a b. (a -> b) -> a -> b
$ \Term abt a
t ->
       case Term abt a
t of
       U.NaryOp_ NaryOp
op' [abt '[] 'U]
args' | NaryOp
op NaryOp -> NaryOp -> Bool
forall a. Eq a => a -> a -> Bool
== NaryOp
op' -> NaryOp -> [abt '[] 'U] -> [abt '[] 'U]
forall (abt :: [Untyped] -> Untyped -> *) (a :: Untyped).
ABT Term abt =>
NaryOp -> [abt '[] a] -> [abt '[] a]
coalesceNaryOp NaryOp
op [abt '[] 'U]
args'
       Term abt a
_                               -> [abt '[] a
ast']


type M  = MetaABT U.SourceSpan U.Term

preserveMetadata
   :: (M xs a -> M xs a)
   -> M xs a
   -> M xs a
preserveMetadata :: (M xs a -> M xs a) -> M xs a -> M xs a
preserveMetadata M xs a -> M xs a
f M xs a
x =
    case M xs a -> Maybe SourceSpan
forall meta k (syn :: ([k] -> k -> *) -> k -> *) (xs :: [k])
       (a :: k).
MetaABT meta syn xs a -> Maybe meta
getMetadata M xs a
x of
      Maybe SourceSpan
Nothing -> M xs a -> M xs a
f M xs a
x
      Just SourceSpan
s  -> SourceSpan -> M xs a -> M xs a
forall k meta (syn :: ([k] -> k -> *) -> k -> *) (xs :: [k])
       (a :: k).
meta -> MetaABT meta syn xs a -> MetaABT meta syn xs a
withMetadata SourceSpan
s (M xs a -> M xs a
f M xs a
x)

cataABT_
    :: (forall    a. U.Term M a -> M '[] a)
    -> (forall xs a. M xs a     -> M xs  a)
cataABT_ :: (forall (a :: Untyped). Term M a -> M '[] a)
-> forall (xs :: [Untyped]) (a :: Untyped). M xs a -> M xs a
cataABT_ forall (a :: Untyped). Term M a -> M '[] a
syn_ = M xs a -> M xs a
forall (xs :: [Untyped]) (a :: Untyped). M xs a -> M xs a
start
    where
    start :: forall xs a. M xs a -> M xs a
    start :: M xs a -> M xs a
start = (M xs a -> M xs a) -> M xs a -> M xs a
forall (xs :: [Untyped]) (a :: Untyped).
(M xs a -> M xs a) -> M xs a -> M xs a
preserveMetadata (View (Term M) xs a -> M xs a
forall (xs :: [Untyped]) (a :: Untyped).
View (Term M) xs a -> M xs a
loop (View (Term M) xs a -> M xs a)
-> (M xs a -> View (Term M) xs a) -> M xs a -> M xs a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. M xs a -> View (Term M) xs a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> View (syn abt) xs a
viewABT)

    loop ::  forall xs a. View (U.Term M) xs a -> M xs a
    loop :: View (Term M) xs a -> M xs a
loop (Syn  Term M a
t)   = Term M a -> M '[] a
forall (a :: Untyped). Term M a -> M '[] a
syn_  ((forall (xs :: [Untyped]) (a :: Untyped). M xs a -> M xs a)
-> Term M a -> Term M a
forall k1 k2 k3 (f :: (k1 -> k2 -> *) -> k3 -> *)
       (a :: k1 -> k2 -> *) (b :: k1 -> k2 -> *) (j :: k3).
Functor21 f =>
(forall (h :: k1) (i :: k2). a h i -> b h i) -> f a j -> f b j
fmap21 forall (xs :: [Untyped]) (a :: Untyped). M xs a -> M xs a
start Term M a
t)
    loop (Var  Variable a
x)   = Variable a -> M '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var  Variable a
x
    loop (Bind Variable a
x View (Term M) xs a
e) = Variable a
-> MetaABT SourceSpan Term xs a
-> MetaABT SourceSpan Term (a : xs) a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) (xs :: [k]) (b :: k).
ABT syn abt =>
Variable a -> abt xs b -> abt (a : xs) b
bind Variable a
x (View (Term M) xs a -> MetaABT SourceSpan Term xs a
forall (xs :: [Untyped]) (a :: Untyped).
View (Term M) xs a -> M xs a
loop View (Term M) xs a
e)