{-|
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 :: forall sym (w :: Natural) (a :: BaseType).
(IsExprBuilder sym, 1 <= w) =>
sym -> NatRepr w -> BaseTypeRepr a -> IO (WordMap sym w a)
emptyWordMap sym
sym NatRepr w
w BaseTypeRepr a
tp = do
  let idxRepr :: Assignment BaseTypeRepr (EmptyCtx ::> 'BaseBVType w)
idxRepr = forall {k} (f :: k -> Type) (tp :: k).
f tp -> Assignment f (EmptyCtx ::> tp)
Ctx.singleton (forall (w :: Natural).
(1 <= w) =>
NatRepr w -> BaseTypeRepr ('BaseBVType w)
BaseBVRepr NatRepr w
w)
  forall sym (w :: Natural) (tp :: BaseType).
SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) BaseBoolType)
-> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) tp)
-> WordMap sym w tp
SimpleWordMap forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall sym (idx :: Ctx BaseType) (tp :: BaseType) (b :: BaseType).
IsExprBuilder sym =>
sym
-> Assignment BaseTypeRepr (idx ::> tp)
-> SymExpr sym b
-> IO (SymArray sym (idx ::> tp) b)
constantArray sym
sym Assignment BaseTypeRepr (EmptyCtx ::> 'BaseBVType w)
idxRepr (forall sym. IsExprBuilder sym => sym -> Pred sym
falsePred sym
sym)
                forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> forall sym (bt :: BaseType).
IsExprBuilder sym =>
sym -> BaseTypeRepr bt -> IO (SymExpr sym bt)
baseDefaultValue sym
sym (forall (idx :: Ctx BaseType) (tp :: BaseType) (xs :: BaseType).
Assignment BaseTypeRepr (idx ::> tp)
-> BaseTypeRepr xs -> BaseTypeRepr ('BaseArrayType (idx ::> tp) xs)
BaseArrayRepr Assignment BaseTypeRepr (EmptyCtx ::> 'BaseBVType w)
idxRepr BaseTypeRepr a
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 :: forall sym (w :: Natural) (a :: BaseType).
IsExprBuilder sym =>
sym
-> NatRepr w
-> BaseTypeRepr a
-> Pred sym
-> WordMap sym w a
-> WordMap sym w a
-> IO (WordMap sym w a)
muxWordMap sym
sym NatRepr w
_w BaseTypeRepr a
_tp Pred sym
p (SimpleWordMap SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) BaseBoolType)
bs1 SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
xs1) (SimpleWordMap SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) BaseBoolType)
bs2 SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
xs2) = do
  forall sym (w :: Natural) (tp :: BaseType).
SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) BaseBoolType)
-> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) tp)
-> WordMap sym w tp
SimpleWordMap forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall sym (idx :: Ctx BaseType) (b :: BaseType).
IsExprBuilder sym =>
sym
-> Pred sym
-> SymArray sym idx b
-> SymArray sym idx b
-> IO (SymArray sym idx b)
arrayIte sym
sym Pred sym
p SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) BaseBoolType)
bs1 SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) BaseBoolType)
bs2
                forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> forall sym (idx :: Ctx BaseType) (b :: BaseType).
IsExprBuilder sym =>
sym
-> Pred sym
-> SymArray sym idx b
-> SymArray sym idx b
-> IO (SymArray sym idx b)
arrayIte sym
sym Pred sym
p SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
xs1 SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
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 :: forall sym (w :: Natural) (a :: BaseType).
IsExprBuilder sym =>
sym
-> NatRepr w
-> BaseTypeRepr a
-> SymBV sym w
-> SymExpr sym a
-> WordMap sym w a
-> IO (WordMap sym w a)
insertWordMap sym
sym NatRepr w
_w BaseTypeRepr a
_ SymBV sym w
idx SymExpr sym a
v (SimpleWordMap SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) BaseBoolType)
bs SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
xs) = do
  let i :: Assignment (SymExpr sym) (EmptyCtx ::> BaseBVType w)
i = forall {k} (f :: k -> Type) (tp :: k).
f tp -> Assignment f (EmptyCtx ::> tp)
Ctx.singleton SymBV sym w
idx
  forall sym (w :: Natural) (tp :: BaseType).
SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) BaseBoolType)
-> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) tp)
-> WordMap sym w tp
SimpleWordMap forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall sym (idx :: Ctx BaseType) (tp :: BaseType) (b :: BaseType).
IsExprBuilder sym =>
sym
-> SymArray sym (idx ::> tp) b
-> Assignment (SymExpr sym) (idx ::> tp)
-> SymExpr sym b
-> IO (SymArray sym (idx ::> tp) b)
arrayUpdate sym
sym SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) BaseBoolType)
bs Assignment (SymExpr sym) (EmptyCtx ::> BaseBVType w)
i (forall sym. IsExprBuilder sym => sym -> Pred sym
truePred sym
sym)
                forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> forall sym (idx :: Ctx BaseType) (tp :: BaseType) (b :: BaseType).
IsExprBuilder sym =>
sym
-> SymArray sym (idx ::> tp) b
-> Assignment (SymExpr sym) (idx ::> tp)
-> SymExpr sym b
-> IO (SymArray sym (idx ::> tp) b)
arrayUpdate sym
sym SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
xs Assignment (SymExpr sym) (EmptyCtx ::> BaseBVType w)
i SymExpr sym a
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 :: forall sym (w :: Natural) (a :: BaseType).
IsExprBuilder sym =>
sym
-> NatRepr w
-> BaseTypeRepr a
-> SymBV sym w
-> WordMap sym w a
-> IO (PartExpr (Pred sym) (SymExpr sym a))
lookupWordMap sym
sym NatRepr w
_w BaseTypeRepr a
_tp SymBV sym w
idx (SimpleWordMap SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) BaseBoolType)
bs SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
xs) = do
  let i :: Assignment (SymExpr sym) (EmptyCtx ::> BaseBVType w)
i = forall {k} (f :: k -> Type) (tp :: k).
f tp -> Assignment f (EmptyCtx ::> tp)
Ctx.singleton SymBV sym w
idx
  Pred sym
p <- forall sym (idx :: Ctx BaseType) (tp :: BaseType) (b :: BaseType).
IsExprBuilder sym =>
sym
-> SymArray sym (idx ::> tp) b
-> Assignment (SymExpr sym) (idx ::> tp)
-> IO (SymExpr sym b)
arrayLookup sym
sym SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) BaseBoolType)
bs Assignment (SymExpr sym) (EmptyCtx ::> BaseBVType w)
i
  case forall (e :: BaseType -> Type).
IsExpr e =>
e BaseBoolType -> Maybe Bool
asConstantPred Pred sym
p of
    Just Bool
False -> forall (m :: Type -> Type) a. Monad m => a -> m a
return forall p v. PartExpr p v
Unassigned
    Maybe Bool
_ -> forall p v. p -> v -> PartExpr p v
PE Pred sym
p forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall sym (idx :: Ctx BaseType) (tp :: BaseType) (b :: BaseType).
IsExprBuilder sym =>
sym
-> SymArray sym (idx ::> tp) b
-> Assignment (SymExpr sym) (idx ::> tp)
-> IO (SymExpr sym b)
arrayLookup sym
sym SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
xs Assignment (SymExpr sym) (EmptyCtx ::> BaseBVType w)
i