{-# LANGUAGE CPP
           , DataKinds
           , FlexibleContexts
           , GADTs
           , GeneralizedNewtypeDeriving
           , KindSignatures
           , MultiParamTypeClasses
           , ScopedTypeVariables
           , TypeOperators
           #-}

{-# OPTIONS_GHC -Wall -fwarn-tabs #-}
----------------------------------------------------------------
--                                                    2017.02.01
-- |
-- Module      :  Language.Hakaru.Syntax.Prune
-- Copyright   :  Copyright (c) 2016 the Hakaru team
-- License     :  BSD3
-- Maintainer  :
-- Stability   :  experimental
-- Portability :  GHC-only
--
--
----------------------------------------------------------------
module Language.Hakaru.Syntax.Prune (prune) where

import           Control.Monad.Reader
import           Data.Maybe

import           Language.Hakaru.Syntax.ABT
import           Language.Hakaru.Syntax.AST
import           Language.Hakaru.Syntax.AST.Eq
import           Language.Hakaru.Syntax.IClasses
import           Language.Hakaru.Syntax.Unroll   (renameInEnv)
import           Language.Hakaru.Types.DataKind

#if __GLASGOW_HASKELL__ < 710
import           Control.Applicative
#endif

-- A Simple pass for pruning the unused let bindings from an AST.

newtype PruneM a = PruneM { PruneM a -> Reader Varmap a
runPruneM :: Reader Varmap a }
  deriving (a -> PruneM b -> PruneM a
(a -> b) -> PruneM a -> PruneM b
(forall a b. (a -> b) -> PruneM a -> PruneM b)
-> (forall a b. a -> PruneM b -> PruneM a) -> Functor PruneM
forall a b. a -> PruneM b -> PruneM a
forall a b. (a -> b) -> PruneM a -> PruneM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> PruneM b -> PruneM a
$c<$ :: forall a b. a -> PruneM b -> PruneM a
fmap :: (a -> b) -> PruneM a -> PruneM b
$cfmap :: forall a b. (a -> b) -> PruneM a -> PruneM b
Functor, Functor PruneM
a -> PruneM a
Functor PruneM
-> (forall a. a -> PruneM a)
-> (forall a b. PruneM (a -> b) -> PruneM a -> PruneM b)
-> (forall a b c.
    (a -> b -> c) -> PruneM a -> PruneM b -> PruneM c)
-> (forall a b. PruneM a -> PruneM b -> PruneM b)
-> (forall a b. PruneM a -> PruneM b -> PruneM a)
-> Applicative PruneM
PruneM a -> PruneM b -> PruneM b
PruneM a -> PruneM b -> PruneM a
PruneM (a -> b) -> PruneM a -> PruneM b
(a -> b -> c) -> PruneM a -> PruneM b -> PruneM c
forall a. a -> PruneM a
forall a b. PruneM a -> PruneM b -> PruneM a
forall a b. PruneM a -> PruneM b -> PruneM b
forall a b. PruneM (a -> b) -> PruneM a -> PruneM b
forall a b c. (a -> b -> c) -> PruneM a -> PruneM b -> PruneM c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: PruneM a -> PruneM b -> PruneM a
$c<* :: forall a b. PruneM a -> PruneM b -> PruneM a
*> :: PruneM a -> PruneM b -> PruneM b
$c*> :: forall a b. PruneM a -> PruneM b -> PruneM b
liftA2 :: (a -> b -> c) -> PruneM a -> PruneM b -> PruneM c
$cliftA2 :: forall a b c. (a -> b -> c) -> PruneM a -> PruneM b -> PruneM c
<*> :: PruneM (a -> b) -> PruneM a -> PruneM b
$c<*> :: forall a b. PruneM (a -> b) -> PruneM a -> PruneM b
pure :: a -> PruneM a
$cpure :: forall a. a -> PruneM a
$cp1Applicative :: Functor PruneM
Applicative, Applicative PruneM
a -> PruneM a
Applicative PruneM
-> (forall a b. PruneM a -> (a -> PruneM b) -> PruneM b)
-> (forall a b. PruneM a -> PruneM b -> PruneM b)
-> (forall a. a -> PruneM a)
-> Monad PruneM
PruneM a -> (a -> PruneM b) -> PruneM b
PruneM a -> PruneM b -> PruneM b
forall a. a -> PruneM a
forall a b. PruneM a -> PruneM b -> PruneM b
forall a b. PruneM a -> (a -> PruneM b) -> PruneM b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: a -> PruneM a
$creturn :: forall a. a -> PruneM a
>> :: PruneM a -> PruneM b -> PruneM b
$c>> :: forall a b. PruneM a -> PruneM b -> PruneM b
>>= :: PruneM a -> (a -> PruneM b) -> PruneM b
$c>>= :: forall a b. PruneM a -> (a -> PruneM b) -> PruneM b
$cp1Monad :: Applicative PruneM
Monad, MonadReader Varmap, Monad PruneM
Monad PruneM
-> (forall a. (a -> PruneM a) -> PruneM a) -> MonadFix PruneM
(a -> PruneM a) -> PruneM a
forall a. (a -> PruneM a) -> PruneM a
forall (m :: * -> *).
Monad m -> (forall a. (a -> m a) -> m a) -> MonadFix m
mfix :: (a -> PruneM a) -> PruneM a
$cmfix :: forall a. (a -> PruneM a) -> PruneM a
$cp1MonadFix :: Monad PruneM
MonadFix)

lookupEnv
  :: forall (a :: Hakaru)
  .  Variable a
  -> Varmap
  -> Variable a
lookupEnv :: Variable a -> Varmap -> Variable a
lookupEnv Variable a
v = Variable a -> Maybe (Variable a) -> Variable a
forall a. a -> Maybe a -> a
fromMaybe Variable a
v (Maybe (Variable a) -> Variable a)
-> (Varmap -> Maybe (Variable a)) -> Varmap -> Variable a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Variable a -> Varmap -> Maybe (Variable a)
forall k (a :: k) (ast :: k -> *).
(Show1 Sing, JmEq1 Sing) =>
Variable a -> Assocs ast -> Maybe (ast a)
lookupAssoc Variable a
v

prune
  :: (ABT Term abt)
  => abt '[] a
  -> abt '[] a
prune :: abt '[] a -> abt '[] a
prune = (Reader Varmap (abt '[] a) -> Varmap -> abt '[] a)
-> Varmap -> Reader Varmap (abt '[] a) -> abt '[] a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Reader Varmap (abt '[] a) -> Varmap -> abt '[] a
forall r a. Reader r a -> r -> a
runReader Varmap
forall k (abt :: k -> *). Assocs abt
emptyAssocs (Reader Varmap (abt '[] a) -> abt '[] a)
-> (abt '[] a -> Reader Varmap (abt '[] a))
-> abt '[] a
-> abt '[] a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PruneM (abt '[] a) -> Reader Varmap (abt '[] a)
forall a. PruneM a -> Reader Varmap a
runPruneM (PruneM (abt '[] a) -> Reader Varmap (abt '[] a))
-> (abt '[] a -> PruneM (abt '[] a))
-> abt '[] a
-> Reader Varmap (abt '[] a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. abt '[] a -> PruneM (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
abt xs a -> PruneM (abt xs a)
prune'

prune'
  :: forall abt xs a . (ABT Term abt)
  => abt xs a
  -> PruneM (abt xs a)
prune' :: abt xs a -> PruneM (abt xs a)
prune' = View (Term abt) xs a -> PruneM (abt xs a)
forall (b :: Hakaru) (ys :: [Hakaru]).
View (Term abt) ys b -> PruneM (abt ys b)
loop (View (Term abt) xs a -> PruneM (abt xs a))
-> (abt xs a -> View (Term abt) xs a)
-> abt xs a
-> PruneM (abt xs a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. abt xs a -> View (Term abt) 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
   where
    loop :: forall (b :: Hakaru) ys . View (Term abt) ys b -> PruneM (abt ys b)
    loop :: View (Term abt) ys b -> PruneM (abt ys b)
loop (Var Variable b
v)    = (Variable b -> abt '[] b
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable b -> abt '[] b)
-> (Varmap -> Variable b) -> Varmap -> abt '[] b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Variable b -> Varmap -> Variable b
forall (a :: Hakaru). Variable a -> Varmap -> Variable a
lookupEnv Variable b
v) (Varmap -> abt '[] b) -> PruneM Varmap -> PruneM (abt '[] b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` PruneM Varmap
forall r (m :: * -> *). MonadReader r m => m r
ask
    loop (Syn Term abt b
s)    = Term abt b -> PruneM (abt '[] b)
forall (a :: Hakaru) (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
Term abt a -> PruneM (abt '[] a)
pruneTerm Term abt b
s
    loop (Bind Variable a
v View (Term abt) xs b
b) = Variable a -> PruneM (abt xs b) -> PruneM (abt (a : xs) b)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (a :: Hakaru)
       (xs :: [Hakaru]) (b :: Hakaru).
(ABT Term abt, MonadReader Varmap m, MonadFix m) =>
Variable a -> m (abt xs b) -> m (abt (a : xs) b)
renameInEnv Variable a
v (View (Term abt) xs b -> PruneM (abt xs b)
forall (b :: Hakaru) (ys :: [Hakaru]).
View (Term abt) ys b -> PruneM (abt ys b)
loop View (Term abt) xs b
b)

pruneTerm
  :: forall a abt
  .  (ABT Term abt)
  => Term abt a
  -> PruneM (abt '[] a)
pruneTerm :: Term abt a -> PruneM (abt '[] a)
pruneTerm (SCon args a
Let_ :$ abt vars a
rhs :* abt vars a
body :* SArgs abt args
End) =
  abt '[a] a
-> (Variable a -> abt '[] a -> PruneM (abt '[] a))
-> PruneM (abt '[] a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt vars a
abt '[a] a
body ((Variable a -> abt '[] a -> PruneM (abt '[] a))
 -> PruneM (abt '[] a))
-> (Variable a -> abt '[] a -> PruneM (abt '[] a))
-> PruneM (abt '[] a)
forall a b. (a -> b) -> a -> b
$ \Variable a
v abt '[] a
body' ->
  let frees :: VarSet (KindOf a)
frees     = abt '[] a -> VarSet (KindOf a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> VarSet (KindOf a)
freeVars abt '[] a
body'
      mklet :: abt '[] a -> abt '[a] a -> abt '[] a
mklet abt '[] a
r abt '[a] a
b = 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 (SCon '[LC a, '( '[a], a)] a
forall (a :: Hakaru) (b :: Hakaru). SCon '[LC a, '( '[a], b)] b
Let_ SCon '[LC a, '( '[a], a)] a
-> SArgs abt '[LC a, '( '[a], a)] -> Term abt a
forall (args :: [([Hakaru], Hakaru)]) (a :: Hakaru)
       (abt :: [Hakaru] -> Hakaru -> *).
SCon args a -> SArgs abt args -> Term abt a
:$ abt '[] a
r abt '[] a
-> SArgs abt '[ '( '[a], a)] -> SArgs abt '[LC a, '( '[a], a)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* abt '[a] a
b abt '[a] a -> SArgs abt '[] -> SArgs abt '[ '( '[a], a)]
forall (abt :: [Hakaru] -> Hakaru -> *) (vars :: [Hakaru])
       (a :: Hakaru) (args :: [([Hakaru], Hakaru)]).
abt vars a -> SArgs abt args -> SArgs abt ('(vars, a) : args)
:* SArgs abt '[]
forall (abt :: [Hakaru] -> Hakaru -> *). SArgs abt '[]
End)
      doRhs :: PruneM (abt vars a)
doRhs     = abt vars a -> PruneM (abt vars a)
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
abt xs a -> PruneM (abt xs a)
prune' abt vars a
rhs
      doBody :: PruneM (abt '[] a)
doBody    = abt '[] a -> PruneM (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
abt xs a -> PruneM (abt xs a)
prune' abt '[] a
body'
      fullExpr :: PruneM (abt '[] a)
fullExpr  = abt '[] a -> abt '[a] a -> abt '[] a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (a :: Hakaru).
ABT Term abt =>
abt '[] a -> abt '[a] a -> abt '[] a
mklet (abt '[] a -> abt '[a] a -> abt '[] a)
-> PruneM (abt '[] a) -> PruneM (abt '[a] a -> abt '[] a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PruneM (abt vars a)
PruneM (abt '[] a)
doRhs PruneM (abt '[a] a -> abt '[] a)
-> PruneM (abt '[a] a) -> PruneM (abt '[] a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Variable a -> PruneM (abt '[] a) -> PruneM (abt '[a] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (m :: * -> *) (a :: Hakaru)
       (xs :: [Hakaru]) (b :: Hakaru).
(ABT Term abt, MonadReader Varmap m, MonadFix m) =>
Variable a -> m (abt xs b) -> m (abt (a : xs) b)
renameInEnv Variable a
v PruneM (abt '[] a)
doBody
  in case abt '[] a -> View (Term abt) '[] 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 abt '[] a
body' of
       Var Variable a
v' | Just TypeEq a a
Refl <- Variable a -> Variable a -> Maybe (TypeEq a a)
forall k (a :: k) (b :: k).
(Show1 Sing, JmEq1 Sing) =>
Variable a -> Variable b -> Maybe (TypeEq a b)
varEq Variable a
v Variable a
v' -> PruneM (abt vars a)
PruneM (abt '[] a)
doRhs
       View (Term abt) '[] a
_      | Variable a -> VarSet (KindOf a) -> Bool
forall k (a :: k) (kproxy :: KProxy k).
(Show1 Sing, JmEq1 Sing) =>
Variable a -> VarSet kproxy -> Bool
memberVarSet Variable a
v VarSet (KindOf a)
frees    -> PruneM (abt '[] a)
fullExpr
              | Bool
otherwise               -> PruneM (abt '[] a)
doBody

pruneTerm Term abt a
term = 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 -> abt '[] a)
-> PruneM (Term abt a) -> PruneM (abt '[] a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall (h :: [Hakaru]) (i :: Hakaru). abt h i -> PruneM (abt h i))
-> Term abt a -> PruneM (Term abt a)
forall k1 k2 k3 (t :: (k1 -> k2 -> *) -> k3 -> *) (f :: * -> *)
       (a :: k1 -> k2 -> *) (b :: k1 -> k2 -> *) (j :: k3).
(Traversable21 t, Applicative f) =>
(forall (h :: k1) (i :: k2). a h i -> f (b h i))
-> t a j -> f (t b j)
traverse21 forall (h :: [Hakaru]) (i :: Hakaru). abt h i -> PruneM (abt h i)
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
abt xs a -> PruneM (abt xs a)
prune' Term abt a
term