{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ImpredicativeTypes #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

module Data.Reify.GADT (
    MuRef (..),
    reifyGraph,
    reifyGraphs,
    module Data.Reify.GADT.Graph,
) where

import Control.Concurrent
import Data.HashMap.Lazy
import qualified Data.HashMap.Lazy as HM
import Data.Hashable
import Data.IntSet
import qualified Data.IntSet as IS
import Data.Kind
import Data.Reify.GADT.Graph
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 :: Type -> Type) where
    type DeRef a :: (Type -> Type) -> (Type -> Type)
    type SubNode a :: Type -> Type
    mapDeRef ::
        (Applicative f, e ~ SubNode a) =>
        (forall t'. (MuRef e) => e t' -> f (u t')) ->
        a t ->
        f ((DeRef a) u t)

-- | '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, SubNode (SubNode s) ~ SubNode s, DeRef (SubNode s) ~ DeRef s) =>
    s a ->
    IO (Graph (DeRef s) a)
reifyGraph :: forall (s :: * -> *) a.
(MuRef s, SubNode (SubNode s) ~ SubNode s,
 DeRef (SubNode s) ~ DeRef s) =>
s a -> IO (Graph (DeRef s) a)
reifyGraph s a
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, SubNode (SubNode s) ~ SubNode s, DeRef (SubNode s) ~ DeRef s, Traversable t) =>
    t (s e) ->
    IO (t (Graph (DeRef s) e))
reifyGraphs :: forall (s :: * -> *) (t :: * -> *) e.
(MuRef s, SubNode (SubNode s) ~ SubNode s,
 DeRef (SubNode s) ~ DeRef s, Traversable t) =>
t (s e) -> IO (t (Graph (DeRef s) e))
reifyGraphs t (s e)
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, SubNode (SubNode s) ~ SubNode s, DeRef (SubNode s) ~ DeRef s) =>
    MVar (HashMap DynStableName Unique) ->
    MVar Unique ->
    s a ->
    IO (Graph (DeRef s) a)
reifyWithContext :: forall (s :: * -> *) a.
(MuRef s, SubNode (SubNode s) ~ SubNode s,
 DeRef (SubNode s) ~ DeRef s) =>
MVar (HashMap DynStableName Unique)
-> MVar Unique -> s a -> IO (Graph (DeRef s) a)
reifyWithContext MVar (HashMap DynStableName Unique)
rt1 MVar Unique
uVar s a
j = do
    rt2 <- [(Unique, Node (DeRef s))] -> IO (MVar [(Unique, Node (DeRef s))])
forall a. a -> IO (MVar a)
newMVar []
    nodeSetVar <- newMVar IS.empty
    Terminal root <- findNodes rt1 rt2 uVar nodeSetVar j
    pairs <- readMVar rt2
    return (Graph pairs root)

-- The workhorse for 'reifyGraph' and 'reifyGraphs'.
findNodes ::
    (MuRef s, SubNode (SubNode s) ~ SubNode s, DeRef (SubNode s) ~ DeRef s) =>
    -- | 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 (HashMap DynStableName 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,DeRef s Unique)]
    MVar [(Unique, Node (DeRef s))] ->
    -- | A supply of unique names.
    MVar Unique ->
    -- | The unique numbers that we have encountered so far.
    --   Invariant: this set is a subset of the range of the map of
    --   stable names.
    MVar IntSet ->
    -- | The value for which we will reify a 'Graph'.
    s t ->
    -- | The unique number for the value above.
    IO (Terminal t)
findNodes :: forall (s :: * -> *) t.
(MuRef s, SubNode (SubNode s) ~ SubNode s,
 DeRef (SubNode s) ~ DeRef s) =>
MVar (HashMap DynStableName Unique)
-> MVar [(Unique, Node (DeRef s))]
-> MVar Unique
-> MVar IntSet
-> s t
-> IO (Terminal t)
findNodes MVar (HashMap DynStableName Unique)
rt1 MVar [(Unique, Node (DeRef s))]
rt2 MVar Unique
uVar MVar IntSet
nodeSetVar !s t
j = do
    st <- s t -> IO DynStableName
forall a. a -> IO DynStableName
makeDynStableName s t
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
                    Terminal t -> IO (Terminal t)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Terminal t -> IO (Terminal t)) -> Terminal t -> IO (Terminal t)
forall a b. (a -> b) -> a -> b
$ Unique -> Terminal t
forall a. Unique -> Terminal a
Terminal Unique
var
                else Unique -> IntSet -> IO (Terminal t)
forall a. Unique -> IntSet -> IO (Terminal a)
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 (Terminal a)
    recurse :: forall a. Unique -> IntSet -> IO (Terminal a)
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 t'. MuRef (SubNode s) => SubNode s t' -> IO (Terminal t'))
-> s t -> IO (DeRef s Terminal t)
forall (a :: * -> *) (f :: * -> *) (e :: * -> *) (u :: * -> *) t.
(MuRef a, Applicative f, e ~ SubNode a) =>
(forall t'. MuRef e => e t' -> f (u t')) -> a t -> f (DeRef a u t)
forall (f :: * -> *) (e :: * -> *) (u :: * -> *) t.
(Applicative f, e ~ SubNode s) =>
(forall t'. MuRef e => e t' -> f (u t')) -> s t -> f (DeRef s u t)
mapDeRef (MVar (HashMap DynStableName Unique)
-> MVar [(Unique, Node (DeRef (SubNode s)))]
-> MVar Unique
-> MVar IntSet
-> SubNode s t'
-> IO (Terminal t')
forall (s :: * -> *) t.
(MuRef s, SubNode (SubNode s) ~ SubNode s,
 DeRef (SubNode s) ~ DeRef s) =>
MVar (HashMap DynStableName Unique)
-> MVar [(Unique, Node (DeRef s))]
-> MVar Unique
-> MVar IntSet
-> s t
-> IO (Terminal t)
findNodes MVar (HashMap DynStableName Unique)
rt1 MVar [(Unique, Node (DeRef s))]
MVar [(Unique, Node (DeRef (SubNode s)))]
rt2 MVar Unique
uVar MVar IntSet
nodeSetVar) s t
j
        tab' <- takeMVar rt2
        putMVar rt2 $ (var, MkNode res) : tab'
        return $ Terminal 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