{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Data.Reify (
        MuRef(..),
        module Data.Reify.Graph,
        reifyGraph,
        reifyGraphs
        ) where

import Control.Concurrent.MVar

import qualified Data.HashMap.Lazy as HM
import Data.HashMap.Lazy (HashMap)
import Data.Hashable as H
import Data.Reify.Graph
import qualified Data.IntSet as IS
import Data.IntSet (IntSet)

import System.Mem.StableName

-- | 'MuRef' is a class that provided a way to reference into a specific type,
-- and a way to map over the deferenced internals.
class MuRef a where
  type DeRef a :: * -> *

  mapDeRef :: (Applicative f) =>
              (forall b . (MuRef b, DeRef a ~ DeRef b) => b -> f u)
                        -> a
                        -> f (DeRef a u)

-- | 'reifyGraph' takes a data structure that admits 'MuRef', and returns a 'Graph' that contains
-- the dereferenced nodes, with their children as 'Unique's rather than recursive values.
reifyGraph :: (MuRef s) => s -> IO (Graph (DeRef s))
reifyGraph :: forall s. MuRef s => s -> IO (Graph (DeRef s))
reifyGraph s
m = do MVar (HashMap DynStableName Unique)
rt1 <- HashMap DynStableName Unique
-> IO (MVar (HashMap DynStableName Unique))
forall a. a -> IO (MVar a)
newMVar HashMap DynStableName Unique
forall k v. HashMap k v
HM.empty
                  MVar Unique
uVar <- Unique -> IO (MVar Unique)
forall a. a -> IO (MVar a)
newMVar Unique
0
                  MVar (HashMap DynStableName Unique)
-> MVar Unique -> s -> IO (Graph (DeRef s))
forall s.
MuRef s =>
MVar (HashMap DynStableName Unique)
-> MVar Unique -> s -> IO (Graph (DeRef s))
reifyWithContext MVar (HashMap DynStableName Unique)
rt1 MVar Unique
uVar s
m

-- | 'reifyGraphs' takes a 'Traversable' container 't s' of a data structure 's'
-- admitting 'MuRef', and returns a 't (Graph (DeRef s))' with the graph nodes
-- resolved within the same context.
--
-- This allows for, e.g., a list of mutually recursive structures.
reifyGraphs :: (MuRef s, Traversable t) => t s -> IO (t (Graph (DeRef s)))
reifyGraphs :: forall s (t :: * -> *).
(MuRef s, Traversable t) =>
t s -> IO (t (Graph (DeRef s)))
reifyGraphs t s
coll = do MVar (HashMap DynStableName Unique)
rt1 <- HashMap DynStableName Unique
-> IO (MVar (HashMap DynStableName Unique))
forall a. a -> IO (MVar a)
newMVar HashMap DynStableName Unique
forall k v. HashMap k v
HM.empty
                      MVar Unique
uVar <- Unique -> IO (MVar Unique)
forall a. a -> IO (MVar a)
newMVar Unique
0
                      (s -> IO (Graph (DeRef s))) -> t s -> IO (t (Graph (DeRef s)))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> t a -> f (t b)
traverse (MVar (HashMap DynStableName Unique)
-> MVar Unique -> s -> IO (Graph (DeRef s))
forall s.
MuRef s =>
MVar (HashMap DynStableName Unique)
-> MVar Unique -> s -> IO (Graph (DeRef s))
reifyWithContext MVar (HashMap DynStableName Unique)
rt1 MVar Unique
uVar) t s
coll
                        -- NB: We deliberately reuse the same map of stable
                        -- names and unique supply across all iterations of the
                        -- traversal to ensure that the same context is used
                        -- when reifying all elements of the container.

-- Reify a data structure's 'Graph' using the supplied map of stable names and
-- unique supply.
reifyWithContext :: (MuRef s)
                 => MVar (HashMap DynStableName Unique)
                 -> MVar Unique
                 -> s
                 -> IO (Graph (DeRef s))
reifyWithContext :: forall s.
MuRef s =>
MVar (HashMap DynStableName Unique)
-> MVar Unique -> s -> IO (Graph (DeRef s))
reifyWithContext MVar (HashMap DynStableName Unique)
rt1 MVar Unique
uVar s
j = do
  MVar [(Unique, DeRef s Unique)]
rt2 <- [(Unique, DeRef s Unique)] -> IO (MVar [(Unique, DeRef s Unique)])
forall a. a -> IO (MVar a)
newMVar []
  MVar IntSet
nodeSetVar <- IntSet -> IO (MVar IntSet)
forall a. a -> IO (MVar a)
newMVar IntSet
IS.empty
  Unique
root <- MVar (HashMap DynStableName Unique)
-> MVar [(Unique, DeRef s Unique)]
-> MVar Unique
-> MVar IntSet
-> s
-> IO Unique
forall s.
MuRef s =>
MVar (HashMap DynStableName Unique)
-> MVar [(Unique, DeRef s Unique)]
-> MVar Unique
-> MVar IntSet
-> s
-> IO Unique
findNodes MVar (HashMap DynStableName Unique)
rt1 MVar [(Unique, DeRef s Unique)]
rt2 MVar Unique
uVar MVar IntSet
nodeSetVar s
j
  [(Unique, DeRef s Unique)]
pairs <- MVar [(Unique, DeRef s Unique)] -> IO [(Unique, DeRef s Unique)]
forall a. MVar a -> IO a
readMVar MVar [(Unique, DeRef s Unique)]
rt2
  Graph (DeRef s) -> IO (Graph (DeRef s))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Unique, DeRef s Unique)] -> Unique -> Graph (DeRef s)
