{-# LANGUAGE CPP
           , Rank2Types
           , MultiParamTypeClasses
           , UndecidableInstances
           , FlexibleInstances
           #-}
{-# OPTIONS_GHC -Wall -fwarn-tabs #-}
----------------------------------------------------------------
--                                                  ~ 2015.03.29
-- |
-- Module      :  Control.Unification.Ranked.STVar
-- Copyright   :  Copyright (c) 2007--2015 wren gayle romano
-- License     :  BSD
-- Maintainer  :  wren@community.haskell.org
-- Stability   :  highly experimental
-- Portability :  semi-portable (Rank2Types, MPTCs,...)
--
-- A ranked variant of "Control.Unification.STVar".
----------------------------------------------------------------
module Control.Unification.Ranked.STVar
    ( STRVar()
    , STRBinding()
    , runSTRBinding
    ) where

import Prelude hiding (mapM, sequence, foldr, foldr1, foldl, foldl1)

import Data.STRef
import Data.Word            (Word8)
#if __GLASGOW_HASKELL__ < 710
import Control.Applicative  (Applicative(..))
#endif
import Control.Monad        (ap)
import Control.Monad.Trans  (lift)
import Control.Monad.ST
import Control.Monad.Reader (ReaderT, runReaderT, ask)
import Control.Unification.Types
----------------------------------------------------------------
----------------------------------------------------------------

-- | A ranked unification variable implemented by 'STRef's. In
-- addition to the @STRef@ for the term itself, we also track the
-- variable's ID (to support visited-sets) and rank (to support
-- weighted path compression).
data STRVar s t =
    STRVar
        {-# UNPACK #-} !Int
        {-# UNPACK #-} !(STRef s Word8)
        {-# UNPACK #-} !(STRef s (Maybe (UTerm t (STRVar s t))))

instance Show (STRVar s t) where
    show :: STRVar s t -> String
show (STRVar Int
i STRef s Word8
_ STRef s (Maybe (UTerm t (STRVar s t)))
_) = String
"STRVar " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
i

instance Eq (STRVar s t) where
    (STRVar Int
i STRef s Word8
_ STRef s (Maybe (UTerm t (STRVar s t)))
_) == :: STRVar s t -> STRVar s t -> Bool
== (STRVar Int
j STRef s Word8
_ STRef s (Maybe (UTerm t (STRVar s t)))
_) = (Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j)
    
instance Variable (STRVar s t) where
    getVarID :: STRVar s t -> Int
getVarID (STRVar Int
i STRef s Word8
_ STRef s (Maybe (UTerm t (STRVar s t)))
_) = Int
i


----------------------------------------------------------------
-- TODO: parameterize this so we can use BacktrackST too. Of course,
-- that means defining another class for STRef-like variables
--
-- TODO: parameterize this so we can share the implementation for STVar and STRVar
--
-- TODO: does MTL still have the overhead that'd make it worthwhile
-- to do this manually instead of using ReaderT?
--
-- | A monad for handling 'STRVar' bindings.
newtype STRBinding s a = STRB { STRBinding s a -> ReaderT (STRef s Int) (ST s) a
unSTRB :: ReaderT (STRef s Int) (ST s) a }


-- | Run the 'ST' ranked binding monad. N.B., because 'STRVar' are
-- rank-2 quantified, this guarantees that the return value has no
-- such references. However, in order to remove the references from
-- terms, you'll need to explicitly apply the bindings.
runSTRBinding :: (forall s. STRBinding s a) -> a
runSTRBinding :: (forall s. STRBinding s a) -> a
runSTRBinding forall s. STRBinding s a
stb =
    (forall s. ST s a) -> a
forall a. (forall s. ST s a) -> a
runST (Int -> ST s (STRef s Int)
forall a s. a -> ST s (STRef s a)
newSTRef Int
forall a. Bounded a => a
minBound ST s (STRef s Int) -> (STRef s Int -> ST s a) -> ST s a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ReaderT (STRef s Int) (ST s) a -> STRef s Int -> ST s a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (STRBinding s a -> ReaderT (STRef s Int) (ST s) a
forall s a. STRBinding s a -> ReaderT (STRef s Int) (ST s) a
unSTRB STRBinding s a
forall s. STRBinding s a
stb))


-- For portability reasons, we're intentionally avoiding
-- -XDeriveFunctor, -XGeneralizedNewtypeDeriving, and the like.

instance Functor (STRBinding s) where
    fmap :: (a -> b) -> STRBinding s a -> STRBinding s b
fmap a -> b
f = ReaderT (STRef s Int) (ST s) b -> STRBinding s b
forall s a. ReaderT (STRef s Int) (ST s) a -> STRBinding s a
STRB (ReaderT (STRef s Int) (ST s) b -> STRBinding s b)
-> (STRBinding s a -> ReaderT (STRef s Int) (ST s) b)
-> STRBinding s a
-> STRBinding s b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b)
-> ReaderT (STRef s Int) (ST s) a -> ReaderT (STRef s Int) (ST s) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f (ReaderT (STRef s Int) (ST s) a -> ReaderT (STRef s Int) (ST s) b)
-> (STRBinding s a -> ReaderT (STRef s Int) (ST s) a)
-> STRBinding s a
-> ReaderT (STRef s Int) (ST s) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STRBinding s a -> ReaderT (STRef s Int) (ST s) a
forall s a. STRBinding s a -> ReaderT (STRef s Int) (ST s) a
unSTRB

instance Applicative (STRBinding s) where
    pure :: a -> STRBinding s a
pure   = a -> STRBinding s a
forall (m :: * -> *) a. Monad m => a -> m a
return
    <*> :: STRBinding s (a -> b) -> STRBinding s a -> STRBinding s b
(<*>)  = STRBinding s (a -> b) -> STRBinding s a -> STRBinding s b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap
    *> :: STRBinding s a -> STRBinding s b -> STRBinding s b
(*>)   = STRBinding s a -> STRBinding s b -> STRBinding s b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
(>>)
    STRBinding s a
x <* :: STRBinding s a -> STRBinding s b -> STRBinding s a
<* STRBinding s b
y = STRBinding s a
x STRBinding s a -> (a -> STRBinding s a) -> STRBinding s a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
a -> STRBinding s b
y STRBinding s b -> STRBinding s a -> STRBinding s a
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> a -> STRBinding s a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a

instance Monad (STRBinding s) where
    return :: a -> STRBinding s a
return    = ReaderT (STRef s Int) (ST s) a -> STRBinding s a
forall s a. ReaderT (STRef s Int) (ST s) a -> STRBinding s a
STRB (ReaderT (STRef s Int) (ST s) a -> STRBinding s a)
-> (a -> ReaderT (STRef s Int) (ST s) a) -> a -> STRBinding s a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ReaderT (STRef s Int) (ST s) a
forall (m :: * -> *) a. Monad m => a -> m a
return
    STRBinding s a
stb >>= :: STRBinding s a -> (a -> STRBinding s b) -> STRBinding s b
>>= a -> STRBinding s b
f = ReaderT (STRef s Int) (ST s) b -> STRBinding s b
forall s a. ReaderT (STRef s Int) (ST s) a -> STRBinding s a
STRB (STRBinding s a -> ReaderT (STRef s Int) (ST s) a
forall s a. STRBinding s a -> ReaderT (STRef s Int) (ST s) a
unSTRB STRBinding s a
stb ReaderT (STRef s Int) (ST s) a
-> (a -> ReaderT (STRef s Int) (ST s) b)
-> ReaderT (STRef s Int) (ST s) b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= STRBinding s b -> ReaderT (STRef s Int) (ST s) b
forall s a. STRBinding s a -> ReaderT (STRef s Int) (ST s) a
unSTRB (STRBinding s b -> ReaderT (STRef s Int) (ST s) b)
-> (a -> STRBinding s b) -> a -> ReaderT (STRef s Int) (ST s) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> STRBinding s b
f)


----------------------------------------------------------------

_newSTRVar
    :: String
    -> Maybe (UTerm t (STRVar s t))
    -> STRBinding s (STRVar s t)
_newSTRVar :: String -> Maybe (UTerm t (STRVar s t)) -> STRBinding s (STRVar s t)
_newSTRVar String
fun Maybe (UTerm t (STRVar s t))
mb = ReaderT (STRef s Int) (ST s) (STRVar s t)
-> STRBinding s (STRVar s t)
forall s a. ReaderT (STRef s Int) (ST s) a -> STRBinding s a
STRB (ReaderT (STRef s Int) (ST s) (STRVar s t)
 -> STRBinding s (STRVar s t))
-> ReaderT (STRef s Int) (ST s) (STRVar s t)
-> STRBinding s (STRVar s t)
forall a b. (a -> b) -> a -> b
$ do
    STRef s Int
nr <- ReaderT (STRef s Int) (ST s) (STRef s Int)
forall r (m :: * -> *). MonadReader r m => m r
ask
    ST s (STRVar s t) -> ReaderT (STRef s Int) (ST s) (STRVar s t)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ST s (STRVar s t) -> ReaderT (STRef s Int) (ST s) (STRVar s t))
-> ST s (STRVar s t) -> ReaderT (STRef s Int) (ST s) (STRVar s t)
forall a b. (a -> b) -> a -> b
$ do
        Int
n <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef STRef s Int
nr
        if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
forall a. Bounded a => a
maxBound
            then String -> ST s (STRVar s t)
forall a. HasCallStack => String -> a
error (String -> ST s (STRVar s t)) -> String -> ST s (STRVar s t)
forall a b. (a -> b) -> a -> b
$ String
fun String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
": no more variables!"
            else do
                STRef s Int -> Int -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s Int
nr (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$! Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1
                -- BUG: no applicative ST
                STRef s Word8
rk  <- Word8 -> ST s (STRef s Word8)
forall a s. a -> ST s (STRef s a)
newSTRef Word8
0
                STRef s (Maybe (UTerm t (STRVar s t)))
ptr <- Maybe (UTerm t (STRVar s t))
-> ST s (STRef s (Maybe (UTerm t (STRVar s t))))
forall a s. a -> ST s (STRef s a)
newSTRef Maybe (UTerm t (STRVar s t))
mb
                STRVar s t -> ST s (STRVar s t)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
-> STRef s Word8
-> STRef s (Maybe (UTerm t (STRVar s t)))
-> STRVar s t
forall s (t :: * -> *).
Int
-> STRef s Word8
-> STRef s (Maybe (UTerm t (STRVar s t)))
-> STRVar s t
STRVar Int
n STRef s Word8
rk STRef s (Maybe (UTerm t (STRVar s t)))
ptr)


instance (Unifiable t) => BindingMonad t (STRVar s t) (STRBinding s) where
    lookupVar :: STRVar s t -> STRBinding s (Maybe (UTerm t (STRVar s t)))
lookupVar (STRVar Int
_ STRef s Word8
_ STRef s (Maybe (UTerm t (STRVar s t)))
p) = ReaderT (STRef s Int) (ST s) (Maybe (UTerm t (STRVar s t)))
-> STRBinding s (Maybe (UTerm t (STRVar s t)))
forall s a. ReaderT (STRef s Int) (ST s) a -> STRBinding s a
STRB (ReaderT (STRef s Int) (ST s) (Maybe (UTerm t (STRVar s t)))
 -> STRBinding s (Maybe (UTerm t (STRVar s t))))
-> (ST s (Maybe (UTerm t (STRVar s t)))
    -> ReaderT (STRef s Int) (ST s) (Maybe (UTerm t (STRVar s t))))
-> ST s (Maybe (UTerm t (STRVar s t)))
-> STRBinding s (Maybe (UTerm t (STRVar s t)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ST s (Maybe (UTerm t (STRVar s t)))
-> ReaderT (STRef s Int) (ST s) (Maybe (UTerm t (STRVar s t)))
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ST s (Maybe (UTerm t (STRVar s t)))
 -> STRBinding s (Maybe (UTerm t (STRVar s t))))
-> ST s (Maybe (UTerm t (STRVar s t)))
-> STRBinding s (Maybe (UTerm t (STRVar s t)))
forall a b. (a -> b) -> a -> b
$ STRef s (Maybe (UTerm t (STRVar s t)))
-> ST s (Maybe (UTerm t (STRVar s t)))
forall s a. STRef s a -> ST s a
readSTRef STRef s (Maybe (UTerm t (STRVar s t)))
p
    
    freeVar :: STRBinding s (STRVar s t)
freeVar  = String -> Maybe (UTerm t (STRVar s t)) -> STRBinding s (STRVar s t)
forall (t :: * -> *) s.
String -> Maybe (UTerm t (STRVar s t)) -> STRBinding s (STRVar s t)
_newSTRVar String
"freeVar" Maybe (UTerm t (STRVar s t))
forall a. Maybe a
Nothing
    
    newVar :: UTerm t (STRVar s t) -> STRBinding s (STRVar s t)
newVar UTerm t (STRVar s t)
t = String -> Maybe (UTerm t (STRVar s t)) -> STRBinding s (STRVar s t)
forall (t :: * -> *) s.
String -> Maybe (UTerm t (STRVar s t)) -> STRBinding s (STRVar s t)
_newSTRVar String
"newVar" (UTerm t (STRVar s t) -> Maybe (UTerm t (STRVar s t))
forall a. a -> Maybe a
Just UTerm t (STRVar s t)
t)
    
    bindVar :: STRVar s t -> UTerm t (STRVar s t) -> STRBinding s ()
bindVar (STRVar Int
_ STRef s Word8
_ STRef s (Maybe (UTerm t (STRVar s t)))
p) UTerm t (STRVar s t)
t = ReaderT (STRef s Int) (ST s) () -> STRBinding s ()
forall s a. ReaderT (STRef s Int) (ST s) a -> STRBinding s a
STRB (ReaderT (STRef s Int) (ST s) () -> STRBinding s ())
-> (ST s () -> ReaderT (STRef s Int) (ST s) ())
-> ST s ()
-> STRBinding s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ST s () -> ReaderT (STRef s Int) (ST s) ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ST s () -> STRBinding s ()) -> ST s () -> STRBinding s ()
forall a b. (a -> b) -> a -> b
$ STRef s (Maybe (UTerm t (STRVar s t)))
-> Maybe (UTerm t (STRVar s t)) -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s (Maybe (UTerm t (STRVar s t)))
p (UTerm t (STRVar s t) -> Maybe (UTerm t (STRVar s t))
forall a. a -> Maybe a
Just UTerm t (STRVar s t)
t)


instance (Unifiable t) =>
    RankedBindingMonad t (STRVar s t) (STRBinding s)
    where
    
    lookupRankVar :: STRVar s t -> STRBinding s (Rank t (STRVar s t))
lookupRankVar (STRVar Int
_ STRef s Word8
r STRef s (Maybe (UTerm t (STRVar s t)))
p) = ReaderT (STRef s Int) (ST s) (Rank t (STRVar s t))
-> STRBinding s (Rank t (STRVar s t))
forall s a. ReaderT (STRef s Int) (ST s) a -> STRBinding s a
STRB (ReaderT (STRef s Int) (ST s) (Rank t (STRVar s t))
 -> STRBinding s (Rank t (STRVar s t)))
-> (ST s (Rank t (STRVar s t))
    -> ReaderT (STRef s Int) (ST s) (Rank t (STRVar s t)))
-> ST s (Rank t (STRVar s t))
-> STRBinding s (Rank t (STRVar s t))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ST s (Rank t (STRVar s t))
-> ReaderT (STRef s Int) (ST s) (Rank t (STRVar s t))
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ST s (Rank t (STRVar s t)) -> STRBinding s (Rank t (STRVar s t)))
-> ST s (Rank t (STRVar s t)) -> STRBinding s (Rank t (STRVar s t))
forall a b. (a -> b) -> a -> b
$ do
        Word8
n  <- STRef s Word8 -> ST s Word8
forall s a. STRef s a -> ST s a
readSTRef STRef s Word8
r
        Maybe (UTerm t (STRVar s t))
mb <- STRef s (Maybe (UTerm t (STRVar s t)))
-> ST s (Maybe (UTerm t (STRVar s t)))
forall s a. STRef s a -> ST s a
readSTRef STRef s (Maybe (UTerm t (STRVar s t)))
p
        Rank t (STRVar s t) -> ST s (Rank t (STRVar s t))
forall (m :: * -> *) a. Monad m => a -> m a
return (Word8 -> Maybe (UTerm t (STRVar s t)) -> Rank t (STRVar s t)
forall (t :: * -> *) v. Word8 -> Maybe (UTerm t v) -> Rank t v
Rank Word8
n Maybe (UTerm t (STRVar s t))
mb)
    
    incrementRank :: STRVar s t -> STRBinding s ()
incrementRank (STRVar Int
_ STRef s Word8
r STRef s (Maybe (UTerm t (STRVar s t)))
_) = ReaderT (STRef s Int) (ST s) () -> STRBinding s ()
forall s a. ReaderT (STRef s Int) (ST s) a -> STRBinding s a
STRB (ReaderT (STRef s Int) (ST s) () -> STRBinding s ())
-> (ST s () -> ReaderT (STRef s Int) (ST s) ())
-> ST s ()
-> STRBinding s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ST s () -> ReaderT (STRef s Int) (ST s) ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ST s () -> STRBinding s ()) -> ST s () -> STRBinding s ()
forall a b. (a -> b) -> a -> b
$ do
        Word8
n <- STRef s Word8 -> ST s Word8
forall s a. STRef s a -> ST s a
readSTRef STRef s Word8
r
        STRef s Word8 -> Word8 -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s Word8
r (Word8 -> ST s ()) -> Word8 -> ST s ()
forall a b. (a -> b) -> a -> b
$! Word8
nWord8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+Word8
1
    
    -- incrementBindVar = default

----------------------------------------------------------------
----------------------------------------------------------- fin.