    {-|
    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 = 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 :: Natural).
    (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 :: Natural) (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)
    forall (idx :: Ctx BaseType) (tp :: BaseType) (b :: BaseType).
    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 a b. IO (a -> b) -> IO a -> IO b
    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 :: 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
      SymExpr
      sym (BaseArrayType (EmptyCtx ::> BaseBVType w) BaseBoolType)
    -> SymExpr sym (BaseArrayType (EmptyCtx ::> BaseBVType w) a)
    -> WordMap sym w a
    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 (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)
    forall (idx :: Ctx BaseType) (b :: BaseType).
    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 a b. IO (a -> b) -> IO a -> IO b
    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)
    forall (idx :: Ctx BaseType) (b :: BaseType).
    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 = 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 :: Natural) (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)
    forall (idx :: Ctx BaseType) (tp :: BaseType) (b :: BaseType).
    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 a b. IO (a -> b) -> IO a -> IO b
    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)
    forall (idx :: Ctx BaseType) (tp :: BaseType) (b :: BaseType).
    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 = 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
      p <- sym
    -> SymExpr
         sym (BaseArrayType (EmptyCtx ::> BaseBVType w) BaseBoolType)
    -> Assignment (SymExpr sym) (EmptyCtx ::> BaseBVType w)
    -> IO (SymExpr sym BaseBoolType)
    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)
    forall (idx :: Ctx BaseType) (tp :: BaseType) (b :: BaseType).
    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 asConstantPred p of
        Just Bool
    False -> PartialWithErr () (SymExpr sym BaseBoolType) (SymExpr sym a)
    -> IO
         (PartialWithErr () (SymExpr sym BaseBoolType) (SymExpr sym a))
    forall a. a -> IO a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return PartialWithErr () (SymExpr sym BaseBoolType) (SymExpr sym a)
    forall p v. PartExpr p v
    Unassigned
        Maybe Bool
    _ -> SymExpr sym BaseBoolType
    -> SymExpr sym a
    -> PartialWithErr () (SymExpr sym BaseBoolType) (SymExpr sym a)
    forall p v. p -> v -> PartExpr p v
    PE SymExpr sym BaseBoolType
    p (SymExpr sym a
     -> PartialWithErr () (SymExpr sym BaseBoolType) (SymExpr sym a))
    -> IO (SymExpr sym a)
    -> IO
         (PartialWithErr () (SymExpr sym BaseBoolType) (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)
    forall (idx :: Ctx BaseType) (tp :: BaseType) (b :: BaseType).
    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