{-|
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 -> 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 = BaseTypeRepr (BaseBVType w)
-> Assignment BaseTypeRepr (EmptyCtx ::> BaseBVType w)
forall k (f :: k -> Type) (tp :: k).
f tp -> Assignment f (EmptyCtx ::> tp)
Ctx.singleton (NatRepr w -> BaseTypeRepr (BaseBVType w)
forall (w :: Nat).
(1 <= w) =>
NatRepr w -> BaseTypeRepr (BaseBVType w)
BaseBVRepr NatRepr w
w)
  SymExpr
  sym ('BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType)
-> SymExpr sym ('BaseArrayType (EmptyCtx ::> BaseBVType w) a)
-> WordMap sym w a
forall sym (w :: Nat) (tp :: BaseType).
SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType)
-> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) tp)
-> WordMap sym w tp
SimpleWordMap (SymExpr
   sym ('BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType)
 -> SymExpr sym ('BaseArrayType (EmptyCtx ::> BaseBVType w) a)
 -> WordMap sym w a)
-> IO
     (SymExpr
        sym ('BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType))
-> IO
     (SymExpr sym ('BaseArrayType (EmptyCtx ::> BaseBVType w) a)
      -> WordMap sym w a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> sym
-> Assignment BaseTypeRepr (EmptyCtx ::> BaseBVType w)
-> SymExpr sym 'BaseBoolType
-> IO
     (SymExpr
        sym ('BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType))
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 (sym -> SymExpr sym 'BaseBoolType
forall sym. IsExprBuilder sym => sym -> Pred sym
falsePred sym
sym)
                IO
  (SymExpr sym ('BaseArrayType (EmptyCtx ::> BaseBVType w) a)
   -> WordMap sym w a)
-> IO (SymExpr sym ('BaseArrayType (EmptyCtx ::> BaseBVType w) a))
-> IO (WordMap sym w a)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> sym
-> BaseTypeRepr ('BaseArrayType (EmptyCtx ::> BaseBVType w) a)
-> IO (SymExpr sym ('BaseArrayType (EmptyCtx ::> BaseBVType w) a))
forall sym (bt :: BaseType).
IsExprBuilder sym =>
sym -> BaseTypeRepr bt -> IO (SymExpr sym bt)
baseDefaultValue sym
sym (Assignment BaseTypeRepr (EmptyCtx ::> BaseBVType w)
-> BaseTypeRepr a
-> BaseTypeRepr ('BaseArrayType (EmptyCtx ::> BaseBVType w) a)
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 :: 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
  SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType)
-> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
-> WordMap sym w a
forall sym (w :: Nat) (tp :: BaseType).
SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType)
-> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) tp)
-> WordMap sym w tp
SimpleWordMap (SymExpr
   sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType)
 -> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
 -> WordMap sym w a)
-> IO
     (SymExpr
        sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType))
-> IO
     (SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
      -> WordMap sym w a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> sym
-> Pred sym
-> SymExpr
     sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType)
-> SymExpr
     sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType)
-> IO
     (SymExpr
        sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType))
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
                IO
  (SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
   -> WordMap sym w a)
-> IO (SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a))
-> IO (WordMap sym w a)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> sym
-> Pred sym
-> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
-> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
-> IO (SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a))
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 :: 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 = SymBV sym w -> Assignment (SymExpr sym) (EmptyCtx ::> BaseBVType w)
forall k (f :: k -> Type) (tp :: k).
f tp -> Assignment f (EmptyCtx ::> tp)
Ctx.singleton SymBV sym w
idx
  SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType)
-> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
-> WordMap sym w a
forall sym (w :: Nat) (tp :: BaseType).
SymExpr
  sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType)
-> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) tp)
-> WordMap sym w tp
SimpleWordMap (SymExpr
   sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType)
 -> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
 -> WordMap sym w a)
-> IO
     (SymExpr
        sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType))
-> IO
     (SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
      -> WordMap sym w a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> sym
-> SymExpr
     sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType)
-> Assignment (SymExpr sym) (EmptyCtx ::> BaseBVType w)
-> SymExpr sym 'BaseBoolType
-> IO
     (SymExpr
        sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType))
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 (sym -> SymExpr sym 'BaseBoolType
forall sym. IsExprBuilder sym => sym -> Pred sym
truePred sym
sym)
                IO
  (SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
   -> WordMap sym w a)
-> IO (SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a))
-> IO (WordMap sym w a)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> sym
-> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
-> Assignment (SymExpr sym) (EmptyCtx ::> BaseBVType w)
-> SymExpr sym a
-> IO (SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a))
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 :: 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 = SymBV sym w -> Assignment (SymExpr sym) (EmptyCtx ::> BaseBVType w)
forall k (f :: k -> Type) (tp :: k).
f tp -> Assignment f (EmptyCtx ::> tp)
Ctx.singleton SymBV sym w
idx
  Pred sym
p <- sym
-> SymExpr
     sym (BaseArrayType (EmptyCtx ::> BaseBVType w) 'BaseBoolType)
-> Assignment (SymExpr sym) (EmptyCtx ::> BaseBVType w)
-> IO (Pred sym)
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 Pred sym -> Maybe Bool
forall (e :: BaseType -> Type).
IsExpr e =>
e 'BaseBoolType -> Maybe Bool
asConstantPred Pred sym
p of
    Just Bool
False -> PartExpr (Pred sym) (SymExpr sym a)
-> IO (PartExpr (Pred sym) (SymExpr sym a))
forall (m :: Type -> Type) a. Monad m => a -> m a
return PartExpr (Pred sym) (SymExpr sym a)
forall p v. PartExpr p v
Unassigned
    Maybe Bool
_ -> Pred sym -> SymExpr sym a -> PartExpr (Pred sym) (SymExpr sym a)
forall p v. p -> v -> PartExpr p v
PE Pred sym
p (SymExpr sym a -> PartExpr (Pred sym) (SymExpr sym a))
-> IO (SymExpr sym a) -> IO (PartExpr (Pred sym) (SymExpr sym a))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> sym
-> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
-> Assignment (SymExpr sym) (EmptyCtx ::> BaseBVType w)
-> IO (SymExpr sym a)
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