content: {-|
  Module           : What4.WordMap
  Description      : Datastructure for mapping bitvectors to values
  Copyright        : (c) Galois, Inc 2014-2020
  License          : BSD3
  Maintainer       : Rob Dockins <rdockins@galois.com>
  -}
  
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE PatternSynonyms #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TypeOperators #-}
  module What4.WordMap
    (
      WordMap(..)
    , emptyWordMap
    , muxWordMap
    , insertWordMap
    , lookupWordMap
    ) where
  
  import           Data.Parameterized.Ctx
  import qualified Data.Parameterized.Context as Ctx
  
  import What4.BaseTypes
  import What4.Interface
  import What4.Partial (PartExpr, pattern PE, pattern Unassigned) -- TODO(langston): use PartialWithErr
  
  -----------------------------------------------------------------------
  -- WordMap operations
  
  -- | A @WordMap@ represents a finite partial map from bitvectors of width @w@
  --   to elements of type @tp@.
  data WordMap sym w tp =
       SimpleWordMap !(SymExpr sym
                         (BaseArrayType (EmptyCtx ::> BaseBVType w) BaseBoolType))
                     !(SymExpr sym
                         (BaseArrayType (EmptyCtx ::> BaseBVType w) tp))
  
  -- | Create a word map where every element is undefined.
  emptyWordMap :: (IsExprBuilder sym, 1 <= w)
               => sym
               -> NatRepr w
               -> BaseTypeRepr a
               -> IO (WordMap sym w a)
  emptyWordMap sym w tp = do
    let idxRepr = Ctx.singleton (BaseBVRepr w)
    SimpleWordMap <$> constantArray sym idxRepr (falsePred sym)
                  <*> baseDefaultValue sym (BaseArrayRepr idxRepr tp)
  
  -- | Compute a pointwise if-then-else operation on the elements of two word maps.
  muxWordMap :: IsExprBuilder sym
             => sym
             -> NatRepr w
             -> BaseTypeRepr a
             -> (Pred sym
                 -> WordMap sym w a
                 -> WordMap sym w a
                 -> IO (WordMap sym w a))
  muxWordMap sym _w _tp p (SimpleWordMap bs1 xs1) (SimpleWordMap bs2 xs2) = do
    SimpleWordMap <$> arrayIte sym p bs1 bs2
                  <*> arrayIte sym p xs1 xs2
  
  -- | Update a word map at the given index.
  insertWordMap :: IsExprBuilder sym
                => sym
                -> NatRepr w
                -> BaseTypeRepr a
                -> SymBV sym w {- ^ index -}
                -> SymExpr sym a {- ^ new value -}
                -> WordMap sym w a {- ^ word map to update -}
                -> IO (WordMap sym w a)
  insertWordMap sym _w _ idx v (SimpleWordMap bs xs) = do
    let i = Ctx.singleton idx
    SimpleWordMap <$> arrayUpdate sym bs i (truePred sym)
                  <*> arrayUpdate sym xs i v
  
  -- | Lookup the value of an index in a word map.
  lookupWordMap :: IsExprBuilder sym
                => sym
                -> NatRepr w
                -> BaseTypeRepr a
                -> SymBV sym w {- ^ index -}
                -> WordMap sym w a
                -> IO (PartExpr (Pred sym) (SymExpr sym a))
  lookupWordMap sym _w _tp idx (SimpleWordMap bs xs) = do
    let i = Ctx.singleton idx
    p <- arrayLookup sym bs i
    case asConstantPred p of
      Just False -> return Unassigned
      _ -> PE p <$> arrayLookup sym xs i
  