{-# LANGUAGE AllowAmbiguousTypes #-} -- For ShowLetName
{-# LANGUAGE BangPatterns #-} -- For makeSharingName
{-# LANGUAGE DataKinds #-} -- For ShowLetName
{-# LANGUAGE ExistentialQuantification #-} -- For SharingName
-- {-# LANGUAGE MagicHash #-} -- For unsafeCoerce#
module Symantic.ObserveSharing 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 Data.String (String)
-- import GHC.Exts (Int(..))
-- import GHC.Prim (unsafeCoerce#)
import GHC.StableName (StableName(..), makeStableName, hashStableName, eqStableName)
-- import Numeric (showHex)
import Prelude ((+), error)
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 Control.Monad.Trans.Writer as MT
import qualified Data.HashMap.Strict as HM
import qualified Data.HashSet as HS

import Symantic.Derive

-- * Class 'Letable'
-- | This class is not for end-users like usual symantic operators,
-- here 'shareable' and 'ref' are introduced by 'observeSharing'.
class Letable letName repr where
  -- | @('ref' isRec letName)@ is a reference to @(letName)@.
  -- @(isRec)@ is 'True' iif. this 'ref'erence is recursive,
  -- ie. is reachable within its 'shareable' definition.
  ref :: Bool -> letName -> repr a
  ref isRec n = liftDerived (ref isRec n)
  default ref ::
    FromDerived (Letable letName) repr =>
    Bool -> letName -> repr a

  -- | @('shareable' letName x)@ let-binds @(letName)@ to be equal to @(x)@.
  shareable :: letName -> repr a -> repr a
  shareable n = liftDerived1 (shareable n)
  default shareable ::
    FromDerived1 (Letable letName) repr =>
    letName -> repr a -> repr a

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

-- ** Type 'ShowLetName'
-- | Useful on golden unit tests because 'StableName's
-- change often when changing unrelated source code
-- or even when changing basic GHC or executable flags.
class ShowLetName (showName::Bool) letName where
  showLetName :: letName -> String
-- | Like 'Show'.
instance Show letName => ShowLetName 'True letName where
  showLetName = show
-- | Always return @"<hidden>"@,
instance ShowLetName 'False letName where
  showLetName _p = "<hidden>"

-- * 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'
newtype ObserveSharing letName repr a = ObserveSharing { unObserveSharing ::
  MT.ReaderT (HashSet SharingName)
             (MT.State (ObserveSharingState letName))
             (FinalizeSharing letName repr a) }

-- | Interpreter detecting some (Haskell embedded) @let@ definitions used at
-- least once and/or recursively, in order to replace them
-- with the 'shareable' and 'ref' combinators.
-- See [Type-safe observable sharing in Haskell](https://doi.org/10.1145/1596638.1596653)
--
-- Beware not to apply 'observeSharing' more than once on the same term
-- otherwise some 'shareable' introduced by the first call
-- would be removed by the second call.
observeSharing ::
  Eq letName =>
  Hashable letName =>
  Show letName =>
  ObserveSharing letName repr a ->
  WithSharing letName repr a
observeSharing (ObserveSharing m) =
  let (fs, st) = MT.runReaderT m mempty `MT.runState`
        ObserveSharingState
          { oss_refs = HM.empty
          , oss_recs = HS.empty
          } in
  let refs = HS.fromList $
        (`foldMap` oss_refs st) $ (\(letName, refCount) ->
          [letName | refCount > 0]) in
  --trace (show refs) $
  MT.runWriter $
    (`MT.runReaderT` refs) $
      unFinalizeSharing fs

-- ** Type 'SomeLet'
data SomeLet repr = forall a. SomeLet (repr a)

-- ** Type 'WithSharing'
type WithSharing letName repr a =
  (repr a, HM.HashMap letName (SomeLet repr))
{-
-- * Type 'WithSharing'
data WithSharing letName repr a = WithSharing
  { lets :: HM.HashMap letName (SomeLet repr)
  , body :: repr a
  }
mapWithSharing ::
  (forall v. repr v -> repr v) ->
  WithSharing letName repr a ->
  WithSharing letName repr a
mapWithSharing f ws = WithSharing
  { lets = (\(SomeLet repr) -> SomeLet (f repr)) <$> lets ws
  , body = f (body ws)
  }
-}

-- ** 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 =>
  Show 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) (shareable letName <$> m)
      else return $ ref False letName

type instance Derived (ObserveSharing letName repr) = FinalizeSharing letName repr
instance
  ( Letable letName repr
  , MakeLetName letName
  , Eq letName
  , Hashable letName
  , Show letName
  ) => LiftDerived (ObserveSharing letName repr) where
  liftDerived = observeSharingNode . ObserveSharing . return
instance
  ( Letable letName repr
  , MakeLetName letName
  , Eq letName
  , Hashable letName
  , Show letName
  ) => LiftDerived1 (ObserveSharing letName repr) where
  liftDerived1 f x = observeSharingNode $ ObserveSharing $
    f <$> unObserveSharing x
instance
  ( Letable letName repr
  , MakeLetName letName
  , Eq letName
  , Hashable letName
  , Show letName
  ) => LiftDerived2 (ObserveSharing letName repr) where
  liftDerived2 f x y = observeSharingNode $ ObserveSharing $
    f <$> unObserveSharing x
      <*> unObserveSharing y
instance
  ( Letable letName repr
  , MakeLetName letName
  , Eq letName
  , Hashable letName
  , Show letName
  ) => LiftDerived3 (ObserveSharing letName repr) where
  liftDerived3 f x y z = observeSharingNode $ ObserveSharing $
    f <$> unObserveSharing x
      <*> unObserveSharing y
      <*> unObserveSharing z
instance Letable letName (ObserveSharing letName repr) where
  shareable = error "[BUG]: observeSharing MUST NOT be applied twice"
  ref = error "[BUG]: observeSharing MUST NOT be applied twice"
instance Letsable letName (ObserveSharing letName repr) where
  lets = error "[BUG]: observeSharing MUST NOT be applied twice"

-- * Type 'FinalizeSharing'
-- | Remove 'shareable' when non-recursive or unused
-- or replace it by 'ref', moving 'shareable's to the top.
newtype FinalizeSharing letName repr a = FinalizeSharing { unFinalizeSharing ::
  MT.ReaderT (HS.HashSet letName)
    (MT.Writer (LetBindings letName repr))
      (repr a) }

-- ** Type 'LetBindings'
type LetBindings letName repr = HM.HashMap letName (SomeLet repr)

type instance Derived (FinalizeSharing _letName repr) = repr
instance
  ( Eq letName
  , Hashable letName
  ) => LiftDerived (FinalizeSharing letName repr) where
  liftDerived = FinalizeSharing . pure
instance
  ( Eq letName
  , Hashable letName
  ) => LiftDerived1 (FinalizeSharing letName repr) where
  liftDerived1 f x = FinalizeSharing $ f <$> unFinalizeSharing x
instance
  ( Eq letName
  , Hashable letName
  ) => LiftDerived2 (FinalizeSharing letName repr) where
  liftDerived2 f x y = FinalizeSharing $
    f <$> unFinalizeSharing x
      <*> unFinalizeSharing y
instance
  ( Eq letName
  , Hashable letName
  ) => LiftDerived3 (FinalizeSharing letName repr) where
  liftDerived3 f x y z = FinalizeSharing $
    f <$> unFinalizeSharing x
      <*> unFinalizeSharing y
      <*> unFinalizeSharing z
instance
  ( Letable letName repr
  , Eq letName
  , Hashable letName
  , Show letName
  ) => Letable letName (FinalizeSharing letName repr) where
  shareable name x = FinalizeSharing $ do
    refs <- MT.ask
    if name `HS.member` refs
      -- This 'shareable' is 'ref'erenced, move it into the result,
      -- to put it in scope even when some 'ref' to it exists outside of 'x'
      -- (which can happen when a sub-expression is shared),
      -- and replace it by a 'ref'.
      then do
        let (repr, defs) = MT.runWriter $ MT.runReaderT (unFinalizeSharing x) refs
        MT.lift $ MT.tell $ HM.insert name (SomeLet repr) defs
        return $ ref False name
      -- Remove 'shareable'.
      else
        unFinalizeSharing x

-- * Class 'Letsable'
class Letsable letName repr where
  -- | @('lets' defs x)@ let-binds @(defs)@ in @(x)@.
  lets :: LetBindings letName repr -> repr a -> repr a
  lets defs = liftDerived1 (lets ((\(SomeLet val) -> SomeLet (derive val)) <$> defs))
  default lets ::
    Derivable repr =>
    FromDerived1 (Letsable letName) repr =>
    LetBindings letName repr -> repr a -> repr a
{-
-- | Not used but can be written nonetheless.
instance
  ( Letsable letName repr
  , Eq letName
  , Hashable letName
  , Show letName
  ) => Letsable letName (FinalizeSharing letName repr) where
  lets defs x = FinalizeSharing $ do
    ds <- traverse (\(SomeLet v) -> do
      r <- unFinalizeSharing v
      return (SomeLet r)
      ) defs
    MT.lift $ MT.tell ds
    unFinalizeSharing x
-}