BoolMap.hs

{-|
Module      : What4.Expr.BoolMap
Description : Datastructure for representing a conjunction of predicates
Copyright   : (c) Galois Inc, 2019-2020
License     : BSD3
Maintainer  : rdockins@galois.com

Declares a datatype for representing n-way conjunctions or disjunctions
in a way that efficiently captures important algebraic
laws like commutativity, associativity and resolution.
-}

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module What4.Expr.BoolMap
  ( BoolMap
  , var
  , addVar
  , fromVars
  , combine
  , Polarity(..)
  , negatePolarity
  , contains
  , isInconsistent
  , isNull
  , BoolMapView(..)
  , viewBoolMap
  , foldMapVars
  , traverseVars
  , reversePolarities
  , removeVar
  , Wrap(..)
    -- * 'ConjMap'
  , ConjMap(..)
  , ConjMapView
  , pattern ConjTrue
  , pattern ConjFalse
  , pattern Conjuncts
  , viewConjMap
  , addConjunct
  , evalConj
  ) where

import           Control.Lens (_1, over)
import           Data.Coerce (coerce)
import           Data.Hashable
import qualified Data.List as List (foldl')
import           Data.List.NonEmpty (NonEmpty(..))
import           Data.Kind (Type)
import           Data.Parameterized.Classes
import           Data.Parameterized.TraversableF

import           What4.BaseTypes
import qualified What4.Utils.AnnotatedMap as AM
import           What4.Utils.IncrHash

-- | Describes the occurrence of a variable or expression, whether it is
--   negated or not.
data Polarity = Positive | Negative
 deriving (Eq,Ord,Show)

instance Hashable Polarity where
  hashWithSalt s Positive = hashWithSalt s (0::Int)
  hashWithSalt s Negative = hashWithSalt s (1::Int)

-- | Swap a polarity value
negatePolarity :: Polarity -> Polarity
negatePolarity Positive = Negative
negatePolarity Negative = Positive

newtype Wrap (f :: k -> Type) (x :: k) = Wrap { unWrap:: f x }

instance TestEquality f => Eq (Wrap f x) where
  Wrap a == Wrap b = isJust $ testEquality a b
instance OrdF f => Ord (Wrap f x) where
  compare (Wrap a) (Wrap b) = toOrdering $ compareF a b
instance (HashableF f, TestEquality f) => Hashable (Wrap f x) where
  hashWithSalt s (Wrap a) = hashWithSaltF s a

-- | A representation of a conjunction or a disjunction.
--
--   This data structure keeps track of a collection of expressions together
--   with their polarities.  The implementation uses a map from expression
--   values to their polarities, and thus automatically implements the
--   associative, commutative and idempotency laws common to both conjunctions
--   and disjunctions.  Moreover, if the same expression occurs in the
--   collection with opposite polarities, the entire collection collapses
--   via a resolution step to an \"inconsistent\" map.  For conjunctions this
--   corresponds to a contradiction and represents false; for disjunction, this
--   corresponds to the law of the excluded middle and represents true.
--
--   The annotation on the 'AM.AnnotatedMap' is an incremental hash ('IncrHash')
--   of the map, used to support a fast 'Hashable' instance.

data BoolMap (f :: BaseType -> Type)
  = InconsistentMap
  | BoolMap !(AM.AnnotatedMap (Wrap f BaseBoolType) IncrHash Polarity)

instance OrdF f => Eq (BoolMap f) where
  InconsistentMap == InconsistentMap = True
  BoolMap m1 == BoolMap m2 = AM.eqBy (==) m1 m2
  _ == _ = False

instance OrdF f => Semigroup (BoolMap f) where
  (<>) = combine

-- | Specialized version of 'foldMapVars'
instance FoldableF BoolMap where
  foldMapF f = foldMapVars f

foldMapVars :: Monoid m => (f BaseBoolType -> m) -> BoolMap f -> m
foldMapVars _ InconsistentMap = mempty
foldMapVars f (BoolMap am) = foldMap (f . unWrap . fst) (AM.toList am)

-- | Traverse the expressions in a bool map, and rebuild the map.
traverseVars :: (Applicative m, HashableF g, OrdF g) =>
  (f BaseBoolType -> m (g (BaseBoolType))) ->
  BoolMap f -> m (BoolMap g)
traverseVars _ InconsistentMap = pure InconsistentMap
traverseVars f (BoolMap m) =
  fromVars <$> traverse (_1 (f . unWrap)) (AM.toList m)

elementHash :: HashableF f => f BaseBoolType -> Polarity -> IncrHash
elementHash x p = mkIncrHash (hashWithSaltF (hash p) x)

instance (OrdF f, HashableF f) => Hashable (BoolMap f) where
  hashWithSalt s InconsistentMap = hashWithSalt s (0::Int)
  hashWithSalt s (BoolMap m) =
    case AM.annotation m of
      Nothing -> hashWithSalt s (1::Int)
      Just h  -> hashWithSalt (hashWithSalt s (1::Int)) h

-- | Represents the state of a 'BoolMap' (either a conjunction or disjunction).
--
-- If you know you are dealing with a 'BoolMap' that represents a conjunction,
-- consider using 'ConjMap' and 'viewConjMap' for the sake of clarity.
data BoolMapView f
  = BoolMapUnit
       -- ^ A bool map with no expressions, represents the unit of the corresponding operation
  | BoolMapDualUnit
       -- ^ An inconsistent bool map, represents the dual of the operation unit
  | BoolMapTerms (NonEmpty (f BaseBoolType, Polarity))
       -- ^ The terms appearing in the bool map, of which there is at least one

-- | Deconstruct the given bool map for later processing
viewBoolMap :: BoolMap f -> BoolMapView f
viewBoolMap InconsistentMap = BoolMapDualUnit
viewBoolMap (BoolMap m) =
  case AM.toList m of
    []  -> BoolMapUnit
    (Wrap x,p):xs -> BoolMapTerms ((x,p):|(map (over _1 unWrap) xs))

-- | Returns true for an inconsistent bool map
isInconsistent :: BoolMap f -> Bool
isInconsistent InconsistentMap = True
isInconsistent _ = False

-- | Returns true for a \"null\" bool map with no terms
isNull :: BoolMap f -> Bool
isNull InconsistentMap = False
isNull (BoolMap m) = AM.null m

-- | Produce a singleton bool map, consisting of just the given term
var :: (HashableF f, OrdF f) => f BaseBoolType -> Polarity -> BoolMap f
var x p = BoolMap (AM.singleton (Wrap x) (elementHash x p) p)

-- | Add a variable to a bool map, performing a resolution step if possible
addVar :: (HashableF f, OrdF f) => f BaseBoolType -> Polarity -> BoolMap f -> BoolMap f
addVar _ _ InconsistentMap = InconsistentMap
addVar x p1 (BoolMap bm) = maybe InconsistentMap BoolMap $ AM.alterF f (Wrap x) bm
 where
 f Nothing = return (Just (elementHash x p1, p1))
 f el@(Just (_,p2)) | p1 == p2  = return el
                    | otherwise = Nothing

-- | Generate a bool map from a list of terms and polarities by repeatedly
--   calling @addVar@.
fromVars :: (HashableF f, OrdF f) => [(f BaseBoolType, Polarity)] -> BoolMap f
fromVars = List.foldl' (\m (x,p) -> addVar x p m) (BoolMap AM.empty)

-- | Merge two bool maps, performing resolution as necessary.
combine :: OrdF f => BoolMap f -> BoolMap f -> BoolMap f
combine InconsistentMap _ = InconsistentMap
combine _ InconsistentMap = InconsistentMap
combine (BoolMap m1) (BoolMap m2) =
    maybe InconsistentMap BoolMap $ AM.mergeA f m1 m2

  where f _k (v,p1) (_,p2)
          | p1 == p2  = Just (v,p1)
          | otherwise = Nothing

-- | Test if the bool map contains the given term, and return the polarity
--   of that term if so.
contains :: OrdF f => BoolMap f -> f BaseBoolType -> Maybe Polarity
contains InconsistentMap _ = Nothing
contains (BoolMap m) x = snd <$> AM.lookup (Wrap x) m

-- | Swap the polarities of the terms in the given bool map.
reversePolarities :: OrdF f => BoolMap f -> BoolMap f
reversePolarities InconsistentMap = InconsistentMap
reversePolarities (BoolMap m) = BoolMap $! fmap negatePolarity m

-- | Remove the given term from the bool map.  The map is unchanged
--   if inconsistent or if the term does not occur.
removeVar :: OrdF f => BoolMap f -> f BaseBoolType -> BoolMap f
removeVar InconsistentMap _ = InconsistentMap
removeVar (BoolMap m) x = BoolMap (AM.delete (Wrap x) m)

--------------------------------------------------------------------------------
-- ConjMap

-- | A 'BoolMap' representing a conjunction.
newtype ConjMap f = ConjMap { getConjMap :: BoolMap f }
  deriving (Eq, FoldableF, Hashable, Semigroup)

-- | Represents the state of a 'ConjMap'. See 'viewConjMap'.
--
-- Like 'BoolMapView', but with more specific patterns for readability.
newtype ConjMapView f = ConjMapView (BoolMapView f)

pattern ConjTrue :: ConjMapView f
pattern ConjTrue = ConjMapView BoolMapUnit

pattern ConjFalse :: ConjMapView f
pattern ConjFalse = ConjMapView BoolMapDualUnit

pattern Conjuncts :: NonEmpty (f BaseBoolType, Polarity) -> ConjMapView f
pattern Conjuncts ts = ConjMapView (BoolMapTerms ts)

{-# COMPLETE ConjTrue, ConjFalse, Conjuncts #-}

-- | Deconstruct the given 'ConjMap' for later processing
viewConjMap :: forall f. ConjMap f -> ConjMapView f
viewConjMap =
  -- The explicit type annotations on `coerce` are likely necessary because of
  -- https://gitlab.haskell.org/ghc/ghc/-/issues/21003
  coerce @(BoolMap f -> BoolMapView f) @(ConjMap f -> ConjMapView f) viewBoolMap
{-# INLINE viewConjMap #-}

-- | Add a conjunct to a 'ConjMap'.
--
-- Wrapper around 'addVar'.
addConjunct ::
  forall f.
  (HashableF f, OrdF f) =>
  f BaseBoolType ->
  Polarity ->
  ConjMap f ->
  ConjMap f
addConjunct =
  -- The explicit type annotations on `coerce` are likely necessary because of
  -- https://gitlab.haskell.org/ghc/ghc/-/issues/21003
  coerce
    @(f BaseBoolType -> Polarity -> BoolMap f -> BoolMap f)
    @(f BaseBoolType -> Polarity -> ConjMap f -> ConjMap f)
    addVar
{-# INLINE addConjunct #-}

-- | Given the means to evaluate the conjuncts of a 'ConjMap' to a concrete
-- 'Bool', evaluate the whole conjunction to a 'Bool'.
evalConj :: Applicative m => (f BaseBoolType -> m Bool) -> ConjMap f -> m Bool
evalConj f cm =
  let pol (x, Positive) = f x
      pol (x, Negative) = not <$> f x
  in
  case viewConjMap cm of
    ConjTrue -> pure True
    ConjFalse -> pure False
    Conjuncts (t:|ts) ->
      List.foldl' (&&) <$> pol t <*> traverse pol ts