content: {-# LANGUAGE DeriveAnyClass #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE DeriveLift #-}
  {-# HLINT ignore "Unused LANGUAGE pragma" #-}
  {-# LANGUAGE DerivingStrategies #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE PatternSynonyms #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE ViewPatterns #-}
  {-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}
  
  -- |
  -- Module      :   Grisette.Internal.Core.Data.Symbol
  -- Copyright   :   (c) Sirui Lu 2024
  -- License     :   BSD-3-Clause (see the LICENSE file)
  --
  -- Maintainer  :   siruilu@cs.washington.edu
  -- Stability   :   Experimental
  -- Portability :   GHC only
  module Grisette.Internal.Core.Data.Symbol
    ( Identifier (..),
      identifier,
      withMetadata,
      withLocation,
      mapMetadata,
      uniqueIdentifier,
      Symbol (..),
      simple,
      indexed,
      symbolIdentifier,
      mapIdentifier,
      AsMetadata (..),
      pattern Metadata,
    )
  where
  
  import Control.DeepSeq (NFData)
  import qualified Data.Binary as Binary
  import Data.Bytes.Serial (Serial (deserialize, serialize))
  import Data.Hashable (Hashable (hashWithSalt))
  import Data.IORef (IORef, atomicModifyIORef', newIORef)
  import qualified Data.Serialize as Cereal
  import Data.String (IsString (fromString))
  import qualified Data.Text as T
  import GHC.Generics (Generic)
  import GHC.IO (unsafePerformIO)
  import Grisette.Internal.Core.Data.SExpr
    ( SExpr (Atom, List, NumberAtom),
      fileLocation,
      showsSExprWithParens,
    )
  import Language.Haskell.TH.Syntax (Lift)
  import Language.Haskell.TH.Syntax.Compat (SpliceQ)
  
  -- $setup
  -- >>> import Grisette.Core
  -- >>> import Grisette.SymPrim
  
  -- | Identifier type used for 'Grisette.Core.GenSym'
  --
  -- The constructor is hidden intentionally.
  -- You can construct an identifier by:
  --
  --   * a raw identifier
  --
  --     The following two expressions will refer to the same identifier (the
  --     solver won't distinguish them and would assign the same value to them).
  --     The user may need to use unique names to avoid unintentional identifier
  --     collision.
  --
  --     >>> identifier "a"
  --     a
  --
  --     >>> "a" :: Identifier -- available when OverloadedStrings is enabled
  --     a
  --
  --   * bundle the identifier with some user provided metadata
  --
  --     Identifiers created with different name or different additional
  --     information will not be the same.
  --
  --     >>> withMetadata "a" (NumberAtom 1)
  --     a:1
  --
  --   * bundle the calling file location with the identifier to ensure global
  --     uniqueness
  --
  --     Identifiers created at different locations will not be the
  --     same. The identifiers created at the same location will be the same.
  --
  --     >>> $$(withLocation "a") -- a sample result could be "a:[grisette-file-location <interactive> 18 (4 18)]"
  --     a:[grisette-file-location <interactive>...]
  data Identifier = Identifier {baseIdent :: T.Text, metadata :: SExpr}
    deriving (Eq, Ord, Generic, Lift)
    deriving anyclass (Hashable, NFData, Serial)
  
  instance Cereal.Serialize Identifier where
    put = serialize
    get = deserialize
  
  instance Binary.Binary Identifier where
    put = serialize
    get = deserialize
  
  instance Show Identifier where
    showsPrec _ (Identifier i (List [])) = showString (T.unpack i)
    showsPrec _ (Identifier i metadata) =
      showString (T.unpack i)
        . showString ":"
        . showsSExprWithParens '[' ']' metadata
  
  instance IsString Identifier where
    fromString i = Identifier (T.pack i) $ List []
  
  -- | A type class for embedding a type into a metadata represented as an
  -- S-expression.
  class AsMetadata a where
    asMetadata :: a -> SExpr
    fromMetadata :: SExpr -> Maybe a
  
  -- | A pattern for extracting a value from a metadata represented as an
  -- S-expression.
  pattern Metadata :: (AsMetadata a) => a -> SExpr
  pattern Metadata m <- (fromMetadata -> Just m)
    where
      Metadata m = asMetadata m
  
  instance AsMetadata SExpr where
    asMetadata = id
    fromMetadata = Just
  
  -- | Simple identifier.
  -- The same identifier refers to the same symbolic variable in the whole
  -- program.
  --
  -- The user may need to use unique identifiers to avoid unintentional identifier
  -- collision.
  identifier :: T.Text -> Identifier
  identifier = flip Identifier $ List []
  
  -- | Identifier with extra metadata.
  --
  -- The same identifier with the same metadata refers to the same symbolic
  -- variable in the whole program.
  --
  -- The user may need to use unique identifiers or additional metadata to
  -- avoid unintentional identifier collision.
  withMetadata :: (AsMetadata a) => T.Text -> a -> Identifier
  withMetadata ident meta = Identifier ident (asMetadata meta)
  
  -- | Identifier with the file location.
  withLocation :: T.Text -> SpliceQ Identifier
  withLocation nm = [||withMetadata nm $$fileLocation||]
  
  -- | Modify the metadata of an identifier.
  mapMetadata ::
    (AsMetadata a) => (SExpr -> a) -> Identifier -> Identifier
  mapMetadata f (Identifier i m) = Identifier i (asMetadata $ f m)
  
  identifierCount :: IORef Int
  identifierCount = unsafePerformIO $ newIORef 0
  {-# NOINLINE identifierCount #-}
  
  -- | Get a globally unique identifier within the 'IO' monad.
  uniqueIdentifier :: T.Text -> IO Identifier
  uniqueIdentifier ident = do
    i <- atomicModifyIORef' identifierCount (\x -> (x + 1, x))
    return $
      withMetadata
        ident
        (List [Atom "grisette-unique", NumberAtom $ toInteger i])
  
  -- | Symbol types for a symbolic variable.
  --
  -- The symbols can be indexed with an integer.
  data Symbol where
    SimpleSymbol :: Identifier -> Symbol
    IndexedSymbol :: Identifier -> Int -> Symbol
    deriving (Eq, Ord, Generic, Lift, NFData, Serial)
  
  instance Cereal.Serialize Symbol where
    put = serialize
    get = deserialize
  
  instance Binary.Binary Symbol where
    put = serialize
    get = deserialize
  
  instance Hashable Symbol where
    hashWithSalt s (SimpleSymbol i) = hashWithSalt s i
    hashWithSalt s (IndexedSymbol i idx) = s `hashWithSalt` i `hashWithSalt` idx
    {-# INLINE hashWithSalt #-}
  
  -- | Get the identifier of a symbol.
  symbolIdentifier :: Symbol -> Identifier
  symbolIdentifier (SimpleSymbol i) = i
  symbolIdentifier (IndexedSymbol i _) = i
  
  -- | Modify the identifier of a symbol.
  mapIdentifier :: (Identifier -> Identifier) -> Symbol -> Symbol
  mapIdentifier f (SimpleSymbol i) = SimpleSymbol (f i)
  mapIdentifier f (IndexedSymbol i idx) = IndexedSymbol (f i) idx
  
  instance Show Symbol where
    show (SimpleSymbol i) = show i
    show (IndexedSymbol i idx) = show i ++ "@" ++ show idx
  
  instance IsString Symbol where
    fromString = SimpleSymbol . fromString
  
  -- | Create a simple symbol.
  simple :: Identifier -> Symbol
  simple = SimpleSymbol
  
  -- | Create an indexed symbol.
  indexed :: Identifier -> Int -> Symbol
  indexed = IndexedSymbol
  