{-|
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 :: forall sym a. IsExprBuilder sym => sym -> a -> MuxTree sym a
toMuxTree sym
sym a
v = Map a (Pred sym) -> MuxTree sym a
forall sym a. Map a (Pred sym) -> MuxTree sym a
MuxTree (a -> Pred sym -> Map a (Pred sym)
forall k a. k -> a -> Map k a
Map.singleton a
v (sym -> Pred sym
forall sym. IsExprBuilder sym => sym -> Pred sym
truePred sym
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 :: forall sym a. MuxTree sym a -> [(a, Pred sym)]
viewMuxTree (MuxTree Map a (Pred sym)
m) = Map a (Pred sym) -> [(a, Pred sym)]
forall k a. Map k a -> [(k, a)]
Map.toList Map a (Pred sym)
m

_conditionMuxTree :: IsExprBuilder sym => sym -> Pred sym -> MuxTree sym a -> IO (MuxTree sym a)
_conditionMuxTree :: forall sym a.
IsExprBuilder sym =>
sym -> Pred sym -> MuxTree sym a -> IO (MuxTree sym a)
_conditionMuxTree sym
sym Pred sym
p (MuxTree Map a (Pred sym)
m) = Map a (Pred sym) -> MuxTree sym a
forall sym a. Map a (Pred sym) -> MuxTree sym a
MuxTree (Map a (Pred sym) -> MuxTree sym a)
-> IO (Map a (Pred sym)) -> IO (MuxTree sym a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> Pred sym -> IO (Maybe (Pred sym)))
-> Map a (Pred sym) -> IO (Map a (Pred sym))
forall (f :: Type -> Type) k a b.
Applicative f =>
(k -> a -> f (Maybe b)) -> Map k a -> f (Map k b)
Map.traverseMaybeWithKey (sym -> Pred sym -> a -> Pred sym -> IO (Maybe (Pred sym))
forall sym a.
IsExprBuilder sym =>
sym -> Pred sym -> a -> Pred sym -> IO (Maybe (Pred sym))
conditionMuxTreeLeaf sym
sym Pred sym
p) Map a (Pred sym)
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 :: forall sym a.
IsExprBuilder sym =>
sym
-> (a -> a -> IO (Pred sym))
-> MuxTree sym a
-> MuxTree sym a
-> IO (Pred sym)
muxTreeCmpOp sym
sym a -> a -> IO (Pred sym)
f MuxTree sym a
xt MuxTree sym a
yt = sym -> Fold [Pred sym] (Pred sym) -> [Pred sym] -> IO (Pred sym)
forall sym s.
IsExprBuilder sym =>
sym -> Fold s (Pred sym) -> s -> IO (Pred sym)
orOneOf sym
sym (Pred sym -> f (Pred sym)) -> [Pred sym] -> f [Pred sym]
Fold [Pred sym] (Pred sym)
forall (f :: Type -> Type) a. Foldable f => IndexedFold Int (f a) a
IndexedFold Int [Pred sym] (Pred sym)
folded ([Pred sym] -> IO (Pred sym)) -> IO [Pred sym] -> IO (Pred sym)
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< [IO (Pred sym)] -> IO [Pred sym]
forall (t :: Type -> Type) (m :: Type -> Type) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: Type -> Type) a. Monad m => [m a] -> m [a]
sequence [IO (Pred sym)]
zs
  where
  zs :: [IO (Pred sym)]
zs = [ do Pred sym
pf <- a -> a -> IO (Pred sym)
f a
x a
y
            sym -> Pred sym -> Pred sym -> IO (Pred sym)
forall sym.
IsExprBuilder sym =>
sym -> Pred sym -> Pred sym -> IO (Pred sym)
andPred sym
sym Pred sym
pf (Pred sym -> IO (Pred sym)) -> IO (Pred sym) -> IO (Pred sym)
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< sym -> Pred sym -> Pred sym -> IO (Pred sym)
forall sym.
IsExprBuilder sym =>
sym -> Pred sym -> Pred sym -> IO (Pred sym)
andPred sym
sym Pred sym
px Pred sym
py
       | (a
x,Pred sym
px) <- [(a, Pred sym)]
xs
       , (a
y,Pred sym
py) <- [(a, Pred sym)]
ys
       ]
  xs :: [(a, Pred sym)]
xs = MuxTree sym a -> [(a, Pred sym)]
forall sym a. MuxTree sym a -> [(a, Pred sym)]
viewMuxTree MuxTree sym a
xt
  ys :: [(a, Pred sym)]
ys = MuxTree sym a -> [(a, Pred sym)]
forall sym a. MuxTree sym a -> [(a, Pred sym)]
viewMuxTree MuxTree sym a
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 :: forall a sym.
(Eq a, IsExprBuilder sym) =>
sym -> MuxTree sym a -> MuxTree sym a -> IO (Pred sym)
muxTreeEq sym
sym = sym
-> (a -> a -> IO (SymExpr sym BaseBoolType))
-> MuxTree sym a
-> MuxTree sym a
-> IO (SymExpr sym BaseBoolType)
forall sym a.
IsExprBuilder sym =>
sym
-> (a -> a -> IO (Pred sym))
-> MuxTree sym a
-> MuxTree sym a
-> IO (Pred sym)
muxTreeCmpOp sym
sym a -> a -> IO (SymExpr sym BaseBoolType)
forall {f :: Type -> Type} {a}.
(Applicative f, Eq a) =>
a -> a -> f (SymExpr sym BaseBoolType)
f
  where f :: a -> a -> f (SymExpr sym BaseBoolType)
f a
x a
y = SymExpr sym BaseBoolType -> f (SymExpr sym BaseBoolType)
forall a. a -> f a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (sym -> Bool -> SymExpr sym BaseBoolType
forall sym. IsExprBuilder sym => sym -> Bool -> Pred sym
backendPred sym
sym (a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
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 :: forall a sym.
(Ord a, IsExprBuilder sym) =>
sym -> MuxTree sym a -> MuxTree sym a -> IO (Pred sym)
muxTreeLt sym
sym = sym
-> (a -> a -> IO (SymExpr sym BaseBoolType))
-> MuxTree sym a
-> MuxTree sym a
-> IO (SymExpr sym BaseBoolType)
forall sym a.
IsExprBuilder sym =>
sym
-> (a -> a -> IO (Pred sym))
-> MuxTree sym a
-> MuxTree sym a
-> IO (Pred sym)
muxTreeCmpOp sym
sym a -> a -> IO (SymExpr sym BaseBoolType)
forall {f :: Type -> Type} {a}.
(Applicative f, Ord a) =>
a -> a -> f (SymExpr sym BaseBoolType)
f
  where f :: a -> a -> f (SymExpr sym BaseBoolType)
f a
x a
y = SymExpr sym BaseBoolType -> f (SymExpr sym BaseBoolType)
forall a. a -> f a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (sym -> Bool -> SymExpr sym BaseBoolType
forall sym. IsExprBuilder sym => sym -> Bool -> Pred sym
backendPred sym
sym (a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
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 :: forall a sym.
(Ord a, IsExprBuilder sym) =>
sym -> MuxTree sym a -> MuxTree sym a -> IO (Pred sym)
muxTreeLe sym
sym = sym
-> (a -> a -> IO (SymExpr sym BaseBoolType))
-> MuxTree sym a
-> MuxTree sym a
-> IO (SymExpr sym BaseBoolType)
forall sym a.
IsExprBuilder sym =>
sym
-> (a -> a -> IO (Pred sym))
-> MuxTree sym a
-> MuxTree sym a
-> IO (Pred sym)
muxTreeCmpOp sym
sym a -> a -> IO (SymExpr sym BaseBoolType)
forall {f :: Type -> Type} {a}.
(Applicative f, Ord a) =>
a -> a -> f (SymExpr sym BaseBoolType)
f
  where f :: a -> a -> f (SymExpr sym BaseBoolType)
f a
x a
y = SymExpr sym BaseBoolType -> f (SymExpr sym BaseBoolType)
forall a. a -> f a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (sym -> Bool -> SymExpr sym BaseBoolType
forall sym. IsExprBuilder sym => sym -> Bool -> Pred sym
backendPred sym
sym (a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
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 :: forall a sym.
(Ord a, IsExprBuilder sym) =>
sym -> MuxTree sym a -> MuxTree sym a -> IO (Pred sym)
muxTreeGt sym
sym = sym
-> (a -> a -> IO (SymExpr sym BaseBoolType))
-> MuxTree sym a
-> MuxTree sym a
-> IO (SymExpr sym BaseBoolType)
forall sym a.
IsExprBuilder sym =>
sym
-> (a -> a -> IO (Pred sym))
-> MuxTree sym a
-> MuxTree sym a
-> IO (Pred sym)
muxTreeCmpOp sym
sym a -> a -> IO (SymExpr sym BaseBoolType)
forall {f :: Type -> Type} {a}.
(Applicative f, Ord a) =>
a -> a -> f (SymExpr sym BaseBoolType)
f
  where f :: a -> a -> f (SymExpr sym BaseBoolType)
f a
x a
y = SymExpr sym BaseBoolType -> f (SymExpr sym BaseBoolType)
forall a. a -> f a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (sym -> Bool -> SymExpr sym BaseBoolType
forall sym. IsExprBuilder sym => sym -> Bool -> Pred sym
backendPred sym
sym (a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
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 :: forall a sym.
(Ord a, IsExprBuilder sym) =>
sym -> MuxTree sym a -> MuxTree sym a -> IO (Pred sym)
muxTreeGe sym
sym = sym
-> (a -> a -> IO (SymExpr sym BaseBoolType))
-> MuxTree sym a
-> MuxTree sym a
-> IO (SymExpr sym BaseBoolType)
forall sym a.
IsExprBuilder sym =>
sym
-> (a -> a -> IO (Pred sym))
-> MuxTree sym a
-> MuxTree sym a
-> IO (Pred sym)
muxTreeCmpOp sym
sym a -> a -> IO (SymExpr sym BaseBoolType)
forall {f :: Type -> Type} {a}.
(Applicative f, Ord a) =>
a -> a -> f (SymExpr sym BaseBoolType)
f
  where f :: a -> a -> f (SymExpr sym BaseBoolType)
f a
x a
y = SymExpr sym BaseBoolType -> f (SymExpr sym BaseBoolType)
forall a. a -> f a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (sym -> Bool -> SymExpr sym BaseBoolType
forall sym. IsExprBuilder sym => sym -> Bool -> Pred sym
backendPred sym
sym (a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
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 :: forall sym a.
IsExprBuilder sym =>
sym -> (Pred sym -> a -> a -> IO a) -> MuxTree sym a -> IO a
collapseMuxTree sym
_sym Pred sym -> a -> a -> IO a
ite MuxTree sym a
xt = [(a, Pred sym)] -> IO a
go (MuxTree sym a -> [(a, Pred sym)]
forall sym a. MuxTree sym a -> [(a, Pred sym)]
viewMuxTree MuxTree sym a
xt)
  where
  go :: [(a, Pred sym)] -> IO a
go []         = String -> [String] -> IO a
forall a. HasCallStack => String -> [String] -> a
panic String
"collapseMuxTree" [String
"empty mux tree"]
  go [(a
x,Pred sym
_p)]   = a -> IO a
forall a. a -> IO a
forall (m :: Type -> Type) a. Monad m => a -> m a
return a
x
  go ((a
x,Pred sym
p):[(a, Pred sym)]
xs) = Pred sym -> a -> a -> IO a
ite Pred sym
p a
x (a -> IO a) -> IO a -> IO a
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< [(a, Pred sym)] -> IO a
go [(a, Pred sym)]
xs

buildMuxTree ::
  (Ord a, IsExprBuilder sym) =>
  sym ->
  [(a, Pred sym)] ->
  IO (MuxTree sym a)
buildMuxTree :: forall a sym.
(Ord a, IsExprBuilder sym) =>
sym -> [(a, Pred sym)] -> IO (MuxTree sym a)
buildMuxTree sym
_sym [] = String -> [String] -> IO (MuxTree sym a)
forall a. HasCallStack => String -> [String] -> a
panic String
"buildMuxTree" [String
"empty mux tree"]
buildMuxTree sym
sym  [(a, Pred sym)]
xs = Map a (Pred sym) -> [(a, Pred sym)] -> IO (MuxTree sym a)
forall {sym} {a}.
(SymExpr sym ~ SymExpr sym, Ord a) =>
Map a (SymExpr sym BaseBoolType)
-> [(a, SymExpr sym BaseBoolType)] -> IO (MuxTree sym a)
go Map a (Pred sym)
forall k a. Map k a
Map.empty [(a, Pred sym)]
xs
  where
  go :: Map a (SymExpr sym BaseBoolType)
-> [(a, SymExpr sym BaseBoolType)] -> IO (MuxTree sym a)
go Map a (SymExpr sym BaseBoolType)
m [] = MuxTree sym a -> IO (MuxTree sym a)
forall a. a -> IO a
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Map a (SymExpr sym BaseBoolType) -> MuxTree sym a
forall sym a. Map a (Pred sym) -> MuxTree sym a
MuxTree Map a (SymExpr sym BaseBoolType)
m)
  go Map a (SymExpr sym BaseBoolType)
m ((a
z,SymExpr sym BaseBoolType
p):[(a, SymExpr sym BaseBoolType)]
zs) =
     case a
-> Map a (SymExpr sym BaseBoolType)
-> Maybe (SymExpr sym BaseBoolType)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
z Map a (SymExpr sym BaseBoolType)
m of
       Maybe (SymExpr sym BaseBoolType)
Nothing -> Map a (SymExpr sym BaseBoolType)
-> [(a, SymExpr sym BaseBoolType)] -> IO (MuxTree sym a)
go (a
-> SymExpr sym BaseBoolType
-> Map a (SymExpr sym BaseBoolType)
-> Map a (SymExpr sym BaseBoolType)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
z SymExpr sym BaseBoolType
p Map a (SymExpr sym BaseBoolType)
m) [(a, SymExpr sym BaseBoolType)]
zs
       Just SymExpr sym BaseBoolType
q -> do SymExpr sym BaseBoolType
pq <- sym -> Pred sym -> Pred sym -> IO (Pred sym)
forall sym.
IsExprBuilder sym =>
sym -> Pred sym -> Pred sym -> IO (Pred sym)
orPred sym
sym Pred sym
SymExpr sym BaseBoolType
p Pred sym
SymExpr sym BaseBoolType
q
                    case SymExpr sym BaseBoolType -> Maybe Bool
forall (e :: BaseType -> Type).
IsExpr e =>
e BaseBoolType -> Maybe Bool
asConstantPred SymExpr sym BaseBoolType
pq of
                      Just Bool
False -> Map a (SymExpr sym BaseBoolType)
-> [(a, SymExpr sym BaseBoolType)] -> IO (MuxTree sym a)
go Map a (SymExpr sym BaseBoolType)
m [(a, SymExpr sym BaseBoolType)]
zs
                      Maybe Bool
_ -> Map a (SymExpr sym BaseBoolType)
-> [(a, SymExpr sym BaseBoolType)] -> IO (MuxTree sym a)
go (a
-> SymExpr sym BaseBoolType
-> Map a (SymExpr sym BaseBoolType)
-> Map a (SymExpr sym BaseBoolType)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
z SymExpr sym BaseBoolType
pq Map a (SymExpr sym BaseBoolType)
m) [(a, SymExpr sym BaseBoolType)]
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 :: forall b sym a.
(Ord b, IsExprBuilder sym) =>
sym -> (a -> IO b) -> MuxTree sym a -> IO (MuxTree sym b)
muxTreeUnaryOp sym
sym a -> IO b
op MuxTree sym a
xt =
  do let xs :: [(a, Pred sym)]
xs = MuxTree sym a -> [(a, Pred sym)]
forall sym a. MuxTree sym a -> [(a, Pred sym)]
viewMuxTree MuxTree sym a
xt
     [(b, Pred sym)]
zs <- [IO (b, Pred sym)] -> IO [(b, Pred sym)]
forall (t :: Type -> Type) (m :: Type -> Type) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: Type -> Type) a. Monad m => [m a] -> m [a]
sequence
            [ do b
z <- a -> IO b
op a
x
                 (b, Pred sym) -> IO (b, Pred sym)
forall a. a -> IO a
forall (m :: Type -> Type) a. Monad m => a -> m a
return (b
z,Pred sym
p)
            | (a
x,Pred sym
p) <- [(a, Pred sym)]
xs
            ]
     sym -> [(b, Pred sym)] -> IO (MuxTree sym b)
forall a sym.
(Ord a, IsExprBuilder sym) =>
sym -> [(a, Pred sym)] -> IO (MuxTree sym a)
buildMuxTree sym
sym [(b, Pred 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 :: forall c sym a b.
(Ord c, IsExprBuilder sym) =>
sym
-> (a -> b -> IO c)
-> MuxTree sym a
-> MuxTree sym b
-> IO (MuxTree sym c)
muxTreeBinOp sym
sym a -> b -> IO c
op MuxTree sym a
xt MuxTree sym b
yt =
  do let xs :: [(a, Pred sym)]
xs = MuxTree sym a -> [(a, Pred sym)]
forall sym a. MuxTree sym a -> [(a, Pred sym)]
viewMuxTree MuxTree sym a
xt
     let ys :: [(b, Pred sym)]
ys = MuxTree sym b -> [(b, Pred sym)]
forall sym a. MuxTree sym a -> [(a, Pred sym)]
viewMuxTree MuxTree sym b
yt
     [(c, Pred sym)]
zs <- [IO (c, Pred sym)] -> IO [(c, Pred sym)]
forall (t :: Type -> Type) (m :: Type -> Type) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: Type -> Type) a. Monad m => [m a] -> m [a]
sequence
           [ do Pred sym
p <- sym -> Pred sym -> Pred sym -> IO (Pred sym)
forall sym.
IsExprBuilder sym =>
sym -> Pred sym -> Pred sym -> IO (Pred sym)
andPred sym
sym Pred sym
px Pred sym
py
                c
z <- a -> b -> IO c
op a
x b
y
                (c, Pred sym) -> IO (c, Pred sym)
forall a. a -> IO a
forall (m :: Type -> Type) a. Monad m => a -> m a
return (c
z,Pred sym
p)
           | (a
x,Pred sym
px) <- [(a, Pred sym)]
xs
           , (b
y,Pred sym
py) <- [(b, Pred sym)]
ys
           ]
     sym -> [(c, Pred sym)] -> IO (MuxTree sym c)
forall a sym.
(Ord a, IsExprBuilder sym) =>
sym -> [(a, Pred sym)] -> IO (MuxTree sym a)
buildMuxTree sym
sym [(c, Pred sym)]
zs


conditionMuxTreeLeaf ::
  IsExprBuilder sym => sym -> Pred sym -> a -> Pred sym -> IO (Maybe (Pred sym))
conditionMuxTreeLeaf :: forall sym a.
IsExprBuilder sym =>
sym -> Pred sym -> a -> Pred sym -> IO (Maybe (Pred sym))
conditionMuxTreeLeaf sym
sym Pred sym
p a
_v Pred sym
pv =
   do Pred sym
p' <- sym -> Pred sym -> Pred sym -> IO (Pred sym)
forall sym.
IsExprBuilder sym =>
sym -> Pred sym -> Pred sym -> IO (Pred sym)
andPred sym
sym Pred sym
p Pred sym
pv
      case Pred sym -> Maybe Bool
forall (e :: BaseType -> Type).
IsExpr e =>
e BaseBoolType -> Maybe Bool
asConstantPred Pred sym
p' of
        Just Bool
False -> Maybe (Pred sym) -> IO (Maybe (Pred sym))
forall a. a -> IO a
forall (m :: Type -> Type) a. Monad m => a -> m a
return Maybe (Pred sym)
forall a. Maybe a
Nothing
        Maybe Bool
_ -> Maybe (Pred sym) -> IO (Maybe (Pred sym))
forall a. a -> IO a
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Pred sym -> Maybe (Pred sym)
forall a. a -> Maybe a
Just Pred sym
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 :: forall a sym.
(Ord a, IsExprBuilder sym) =>
sym
-> Pred sym -> MuxTree sym a -> MuxTree sym a -> IO (MuxTree sym a)
mergeMuxTree sym
sym Pred sym
p (MuxTree Map a (Pred sym)
mx) (MuxTree Map a (Pred sym)
my) =
   do Pred sym
np <- sym -> Pred sym -> IO (Pred sym)
forall sym. IsExprBuilder sym => sym -> Pred sym -> IO (Pred sym)
notPred sym
sym Pred sym
p
      Map a (Pred sym) -> MuxTree sym a
forall sym a. Map a (Pred sym) -> MuxTree sym a
MuxTree (Map a (Pred sym) -> MuxTree sym a)
-> IO (Map a (Pred sym)) -> IO (MuxTree sym a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Pred sym
-> Map a (Pred sym) -> Map a (Pred sym) -> IO (Map a (Pred sym))
forall {k}.
Ord k =>
Pred sym
-> Map k (Pred sym) -> Map k (Pred sym) -> IO (Map k (Pred sym))
doMerge Pred sym
np Map a (Pred sym)
mx Map a (Pred sym)
my

  where
  f :: p -> Pred sym -> Pred sym -> IO (Maybe (Pred sym))
f p
_v Pred sym
px Pred sym
py =
    do Pred sym
p' <- sym -> Pred sym -> Pred sym -> Pred sym -> IO (Pred sym)
forall sym.
IsExprBuilder sym =>
sym -> Pred sym -> Pred sym -> Pred sym -> IO (Pred sym)
itePred sym
sym Pred sym
p Pred sym
px Pred sym
py
       case Pred sym -> Maybe Bool
forall (e :: BaseType -> Type).
IsExpr e =>
e BaseBoolType -> Maybe Bool
asConstantPred Pred sym
p' of
         Just Bool
False -> Maybe (Pred sym) -> IO (Maybe (Pred sym))
forall a. a -> IO a
forall (m :: Type -> Type) a. Monad m => a -> m a
return Maybe (Pred sym)
forall a. Maybe a
Nothing
         Maybe Bool
_ -> Maybe (Pred sym) -> IO (Maybe (Pred sym))
forall a. a -> IO a
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Pred sym -> Maybe (Pred sym)
forall a. a -> Maybe a
Just Pred sym
p')

  doMerge :: Pred sym
-> Map k (Pred sym) -> Map k (Pred sym) -> IO (Map k (Pred sym))
doMerge Pred sym
np = WhenMissing IO k (Pred sym) (Pred sym)
-> WhenMissing IO k (Pred sym) (Pred sym)
-> WhenMatched IO k (Pred sym) (Pred sym) (Pred sym)
-> Map k (Pred sym)
-> Map k (Pred sym)
-> IO (Map k (Pred sym))
forall (f :: Type -> Type) k a c b.
(Applicative f, Ord k) =>
WhenMissing f k a c
-> WhenMissing f k b c
-> WhenMatched f k a b c
-> Map k a
-> Map k b
-> f (Map k c)
Map.mergeA ((k -> Pred sym -> IO (Maybe (Pred sym)))
-> WhenMissing IO k (Pred sym) (Pred sym)
forall (f :: Type -> Type) k x y.
Applicative f =>
(k -> x -> f (Maybe y)) -> WhenMissing f k x y
Map.traverseMaybeMissing (sym -> Pred sym -> k -> Pred sym -> IO (Maybe (Pred sym))
forall sym a.
IsExprBuilder sym =>
sym -> Pred sym -> a -> Pred sym -> IO (Maybe (Pred sym))
conditionMuxTreeLeaf sym
sym Pred sym
p))
                          ((k -> Pred sym -> IO (Maybe (Pred sym)))
-> WhenMissing IO k (Pred sym) (Pred sym)
forall (f :: Type -> Type) k x y.
Applicative f =>
(k -> x -> f (Maybe y)) -> WhenMissing f k x y
Map.traverseMaybeMissing (sym -> Pred sym -> k -> Pred sym -> IO (Maybe (Pred sym))
forall sym a.
IsExprBuilder sym =>
sym -> Pred sym -> a -> Pred sym -> IO (Maybe (Pred sym))
conditionMuxTreeLeaf sym
sym Pred sym
np))
                          ((k -> Pred sym -> Pred sym -> IO (Maybe (Pred sym)))
-> WhenMatched IO k (Pred sym) (Pred sym) (Pred sym)
forall (f :: Type -> Type) k x y z.
Applicative f =>
(k -> x -> y -> f (Maybe z)) -> WhenMatched f k x y z
Map.zipWithMaybeAMatched k -> Pred sym -> Pred sym -> IO (Maybe (Pred sym))
forall {p}. p -> Pred sym -> Pred sym -> IO (Maybe (Pred sym))
f)