{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE ExistentialQuantification #-} -- For SharingName
{-# LANGUAGE MagicHash #-} -- For unsafeCoerce#
module Symantic.Univariant.Letable where

import Control.Applicative (Applicative(..))
import Control.Monad (Monad(..))
import Data.Bool (Bool(..))
import Data.Eq (Eq(..))
import Data.Function (($), (.))
import Data.Functor ((<$>))
import Data.Functor.Compose (Compose(..))
import Data.HashMap.Strict (HashMap)
import Data.HashSet (HashSet)
import Data.Hashable (Hashable, hashWithSalt, hash)
import Data.Maybe (Maybe(..), isNothing)
import Data.Monoid (Monoid(..))
import Data.Ord (Ord(..))
import Data.Tuple (fst)
import GHC.Exts (Int(..))
import GHC.Prim (unsafeCoerce#)
import GHC.StableName (StableName(..), makeStableName, hashStableName, eqStableName)
import Numeric (showHex)
import Prelude ((+))
import System.IO (IO)
import Text.Show (Show(..))
import qualified Control.Monad.Trans.Class as MT
import qualified Control.Monad.Trans.Reader as MT
import qualified Control.Monad.Trans.State as MT
import qualified Data.HashMap.Strict as HM
import qualified Data.HashSet as HS
import qualified Data.List as List

import Symantic.Univariant.Liftable

-- * Class 'Letable'
-- | This class is not for manual usage like usual symantic operators, here 'def' and 'ref' are introduced by 'observeSharing'.
class Letable letName repr where
  -- | @('def' letName x)@ let-binds @(letName)@ to be equal to @(x)@.
  def :: letName -> repr a -> repr a
  -- | @('ref' isRec letName)@ is a reference to @(letName)@. @(isRec)@ is 'True' iif. this 'ref'erence is recursive, ie. is reachable within its 'def'inition.
  ref :: Bool -> letName -> repr a
  default def ::
    Liftable repr => Letable letName (Unlift repr) =>
    letName -> repr a -> repr a
  default ref ::
    Liftable repr => Letable letName (Unlift repr) =>
    Bool -> letName -> repr a
  def n = lift1 (def n)
  ref r n = lift (ref r n)

-- * Class 'MakeLetName'
class MakeLetName letName where
  makeLetName :: SharingName -> IO letName

-- * Type 'SharingName'
-- | Note that the observable sharing enabled by 'StableName'
-- is not perfect as it will not observe all the sharing explicitely done.
data SharingName = forall a. SharingName (StableName a)
-- | @('makeSharingName' x)@ is like @('makeStableName' x)@ but it also forces evaluation of @(x)@ to ensure that the 'StableName' is correct first time, which avoids to produce a tree bigger than needed.
makeSharingName :: a -> IO SharingName
makeSharingName !x = SharingName <$> makeStableName x
instance Eq SharingName where
  SharingName n == SharingName m = eqStableName n m
instance Hashable SharingName where
  hash (SharingName n) = hashStableName n
  hashWithSalt salt (SharingName n) = hashWithSalt salt n
instance Show SharingName where
  showsPrec _ (SharingName n) = showHex (I# (unsafeCoerce# n))

-- * Type 'ObserveSharing'
-- | Interpreter detecting some (Haskell embedded) @let@ definitions used at least once and/or recursively, in order to replace them with the 'def' and 'ref' combinators.
-- See [Type-safe observable sharing in Haskell](https://doi.org/10.1145/1596638.1596653)
newtype ObserveSharing letName repr a = ObserveSharing { unObserveSharing ::
  MT.ReaderT (HashSet SharingName)
             (MT.StateT (ObserveSharingState letName) IO)
             (CleanDefs letName repr a) }

observeSharing ::
  Eq letName =>
  Hashable letName =>
  ObserveSharing letName repr a -> IO (repr a)
observeSharing (ObserveSharing m) = do
  (a, st) <- MT.runReaderT m mempty `MT.runStateT`
    ObserveSharingState
      { oss_refs = HM.empty
      , oss_recs = HS.empty
      }
  let refs = HS.fromList $
        (fst <$>) $
        List.filter (\(_letName, refCount) -> refCount > 0) $
        HM.elems $ oss_refs st
  return $
    -- trace (show refs) $
    unCleanDefs a refs

-- ** Type 'ObserveSharingState'
data ObserveSharingState letName = ObserveSharingState
  { oss_refs :: HashMap SharingName (letName, Int)
  , oss_recs :: HashSet SharingName
  }

observeSharingNode ::
  Eq letName =>
  Hashable letName =>
  Letable letName repr =>
  MakeLetName letName =>
  ObserveSharing letName repr a -> ObserveSharing letName repr a
observeSharingNode node@(ObserveSharing m) = ObserveSharing $ do
  nodeName <- MT.lift $ MT.lift $ makeSharingName node
  st <- MT.lift MT.get
  ((letName, before), preds) <- getCompose $ HM.alterF (\before ->
    Compose $ case before of
      Nothing -> do
        letName <- MT.lift $ MT.lift $ makeLetName nodeName
        return ((letName, before), Just (letName, 0))
      Just (letName, refCount) -> do
        return ((letName, before), Just (letName, refCount + 1))
    ) nodeName (oss_refs st)
  parentNames <- MT.ask
  if nodeName `HS.member` parentNames
  then do
    MT.lift $ MT.put st
      { oss_refs = preds
      , oss_recs = HS.insert nodeName (oss_recs st)
      }
    return $ ref True letName
  else do
    MT.lift $ MT.put st{ oss_refs = preds }
    if isNothing before
     then MT.local (HS.insert nodeName) (def letName <$> m)
     else return $ ref False letName

type instance Unlift (ObserveSharing letName repr) = CleanDefs letName repr
instance
  ( Letable letName repr
  , MakeLetName letName
  , Eq letName
  , Hashable letName
  ) => Liftable (ObserveSharing letName repr) where
  lift x = observeSharingNode (ObserveSharing (return x))
  lift1 f x = observeSharingNode (ObserveSharing (f <$> unObserveSharing x))
  lift2 f x y = observeSharingNode $ ObserveSharing $
    f <$> unObserveSharing x
      <*> unObserveSharing y
  lift3 f x y z = observeSharingNode $ ObserveSharing $
    f <$> unObserveSharing x
      <*> unObserveSharing y
      <*> unObserveSharing z

-- * Type 'CleanDefs'
-- | Remove 'def' when non-recursive or unused.
newtype CleanDefs letName repr a = CleanDefs { unCleanDefs ::
  HS.HashSet letName -> repr a }

type instance Unlift (CleanDefs letName repr) = repr
instance Liftable (CleanDefs letName repr) where
  lift = CleanDefs . pure
  lift1 f x = CleanDefs $
    f <$> unCleanDefs x
  lift2 f x y = CleanDefs $
    f <$> unCleanDefs x
      <*> unCleanDefs y
  lift3 f x y z = CleanDefs $
    f <$> unCleanDefs x
      <*> unCleanDefs y
      <*> unCleanDefs z
instance
  ( Letable letName repr
  , Eq letName
  , Hashable letName
  ) => Letable letName (CleanDefs letName repr) where
  def name x = CleanDefs $ \refs ->
    if name `HS.member` refs
    then -- Perserve 'def'
      def name $ unCleanDefs x refs
    else -- Remove 'def'
      unCleanDefs x refs