-----------------------------------------------------------------------
-- |
-- Module           : What4.CachedArray
-- Description      : What4 array storage with a concrete backing supporting symbolic indexes
-- Copyright        : (c) Galois, Inc 2020
-- License          : BSD3
-- Maintainer       : Daniel Matichuk <dmatichuk@galois.com>
-- Stability        : provisional
--
--
-- This module provides a storage structure that supports arrays that have reads
-- from and writes to a mix of concrete and symbolic indexes. It can be thought
-- of as a multi-dimensional array that supports reading and writing contiguous
-- "chunks".  It is built on the 'Data.Parameterized.IntervalsMap' structure,
-- which computes an abstract domain over indexes (supporting symbolic
-- reads/writes).
------------------------------------------------------------------------

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleInstances, FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE UndecidableInstances #-}

module What4.CachedArray
  (
    CachedArray
  , ArrayChunk
  , mkArrayChunk
  , evalChunk
  , writeChunk
  , writeSingle
  , readSingle
  , readChunk
  , arrayToChunk
  , chunkToArray
  , muxArrays
  , initArrayConcrete
  , initArray
  ) where

import           Control.Lens ( (.~), (&) )
import           Control.Monad ( foldM, join )
import           Control.Monad.Trans ( lift )
import           Data.Functor.Const
import           Data.Maybe ( catMaybes )
import qualified Data.Map as Map
import           Data.Maybe (mapMaybe)
import qualified Data.IORef as IO

import qualified Data.Parameterized.TraversableFC as FC
import qualified Data.Parameterized.Context as Ctx
import           Data.Parameterized.Classes
import           Data.Parameterized.NatRepr ( type (<=) )
import qualified Data.Parameterized.Nonce as PN
import qualified Data.BitVector.Sized as BV

import qualified Lang.Crucible.Utils.MuxTree as MT

import qualified What4.Interface as W4
import qualified What4.Partial as W4
import qualified What4.Concrete as W4
import qualified What4.Utils.AbstractDomains as W4
import qualified What4.Utils.BVDomain as BVD

import qualified Data.Parameterized.IntervalsMap as IM
import           Data.Parameterized.IntervalsMap ( AsOrd(..) )

------------------------------------------------
-- Interface

-- TODO: add coalescing function for merging adjacent entries

newtype ArrayChunk sym idx tp =
  ArrayChunk { evalChunk :: (W4.SymExpr sym idx -> IO (W4.SymExpr sym tp)) }

mkArrayChunk ::
  forall sym idx tp.
  W4.IsSymExprBuilder sym =>
  sym ->
  (W4.SymExpr sym idx -> IO (W4.SymExpr sym tp)) ->
  IO (ArrayChunk sym idx tp)
mkArrayChunk _sym f = do
  ref <- IO.newIORef Map.empty
  let f' idx = do
        m <- IO.readIORef ref
        case Map.lookup (AsOrd idx) m of
          Just v -> return v
          Nothing -> do
            v <- f idx
            IO.modifyIORef ref (Map.insert (AsOrd idx) v)
            return v
  return $ ArrayChunk f'

writeChunk ::
  forall sym ctx tp.
  NonEmptyCtx ctx =>
  W4.IsSymExprBuilder sym =>
  sym ->
  -- | base address to write to
  Ctx.Assignment (W4.SymExpr sym) ctx ->
  -- | size of write
  W4.SymExpr sym (CtxFirst ctx) ->
  -- | symbolic value to write
  ArrayChunk sym (CtxFirst ctx) tp ->
  CachedArray sym ctx tp ->
  IO (CachedArray sym ctx tp)
