what4-1.8: src/What4/Expr/WeightedSum.hs
{-|
Module : What4.Expr.WeightedSum
Description : Representations for weighted sums and products in semirings
Copyright : (c) Galois Inc, 2015-2020
License : BSD3
Maintainer : jhendrix@galois.com
Declares a weighted sum type used for representing sums over variables and
an offset in one of the supported semirings. This module also implements a
representation of semiring products. See "What4.Expr.App" for an overview of
normalizing data structures such as these.
-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -Wwarn #-}
module What4.Expr.WeightedSum
( -- * Utilities
Tm
-- * Weighted sums
, WeightedSum
, sumRepr
, sumOffset
, sumAbsValue
, constant
, var
, scaledVar
, asConstant
, asVar
, asWeightedVar
, asAffineVar
, isZero
, traverseVars
, traverseCoeffs
, add
, addVar
, addVars
, addConstant
, scale
, eval
, evalM
, extractCommon
, fromTerms
, transformSum
, reduceIntSumMod
-- * Ring products
, SemiRingProduct
, traverseProdVars
, nullProd
, asProdVar
, prodRepr
, prodVar
, prodAbsValue
, prodMul
, prodEval
, prodEvalM
, prodContains
) where
import qualified Control.Exception as Ex
import Data.Functor.Identity (Identity(..))
import Lens.Micro
import Control.Monad (unless)
import qualified Data.BitVector.Sized as BV
import Data.Hashable
import Data.Kind
import qualified Data.List as List (foldl')
import Data.Maybe
import Data.Parameterized.Classes
import What4.BaseTypes
import qualified What4.Domains.BV as BVD
import qualified What4.Domains.BV.Arith as A
import qualified What4.Domains.BV.XOR as X
import qualified What4.SemiRing as SR
import What4.Utils.AnnotatedMap (AnnotatedMap)
import qualified What4.Utils.AnnotatedMap as AM
import qualified What4.Utils.AbstractDomains as AD
import What4.Utils.IncrHash
--------------------------------------------------------------------------------
data SRAbsValue :: SR.SemiRing -> Type where
SRAbsIntAdd :: !(AD.ValueRange Integer) -> SRAbsValue SR.SemiRingInteger
SRAbsRealAdd :: !AD.RealAbstractValue -> SRAbsValue SR.SemiRingReal
SRAbsBVAdd :: (1 <= w) => !(A.Domain w) -> SRAbsValue (SR.SemiRingBV SR.BVArith w)
SRAbsBVXor :: (1 <= w) => !(X.Domain w) -> SRAbsValue (SR.SemiRingBV SR.BVBits w)
instance Semigroup (SRAbsValue sr) where
SRAbsIntAdd x <> SRAbsIntAdd y = SRAbsIntAdd (AD.addRange x y)
SRAbsRealAdd x <> SRAbsRealAdd y = SRAbsRealAdd (AD.ravAdd x y)
SRAbsBVAdd x <> SRAbsBVAdd y = SRAbsBVAdd (A.add x y)
SRAbsBVXor x <> SRAbsBVXor y = SRAbsBVXor (X.xor x y)
(.**) :: SRAbsValue sr -> SRAbsValue sr -> SRAbsValue sr
SRAbsIntAdd x .** SRAbsIntAdd y = SRAbsIntAdd (AD.mulRange x y)
SRAbsRealAdd x .** SRAbsRealAdd y = SRAbsRealAdd (AD.ravMul x y)
SRAbsBVAdd x .** SRAbsBVAdd y = SRAbsBVAdd (A.mul x y)
SRAbsBVXor x .** SRAbsBVXor y = SRAbsBVXor (X.and x y)
abstractTerm ::
AD.HasAbsValue f =>
SR.SemiRingRepr sr -> SR.Coefficient sr -> f (SR.SemiRingBase sr) -> SRAbsValue sr
abstractTerm sr c e =
case sr of
SR.SemiRingIntegerRepr -> SRAbsIntAdd (AD.rangeScalarMul c (AD.getAbsValue e))
SR.SemiRingRealRepr -> SRAbsRealAdd (AD.ravScalarMul c (AD.getAbsValue e))
SR.SemiRingBVRepr fv w ->
case fv of
SR.BVArithRepr ->
-- A.scale expects a signed integer coefficient
SRAbsBVAdd (A.scale (BV.asSigned w c) (BVD.asArithDomain (AD.getAbsValue e)))
SR.BVBitsRepr -> SRAbsBVXor (X.and_scalar (BV.asUnsigned c) (BVD.asXorDomain (AD.getAbsValue e)))
abstractVal :: AD.HasAbsValue f => SR.SemiRingRepr sr -> f (SR.SemiRingBase sr) -> SRAbsValue sr
abstractVal sr e =
case sr of
SR.SemiRingIntegerRepr -> SRAbsIntAdd (AD.getAbsValue e)
SR.SemiRingRealRepr -> SRAbsRealAdd (AD.getAbsValue e)
SR.SemiRingBVRepr fv _w ->
case fv of
SR.BVArithRepr -> SRAbsBVAdd (BVD.asArithDomain (AD.getAbsValue e))
SR.BVBitsRepr -> SRAbsBVXor (BVD.asXorDomain (AD.getAbsValue e))
abstractScalar ::
SR.SemiRingRepr sr -> SR.Coefficient sr -> SRAbsValue sr
abstractScalar sr c =
case sr of
SR.SemiRingIntegerRepr -> SRAbsIntAdd (AD.SingleRange c)
SR.SemiRingRealRepr -> SRAbsRealAdd (AD.ravSingle c)
SR.SemiRingBVRepr fv w ->
case fv of
SR.BVArithRepr -> SRAbsBVAdd (A.singleton w (BV.asUnsigned c))
SR.BVBitsRepr -> SRAbsBVXor (X.singleton w (BV.asUnsigned c))
fromSRAbsValue ::
SRAbsValue sr -> AD.AbstractValue (SR.SemiRingBase sr)
fromSRAbsValue v =
case v of
SRAbsIntAdd x -> x
SRAbsRealAdd x -> x
SRAbsBVAdd x -> BVD.BVDArith x
SRAbsBVXor x -> BVD.fromXorDomain x
-- | Returns 'Just' when the abstract value is a singleton.
asCoeff ::
AD.HasAbsValue f =>
SR.SemiRingRepr sr ->
f (SR.SemiRingBase sr) ->
Maybe (SR.Coefficient sr)
asCoeff =
\case
SR.SemiRingIntegerRepr -> AD.asSingleRange . AD.getAbsValue
SR.SemiRingRealRepr -> AD.asSingleRange . AD.ravRange . AD.getAbsValue
SR.SemiRingBVRepr _ w -> fmap (BV.mkBV w) . BVD.asSingleton . AD.getAbsValue
varIsConst ::
AD.HasAbsValue f =>
SR.SemiRingRepr sr ->
f (SR.SemiRingBase sr) ->
Bool
varIsConst sr = isJust . asCoeff sr
--------------------------------------------------------------------------------
type Tm f = (HashableF f, OrdF f, AD.HasAbsValue f)
newtype WrapF (f :: BaseType -> Type) (i :: SR.SemiRing) = WrapF (f (SR.SemiRingBase i))
instance OrdF f => Ord (WrapF f i) where
compare (WrapF x) (WrapF y) = toOrdering $ compareF x y
instance TestEquality f => Eq (WrapF f i) where
(WrapF x) == (WrapF y) = isJust $ testEquality x y
instance (HashableF f, TestEquality f) => Hashable (WrapF f i) where
hashWithSalt s (WrapF x) = hashWithSaltF s x
traverseWrap :: Functor m => (f (SR.SemiRingBase i) -> m (g (SR.SemiRingBase i))) -> WrapF f i -> m (WrapF g i)
traverseWrap f (WrapF x) = WrapF <$> f x
-- | The annotation type used for the annotated map. It consists of
-- the hash value and the abstract domain representation of type @d@
-- for each submap.
data Note sr = Note !IncrHash !(SRAbsValue sr)
instance Semigroup (Note sr) where
Note h1 d1 <> Note h2 d2 = Note (h1 <> h2) (d1 <> d2)
data ProdNote sr = ProdNote !IncrHash !(SRAbsValue sr)
-- | The annotation type used for the annotated map for products.
-- It consists of the hash value and the abstract domain representation
-- of type @d@ for each submap. NOTE! that the multiplication operation
-- on abstract values is not always associative. This, however, is
-- acceptable because all associative groupings lead to sound (but perhaps not best)
-- approximate values.
instance Semigroup (ProdNote sr) where
ProdNote h1 d1 <> ProdNote h2 d2 = ProdNote (h1 <> h2) (d1 .** d2)
-- | Construct the annotation for a single map entry.
mkNote ::
(HashableF f, AD.HasAbsValue f) =>
SR.SemiRingRepr sr -> SR.Coefficient sr -> f (SR.SemiRingBase sr) -> Note sr
mkNote sr c t = Note (mkIncrHash h) d
where
h = SR.sr_hashWithSalt sr (hashF t) c
d = abstractTerm sr c t
mkProdNote ::
(HashableF f, AD.HasAbsValue f) =>
SR.SemiRingRepr sr ->
SR.Occurrence sr ->
f (SR.SemiRingBase sr) ->
ProdNote sr
mkProdNote sr occ t = ProdNote (mkIncrHash h) d
where
h = SR.occ_hashWithSalt sr (hashF t) occ
v = abstractVal sr t
power = fromIntegral (SR.occ_count sr occ)
d = go (power - 1) v
go (n::Integer) x
| n > 0 = go (n-1) (v .** x)
| otherwise = x
type SumMap f sr = AnnotatedMap (WrapF f sr) (Note sr) (SR.Coefficient sr)
type ProdMap f sr = AnnotatedMap (WrapF f sr) (ProdNote sr) (SR.Occurrence sr)
insertSumMap ::
Tm f =>
SR.SemiRingRepr sr ->
SR.Coefficient sr -> f (SR.SemiRingBase sr) -> SumMap f sr -> SumMap f sr
insertSumMap sr c t = AM.alter f (WrapF t)
where
f Nothing = Just (mkNote sr c t, c)
f (Just (_, c0))
| SR.eq sr (SR.zero sr) c' = Nothing
| otherwise = Just (mkNote sr c' t, c')
where c' = SR.add sr c0 c
singletonSumMap ::
Tm f =>
SR.SemiRingRepr sr ->
SR.Coefficient sr -> f (SR.SemiRingBase sr) -> SumMap f sr
singletonSumMap sr c t = AM.singleton (WrapF t) (mkNote sr c t) c
singletonProdMap ::
Tm f =>
SR.SemiRingRepr sr ->
SR.Occurrence sr ->
f (SR.SemiRingBase sr) ->
ProdMap f sr
singletonProdMap sr occ t = AM.singleton (WrapF t) (mkProdNote sr occ t) occ
-- | Returns a 'SumMap' together with the sum of all 'SR.Cofficient'-only terms.
fromListSumMap ::
Tm f =>
SR.SemiRingRepr sr ->
[(f (SR.SemiRingBase sr), SR.Coefficient sr)] ->
(SumMap f sr, SR.Coefficient sr)
fromListSumMap sr = foldr go (AM.empty, SR.zero sr)
where
go (t, c) (m, acc) =
if SR.eq sr (SR.zero sr) c
then (m, acc)
else
case asCoeff sr t of
Just c' -> (m, SR.add sr (SR.mul sr c c') acc)
Nothing -> (insertSumMap sr c t m, acc)
toListSumMap :: SumMap f sr -> [(f (SR.SemiRingBase sr), SR.Coefficient sr)]
toListSumMap am = [ (t, c) | (WrapF t, c) <- AM.toList am ]
-- | A weighted sum of semiring values. Mathematically, this represents
-- an affine operation on the underlying expressions.
data WeightedSum (f :: BaseType -> Type) (sr :: SR.SemiRing)
= WeightedSum { _sumMap :: !(SumMap f sr)
-- ^ Map from terms to their coefficients
--
-- INVARIANT: The terms in the map should not be constant
-- (i.e., according to 'varIsConst') and the coefficients
-- should not be 'SR.zero'. This is not a safety invariant,
-- but helps ensure normalized terms.
, _sumOffset :: !(SR.Coefficient sr)
, sumRepr :: !(SR.SemiRingRepr sr)
-- ^ Runtime representation of the semiring for this sum.
}
-- | A product of semiring values.
data SemiRingProduct (f :: BaseType -> Type) (sr :: SR.SemiRing)
= SemiRingProduct { _prodMap :: !(ProdMap f sr)
, prodRepr :: !(SR.SemiRingRepr sr)
-- ^ Runtime representation of the semiring for this product
}
-- | Return the hash of the 'SumMap' part of the 'WeightedSum'.
sumMapHash :: OrdF f => WeightedSum f sr -> IncrHash
sumMapHash x =
case AM.annotation (_sumMap x) of
Nothing -> mempty
Just (Note h _) -> h
prodMapHash :: OrdF f => SemiRingProduct f sr -> IncrHash
prodMapHash pd =
case AM.annotation (_prodMap pd) of
Nothing -> mempty
Just (ProdNote h _) -> h
sumAbsValue :: OrdF f => WeightedSum f sr -> AD.AbstractValue (SR.SemiRingBase sr)
sumAbsValue wsum =
fromSRAbsValue $
case AM.annotation (_sumMap wsum) of
Nothing -> absOffset
Just (Note _ v) -> absOffset <> v
where
absOffset = abstractScalar (sumRepr wsum) (_sumOffset wsum)
instance OrdF f => TestEquality (SemiRingProduct f) where
testEquality x y
| prodMapHash x /= prodMapHash y = Nothing
| otherwise =
do Refl <- testEquality (prodRepr x) (prodRepr y)
unless (AM.eqBy (SR.occ_eq (prodRepr x)) (_prodMap x) (_prodMap y)) Nothing
return Refl
instance OrdF f => Eq (SemiRingProduct f sr) where
x == y = isJust (testEquality x y)
instance OrdF f => TestEquality (WeightedSum f) where
testEquality x y
| sumMapHash x /= sumMapHash y = Nothing
| otherwise =
do Refl <- testEquality (sumRepr x) (sumRepr y)
unless (SR.eq (sumRepr x) (_sumOffset x) (_sumOffset y)) Nothing
unless (AM.eqBy (SR.eq (sumRepr x)) (_sumMap x) (_sumMap y)) Nothing
return Refl
instance OrdF f => Eq (WeightedSum f sr) where
x == y = isJust (testEquality x y)
-- | Created a weighted sum directly from a map and constant.
--
-- When calling this, one should ensure values with coefficients equal to
-- @'SR.zero'@ have been removed and none of the terms of the map satisfy
-- 'varIsConst'. See INVARIANT on '_sumMap'.
unfilteredSum ::
AD.HasAbsValue f =>
SR.SemiRingRepr sr ->
SumMap f sr ->
SR.Coefficient sr ->
WeightedSum f sr
unfilteredSum sr m c =
Ex.assert (all (uncurry notConst) (toListSumMap m)) $
WeightedSum m c sr
where
notConst v coeff =
not (varIsConst sr v) && not (SR.eq sr (SR.zero sr) coeff)
-- | Retrieve the mapping from terms to coefficients.
sumMap :: Lens' (WeightedSum f sr) (SumMap f sr)
sumMap = lens _sumMap (\w m -> w{ _sumMap = m })
-- | Retrieve the constant addend of the weighted sum.
sumOffset :: Lens' (WeightedSum f sr) (SR.Coefficient sr)
sumOffset = lens _sumOffset (\s v -> s { _sumOffset = v })
instance OrdF f => Hashable (WeightedSum f sr) where
hashWithSalt s0 w =
hashWithSalt (SR.sr_hashWithSalt (sumRepr w) s0 (_sumOffset w)) (sumMapHash w)
instance OrdF f => Hashable (SemiRingProduct f sr) where
hashWithSalt s0 w = hashWithSalt s0 (prodMapHash w)
-- | Attempt to parse a weighted sum as a constant.
asConstant :: WeightedSum f sr -> Maybe (SR.Coefficient sr)
asConstant w
| AM.null (_sumMap w) = Just (_sumOffset w)
| otherwise = Nothing
-- | Return true if a weighted sum is equal to constant 0.
isZero :: SR.SemiRingRepr sr -> WeightedSum f sr -> Bool
isZero sr s =
case asConstant s of
Just c -> SR.sr_compare sr (SR.zero sr) c == EQ
Nothing -> False
-- | Attempt to parse a weighted sum as a single expression with a coefficient and offset.
-- @asAffineVar w = Just (c,r,o)@ when @denotation(w) = c*r + o@.
asAffineVar :: WeightedSum f sr -> Maybe (SR.Coefficient sr, f (SR.SemiRingBase sr), SR.Coefficient sr)
asAffineVar w
| [(WrapF r, c)] <- AM.toList (_sumMap w)
= Just (c,r,_sumOffset w)
| otherwise
= Nothing
-- | Attempt to parse weighted sum as a single expression with a coefficient.
-- @asWeightedVar w = Just (c,r)@ when @denotation(w) = c*r@.
asWeightedVar :: WeightedSum f sr -> Maybe (SR.Coefficient sr, f (SR.SemiRingBase sr))
asWeightedVar w
| [(WrapF r, c)] <- AM.toList (_sumMap w)
, let sr = sumRepr w
, SR.eq sr (SR.zero sr) (_sumOffset w)
= Just (c,r)
| otherwise
= Nothing
-- | Attempt to parse a weighted sum as a single expression.
-- @asVar w = Just r@ when @denotation(w) = r@
asVar :: WeightedSum f sr -> Maybe (f (SR.SemiRingBase sr))
asVar w
| [(WrapF r, c)] <- AM.toList (_sumMap w)
, let sr = sumRepr w
, SR.eq sr (SR.one sr) c
, SR.eq sr (SR.zero sr) (_sumOffset w)
= Just r
| otherwise
= Nothing
-- | Create a sum from a constant coefficient value.
constant :: Tm f => SR.SemiRingRepr sr -> SR.Coefficient sr -> WeightedSum f sr
constant sr c = unfilteredSum sr AM.empty c
-- | Create a sum from a single affine variable (@st + c@).
affineVar ::
Tm f =>
SR.SemiRingRepr sr ->
SR.Coefficient sr ->
f (SR.SemiRingBase sr) ->
SR.Coefficient sr ->
WeightedSum f sr
affineVar sr s t c
| SR.eq sr (SR.zero sr) s = constant sr c
| Just s' <- asCoeff sr t = constant sr (SR.add sr (SR.mul sr s s') c)
| otherwise = unfilteredSum sr (singletonSumMap sr s t) c
-- | Traverse the expressions in a weighted sum.
traverseVars :: forall k j m sr.
(Applicative m, Tm k) =>
(j (SR.SemiRingBase sr) -> m (k (SR.SemiRingBase sr))) ->
WeightedSum j sr ->
m (WeightedSum k sr)
traverseVars f w =
(\tms -> fromTerms sr tms (_sumOffset w)) <$>
traverse (_1 f) (toListSumMap (_sumMap w))
where sr = sumRepr w
-- This is the type at its use in 'evalBoundVars'
{-# SPECIALIZE traverseVars :: forall k sr.
Tm k =>
(k (SR.SemiRingBase sr) -> IO (k (SR.SemiRingBase sr))) ->
WeightedSum k sr ->
IO (WeightedSum k sr)
#-}
-- | Traverse the coefficients in a weighted sum.
traverseCoeffs :: forall m f sr.
(Applicative m, Tm f) =>
(SR.Coefficient sr -> m (SR.Coefficient sr)) ->
WeightedSum f sr ->
m (WeightedSum f sr)
traverseCoeffs f w =
unfilteredSum sr <$> AM.traverseMaybeWithKey g (_sumMap w) <*> f (_sumOffset w)
where
sr = sumRepr w
g (WrapF t) _ c = mk t <$> f c
mk t c = if SR.eq sr (SR.zero sr) c then Nothing else Just (mkNote sr c t, c)
-- | Traverse the expressions in a product.
traverseProdVars :: forall k j m sr.
(Applicative m, Tm k) =>
(j (SR.SemiRingBase sr) -> m (k (SR.SemiRingBase sr))) ->
SemiRingProduct j sr ->
m (SemiRingProduct k sr)
traverseProdVars f pd =
mkProd sr . rebuild <$>
traverse (_1 (traverseWrap f)) (AM.toList (_prodMap pd))
where
sr = prodRepr pd
rebuild = List.foldl' (\m (WrapF t, occ) -> AM.insert (WrapF t) (mkProdNote sr occ t) occ m) AM.empty
-- This is the type at its use in 'evalBoundVars'
{-# SPECIALIZE traverseProdVars :: forall k sr.
Tm k =>
(k (SR.SemiRingBase sr) -> IO (k (SR.SemiRingBase sr))) ->
SemiRingProduct k sr ->
IO (SemiRingProduct k sr)
#-}
-- | This returns a variable times a constant.
scaledVar :: Tm f => SR.SemiRingRepr sr -> SR.Coefficient sr -> f (SR.SemiRingBase sr) -> WeightedSum f sr
scaledVar sr s t = affineVar sr s t (SR.zero sr)
-- | Create a weighted sum corresponding to the given variable.
var :: Tm f => SR.SemiRingRepr sr -> f (SR.SemiRingBase sr) -> WeightedSum f sr
var sr t = scaledVar sr (SR.one sr) t
-- | Add two sums, collecting terms as necessary and deleting terms whose
-- coefficients sum to 0.
add ::
Tm f =>
SR.SemiRingRepr sr ->
WeightedSum f sr ->
WeightedSum f sr ->
WeightedSum f sr
add sr x y = unfilteredSum sr zm zc
where
merge (WrapF k) u v | SR.eq sr r (SR.zero sr) = Nothing
| otherwise = Just (mkNote sr r k, r)
where r = SR.add sr u v
zm = AM.unionWithKeyMaybe merge (_sumMap x) (_sumMap y)
zc = SR.add sr (x^.sumOffset) (y^.sumOffset)
-- | Create a weighted sum that represents the sum of two terms.
addVars ::
Tm f =>
SR.SemiRingRepr sr ->
f (SR.SemiRingBase sr) ->
f (SR.SemiRingBase sr) ->
WeightedSum f sr
addVars sr x y = fromTerms sr [(x, SR.one sr), (y, SR.one sr)] (SR.zero sr)
-- | Add a variable to the sum.
addVar ::
Tm f =>
SR.SemiRingRepr sr ->
WeightedSum f sr -> f (SR.SemiRingBase sr) -> WeightedSum f sr
addVar sr wsum x = wsum { _sumMap = m' }
where m' = insertSumMap sr (SR.one sr) x (_sumMap wsum)
-- | Add a constant to the sum.
addConstant :: SR.SemiRingRepr sr -> WeightedSum f sr -> SR.Coefficient sr -> WeightedSum f sr
addConstant sr x r = x & sumOffset %~ SR.add sr r
-- | Multiply a sum by a constant coefficient.
scale :: Tm f => SR.SemiRingRepr sr -> SR.Coefficient sr -> WeightedSum f sr -> WeightedSum f sr
scale sr c wsum
| SR.eq sr c (SR.zero sr) = constant sr (SR.zero sr)
| otherwise = unfilteredSum sr m' (SR.mul sr c (wsum^.sumOffset))
where
m' = AM.mapMaybeWithKey f (wsum^.sumMap)
f (WrapF t) _ x =
Ex.assert (not (varIsConst sr t)) $ -- INVARIANT
-- Filter out terms that become zero after scaling. This can happen in
-- bitvector arthmetic. For example, for 8-bit bitvectors c and x where
-- c = 2 and x = 128, c * x mod 256 = 0.
--
-- Necessary to uphold INVARIANT.
if SR.eq sr (SR.zero sr) cx
then Nothing
else Just (mkNote sr cx t, cx)
where cx = SR.mul sr c x
-- | Produce a weighted sum from a list of terms and an offset.
fromTerms ::
Tm f =>
SR.SemiRingRepr sr ->
[(f (SR.SemiRingBase sr), SR.Coefficient sr)] ->
SR.Coefficient sr ->
WeightedSum f sr
fromTerms sr tms offset =
let (m, offset') = fromListSumMap sr tms in
unfilteredSum sr m (SR.add sr offset offset')
-- | Apply update functions to the terms and coefficients of a weighted sum.
transformSum :: (Applicative m, Tm g) =>
SR.SemiRingRepr sr' ->
(SR.Coefficient sr -> m (SR.Coefficient sr')) ->
(f (SR.SemiRingBase sr) -> m (g (SR.SemiRingBase sr'))) ->
WeightedSum f sr ->
m (WeightedSum g sr')
transformSum sr' transCoef transTm s = fromTerms sr' <$> tms <*> c
where
f (t, x) = (,) <$> transTm t <*> transCoef x
tms = traverse f (toListSumMap (_sumMap s))
c = transCoef (_sumOffset s)
-- | Evaluate a sum given interpretations of addition, scalar
-- multiplication, and a constant. This evaluation is threaded through
-- a monad. The addition function is associated to the left, as in
-- 'foldlM'.
evalM :: Monad m =>
(r -> r -> m r) {- ^ Addition function -} ->
(SR.Coefficient sr -> f (SR.SemiRingBase sr) -> m r) {- ^ Scalar multiply -} ->
(SR.Coefficient sr -> m r) {- ^ Constant evaluation -} ->
WeightedSum f sr ->
m r
evalM addFn smul cnst sm
| SR.eq sr (_sumOffset sm) (SR.zero sr) =
case toListSumMap (_sumMap sm) of
[] -> cnst (SR.zero sr)
((e, s) : tms) -> go tms =<< smul s e
| otherwise =
go (toListSumMap (_sumMap sm)) =<< cnst (_sumOffset sm)
where
sr = sumRepr sm
go [] x = return x
go ((e, s) : tms) x = go tms =<< addFn x =<< smul s e
-- | Evaluate a sum given interpretations of addition, scalar multiplication, and
-- a constant rational.
eval ::
(r -> r -> r) {- ^ Addition function -} ->
(SR.Coefficient sr -> f (SR.SemiRingBase sr) -> r) {- ^ Scalar multiply -} ->
(SR.Coefficient sr -> r) {- ^ Constant evaluation -} ->
WeightedSum f sr ->
r
eval addFn smul cnst w
| SR.eq sr (_sumOffset w) (SR.zero sr) =
case toListSumMap (_sumMap w) of
[] -> cnst (SR.zero sr)
((e, s) : tms) -> go tms (smul s e)
| otherwise =
go (toListSumMap (_sumMap w)) (cnst (_sumOffset w))
where
sr = sumRepr w
go [] x = x
go ((e, s) : tms) x = go tms (addFn (smul s e) x)
{-# INLINABLE eval #-}
-- | Reduce a weighted sum of integers modulo a concrete integer.
-- This reduces each of the coefficients modulo the given integer,
-- removing any that are congruent to 0; the offset value is
-- also reduced.
reduceIntSumMod ::
Tm f =>
WeightedSum f SR.SemiRingInteger {- ^ The sum to reduce -} ->
Integer {- ^ The modulus, must not be 0 -} ->
WeightedSum f SR.SemiRingInteger
reduceIntSumMod ws k = unfilteredSum SR.SemiRingIntegerRepr m (ws^.sumOffset `mod` k)
where
sr = sumRepr ws
m = runIdentity (AM.traverseMaybeWithKey f (ws^.sumMap))
f (WrapF t) _ x
| x' == 0 = return Nothing
| otherwise = return (Just (mkNote sr x' t, x'))
where x' = x `mod` k
{-# INLINABLE extractCommon #-}
-- | Given two weighted sums @x@ and @y@, this returns a triple @(z,x',y')@
-- where @x = z + x'@ and @y = z + y'@ and @z@ contains the "common"
-- parts of @x@ and @y@. We only extract common terms when both
-- terms occur with the same coefficient in each sum.
--
-- This is primarily used to simplify if-then-else expressions to
-- preserve shared subterms.
extractCommon ::
Tm f =>
WeightedSum f sr ->
WeightedSum f sr ->
(WeightedSum f sr, WeightedSum f sr, WeightedSum f sr)
extractCommon (WeightedSum xm xc sr) (WeightedSum ym yc _) = (z, x', y')
where
mergeCommon (WrapF t) (_, xv) (_, yv)
| SR.eq sr xv yv = Just (mkNote sr xv t, xv)
| otherwise = Nothing
zm = AM.mergeWithKey mergeCommon (const AM.empty) (const AM.empty) xm ym
(zc, xc', yc')
| SR.eq sr xc yc = (xc, SR.zero sr, SR.zero sr)
| otherwise = (SR.zero sr, xc, yc)
z = unfilteredSum sr zm zc
x' = unfilteredSum sr (xm `AM.difference` zm) xc'
y' = unfilteredSum sr (ym `AM.difference` zm) yc'
-- | Returns true if the product is trivial (contains no terms).
nullProd :: SemiRingProduct f sr -> Bool
nullProd pd = AM.null (_prodMap pd)
-- | If the product consists of exactly on term, return it.
asProdVar :: SemiRingProduct f sr -> Maybe (f (SR.SemiRingBase sr))
asProdVar pd
| [(WrapF x, SR.occ_count sr -> 1)] <- AM.toList (_prodMap pd) = Just x
| otherwise = Nothing
where
sr = prodRepr pd
prodAbsValue :: OrdF f => SemiRingProduct f sr -> AD.AbstractValue (SR.SemiRingBase sr)
prodAbsValue pd =
fromSRAbsValue $
case AM.annotation (_prodMap pd) of
Nothing -> abstractScalar (prodRepr pd) (SR.one (prodRepr pd))
Just (ProdNote _ v) -> v
-- | Returns true if the product contains at least on occurrence of the given term.
prodContains :: OrdF f => SemiRingProduct f sr -> f (SR.SemiRingBase sr) -> Bool
prodContains pd x = isJust $ AM.lookup (WrapF x) (_prodMap pd)
-- | Produce a product map from a raw map of terms to occurrences.
-- PRECONDITION: the occurrence value for each term should be non-zero.
mkProd :: SR.SemiRingRepr sr -> ProdMap f sr -> SemiRingProduct f sr
mkProd sr m = SemiRingProduct m sr
-- | Produce a product representing the single given term.
prodVar :: Tm f => SR.SemiRingRepr sr -> f (SR.SemiRingBase sr) -> SemiRingProduct f sr
prodVar sr x = mkProd sr (singletonProdMap sr (SR.occ_one sr) x)
-- | Multiply two products, collecting terms and adding occurrences.
prodMul :: Tm f => SemiRingProduct f sr -> SemiRingProduct f sr -> SemiRingProduct f sr
prodMul x y = mkProd sr m
where
sr = prodRepr x
mergeCommon (WrapF k) (_,a) (_,b) = Just (mkProdNote sr c k, c)
where c = SR.occ_add sr a b
m = AM.mergeWithKey mergeCommon id id (_prodMap x) (_prodMap y)
-- | Evaluate a product, given a function representing multiplication
-- and a function to evaluate terms.
prodEval ::
(r -> r -> r) {-^ multiplication evalation -} ->
(f (SR.SemiRingBase sr) -> r) {-^ term evaluation -} ->
SemiRingProduct f sr ->
Maybe r
prodEval mul tm om =
runIdentity (prodEvalM (\x y -> Identity (mul x y)) (Identity . tm) om)
-- | Evaluate a product, given a function representing multiplication
-- and a function to evaluate terms, where both functions are threaded
-- through a monad.
prodEvalM :: Monad m =>
(r -> r -> m r) {-^ multiplication evalation -} ->
(f (SR.SemiRingBase sr) -> m r) {-^ term evaluation -} ->
SemiRingProduct f sr ->
m (Maybe r)
prodEvalM mul tm om = f (AM.toList (_prodMap om))
where
sr = prodRepr om
-- we have not yet encountered a term with non-zero occurrences
f [] = return Nothing
f ((WrapF x, SR.occ_count sr -> n):xs)
| n == 0 = f xs
| otherwise =
do t <- tm x
t' <- go (n-1) t t
g xs t'
-- we have a partial product @z@ already computed and need to multiply
-- in the remaining terms in the list
g [] z = return (Just z)
g ((WrapF x, SR.occ_count sr -> n):xs) z
| n == 0 = g xs z
| otherwise =
do t <- tm x
t' <- go n t z
g xs t'
-- compute: z * t^n
go n t z
| n > 0 = go (n-1) t =<< mul z t
| otherwise = return z