{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE MagicHash #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Symantic.Parser.Grammar.ObserveSharing where

import Control.Applicative (Applicative(..))
import Control.Monad (Monad(..), mapM)
import Data.Bool (Bool(..))
import Data.Eq (Eq(..))
import Data.Function (($), (.))
import Data.Functor (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 Debug.Trace (trace)
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.List as List
import qualified Language.Haskell.TH.Syntax as TH

import qualified Data.HashMap.Strict as HM
import qualified Data.HashSet        as HS

import Symantic.Base.Univariant
import qualified Symantic.Parser.Grammar.Combinators as P

-- * Type 'SharingName'
data SharingName = forall a. SharingName (StableName a)
-- Force evaluation of p to ensure that the StableName is correct first time, which avoid to produce a tree bigger than needed
makeSharingName :: repr a -> IO SharingName
makeSharingName !p = fmap SharingName (makeStableName p)
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'
-- | Combinator interpreter detecting (Haskell embedded) @let@ definitions used at least twice 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 repr a = ObserveSharing { unObserveSharing :: MT.ReaderT (HashSet SharingName) (MT.StateT ObserveSharingState IO) (CleanDefs repr a) }

observeSharing :: ObserveSharing repr a -> IO (repr a)
observeSharing (ObserveSharing m) = do
  (a, st) <- MT.runStateT (MT.runReaderT m mempty) emptyObserveSharingState
  let refs = HM.fromList $ List.filter (\(_n,c) -> c > 1) $ HM.elems $ defref_shared st
  return $
    trace (show refs) $
    unCleanDefs a refs

instance Hashable TH.Name where
  hashWithSalt s = hashWithSalt s . show

-- * Type 'CleanDefs'
-- | Remove 'def' non-recursive or only used once.
newtype CleanDefs repr a = CleanDefs { unCleanDefs :: HM.HashMap TH.Name Int -> repr a }

type instance Unlift (CleanDefs repr) = repr
instance Liftable (CleanDefs 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 repr => Letable (CleanDefs repr) where
  def name x = CleanDefs $ \refs ->
    case HM.lookup name refs of
      Just c | c > 1 -> def name $ unCleanDefs x refs
      _ -> unCleanDefs x refs
instance P.Applicable repr => P.Applicable (CleanDefs repr)
instance P.Alternable repr => P.Alternable (CleanDefs repr)
instance P.Charable repr => P.Charable (CleanDefs repr)
instance P.Selectable repr => P.Selectable (CleanDefs repr)
instance P.Matchable repr => P.Matchable (CleanDefs repr) where
  conditional cs bs a b =
    CleanDefs (P.conditional cs <$> mapM unCleanDefs bs <*> unCleanDefs a <*> unCleanDefs b)
instance P.Lookable repr => P.Lookable (CleanDefs repr)
instance P.Foldable repr => P.Foldable (CleanDefs repr)

-- ** Type 'ObserveSharingState'
data ObserveSharingState = ObserveSharingState
 { defref_shared :: HashMap SharingName (TH.Name, Int)
 , defref_recs :: HashSet SharingName
 } deriving (Show)

emptyObserveSharingState :: ObserveSharingState
emptyObserveSharingState = ObserveSharingState
 { defref_shared = HM.empty
 , defref_recs = HS.empty
 }

-- ** Class 'Letable'
class Letable repr where
  def :: Pointer -> repr a -> repr a
  ref :: Bool -> Pointer -> repr a
  default def ::
    Liftable repr => Letable (Unlift repr) =>
    Pointer -> repr a -> repr a
  default ref ::
    Liftable repr => Letable (Unlift repr) =>
    Bool -> Pointer -> repr a
  def n = lift1 (def n)
  ref r n = lift (ref r n)

-- *** Type 'Pointer'
type Pointer = TH.Name

observeSharingNode :: Letable repr => ObserveSharing repr a -> ObserveSharing repr a
observeSharingNode node@(ObserveSharing m) = ObserveSharing $ do
  pName <- MT.lift $ MT.lift $ makeSharingName node
  -- let pName = showHex (I# (unsafeCoerce# name)) ""
  -- let pName = SharingName name
  st <- MT.lift MT.get
  ((before, qName), preds) <- getCompose $ HM.alterF (\v ->
    Compose $ case v of
      Nothing -> do
        qName <- MT.lift $ MT.lift $ TH.qNewName ("let"{-<>show pName-})
        return ((v, qName), Just (qName, 1))
      Just (qName, c) -> do
        return ((v, qName), Just (qName, c + 1))
    ) pName (defref_shared st)
  seen <- MT.ask
  if HS.member pName seen
  then do
    MT.lift $ MT.put st
     { defref_shared = preds
     , defref_recs = HS.insert pName (defref_recs st)
     }
    return $ ref True qName
  else do
    MT.lift $ MT.put st{ defref_shared = preds }
    if isNothing before
     then MT.local (HS.insert pName) (def qName <$> m)
     else return $ ref False qName

type instance Unlift (ObserveSharing repr) = CleanDefs repr
instance Letable repr => Liftable (ObserveSharing 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))
instance (Letable repr, P.Charable repr) => P.Charable (ObserveSharing repr)
instance (Letable repr, P.Alternable repr) => P.Alternable (ObserveSharing repr)
instance (Letable repr, P.Applicable repr) => P.Applicable (ObserveSharing repr)
instance (Letable repr, P.Selectable repr) => P.Selectable (ObserveSharing repr)
instance (Letable repr, P.Matchable repr) => P.Matchable (ObserveSharing repr) where
  -- Here the default definition does not fit since there is no lift* for its type and thus handles 'bs' itself, which is not the transformation wanted.
  conditional cs bs a b =
    observeSharingNode (ObserveSharing (P.conditional cs <$> mapM unObserveSharing bs <*> unObserveSharing a <*> unObserveSharing b))
instance (Letable repr, P.Foldable repr) => P.Foldable (ObserveSharing repr)
instance (Letable repr, P.Lookable repr) => P.Lookable (ObserveSharing repr)