{-# 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.Foldable (foldMap)
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.Int (Int)
import Data.Maybe (Maybe(..), isNothing)
import Data.Monoid (Monoid(..))
import Data.Ord (Ord(..))
-- 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 System.IO.Unsafe (unsafePerformIO)
-- 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 Symantic.Univariant.Trans

-- import Debug.Trace (trace)

-- * 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 ::
    Liftable1 repr => Letable letName (Output repr) =>
    letName -> repr a -> repr a
  default ref ::
    Liftable repr => Letable letName (Output 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.
--
-- Note also that the observed sharing could be different between ghc and ghci.
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.
--
-- Note that this function uses 'unsafePerformIO' instead of returning in 'IO',
-- this is apparently required to avoid infinite loops due to unstable 'StableName'
-- in compiled code, and sometimes also in ghci.
--
-- Note that maybe [pseq should be used here](https://gitlab.haskell.org/ghc/ghc/-/issues/2916).
makeSharingName :: a -> SharingName
makeSharingName !x = SharingName $ unsafePerformIO $ makeStableName x

instance Eq SharingName where
  SharingName x == SharingName y = eqStableName x y
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.State (ObserveSharingState letName))
             (CleanDefs letName repr a) }

observeSharing ::
  Eq letName =>
  Hashable letName =>
  ObserveSharing letName repr a -> repr a
observeSharing (ObserveSharing m) = do
  let (a, st) = MT.runReaderT m mempty `MT.runState`
        ObserveSharingState
          { oss_refs = HM.empty
          , oss_recs = HS.empty
          }
  let refs = HS.fromList $
        foldMap (\(letName, refCount) -> if refCount > 0 then [letName] else []) $
        HM.elems $ oss_refs st
  -- trace (show refs) $
  unCleanDefs a refs

-- ** Type 'ObserveSharingState'
data ObserveSharingState letName = ObserveSharingState
  { oss_refs :: HashMap SharingName (letName, Int)
  , oss_recs :: HashSet SharingName
    -- ^ TODO: unused so far, will it be useful somewhere at a later stage?
  }

observeSharingNode ::
  Eq letName =>
  Hashable letName =>
  Letable letName repr =>
  MakeLetName letName =>
  ObserveSharing letName repr a -> ObserveSharing letName repr a
observeSharingNode (ObserveSharing m) = ObserveSharing $ do
  let nodeName = makeSharingName m
  st <- MT.lift MT.get
  ((letName, before), preds) <- getCompose $ HM.alterF (\before ->
    Compose $ case before of
      Nothing -> do
        let letName = unsafePerformIO $ 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 Output (ObserveSharing letName repr) = CleanDefs letName repr
instance
  ( Letable letName repr
  , MakeLetName letName
  , Eq letName
  , Hashable letName
  ) => Trans (CleanDefs letName repr) (ObserveSharing letName repr) where
  trans = observeSharingNode . ObserveSharing . return
instance
  ( Letable letName repr
  , MakeLetName letName
  , Eq letName
  , Hashable letName
  ) => Trans1 (CleanDefs letName repr) (ObserveSharing letName repr) where
  trans1 f x = observeSharingNode $ ObserveSharing $
    f <$> unObserveSharing x
instance
  ( Letable letName repr
  , MakeLetName letName
  , Eq letName
  , Hashable letName
  ) => Trans2 (CleanDefs letName repr) (ObserveSharing letName repr) where
  trans2 f x y = observeSharingNode $ ObserveSharing $
    f <$> unObserveSharing x
      <*> unObserveSharing y
instance
  ( Letable letName repr
  , MakeLetName letName
  , Eq letName
  , Hashable letName
  ) => Trans3 (CleanDefs letName repr) (ObserveSharing letName repr) where
  trans3 f x y z = observeSharingNode $ ObserveSharing $
    f <$> unObserveSharing x
      <*> unObserveSharing y
      <*> unObserveSharing z
instance
  ( Letable letName repr
  , MakeLetName letName
  , Eq letName
  , Hashable letName
  ) => Letable letName (ObserveSharing letName repr)

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

type instance Output (CleanDefs _letName repr) = repr
instance Trans repr (CleanDefs letName repr) where
  trans = CleanDefs . pure
instance Trans1 repr (CleanDefs letName repr) where
  trans1 f x = CleanDefs $ f <$> unCleanDefs x
instance Trans2 repr (CleanDefs letName repr) where
  trans2 f x y = CleanDefs $
    f <$> unCleanDefs x
      <*> unCleanDefs y
instance Trans3 repr (CleanDefs letName repr) where
  trans3 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