content: {-# LANGUAGE Trustworthy #-}
  
  module Language.Hasmtlib.Internal.Sharing
    ( Sharing(..)
    , SharingMode(..)
    , runSharing, share
    )
  where
  
  import Language.Hasmtlib.Internal.Uniplate1
  import Language.Hasmtlib.Type.MonadSMT
  import Language.Hasmtlib.Type.SMTSort
  import Language.Hasmtlib.Type.Expr
  import Data.GADT.Compare
  import Data.Some.Constraint
  import Data.HashMap.Lazy
  import Data.Default
  import Data.Kind
  import Control.Monad.State
  import Control.Lens
  import System.Mem.StableName
  import System.IO.Unsafe
  import Unsafe.Coerce
  
  -- | Mode used for sharing.
  data SharingMode =
      None            -- ^ Common expressions are not shared at all
    | StableNames     -- ^ Expressions that resolve to the same 'StableName' are shared
    deriving Show
  
  instance Default SharingMode where
    def = None
  
  -- | States that can share expressions by comparing their 'StableName's.
  class Sharing s where
    -- | A constraint on the monad used when asserting the shared node in 'assertSharedNode'.
    type SharingMonad s :: (Type -> Type) -> Constraint
  
    -- | A 'Lens'' on a mapping between a 'StableName' and it's 'Expr' we may share.
    stableMap :: Lens' s (HashMap (StableName ()) (SomeKnownSMTSort Expr))
  
    -- | Asserts that a node-expression is represented by it's auxiliary node-variable: @nodeExpr :: Expr t === nodeVar@.
    --   Also gives access to the 'StableName' of the original expression.
    assertSharedNode :: (MonadState s m, SharingMonad s m) => StableName () -> Expr BoolSort -> m ()
  
    -- | Sets the mode used for sharing common expressions. Defaults to 'StableNames'.
    setSharingMode :: MonadState s m => SharingMode -> m ()
  
  -- | Shares all possible sub-expressions in given expression.
  --   Replaces each node in the expression-tree with an auxiliary variable.
  --   All nodes @x@ @y@ where @makeStableName x == makeStableName y@ are replaced with the same auxiliary variable.
  --   Therefore this creates a DAG.
  runSharing :: (KnownSMTSort t, MonadSMT s m, Sharing s, SharingMonad s m) => SharingMode -> Expr t -> m (Expr t)
  runSharing None = return
  runSharing StableNames = lazyParaM1 (
      \origExpr expr ->
        if isLeaf origExpr
        then return origExpr
        else case sortSing' origExpr of   -- scopes Equatable (Expr t) for specific t
          SBoolSort      -> share origExpr expr
          SIntSort       -> share origExpr expr
          SRealSort      -> share origExpr expr
          SBvSort _ _    -> share origExpr expr
          SArraySort _ _ -> share origExpr expr
          SStringSort    -> share origExpr expr)
  
  -- | Returns an auxiliary variable representing this expression node.
  --   If such a shared auxiliary variable exists already, returns that.
  --   Otherwise creates one and returns it.
  share :: (Equatable (Expr t), KnownSMTSort t, MonadSMT s m, Sharing s, SharingMonad s m) => Expr t -> m (Expr t) -> m (Expr t)
  share expr@(ForAll _ _) _ = return expr     -- sharing quantified expression would out-scope quantified var
  share expr@(Exists _ _) _ = return expr
  share origExpr expr = do
    let sn = unsafePerformIO (makeStableName' origExpr)
     in use (stableMap.at sn) >>= \mexpr' -> case mexpr' of
          Just (Some1 expr') -> case geq (sortSing' origExpr) (sortSing' expr') of
            Nothing -> expr >>= makeNode sn
            Just Refl -> return expr'
          Nothing -> expr >>= makeNode sn
  
  makeNode :: (Equatable (Expr t), KnownSMTSort t, MonadSMT s m, Sharing s, SharingMonad s m) => StableName () -> Expr t -> m (Expr t)
  makeNode sn nodeExpr = do
    nodeVar <- var
    assertSharedNode sn $ nodeVar === nodeExpr
    stableMap.at sn ?= Some1 nodeVar
    return nodeVar
  
  makeStableName' :: a -> IO (StableName ())
  makeStableName' x = x `seq` fmap unsafeCoerce (makeStableName x)
  