writeChunk sym loExpr offExpr chunk arr | NonEmptyCtxRepr <- nonEmptyCtxRepr @_ @ctx =
  arrConstraints arr $ do
  rng <- mkSymRangeOff sym loExpr offExpr
  arr' <- invalidateEntries sym rng arr
  -- offset the incoming function so that its value at zero becomes the value at
  -- the base address
  let
    off = indexToOffset $ symRangeLo rng
    vals :: SymIndex sym ctx -> IO (W4.PartExpr (W4.Pred sym) (W4.SymExpr sym tp))
    vals idx' = do
        p <- isInRange sym rng idx'
        (SymOffset idxOffsetExpr) <- indexToOffset <$> subSymOffset sym idx' off
        v <- evalChunk chunk idxOffsetExpr
        return $ W4.mkPE p v
  entry <- mkMultiEntry vals
  arr'' <- IM.insertWithM (mergeEntriesMux sym (isInRange sym rng)) (symRangeToAbs rng) (toPMuxTree sym entry)  (arrMap arr')
  incNonce $ arr { arrMap = arr''}


writeSingle ::
  forall sym ctx tp.
  W4.IsSymExprBuilder sym =>
  sym ->
  Ctx.Assignment (W4.SymExpr sym) ctx ->
  W4.SymExpr sym  tp ->
  CachedArray sym ctx tp ->
  IO (CachedArray sym ctx tp)
writeSingle sym symIdxExpr val arr = arrConstraints arr $ do
  arr' <- invalidateEntries sym (SymRangeSingle symIdx) arr
  entry <- mkValEntry sym symIdx val

  arr'' <- IM.insertWithM (mergeEntriesMux sym (isEqIndex sym symIdx)) (symIdxToAbs symIdx) (toPMuxTree sym entry)  (arrMap arr')
  incNonce $ arr { arrMap = arr'' }
  where
    symIdx = mkSymIndex symIdxExpr


readSingle ::
  forall sym idx tp.
  W4.IsSymExprBuilder sym =>
  sym ->
  Ctx.Assignment (W4.SymExpr sym) idx ->
  CachedArray sym idx tp ->
  IO (W4.SymExpr sym tp)
readSingle sym symIdxExpr arr = readArrayBase sym symIdx arr
  where
    symIdx = SymIndex symIdxExpr Nothing

readChunk ::
  forall sym ctx tp.
  NonEmptyCtx ctx =>
  W4.IsSymExprBuilder sym =>
  sym ->
  -- | base address to read from
  Ctx.Assignment (W4.SymExpr sym) ctx ->
  -- | size of read
  W4.SymExpr sym (CtxFirst ctx) ->
  CachedArray sym ctx tp ->
  IO (ArrayChunk sym (CtxFirst ctx) tp)
readChunk sym loExpr offExpr arr | NonEmptyCtxRepr <- nonEmptyCtxRepr @_ @ctx = do
  rng <- mkSymRangeOff sym loExpr offExpr
  let absIdx = symRangeToAbs rng
  -- offset the outgoing array so that its value at zero is the value at
  -- the base address
  return $ ArrayChunk $ \idxExpr -> do
    let off = SymOffset idxExpr
    offsetIdx <- addSymOffset sym (symRangeLo rng) off
    readArrayBase sym (offsetIdx { symIdxAbs = Just absIdx}) arr

chunkToArray ::
  forall sym idx tp.
  W4.IsSymExprBuilder sym =>
  sym ->
  W4.BaseTypeRepr idx ->
  ArrayChunk sym idx tp ->
  IO (W4.SymArray sym (Ctx.EmptyCtx Ctx.::> idx) tp)
chunkToArray sym repr chunk = do
  var <- W4.freshBoundVar sym W4.emptySymbol repr
  body <- evalChunk chunk (W4.varExpr sym var)
  fn <- W4.definedFn sym (W4.safeSymbol "readRange") (Ctx.empty Ctx.:> var) body W4.AlwaysUnfold
  W4.arrayFromFn sym fn

arrayToChunk ::
  forall sym idx tp.
  W4.IsSymExprBuilder sym =>
  sym ->
  (W4.SymArray sym (Ctx.EmptyCtx Ctx.::> idx) tp) ->
  IO (ArrayChunk sym idx tp)
arrayToChunk sym arr = mkArrayChunk sym $ \idx -> W4.arrayLookup sym arr (Ctx.empty Ctx.:> idx)


muxArrays ::
  forall sym idx tp.
  W4.IsSymExprBuilder sym =>
  sym ->
  W4.Pred sym ->
  CachedArray sym idx tp ->
  CachedArray sym idx tp ->
  IO (CachedArray sym idx tp)
muxArrays sym p arr1 arr2 = case arr1 == arr2 of
  True -> return arr1
  False -> arrConstraints arr1 $ do
    notp <- W4.notPred sym p
    arr' <- IM.mergeWithM
              (pmuxTreeAddCondition sym p)
              (pmuxTreeAddCondition sym notp)
              (muxEntries sym p)
              (arrMap arr1)
              (arrMap arr2)
    incNonce $ arr1 { arrMap = arr' }

-- | Initialize an array with symbolic contents at concrete locations
initArrayConcrete ::
  forall sym idx tp idx' tp'.
  W4.IsSymExprBuilder sym =>
  idx ~ (idx' Ctx.::> tp') =>
  sym ->
  W4.BaseTypeRepr tp ->
  [(Ctx.Assignment W4.ConcreteVal idx, W4.SymExpr sym tp)] ->
  IO (CachedArray sym idx tp)
initArrayConcrete sym repr m = do
  nonce <- freshArrayNonce
  im <- IM.fromList <$> mapM go m
  return $ CachedArray im (\x -> x) repr nonce
  where
    go ::
      (Ctx.Assignment W4.ConcreteVal idx, W4.SymExpr sym tp) ->
      IO (AbsIndex idx, PMuxTree sym (ArrayEntry sym idx tp))
    go (cidx, v) = do
      symIdx <- concreteIdxToSym sym cidx
      entry <- mkValEntry sym symIdx v
      return $ (symIdxToAbs symIdx, toPMuxTree sym entry)

-- | Initialize an array with symbolic contents at symbolic locations
initArray ::
  forall sym idx tp idx' tp'.
  W4.IsSymExprBuilder sym =>
  idx ~ (idx' Ctx.::> tp') =>
  sym ->
  W4.BaseTypeRepr tp ->
  [(Ctx.Assignment (W4.SymExpr sym) idx, W4.SymExpr sym tp)] ->
  IO (CachedArray sym idx tp)
initArray sym repr m = do
  nonce <- freshArrayNonce
  im <- IM.fromList <$> mapM go m
  return $ CachedArray im (\x -> x) repr nonce
  where
    go ::
      (Ctx.Assignment (W4.SymExpr sym) idx, W4.SymExpr sym tp) ->
      IO (AbsIndex idx, PMuxTree sym (ArrayEntry sym idx tp))
    go (symIdxExpr, v) = do
      let
        symIdx = SymIndex symIdxExpr Nothing
      entry <- mkValEntry sym symIdx v
      return $ (symIdxToAbs symIdx, toPMuxTree sym entry)

---------------------------------------------------
-- Implementation

-- | A sentinel nonce that is refreshed every time the array is updated.
newtype ArrayNonce = ArrayNonce (PN.Nonce PN.GlobalNonceGenerator IO)

instance Eq ArrayNonce where
  (ArrayNonce i1) == (ArrayNonce i2) | Just Refl <- testEquality i1 i2 = True
  _ == _ = False

instance Ord ArrayNonce where
  compare (ArrayNonce i1) (ArrayNonce i2) = toOrdering $ compareF i1 i2

freshArrayNonce :: IO ArrayNonce
freshArrayNonce = ArrayNonce <$> PN.freshNonce PN.globalNonceGenerator

-- | An array that supports reading from a stack of mixed concrete/symbolic writes efficiently.
--
-- The primary interface is intended to be 'readChunk' and 'writeChunk', which
-- allow writing contiguous subsequences of data to the array.
--
-- Note that the equality instances is based on a unique nonce (see
-- 'ArrayNonce') that is incremented each time the array is updated, and is thus
-- an identity test rather than a structural equality test.
data CachedArray sym (ctx :: Ctx.Ctx W4.BaseType) (tp :: W4.BaseType) where
  CachedArray ::
    {
      arrMap :: IM.IntervalsMap AbsIntervalEnd ctx (PMuxTree sym (ArrayEntry sym ctx tp))
    , arrConstraints :: forall a. (NonEmptyCtx ctx => a) -> a
    , arrTypeRepr :: W4.BaseTypeRepr tp
    , _arrNonce :: ArrayNonce
    } -> CachedArray sym ctx tp

instance Eq (CachedArray sym idx tp) where
  (CachedArray _ _ _ nonce1) == (CachedArray _ _ _ nonce2) = nonce1 == nonce2

incNonce ::
  CachedArray sym idx tp ->
  IO (CachedArray sym idx tp)
incNonce (CachedArray am ac tr _) = do
  nonce <- freshArrayNonce
  return $ CachedArray am ac tr nonce

-- | An array entry defines a set of possible values for a given
-- abstract domain. Entries may overlap, and so as an invariant we
-- preserve the fact that at each logical index, exactly one entry is valid
data ArrayEntry sym ctx tp where
  ArrayEntry ::
    { -- TODO: should we cache these results?
      entryVals :: (SymIndex sym ctx -> IO (W4.PartExpr (W4.Pred sym) (W4.SymExpr sym tp)))
    , entryNonce :: ArrayNonce
    } -> ArrayEntry sym ctx tp


incNonceEntry ::
  ArrayEntry sym ctx tp ->
  IO (ArrayEntry sym ctx tp)
incNonceEntry (ArrayEntry vals _) = do
  nonce <- freshArrayNonce
  return $ ArrayEntry vals nonce

instance Eq (ArrayEntry sym ctx tp) where
  e1 == e2 = entryNonce e1 == entryNonce e2

instance Ord (ArrayEntry sym ctx tp) where
  compare e1 e2 = compare (entryNonce e1) (entryNonce e2)

-- | A symbolic index into the array. It represents the index for a single array element,
-- although its value may be symbolic
data SymIndex sym ctx =
  SymIndex
    { -- | the symbolic index
      _symIdxExpr :: Ctx.Assignment (W4.SymExpr sym) ctx
      -- | an optional override for the abstract domain of the index
    , symIdxAbs :: Maybe (AbsIndex ctx)
    }

deriving instance W4.IsSymExprBuilder sym => Eq (SymIndex sym ctx)
deriving instance W4.IsSymExprBuilder sym => Ord (SymIndex sym ctx)

-- | An offset is an index into the last element of the array index
-- A value range is always representable as a base + offset
newtype SymOffset sym ctx where
  SymOffset :: W4.SymExpr sym (CtxFirst ctx) -> SymOffset sym ctx

newtype FirstIndex ctx where
  FirstIndex :: Ctx.Index ctx (CtxFirst ctx) -> FirstIndex ctx

skipFirst ::
  FirstIndex (ctx Ctx.::> tp1) -> FirstIndex (ctx Ctx.::> tp1 Ctx.::> tp2)
skipFirst (FirstIndex idx) = FirstIndex (Ctx.skipIndex idx)

firstIndex ::
  forall ctx.
  NonEmptyCtx ctx =>
  Ctx.Size ctx ->
  FirstIndex ctx
firstIndex sz | NonEmptyCtxRepr <- nonEmptyCtxRepr @_ @ctx =
  case Ctx.viewSize (Ctx.decSize sz) of
    Ctx.ZeroSize -> FirstIndex (Ctx.baseIndex)
    Ctx.IncSize _ -> skipFirst (firstIndex (Ctx.decSize sz))

indexToOffset ::
  forall sym ctx.
  NonEmptyCtx ctx =>
  W4.IsSymExprBuilder sym =>
  SymIndex sym ctx ->
  SymOffset sym ctx
indexToOffset (SymIndex eCtx _) =
  let
    FirstIndex idx = firstIndex (Ctx.size eCtx)
    e = eCtx Ctx.! idx
  in SymOffset e

addSymOffset ::
  forall sym ctx.
  W4.IsSymExprBuilder sym =>
  NonEmptyCtx ctx =>
  sym ->
  SymIndex sym ctx ->
  SymOffset sym ctx ->
  IO (SymIndex sym ctx)
addSymOffset sym (SymIndex eCtx _) (SymOffset off) = do
  let
    FirstIndex idx = firstIndex (Ctx.size eCtx)
    e = eCtx Ctx.! idx
  e' <- case W4.exprType off of
    W4.BaseIntegerRepr -> W4.intAdd sym e off
    W4.BaseBVRepr _ -> W4.bvAdd sym e off
    _ -> fail $ "Unsupported type"
  return $ SymIndex (eCtx & (ixF idx) .~ e') Nothing

negateSymOffset ::
  W4.IsSymExprBuilder sym =>
  sym ->
  SymOffset sym ctx ->
  IO (SymOffset sym ctx)
negateSymOffset sym (SymOffset off) = do
  e' <- case W4.exprType off of
    W4.BaseIntegerRepr -> W4.intNeg sym off
    W4.BaseBVRepr _ -> W4.bvNeg sym off
    _ -> fail $ "Unsupported type"
  return $ SymOffset e'

-- | Previous offset from the given one, to create an exclusive upper bound
prevSymOffset ::
  W4.IsSymExprBuilder sym =>
  sym ->
  SymOffset sym ctx ->
  IO (SymOffset sym ctx)
prevSymOffset sym (SymOffset off) = do
  e' <- case W4.exprType off of
    W4.BaseIntegerRepr -> do
      one <- W4.intLit sym 1
      W4.intSub sym off one
    W4.BaseBVRepr w -> do
      one <- W4.bvLit sym w (BV.mkBV w 1)
      W4.bvSub sym off one
    _ -> fail $ "Unsupported type"
  return $ SymOffset e'

subSymOffset ::
  W4.IsSymExprBuilder sym =>
  NonEmptyCtx ctx =>
  sym ->
  SymIndex sym ctx ->
  SymOffset sym ctx ->
  IO (SymIndex sym ctx)
subSymOffset sym idx off = do
  negoff <- negateSymOffset sym off
  addSymOffset sym idx negoff

mkSymIndex ::
  forall sym ctx.
  Ctx.Assignment (W4.SymExpr sym) ctx ->
  SymIndex sym ctx
mkSymIndex e = SymIndex e Nothing

-- | Represents a symbolic range, where equality and ordering is defined on the
-- abstract domain of the underlying expression
data SymRange sym ctx =
    SymRangeSingle (SymIndex sym ctx)
  | SymRangeMulti (SymIndex sym ctx) (SymIndex sym ctx)

symRangeLo :: SymRange sym ctx -> SymIndex sym ctx
symRangeLo (SymRangeSingle symIdx) = symIdx
symRangeLo (SymRangeMulti loIdx _) = loIdx


symRangeToAbs ::
  W4.IsSymExprBuilder sym =>
  SymRange sym ctx ->
  AbsIndex ctx
symRangeToAbs (SymRangeSingle symIdx) = symIdxToAbs symIdx
symRangeToAbs (SymRangeMulti loIdx hiIdx) =
  joinAbsIndex (symIdxToAbs loIdx) (symIdxToAbs hiIdx)


-- | Create a range that is exclusive of the given offset
-- i.e. 2 + 4 --> (2, 5)
mkSymRangeOff ::
  forall sym ctx .
  W4.IsSymExprBuilder sym =>
  NonEmptyCtx ctx =>
  sym ->
  Ctx.Assignment (W4.SymExpr sym) ctx ->
  W4.SymExpr sym (CtxFirst ctx) ->
  IO (SymRange sym ctx)
mkSymRangeOff sym loExpr offExpr = do
  let
    lo = mkSymIndex @sym loExpr
    off = SymOffset offExpr
  offPrev <- prevSymOffset sym off
  hi <- addSymOffset sym lo offPrev
  return $ (SymRangeMulti lo hi)

data NonEmptyCtxRepr (ctx :: Ctx.Ctx k) where
  NonEmptyCtxRepr :: NonEmptyCtxRepr (ctx Ctx.::> x)

type family CtxFirst (ctx :: Ctx.Ctx k) where
  CtxFirst (Ctx.EmptyCtx Ctx.::> a) = a
  CtxFirst (ctx Ctx.::> _) = CtxFirst ctx

class NonEmptyCtx (ctx :: Ctx.Ctx k) where
  type CtxHead ctx :: k
  type CtxTail ctx :: Ctx.Ctx k

  nonEmptyCtxRepr :: NonEmptyCtxRepr ctx

instance NonEmptyCtx (ctx Ctx.::> tp) where
  type CtxHead (ctx Ctx.::> tp) = tp
  type CtxTail (ctx Ctx.::> tp) = ctx
  nonEmptyCtxRepr = NonEmptyCtxRepr

data AbsIntervalEnd tp where
  AbsIntervalEndInt :: W4.ValueBound Integer -> AbsIntervalEnd W4.BaseIntegerType
  AbsIntervalEndBV :: (1 <= w) => W4.NatRepr w -> W4.ValueBound Integer -> AbsIntervalEnd (W4.BaseBVType w)

instance Ord (AbsIntervalEnd tp) where
  compare a1 a2 = toOrdering $ compareF a1 a2

instance Eq (AbsIntervalEnd tp) where
  a1 == a2 = (compare a1 a2) == EQ

instance TestEquality AbsIntervalEnd where
  testEquality a1 a2 = case compareF a1 a2 of
    EQF -> Just Refl
    _ -> Nothing

instance OrdF AbsIntervalEnd where
  compareF a1 a2 = case (a1, a2) of
    (AbsIntervalEndInt n1, AbsIntervalEndInt n2) -> fromOrdering $ compare n1 n2
    (AbsIntervalEndBV w1 i1, AbsIntervalEndBV w2 i2) ->
      lexCompareF w1 w2 $ fromOrdering $ compare i1 i2
    (AbsIntervalEndInt{}, AbsIntervalEndBV{}) -> LTF
    (AbsIntervalEndBV{}, AbsIntervalEndInt{}) -> GTF


type AbsIndex (idx :: Ctx.Ctx W4.BaseType) = IM.Intervals AbsIntervalEnd idx
type AbsInterval tp = IM.IntervalF AbsIntervalEnd tp


bvDomainRange ::
  1 <= w =>
  W4.NatRepr w ->
  BVD.BVDomain w ->
  AbsInterval (W4.BaseBVType w)
bvDomainRange w d = case BVD.ubounds d of
  (i1, i2) -> IM.mkIntervalF $ IM.ClosedInterval (AbsIntervalEndBV w (W4.Inclusive i1)) (AbsIntervalEndBV w (W4.Inclusive i2))

exprToAbsInterval ::
  forall sym tp.
  W4.IsSymExprBuilder sym =>
  W4.SymExpr sym tp ->
  AbsInterval tp
exprToAbsInterval e = absToInterval (W4.exprType e) (W4.getAbsValue e)

absToInterval ::
  W4.BaseTypeRepr tp ->
  W4.AbstractValue tp ->
  AbsInterval tp
absToInterval repr v = case repr of
  W4.BaseIntegerRepr -> case v of
    W4.SingleRange x -> IM.mkIntervalF $ IM.ClosedInterval (AbsIntervalEndInt (W4.Inclusive x)) (AbsIntervalEndInt (W4.Inclusive x))
    W4.MultiRange lo hi -> IM.mkIntervalF $ IM.ClosedInterval (AbsIntervalEndInt lo) (AbsIntervalEndInt hi)
  W4.BaseBVRepr w -> bvDomainRange w v
  _ -> error "Unsupported type"


readArrayBase ::
  forall sym idx tp.
  W4.IsSymExprBuilder sym =>
  sym ->
  SymIndex sym idx ->
  CachedArray sym idx tp ->
  IO (W4.SymExpr sym tp)
readArrayBase sym symIdx arr = do
  let
    intersecting = IM.toList $ IM.intersecting (arrMap arr) (symIdxToAbs symIdx)
  entries <- mapM expandEntry $ concat $ map (viewPMuxTree . snd) intersecting
  case entries of
    [(W4.PE p (AsOrd e), path_cond)]
      | Just True <- W4.asConstantPred path_cond
      , Just True <- W4.asConstantPred p -> return e
    entryExprs -> arrConstraints arr $ do

      muxTree <- mkPMuxTreePartial sym entryExprs
      MT.collapseMuxTree sym ite muxTree >>= \case
        Just (AsOrd e) -> return e
        -- garbage result
        Nothing -> W4.freshConstant sym W4.emptySymbol (arrTypeRepr arr)

  where
    ite ::
      W4.Pred sym ->
      Maybe (AsOrd (W4.SymExpr sym) tp) ->
      Maybe (AsOrd (W4.SymExpr sym) tp) ->
      IO (Maybe (AsOrd (W4.SymExpr sym) tp))
    ite p (Just (AsOrd e1)) (Just (AsOrd e2)) = (Just . AsOrd) <$> W4.baseTypeIte sym p e1 e2
    ite _ Nothing (Just e2) = return $ Just e2
    ite _ (Just e1) Nothing = return $ Just e1
    ite _ Nothing Nothing = return Nothing

    expandEntry ::
      (ArrayEntry sym idx tp, W4.Pred sym) ->
      IO (W4.PartExpr (W4.Pred sym) (AsOrd (W4.SymExpr sym) tp), W4.Pred sym)
    expandEntry (entry, path_cond) = do
      val <- entryVals entry symIdx
      return $ (fmap AsOrd val, path_cond)

mkValEntry ::
  W4.IsSymExprBuilder sym =>
  sym ->
  SymIndex sym ctx ->
  W4.SymExpr sym tp ->
  IO (ArrayEntry sym ctx tp)
mkValEntry sym idx v = do
  let vals idx' = do
        p <- isEqIndex sym idx idx'
        return $ W4.mkPE p v
  mkMultiEntry vals

mkMultiEntry ::
  (SymIndex sym ctx -> IO (W4.PartExpr (W4.Pred sym) (W4.SymExpr sym tp))) ->
  IO (ArrayEntry sym ctx tp)
mkMultiEntry vals = do
  nonce <- freshArrayNonce
  return $ ArrayEntry vals nonce

symIdxToAbs ::
  forall sym ctx.
  W4.IsSymExprBuilder sym =>
  SymIndex sym ctx -> AbsIndex ctx
symIdxToAbs (SymIndex symIdxExpr Nothing) = IM.Intervals $ FC.fmapFC (exprToAbsInterval @sym) symIdxExpr
symIdxToAbs (SymIndex _ (Just absIdx)) = absIdx

concreteIdxToSym ::
  forall sym ctx.
  W4.IsSymExprBuilder sym =>
  sym ->
  Ctx.Assignment W4.ConcreteVal ctx ->
  IO (SymIndex sym ctx)
concreteIdxToSym sym conc = do
 symIdxExpr <- FC.traverseFC (W4.concreteToSym sym) conc
 return $ SymIndex symIdxExpr Nothing


-- | Invalidate all entries within the given range
-- TODO: delete entries which are statically invalid
invalidateRange ::
  forall sym ctx tp.
  W4.IsSymExprBuilder sym =>
  sym ->
  -- | range to invalidate
  SymRange sym ctx ->
  ArrayEntry sym ctx tp ->
  IO (Maybe (ArrayEntry sym ctx tp))
invalidateRange sym invalid_rng entry = do
  let vals symIdx' = do
        notThis <- W4.notPred sym =<< isInRange sym invalid_rng symIdx'
        val <- entryVals entry symIdx'
        W4.runPartialT sym notThis $ W4.returnPartial val
  entry' <- incNonceEntry $ entry { entryVals = vals }
  return $ Just entry'


isInRange ::
  forall sym ctx.
  W4.IsSymExprBuilder sym =>
  sym ->
  SymRange sym ctx ->
  SymIndex sym ctx ->
  IO (W4.Pred sym)
isInRange sym rng symIdx2@(SymIndex symIdxExpr _) = case rng of
  SymRangeSingle symIdx1 -> isEqIndex sym symIdx1 symIdx2
  SymRangeMulti (SymIndex loIdxExpr _) (SymIndex hiIdxExpr _) -> do
    lo <- FC.toListFC getConst <$> Ctx.zipWithM doLe loIdxExpr symIdxExpr
    hi <- FC.toListFC getConst <$> Ctx.zipWithM doLe symIdxExpr hiIdxExpr
    foldM (W4.andPred sym) (W4.truePred sym) $ lo ++ hi
  where
    doLe ::
      forall tp.
      W4.SymExpr sym tp ->
      W4.SymExpr sym tp ->
      IO (Const (W4.Pred sym) tp)
    doLe e1 e2 = Const <$> case W4.exprType e1 of
      W4.BaseBVRepr _ -> W4.bvUle sym e1 e2
      W4.BaseIntegerRepr -> W4.intLe sym e1 e2
      _ -> fail "isInRange: unsupported type"

isEqIndex ::
  forall sym ctx.
  W4.IsSymExprBuilder sym =>
  sym ->
  SymIndex sym ctx ->
  SymIndex sym ctx ->
  IO (W4.Pred sym)
isEqIndex sym (SymIndex symIdxExpr1 _) (SymIndex symIdxExpr2 _) = do
  preds <- FC.toListFC getConst <$> Ctx.zipWithM (\e1 e2 -> Const <$> W4.isEq sym e1 e2) symIdxExpr1 symIdxExpr2
  foldM (W4.andPred sym) (W4.truePred sym) preds

-- | Invalidate all existing symbolic entries at exactly this index
invalidateEntries ::
  forall sym ctx tp.
  W4.IsSymExprBuilder sym =>
  sym ->
  SymRange sym ctx ->
  CachedArray sym ctx tp ->
  IO (CachedArray sym ctx tp)
invalidateEntries sym symRange arr = arrConstraints arr $ do
  NonEmptyCtxRepr <- return $ nonEmptyCtxRepr @_ @ctx
  cmap <- IM.mapMIntersecting absIndex (\v -> getMaybe <$> pmuxTreeMaybeOp sym (invalidateRange sym symRange) v) (arrMap arr)
  return $ arr { arrMap = cmap }
  where
    absIndex = symRangeToAbs symRange
    getMaybe :: PMuxTree sym (ArrayEntry sym ctx tp) -> Maybe (PMuxTree sym (ArrayEntry sym ctx tp))
    getMaybe mt | isEmptyPMuxTree mt = Nothing
    getMaybe mt = Just mt

buildMuxTree :: (W4.IsExprBuilder sym, Ord a) => sym -> a -> [(a, W4.Pred sym)] -> IO (MT.MuxTree sym a)
buildMuxTree sym a as =
  foldM (\mt (a',p) -> MT.mergeMuxTree sym p (MT.toMuxTree sym a') mt) (MT.toMuxTree sym a) as


joinAbsIndex ::
  AbsIndex ctx ->
  AbsIndex ctx ->
  AbsIndex ctx
joinAbsIndex (IM.Intervals idx1) (IM.Intervals idx2) = IM.Intervals $ Ctx.zipWith IM.mergeIntervalsF idx1 idx2


muxEntries ::
  W4.IsSymExprBuilder sym =>
  sym ->
  W4.Pred sym ->
  PMuxTree sym (ArrayEntry sym ctx tp) ->
  PMuxTree sym (ArrayEntry sym ctx tp) ->
  IO (PMuxTree sym (ArrayEntry sym ctx tp))
muxEntries sym p mtT mtF = MT.mergeMuxTree sym p mtT mtF

mergeEntries ::
  forall sym ctx tp.
  W4.IsSymExprBuilder sym =>
  NonEmptyCtx ctx =>
  sym ->
  (SymIndex sym ctx -> IO (W4.Pred sym)) ->
  ArrayEntry sym ctx tp ->
  ArrayEntry sym ctx tp ->
  IO (ArrayEntry sym ctx tp)
mergeEntries sym pickLeftFn e1 e2 = do
  let vals symIdx' = do
        pickLeft <- pickLeftFn symIdx'
        val1 <- entryVals e1 symIdx'
        val2 <- entryVals e2 symIdx'
        W4.mergePartial sym (\p a b -> lift $ W4.baseTypeIte sym p a b)
          pickLeft val1 val2
  incNonceEntry $ e1 { entryVals = vals }

mergeEntriesMux ::
  forall sym ctx tp.
  W4.IsSymExprBuilder sym =>
  NonEmptyCtx ctx =>
  sym ->
  (SymIndex sym ctx -> IO (W4.Pred sym)) ->
  PMuxTree sym (ArrayEntry sym ctx tp) ->
  PMuxTree sym (ArrayEntry sym ctx tp) ->
  IO (PMuxTree sym (ArrayEntry sym ctx tp))
mergeEntriesMux sym pickLeftFn = pmuxTreeBinOp sym (mergeEntries sym pickLeftFn)

-- | A partial mux tree
type PMuxTree sym tp = MT.MuxTree sym (Maybe tp)

viewPMuxTree :: forall sym a. PMuxTree sym a -> [(a, W4.Pred sym)]
viewPMuxTree mt = mapMaybe go $ MT.viewMuxTree mt
  where
    go :: (Maybe a, W4.Pred sym) -> Maybe (a, W4.Pred sym)
    go (Just a, p) = Just (a, p)
    go _ = Nothing

isEmptyPMuxTree :: PMuxTree sym tp -> Bool
isEmptyPMuxTree mt = case MT.viewMuxTree mt of
  [(Nothing, _)] -> True
  _ -> False

mkPMuxTree ::
  (W4.IsExprBuilder sym, Ord a) =>
  sym ->
  [(a, W4.Pred sym)] ->
  IO (PMuxTree sym a)
mkPMuxTree sym ls = buildMuxTree sym Nothing (map (\(a, p) -> (Just a, p)) ls)

mkPMuxTreePartial ::
  forall sym a.
  (W4.IsExprBuilder sym, Ord a) =>
  sym ->
  [(W4.PartExpr (W4.Pred sym) a, W4.Pred sym)] ->
  IO (PMuxTree sym a)
mkPMuxTreePartial sym ls = mkPMuxTree sym =<< (catMaybes <$> mapM go ls)
  where
    go :: (W4.PartExpr (W4.Pred sym) a, W4.Pred sym) -> IO (Maybe (a, W4.Pred sym))
    go (W4.PE p a, cond) = do
      p' <- W4.andPred sym p cond
      return $ Just (a, p')
    go (W4.Unassigned, _) = return Nothing

pmuxTreeAddCondition ::
  forall sym a.
  W4.IsExprBuilder sym =>
  Ord a =>
  sym ->
  W4.Pred sym ->
  PMuxTree sym a ->
  IO (PMuxTree sym a)
pmuxTreeAddCondition sym cond mt = mkPMuxTree sym =<< mapM addCond (viewPMuxTree mt)
  where
    addCond :: (a, W4.Pred sym) -> IO (a, W4.Pred sym)
    addCond (a, cond') = do
      cond'' <- W4.andPred sym cond cond'
      return $ (a, cond'')


pmuxTreeMaybeOp ::
  (W4.IsExprBuilder sym, Ord a, Ord b) =>
  sym ->
  (a -> IO (Maybe b)) ->
  PMuxTree sym a ->
  IO (PMuxTree sym b)
pmuxTreeMaybeOp sym f mt = MT.muxTreeUnaryOp sym (\a -> join <$> mapM f a) mt

_pmuxTreeUnaryOp ::
  (W4.IsExprBuilder sym, Ord b) =>
  sym ->
  (a -> IO b) ->
  PMuxTree sym a ->
  IO (PMuxTree sym b)
_pmuxTreeUnaryOp sym f mt = MT.muxTreeUnaryOp sym (\a -> mapM f a) mt


pmuxTreeBinOp ::
  forall sym a b c.
  (W4.IsExprBuilder sym, Ord c) =>
  sym ->
  (a -> b -> IO c) ->
  PMuxTree sym a ->
  PMuxTree sym b ->
  IO (PMuxTree sym c)
pmuxTreeBinOp sym f mt1 mt2 = MT.muxTreeBinOp sym g mt1 mt2
  where
    g :: Maybe a -> Maybe b -> IO (Maybe c)
    g (Just a) (Just b) = Just <$> f a b
    g _ _ = return Nothing

toPMuxTree :: W4.IsExprBuilder sym => sym -> a -> PMuxTree sym a
toPMuxTree sym a = MT.toMuxTree sym (Just a)