content: {-|
  Module           : Lang.Crucible.Utils.MuxTree
  Copyright        : (c) Galois, Inc 2018
  License          : BSD3
  Maintainer       : Rob Dockins <rdockins@galois.com>
  
  This module defines a @MuxTree@ type that notionally represents
  a collection of values organized into an if-then-else tree.  This
  data structure allows values that otherwise do not have a useful notion
  of symbolic values to nonetheless be merged as control flow merge points
  by simply remembering which concrete values were obtained, and the
  logical conditions under which they were found.
  
  Note that we require an @Ord@ instance on the type @a@ over which we are
  building the mux trees.  It is sufficent that this operation be merely
  syntactic equality; it is not necessary for correctness that terms with
  the same semantics compare equal.
  -}
  
  {-# LANGUAGE FlexibleContexts #-}
  module Lang.Crucible.Utils.MuxTree
    ( MuxTree
    , toMuxTree
    , mergeMuxTree
    , viewMuxTree
    , muxTreeUnaryOp
    , muxTreeBinOp
    , muxTreeCmpOp
    , collapseMuxTree
    , muxTreeEq
    , muxTreeLe
    , muxTreeLt
    , muxTreeGe
    , muxTreeGt
    ) where
  
  import           Control.Lens (folded)
  
  import           Data.Map.Strict (Map)
  import qualified Data.Map.Strict as Map
  import qualified Data.Map.Merge.Strict as Map
  
  import           What4.Interface
  import           Lang.Crucible.Panic
  
  -- | A mux tree represents a collection of if-then-else branches over
  --   a collection of values.  Generally, a mux tree is used to provide
  --   a way to conditionally merge values that otherwise do not
  --   naturally have a merge operation.
  newtype MuxTree sym a = MuxTree (Map a (Pred sym))
  {- INVARIANT: The map inside a mux tree is non-empty! -}
  
  -- Turn a single value into a trivial mux tree
  toMuxTree :: IsExprBuilder sym => sym -> a -> MuxTree sym a
  toMuxTree sym v = MuxTree (Map.singleton v (truePred sym))
  
  -- View all the leaf values of the mux tree, along with the
  -- conditions that lead to those values.
  viewMuxTree :: MuxTree sym a -> [(a, Pred sym)]
  viewMuxTree (MuxTree m) = Map.toList m
  
  _conditionMuxTree :: IsExprBuilder sym => sym -> Pred sym -> MuxTree sym a -> IO (MuxTree sym a)
  _conditionMuxTree sym p (MuxTree m) = MuxTree <$> Map.traverseMaybeWithKey (conditionMuxTreeLeaf sym p) m
  
  -- | Compute a binary boolean predicate between two mux trees.
  --   This operation decomposes the mux trees and compares
  --   all combinations of the underlying values, conditional on
  --   the path conditions leading to those values.
  muxTreeCmpOp ::
    IsExprBuilder sym =>
    sym ->
    (a -> a -> IO (Pred sym)) {- ^ compute the predicate on the underlying type -} ->
    MuxTree sym a ->
    MuxTree sym a ->
    IO (Pred sym)
  muxTreeCmpOp sym f xt yt = orOneOf sym folded =<< sequence zs
    where
    zs = [ do pf <- f x y
              andPred sym pf =<< andPred sym px py
         | (x,px) <- xs
         , (y,py) <- ys
         ]
    xs = viewMuxTree xt
    ys = viewMuxTree yt
  
  
  -- | Compute an equality predicate on mux trees.
  --
  --   NOTE! This assumes the equality relation
  --   defined by `Eq` is the semantic equality
  --   relation on @a@.
  muxTreeEq ::
    (Eq a, IsExprBuilder sym) =>
    sym ->
    MuxTree sym a ->
    MuxTree sym a ->
    IO (Pred sym)
  muxTreeEq sym = muxTreeCmpOp sym f
    where f x y = pure (backendPred sym (x == y))
  
  -- | Compute a less-than predicate on mux trees.
  --
  --   NOTE! This assumes the order relation
  --   defined by `Ord` is the semantic order
  --   relation on @a@.
  muxTreeLt ::
    (Ord a, IsExprBuilder sym) =>
    sym ->
    MuxTree sym a ->
    MuxTree sym a ->
    IO (Pred sym)
  muxTreeLt sym = muxTreeCmpOp sym f
    where f x y = pure (backendPred sym (x < y))
  
  -- | Compute a less-than-or-equal predicate on mux trees.
  --
  --   NOTE! This assumes the order relation
  --   defined by `Ord` is the semantic order
  --   relation on @a@.
  muxTreeLe ::
    (Ord a, IsExprBuilder sym) =>
    sym ->
    MuxTree sym a ->
    MuxTree sym a ->
    IO (Pred sym)
  muxTreeLe sym = muxTreeCmpOp sym f
    where f x y = pure (backendPred sym (x <= y))
  
  -- | Compute a greater-than predicate on mux trees.
  --
  --   NOTE! This assumes the order relation
  --   defined by `Ord` is the semantic order
  --   relation on @a@.
  muxTreeGt ::
    (Ord a, IsExprBuilder sym) =>
    sym ->
    MuxTree sym a ->
    MuxTree sym a ->
    IO (Pred sym)
  muxTreeGt sym = muxTreeCmpOp sym f
    where f x y = pure (backendPred sym (x > y))
  
  -- | Compute a greater-than-or-equal predicate on mux trees.
  --
  --   NOTE! This assumes the order relation
  --   defined by `Ord` is the semantic order
  --   relation on @a@.
  muxTreeGe ::
    (Ord a, IsExprBuilder sym) =>
    sym ->
    MuxTree sym a ->
    MuxTree sym a ->
    IO (Pred sym)
  muxTreeGe sym = muxTreeCmpOp sym f
    where f x y = pure (backendPred sym (x >= y))
  
  
  -- | Use the provided if-then-else operation to collapse the given mux tree
  --   into its underlying type.
  collapseMuxTree ::
    IsExprBuilder sym =>
    sym ->
    (Pred sym -> a -> a -> IO a) ->
    MuxTree sym a ->
    IO a
  collapseMuxTree _sym ite xt = go (viewMuxTree xt)
    where
    go []         = panic "collapseMuxTree" ["empty mux tree"]
    go [(x,_p)]   = return x
    go ((x,p):xs) = ite p x =<< go xs
  
  buildMuxTree ::
    (Ord a, IsExprBuilder sym) =>
    sym ->
    [(a, Pred sym)] ->
    IO (MuxTree sym a)
  buildMuxTree _sym [] = panic "buildMuxTree" ["empty mux tree"]
  buildMuxTree sym  xs = go Map.empty xs
    where
    go m [] = return (MuxTree m)
    go m ((z,p):zs) =
       case Map.lookup z m of
         Nothing -> go (Map.insert z p m) zs
         Just q -> do pq <- orPred sym p q
                      case asConstantPred pq of
                        Just False -> go m zs
                        _ -> go (Map.insert z pq m) zs
  
  -- | Apply a unary operation through a mux tree.  The provided operation
  --   is applied to each leaf of the tree.
  muxTreeUnaryOp ::
    (Ord b, IsExprBuilder sym) =>
    sym ->
    (a -> IO b) ->
    MuxTree sym a ->
    IO (MuxTree sym b)
  muxTreeUnaryOp sym op xt =
    do let xs = viewMuxTree xt
       zs <- sequence
              [ do z <- op x
                   return (z,p)
              | (x,p) <- xs
              ]
       buildMuxTree sym zs
  
  -- | Apply a binary operation through two mux trees.  The provided operation
  --   is applied pairwise to each leaf of the two trees, and appropriate path
  --   conditions are computed for the resulting values.
  muxTreeBinOp ::
    (Ord c, IsExprBuilder sym) =>
    sym ->
    (a -> b -> IO c) ->
    MuxTree sym a ->
    MuxTree sym b ->
    IO (MuxTree sym c)
  muxTreeBinOp sym op xt yt =
    do let xs = viewMuxTree xt
       let ys = viewMuxTree yt
       zs <- sequence
             [ do p <- andPred sym px py
                  z <- op x y
                  return (z,p)
             | (x,px) <- xs
             , (y,py) <- ys
             ]
       buildMuxTree sym zs
  
  
  conditionMuxTreeLeaf ::
    IsExprBuilder sym => sym -> Pred sym -> a -> Pred sym -> IO (Maybe (Pred sym))
  conditionMuxTreeLeaf sym p _v pv =
     do p' <- andPred sym p pv
        case asConstantPred p' of
          Just False -> return Nothing
          _ -> return (Just p')
  
  -- | Compute the if-then-else operation on mux trees.
  mergeMuxTree ::
    (Ord a, IsExprBuilder sym) =>
    sym ->
    Pred sym ->
    MuxTree sym a ->
    MuxTree sym a ->
    IO (MuxTree sym a)
  mergeMuxTree sym p (MuxTree mx) (MuxTree my) =
     do np <- notPred sym p
        MuxTree <$> doMerge np mx my
  
    where
    f _v px py =
      do p' <- itePred sym p px py
         case asConstantPred p' of
           Just False -> return Nothing
           _ -> return (Just p')
  
    doMerge np = Map.mergeA (Map.traverseMaybeMissing (conditionMuxTreeLeaf sym p))
                            (Map.traverseMaybeMissing (conditionMuxTreeLeaf sym np))
                            (Map.zipWithMaybeAMatched f)
  