SymBiMap.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Strict #-}

-- |
-- Module      :   Grisette.Internal.Backend.SymBiMap
-- Copyright   :   (c) Sirui Lu 2021-2024
-- License     :   BSD-3-Clause (see the LICENSE file)
--
-- Maintainer  :   siruilu@cs.washington.edu
-- Stability   :   Experimental
-- Portability :   GHC only
module Grisette.Internal.Backend.SymBiMap
  ( SymBiMap (..),
    emptySymBiMap,
    sizeBiMap,
    addBiMap,
    addBiMapIntermediate,
    findStringToSymbol,
    lookupTerm,
    attachNextQuantifiedSymbolInfo,
  )
where

import Data.Dynamic (Dynamic)
import qualified Data.HashMap.Strict as M
import GHC.Stack (HasCallStack)
import Grisette.Internal.Backend.QuantifiedStack (QuantifiedStack)
import Grisette.Internal.Core.Data.SExpr (SExpr (Atom, List, NumberAtom))
import Grisette.Internal.Core.Data.Symbol
  ( mapIdentifier,
    mapMetadata,
  )
import Grisette.Internal.SymPrim.Prim.SomeTerm
  ( SomeTerm,
  )
import Grisette.Internal.SymPrim.Prim.Term
  ( IsSymbolKind,
    SomeTypedAnySymbol,
    SomeTypedSymbol,
    TypedConstantSymbol,
    TypedSymbol (unTypedSymbol),
    castSomeTypedSymbol,
    typedConstantSymbol,
    pattern SupportedConstantTypedSymbol,
  )

-- | A bidirectional map between symbolic Grisette terms and sbv terms.
data SymBiMap = SymBiMap
  { biMapToSBV :: M.HashMap SomeTerm (QuantifiedStack -> Dynamic),
    biMapSize :: Int,
    biMapFromSBV :: M.HashMap String SomeTypedAnySymbol,
    quantifiedSymbolNum :: Int
  }

instance Show SymBiMap where
  show (SymBiMap t s f _) =
    "SymBiMap { size: "
      ++ show s
      ++ ", toSBV: "
      ++ show (M.keys t)
      ++ ", fromSBV: "
      ++ show (M.toList f)
      ++ " }"

-- | Information about a quantified symbol.
-- newtype QuantifiedSymbolInfo = QuantifiedSymbolInfo Int
--   deriving (Generic, Ord, Eq, Show, Hashable, Lift, NFData)
nextQuantifiedSymbolInfo :: SymBiMap -> (SymBiMap, SExpr -> SExpr)
nextQuantifiedSymbolInfo (SymBiMap t s f num) =
  ( SymBiMap t s f (num + 1),
    \meta ->
      List
        [ Atom "grisette-quantified",
          NumberAtom $ fromIntegral num,
          meta
        ]
  )

attachQuantifiedSymbolInfo ::
  (SExpr -> SExpr) -> TypedConstantSymbol a -> TypedConstantSymbol a
attachQuantifiedSymbolInfo info tsym@SupportedConstantTypedSymbol =
  typedConstantSymbol $
    mapIdentifier (mapMetadata info) $
      unTypedSymbol tsym
attachQuantifiedSymbolInfo _ _ = error "Should not happen"

-- | Attach the next quantified symbol info to a symbol.
attachNextQuantifiedSymbolInfo ::
  SymBiMap -> TypedConstantSymbol a -> (SymBiMap, TypedConstantSymbol a)
attachNextQuantifiedSymbolInfo m s =
  let (m', info) = nextQuantifiedSymbolInfo m
   in (m', attachQuantifiedSymbolInfo info s)

-- | An empty bidirectional map.
emptySymBiMap :: SymBiMap
emptySymBiMap = SymBiMap M.empty 0 M.empty 0

-- | The size of the bidirectional map.
sizeBiMap :: SymBiMap -> Int
sizeBiMap = biMapSize

-- | Add a new entry to the bidirectional map.
addBiMap ::
  (HasCallStack) =>
  SomeTerm ->
  Dynamic ->
  String ->
  SomeTypedSymbol knd ->
  SymBiMap ->
  SymBiMap
addBiMap s d n sb (SymBiMap t sz f num) =
  case castSomeTypedSymbol sb of
    Just sb' -> SymBiMap (M.insert s (const d) t) (sz + 1) (M.insert n sb' f) num
    _ -> error "Casting to AnySymbol, should not fail"

-- | Add a new entry to the bidirectional map for intermediate values.
addBiMapIntermediate ::
  (HasCallStack) => SomeTerm -> (QuantifiedStack -> Dynamic) -> SymBiMap -> SymBiMap
addBiMapIntermediate s d (SymBiMap t sz f num) =
  SymBiMap (M.insert s d t) (sz + 1) f num

-- | Find a symbolic Grisette term from a string.
findStringToSymbol :: (IsSymbolKind knd) => String -> SymBiMap -> Maybe (SomeTypedSymbol knd)
findStringToSymbol s (SymBiMap _ _ f _) = do
  r <- M.lookup s f
  castSomeTypedSymbol r

-- | Look up an sbv value with a symbolic Grisette term in the bidirectional
-- map.
lookupTerm :: (HasCallStack) => SomeTerm -> SymBiMap -> Maybe (QuantifiedStack -> Dynamic)
lookupTerm t m = M.lookup t (biMapToSBV m)