forall (e :: * -> *). [(Unique, e Unique)] -> Unique -> Graph e
Graph [(Unique, DeRef s Unique)]
pairs Unique
root)

-- The workhorse for 'reifyGraph' and 'reifyGraphs'.
findNodes :: (MuRef s)
          => MVar (HashMap DynStableName Unique)
             -- ^ A map of stable names to unique numbers.
             --   Invariant: all 'Uniques' that appear in the range are less
             --   than the current value in the unique name supply.
          -> MVar [(Unique,DeRef s Unique)]
             -- ^ The key-value pairs in the 'Graph' that is being built.
             --   Invariant 1: the domain of this association list is a subset
             --   of the range of the map of stable names.
             --   Invariant 2: the domain of this association list will never
             --   contain duplicate keys.
          -> MVar Unique
             -- ^ A supply of unique names.
          -> MVar IntSet
             -- ^ The unique numbers that we have encountered so far.
             --   Invariant: this set is a subset of the range of the map of
             --   stable names.
          -> s
             -- ^ The value for which we will reify a 'Graph'.
          -> IO Unique
             -- ^ The unique number for the value above.
findNodes :: forall s.
MuRef s =>
MVar (HashMap DynStableName Unique)
-> MVar [(Unique, DeRef s Unique)]
-> MVar Unique
-> MVar IntSet
-> s
-> IO Unique
findNodes MVar (HashMap DynStableName Unique)
rt1 MVar [(Unique, DeRef s Unique)]
rt2 MVar Unique
uVar MVar IntSet
nodeSetVar !s
j = do
        DynStableName
st <- s -> IO DynStableName
forall a. a -> IO DynStableName
makeDynStableName s
j
        HashMap DynStableName Unique
tab <- MVar (HashMap DynStableName Unique)
-> IO (HashMap DynStableName Unique)
forall a. MVar a -> IO a
takeMVar MVar (HashMap DynStableName Unique)
rt1
        IntSet
nodeSet <- MVar IntSet -> IO IntSet
forall a. MVar a -> IO a
takeMVar MVar IntSet
nodeSetVar
        case DynStableName -> HashMap DynStableName Unique -> Maybe Unique
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup DynStableName
st HashMap DynStableName Unique
tab of
          Just Unique
var -> do MVar (HashMap DynStableName Unique)
-> HashMap DynStableName Unique -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (HashMap DynStableName Unique)
rt1 HashMap DynStableName Unique
tab
                         if Unique
var Unique -> IntSet -> Bool
`IS.member` IntSet
nodeSet
                           then do MVar IntSet -> IntSet -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar IntSet
nodeSetVar IntSet
nodeSet
                                   Unique -> IO Unique
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Unique
var
                           else Unique -> IntSet -> IO Unique
recurse Unique
var IntSet
nodeSet
          Maybe Unique
Nothing -> do Unique
var <- MVar Unique -> IO Unique
newUnique MVar Unique
uVar
                        MVar (HashMap DynStableName Unique)
-> HashMap DynStableName Unique -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (HashMap DynStableName Unique)
rt1 (HashMap DynStableName Unique -> IO ())
-> HashMap DynStableName Unique -> IO ()
forall a b. (a -> b) -> a -> b
$ DynStableName
-> Unique
-> HashMap DynStableName Unique
-> HashMap DynStableName Unique
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert DynStableName
st Unique
var HashMap DynStableName Unique
tab
                        Unique -> IntSet -> IO Unique
