{-# 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 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
                  uVar <- newMVar 0
                  reifyWithContext rt1 uVar 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 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
                      uVar <- newMVar 0
                      traverse (reifyWithContext rt1 uVar) 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
  rt2 <- [(Unique, DeRef s Unique)] -> IO (MVar [(Unique, DeRef s Unique)])
forall a. a -> IO (MVar a)
newMVar []
  nodeSetVar <- newMVar IS.empty
  root <- findNodes rt1 rt2 uVar nodeSetVar j
  pairs <- readMVar rt2
  return (Graph pairs 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
        st <- s -> IO DynStableName
forall a. a -> IO DynStableName
makeDynStableName s
j
        tab <- takeMVar rt1
        nodeSet <- takeMVar nodeSetVar
        case HM.lookup st 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 var <- MVar Unique -> IO Unique
newUnique MVar Unique
uVar
                        putMVar rt1 $ HM.insert st var tab
                        recurse var 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
      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
      tab' <- takeMVar rt2
      putMVar rt2 $ (var,res) : tab'
      return var

newUnique :: MVar Unique -> IO Unique
newUnique :: MVar Unique -> IO Unique
newUnique MVar Unique
var = do
  v <- MVar Unique -> IO Unique
forall a. MVar a -> IO a
takeMVar MVar Unique
var
  let v' = Unique -> Unique
forall a. Enum a => a -> a
succ Unique
v
  putMVar var v'
  return 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
    st <- a -> IO (StableName a)
forall a. a -> IO (StableName a)
makeStableName a
a
    return $ DynStableName st