{-# LANGUAGE CPP
           , DataKinds
           , EmptyCase
           , ExistentialQuantification
           , FlexibleContexts
           , GADTs
           , GeneralizedNewtypeDeriving
           , KindSignatures
           , MultiParamTypeClasses
           , OverloadedStrings
           , PolyKinds
           , ScopedTypeVariables
           , TypeFamilies
           , TypeOperators
           #-}

{-# OPTIONS_GHC -Wall -fwarn-tabs #-}
----------------------------------------------------------------
--                                                    2017.02.01
-- |
-- Module      :  Language.Hakaru.Syntax.Uniquify
-- Copyright   :  Copyright (c) 2016 the Hakaru team
-- License     :  BSD3
-- Maintainer  :
-- Stability   :  experimental
-- Portability :  GHC-only
--
-- Performs renaming of Hakaru expressions to ensure globally unique variable
-- identifiers.
--
----------------------------------------------------------------
module Language.Hakaru.Syntax.Uniquify (uniquify) where

import           Control.Monad.Reader
import           Control.Monad.State
import           Data.Maybe                      (fromMaybe)
import           Data.Number.Nat

import           Language.Hakaru.Syntax.ABT
import           Language.Hakaru.Syntax.AST
import           Language.Hakaru.Syntax.AST.Eq   (Varmap)
import           Language.Hakaru.Syntax.Gensym
import           Language.Hakaru.Syntax.IClasses

#if __GLASGOW_HASKELL__ < 710
import           Control.Applicative
#endif

newtype Uniquifier a = Uniquifier { Uniquifier a -> StateT Nat (Reader Varmap) a
runUniquifier :: StateT Nat (Reader Varmap) a }
  deriving (a -> Uniquifier b -> Uniquifier a
(a -> b) -> Uniquifier a -> Uniquifier b
(forall a b. (a -> b) -> Uniquifier a -> Uniquifier b)
-> (forall a b. a -> Uniquifier b -> Uniquifier a)
-> Functor Uniquifier
forall a b. a -> Uniquifier b -> Uniquifier a
forall a b. (a -> b) -> Uniquifier a -> Uniquifier b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> Uniquifier b -> Uniquifier a
$c<$ :: forall a b. a -> Uniquifier b -> Uniquifier a
fmap :: (a -> b) -> Uniquifier a -> Uniquifier b
$cfmap :: forall a b. (a -> b) -> Uniquifier a -> Uniquifier b
Functor, Functor Uniquifier
a -> Uniquifier a
Functor Uniquifier
-> (forall a. a -> Uniquifier a)
-> (forall a b.
    Uniquifier (a -> b) -> Uniquifier a -> Uniquifier b)
-> (forall a b c.
    (a -> b -> c) -> Uniquifier a -> Uniquifier b -> Uniquifier c)
-> (forall a b. Uniquifier a -> Uniquifier b -> Uniquifier b)
-> (forall a b. Uniquifier a -> Uniquifier b -> Uniquifier a)
-> Applicative Uniquifier
Uniquifier a -> Uniquifier b -> Uniquifier b
Uniquifier a -> Uniquifier b -> Uniquifier a
Uniquifier (a -> b) -> Uniquifier a -> Uniquifier b
(a -> b -> c) -> Uniquifier a -> Uniquifier b -> Uniquifier c
forall a. a -> Uniquifier a
forall a b. Uniquifier a -> Uniquifier b -> Uniquifier a
forall a b. Uniquifier a -> Uniquifier b -> Uniquifier b
forall a b. Uniquifier (a -> b) -> Uniquifier a -> Uniquifier b
forall a b c.
(a -> b -> c) -> Uniquifier a -> Uniquifier b -> Uniquifier 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
<* :: Uniquifier a -> Uniquifier b -> Uniquifier a
$c<* :: forall a b. Uniquifier a -> Uniquifier b -> Uniquifier a
*> :: Uniquifier a -> Uniquifier b -> Uniquifier b
$c*> :: forall a b. Uniquifier a -> Uniquifier b -> Uniquifier b
liftA2 :: (a -> b -> c) -> Uniquifier a -> Uniquifier b -> Uniquifier c
$cliftA2 :: forall a b c.
(a -> b -> c) -> Uniquifier a -> Uniquifier b -> Uniquifier c
<*> :: Uniquifier (a -> b) -> Uniquifier a -> Uniquifier b
$c<*> :: forall a b. Uniquifier (a -> b) -> Uniquifier a -> Uniquifier b
pure :: a -> Uniquifier a
$cpure :: forall a. a -> Uniquifier a
$cp1Applicative :: Functor Uniquifier
Applicative, Applicative Uniquifier
a -> Uniquifier a
Applicative Uniquifier
-> (forall a b.
    Uniquifier a -> (a -> Uniquifier b) -> Uniquifier b)
-> (forall a b. Uniquifier a -> Uniquifier b -> Uniquifier b)
-> (forall a. a -> Uniquifier a)
-> Monad Uniquifier
Uniquifier a -> (a -> Uniquifier b) -> Uniquifier b
Uniquifier a -> Uniquifier b -> Uniquifier b
forall a. a -> Uniquifier a
forall a b. Uniquifier a -> Uniquifier b -> Uniquifier b
forall a b. Uniquifier a -> (a -> Uniquifier b) -> Uniquifier 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 -> Uniquifier a
$creturn :: forall a. a -> Uniquifier a
>> :: Uniquifier a -> Uniquifier b -> Uniquifier b
$c>> :: forall a b. Uniquifier a -> Uniquifier b -> Uniquifier b
>>= :: Uniquifier a -> (a -> Uniquifier b) -> Uniquifier b
$c>>= :: forall a b. Uniquifier a -> (a -> Uniquifier b) -> Uniquifier b
$cp1Monad :: Applicative Uniquifier
Monad, MonadState Nat, MonadReader Varmap)

uniquify :: (ABT Term abt) => abt '[] a -> abt '[] a
uniquify :: abt '[] a -> abt '[] a
uniquify abt '[] a
abt = (abt '[] a, Nat) -> abt '[] a
forall a b. (a, b) -> a
fst ((abt '[] a, Nat) -> abt '[] a) -> (abt '[] a, Nat) -> abt '[] a
forall a b. (a -> b) -> a -> b
$ Reader Varmap (abt '[] a, Nat) -> Varmap -> (abt '[] a, Nat)
forall r a. Reader r a -> r -> a
runReader (StateT Nat (Reader Varmap) (abt '[] a)
-> Nat -> Reader Varmap (abt '[] a, Nat)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT Nat (Reader Varmap) (abt '[] a)
unique Nat
seed) Varmap
forall k (abt :: k -> *). Assocs abt
emptyAssocs
  where
    unique :: StateT Nat (Reader Varmap) (abt '[] a)
unique = Uniquifier (abt '[] a) -> StateT Nat (Reader Varmap) (abt '[] a)
forall a. Uniquifier a -> StateT Nat (Reader Varmap) a
runUniquifier (abt '[] a -> Uniquifier (abt '[] a)
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
abt xs a -> Uniquifier (abt xs a)
uniquify' abt '[] a
abt)
    seed :: Nat
seed   = abt '[] a -> Nat
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> Nat
nextFreeOrBind abt '[] a
abt

uniquify'
  :: forall abt xs a . (ABT Term abt)
  => abt xs a
  -> Uniquifier (abt xs a)
uniquify' :: abt xs a -> Uniquifier (abt xs a)
uniquify' = abt xs a -> Uniquifier (abt xs a)
forall (ys :: [Hakaru]) (b :: Hakaru).
abt ys b -> Uniquifier (abt ys b)
start
  where
    start :: abt ys b -> Uniquifier (abt ys b)
    start :: abt ys b -> Uniquifier (abt ys b)
start = View (Term abt) ys b -> Uniquifier (abt ys b)
forall (ys :: [Hakaru]) (b :: Hakaru).
View (Term abt) ys b -> Uniquifier (abt ys b)
loop (View (Term abt) ys b -> Uniquifier (abt ys b))
-> (abt ys b -> View (Term abt) ys b)
-> abt ys b
-> Uniquifier (abt ys b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. abt ys b -> View (Term abt) ys b
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 :: View (Term abt) ys b -> Uniquifier (abt ys b)
    loop :: View (Term abt) ys b -> Uniquifier (abt ys b)
loop (Var Variable b
v)    = Variable b -> Uniquifier (abt '[] b)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Variable a -> Uniquifier (abt '[] a)
uniquifyVar Variable b
v
    loop (Syn Term abt b
s)    = (Term abt b -> abt '[] b)
-> Uniquifier (Term abt b) -> Uniquifier (abt '[] b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Term abt b -> abt '[] b
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
syn abt a -> abt '[] a
syn ((forall (ys :: [Hakaru]) (b :: Hakaru).
 abt ys b -> Uniquifier (abt ys b))
-> Term abt b -> Uniquifier (Term abt b)
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 (ys :: [Hakaru]) (b :: Hakaru).
abt ys b -> Uniquifier (abt ys b)
start Term abt b
s)
    loop (Bind Variable a
v View (Term abt) xs b
b) = do
      Variable a
fresh <- Variable a -> Uniquifier (Variable a)
forall (m :: * -> *) (a :: Hakaru).
(Functor m, Gensym m) =>
Variable a -> m (Variable a)
freshVar Variable a
v
      let assoc :: Assoc Variable
assoc = Variable a -> Variable a -> Assoc Variable
forall k (ast :: k -> *) (a :: k). Variable a -> ast a -> Assoc ast
Assoc Variable a
v Variable a
fresh
      -- Process the body with the updated Varmap and wrap the
      -- result in a bind form
      Variable a -> abt xs b -> abt (a : xs) b
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
fresh (abt xs b -> abt (a : xs) b)
-> Uniquifier (abt xs b) -> Uniquifier (abt (a : xs) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Varmap -> Varmap)
-> Uniquifier (abt xs b) -> Uniquifier (abt xs b)
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Assoc Variable -> Varmap -> Varmap
forall k (ast :: k -> *). Assoc ast -> Assocs ast -> Assocs ast
insertAssoc Assoc Variable
assoc) (View (Term abt) xs b -> Uniquifier (abt xs b)
forall (ys :: [Hakaru]) (b :: Hakaru).
View (Term abt) ys b -> Uniquifier (abt ys b)
loop View (Term abt) xs b
b)

uniquifyVar
  :: (ABT Term abt)
  => Variable a
  -> Uniquifier (abt '[] a)
uniquifyVar :: Variable a -> Uniquifier (abt '[] a)
uniquifyVar Variable a
v = (Variable a -> abt '[] a
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k).
ABT syn abt =>
Variable a -> abt '[] a
var (Variable a -> abt '[] a)
-> (Varmap -> Variable a) -> Varmap -> abt '[] a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) (Varmap -> abt '[] a)
-> Uniquifier Varmap -> Uniquifier (abt '[] a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Uniquifier Varmap
forall r (m :: * -> *). MonadReader r m => m r
ask