    {-# 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