recurse Unique
var IntSet
nodeSet
  where
    recurse :: Unique -> IntSet -> IO Unique
    recurse :: Unique -> IntSet -> IO Unique
recurse Unique
var IntSet
nodeSet = do
      MVar IntSet -> IntSet -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar IntSet
nodeSetVar (IntSet -> IO ()) -> IntSet -> IO ()
forall a b. (a -> b) -> a -> b
$ Unique -> IntSet -> IntSet
IS.insert Unique
var IntSet
nodeSet
      DeRef s Unique
res <- (forall b. (MuRef b, DeRef s ~ DeRef b) => b -> IO Unique)
-> s -> IO (DeRef s Unique)
forall a (f :: * -> *) u.
(MuRef a, Applicative f) =>
(forall b. (MuRef b, DeRef a ~ DeRef b) => b -> f u)
-> a -> f (DeRef a u)
forall (f :: * -> *) u.
Applicative f =>
(forall b. (MuRef b, DeRef s ~ DeRef b) => b -> f u)
-> s -> f (DeRef s u)
mapDeRef (MVar (HashMap DynStableName Unique)
-> MVar [(Unique, DeRef b Unique)]
-> MVar Unique
-> MVar IntSet
-> b
-> IO Unique
forall s.
MuRef s =>
MVar (HashMap DynStableName Unique)
-> MVar [(Unique, DeRef s Unique)]
-> MVar Unique
-> MVar IntSet
-> s
-> IO Unique
findNodes MVar (HashMap DynStableName Unique)
rt1 MVar [(Unique, DeRef s Unique)]
MVar [(Unique, DeRef b Unique)]
rt2 MVar Unique
uVar MVar IntSet
nodeSetVar) s
j
      [(Unique, DeRef s Unique)]
tab' <- MVar [(Unique, DeRef s Unique)] -> IO [(Unique, DeRef s Unique)]
forall a. MVar a -> IO a
takeMVar MVar [(Unique, DeRef s Unique)]
rt2
      MVar [(Unique, DeRef s Unique)]
-> [(Unique, DeRef s Unique)] -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar [(Unique, DeRef s Unique)]
rt2 ([(Unique, DeRef s Unique)] -> IO ())
-> [(Unique, DeRef s Unique)] -> IO ()
forall a b. (a -> b) -> a -> b
$ (Unique
var,DeRef s Unique
res) (Unique, DeRef s Unique)
-> [(Unique, DeRef s Unique)] -> [(Unique, DeRef s Unique)]
forall a. a -> [a] -> [a]
: [(Unique, DeRef s Unique)]
tab'
      Unique -> IO Unique
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Unique
var

newUnique :: MVar Unique -> IO Unique
newUnique :: MVar Unique -> IO Unique
newUnique MVar Unique
var = do
  Unique
v <- MVar Unique -> IO Unique
forall a. MVar a -> IO a
takeMVar MVar Unique
var
  let v' :: Unique
v' = Unique -> Unique
forall a. Enum a => a -> a
succ Unique
v
  MVar Unique -> Unique -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar Unique
var Unique
v'
  Unique -> IO Unique
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Unique
v'

-- Stable names that do not use phantom types.
-- As suggested by Ganesh Sittampalam.
-- Note: GHC can't unpack these because of the existential
-- quantification, but there doesn't seem to be much
-- potential to unpack them anyway.
data DynStableName = forall a. DynStableName !(StableName a)

instance Hashable DynStableName where
  hashWithSalt :: Unique -> DynStableName -> Unique
hashWithSalt Unique
s (DynStableName StableName a
n) = Unique -> StableName a -> Unique
forall a. Hashable a => Unique -> a -> Unique
hashWithSalt Unique
s StableName a
n

instance Eq DynStableName where
  DynStableName StableName a
m == :: DynStableName -> DynStableName -> Bool
== DynStableName StableName a
n =
    StableName a -> StableName a -> Bool
forall a b. StableName a -> StableName b -> Bool
eqStableName StableName a
m StableName a
n

makeDynStableName :: a -> IO DynStableName
makeDynStableName :: forall a. a -> IO DynStableName
makeDynStableName a
a = do
    StableName a
st <- a -> IO (StableName a)
forall a. a -> IO (StableName a)
makeStableName a
a
    DynStableName -> IO DynStableName
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (DynStableName -> IO DynStableName)
-> DynStableName -> IO DynStableName
forall a b. (a -> b) -> a -> b
$ StableName a -> DynStableName
forall a. StableName a -> DynStableName
DynStableName StableName a
st