{-|
Copyright  :  (C) 2015-2016, University of Twente,
                  2017     , QBayLogic B.V.
License    :  BSD2 (see the file LICENSE)
Maintainer :  Christiaan Baaij <christiaan.baaij@gmail.com>
-}

{-# LANGUAGE CPP                        #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MagicHash                  #-}
{-# LANGUAGE RecordWildCards            #-}

{-# OPTIONS_GHC -fno-warn-unused-imports #-}
#if __GLASGOW_HASKELL__ < 801
#define nonDetCmpType cmpType
#endif

module GHC.TypeLits.Normalise.Unify
  ( -- * 'Nat' expressions \<-\> 'SOP' terms
    CType (..)
  , CoreSOP
  , normaliseNat
  , normaliseNatEverywhere
  , normaliseSimplifyNat
  , reifySOP
    -- * Substitution on 'SOP' terms
  , UnifyItem (..)
  , CoreUnify
  , substsSOP
  , substsSubst
    -- * Find unifiers
  , UnifyResult (..)
  , unifyNats
  , unifiers
    -- * Free variables in 'SOP' terms
  , fvSOP
    -- * Inequalities
  , subtractIneq
  , solveIneq
  , ineqToSubst
  , subtractionToPred
  , instantSolveIneq
  , solvedInEqSmallestConstraint
    -- * Properties
  , isNatural
  )
where

-- External
import Control.Arrow (first, second)
import Control.Monad.Trans.Writer.Strict
import Data.Function (on)
import Data.List     ((\\), intersect, nub)
import Data.Maybe    (fromMaybe, mapMaybe, isJust)
import Data.Set      (Set)
import qualified Data.Set as Set

import GHC.Base               (isTrue#,(==#))
import GHC.Integer            (smallInteger)
import GHC.Integer.Logarithms (integerLogBase#)

-- GHC API
#if MIN_VERSION_ghc(9,0,0)
import GHC.Builtin.Types (boolTy, promotedTrueDataCon)
import GHC.Builtin.Types.Literals
  (typeNatAddTyCon, typeNatExpTyCon, typeNatMulTyCon, typeNatSubTyCon)
#if MIN_VERSION_ghc(9,2,0)
import GHC.Builtin.Types (naturalTy, promotedFalseDataCon)
import GHC.Builtin.Types.Literals (typeNatCmpTyCon)
#else
import GHC.Builtin.Types (typeNatKind)
import GHC.Builtin.Types.Literals (typeNatLeqTyCon)
#endif
import GHC.Core.Predicate (EqRel (NomEq), Pred (EqPred), classifyPredType, mkPrimEqPred)
import GHC.Core.TyCon (TyCon)
#if MIN_VERSION_ghc(9,6,0)
import GHC.Core.Type
  (PredType, TyVar, coreView, mkNumLitTy, mkTyConApp, mkTyVarTy, typeKind)
import GHC.Core.TyCo.Compare
  (eqType, nonDetCmpType)
#else
import GHC.Core.Type
  (PredType, TyVar, coreView, eqType, mkNumLitTy, mkTyConApp, mkTyVarTy, nonDetCmpType, typeKind)
#endif
import GHC.Core.TyCo.Rep (Kind, Type (..), TyLit (..))
import GHC.Tc.Plugin (TcPluginM, tcPluginTrace)
import GHC.Tc.Types.Constraint (Ct, ctEvidence, ctEvId, ctEvPred, isGiven)
import GHC.Types.Unique.Set
  (UniqSet, unionManyUniqSets, emptyUniqSet, unionUniqSets, unitUniqSet)
import GHC.Utils.Outputable (Outputable (..), (<+>), ($$), text)
#else
import Outputable    (Outputable (..), (<+>), ($$), text)
import TcPluginM     (TcPluginM, tcPluginTrace)
import TcTypeNats    (typeNatAddTyCon, typeNatExpTyCon, typeNatMulTyCon,
                      typeNatSubTyCon, typeNatLeqTyCon)
import TyCon         (TyCon)
import Type          (TyVar,
                      coreView, eqType, mkNumLitTy, mkTyConApp, mkTyVarTy,
                      nonDetCmpType, PredType, typeKind)
import TyCoRep       (Kind, Type (..), TyLit (..))
import TysWiredIn    (boolTy, promotedTrueDataCon, typeNatKind)
import UniqSet       (UniqSet, unionManyUniqSets, emptyUniqSet, unionUniqSets,
                      unitUniqSet)

#if MIN_VERSION_ghc(8,10,0)
import Constraint (Ct,  ctEvidence, ctEvId, ctEvPred, isGiven)
import Predicate  (EqRel (NomEq), Pred (EqPred), classifyPredType, mkPrimEqPred)
#else
import TcRnMonad  (Ct, ctEvidence, isGiven)
import TcRnTypes  (ctEvPred)
import Type       (EqRel (NomEq), PredTree (EqPred), classifyPredType, mkPrimEqPred)
#endif
#endif

-- Internal
import GHC.TypeLits.Normalise.SOP

-- Used for haddock
import GHC.TypeLits (Nat)

#if MIN_VERSION_ghc(9,2,0)
typeNatKind :: Type
typeNatKind :: Type
typeNatKind = Type
naturalTy
#endif

newtype CType = CType { CType -> Type
unCType :: Type }
  deriving CType -> SDoc
forall a. (a -> SDoc) -> Outputable a
ppr :: CType -> SDoc
$cppr :: CType -> SDoc
Outputable

instance Eq CType where
  (CType Type
ty1) == :: CType -> CType -> Bool
== (CType Type
ty2) = Type -> Type -> Bool
eqType Type
ty1 Type
ty2

instance Ord CType where
  compare :: CType -> CType -> Ordering
compare (CType Type
ty1) (CType Type
ty2) = Type -> Type -> Ordering
nonDetCmpType Type
ty1 Type
ty2

-- | 'SOP' with 'TyVar' variables
type CoreSOP     = SOP TyVar CType
type CoreProduct = Product TyVar CType
type CoreSymbol  = Symbol TyVar CType

-- | Convert a type of /kind/ 'GHC.TypeLits.Nat' to an 'SOP' term, but
-- only when the type is constructed out of:
--
-- * literals
-- * type variables
-- * Applications of the arithmetic operators @(+,-,*,^)@
normaliseNat :: Type -> Writer [(Type,Type)] CoreSOP
normaliseNat :: Type -> Writer [(Type, Type)] (SOP TyVar CType)
normaliseNat Type
ty | Just Type
ty1 <- Type -> Maybe Type
coreView Type
ty = Type -> Writer [(Type, Type)] (SOP TyVar CType)
normaliseNat Type
ty1
normaliseNat (TyVarTy TyVar
v)          = forall (m :: * -> *) a. Monad m => a -> m a
return (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. v -> Symbol v c
V TyVar
v]])
normaliseNat (LitTy (NumTyLit Integer
i)) = forall (m :: * -> *) a. Monad m => a -> m a
return (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
i]])
normaliseNat (TyConApp TyCon
tc [Type
x,Type
y])
  | TyCon
tc forall a. Eq a => a -> a -> Bool
== TyCon
typeNatAddTyCon = forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPAdd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> Writer [(Type, Type)] (SOP TyVar CType)
normaliseNat Type
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> Writer [(Type, Type)] (SOP TyVar CType)
normaliseNat Type
y
  | TyCon
tc forall a. Eq a => a -> a -> Bool
== TyCon
typeNatSubTyCon = do
    forall (m :: * -> *) w. Monad m => w -> WriterT w m ()
tell [(Type
x,Type
y)]
    forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPAdd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> Writer [(Type, Type)] (SOP TyVar CType)
normaliseNat Type
x
                forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I (-Integer
1)]]) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> Writer [(Type, Type)] (SOP TyVar CType)
normaliseNat Type
y)
  | TyCon
tc forall a. Eq a => a -> a -> Bool
== TyCon
typeNatMulTyCon = forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> Writer [(Type, Type)] (SOP TyVar CType)
normaliseNat Type
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> Writer [(Type, Type)] (SOP TyVar CType)
normaliseNat Type
y
  | TyCon
tc forall a. Eq a => a -> a -> Bool
== TyCon
typeNatExpTyCon = forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
normaliseExp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> Writer [(Type, Type)] (SOP TyVar CType)
normaliseNat Type
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> Writer [(Type, Type)] (SOP TyVar CType)
normaliseNat Type
y
normaliseNat Type
t = forall (m :: * -> *) a. Monad m => a -> m a
return (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. c -> Symbol v c
C (Type -> CType
CType Type
t)]])

-- | Runs writer action. If the result /Nothing/ writer actions will be
-- discarded.
maybeRunWriter
  :: Monoid a
  => Writer a (Maybe b)
  -> Writer a (Maybe b)
maybeRunWriter :: forall a b. Monoid a => Writer a (Maybe b) -> Writer a (Maybe b)
maybeRunWriter Writer a (Maybe b)
w =
  case forall w a. Writer w a -> (a, w)
runWriter Writer a (Maybe b)
w of
    (Maybe b
Nothing, a
_) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
    (Maybe b
b, a
a) -> forall (m :: * -> *) w. Monad m => w -> WriterT w m ()
tell a
a forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe b
b

-- | Applies 'normaliseNat' and 'simplifySOP' to type or predicates to reduce
-- any occurrences of sub-terms of /kind/ 'GHC.TypeLits.Nat'. If the result is
-- the same as input, returns @'Nothing'@.
normaliseNatEverywhere :: Type -> Writer [(Type, Type)] (Maybe Type)
normaliseNatEverywhere :: Type -> Writer [(Type, Type)] (Maybe Type)
normaliseNatEverywhere Type
ty0
  | TyConApp TyCon
tc [Type]
_fields <- Type
ty0
  , TyCon
tc forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TyCon]
knownTyCons = do
    -- Normalize under current type constructor application. 'go' skips all
    -- known type constructors.
    Maybe Type
ty1M <- forall a b. Monoid a => Writer a (Maybe b) -> Writer a (Maybe b)
maybeRunWriter (Type -> Writer [(Type, Type)] (Maybe Type)
go Type
ty0)
    let ty1 :: Type
ty1 = forall a. a -> Maybe a -> a
fromMaybe Type
ty0 Maybe Type
ty1M

    -- Normalize (subterm-normalized) type given to 'normaliseNatEverywhere'
    Type
ty2 <- Type -> Writer [(Type, Type)] Type
normaliseSimplifyNat Type
ty1
    -- TODO: 'normaliseNat' could keep track whether it changed anything. That's
    -- TODO: probably cheaper than checking for equality here.
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Type
ty2 Type -> Type -> Bool
`eqType` Type
ty1 then Maybe Type
ty1M else forall a. a -> Maybe a
Just Type
ty2)
  | Bool
otherwise = Type -> Writer [(Type, Type)] (Maybe Type)
go Type
ty0
 where
  knownTyCons :: [TyCon]
  knownTyCons :: [TyCon]
knownTyCons = [TyCon
typeNatExpTyCon, TyCon
typeNatMulTyCon, TyCon
typeNatSubTyCon, TyCon
typeNatAddTyCon]

  -- Normalize given type, but ignore all top-level
  go :: Type -> Writer [(Type, Type)] (Maybe Type)
  go :: Type -> Writer [(Type, Type)] (Maybe Type)
go (TyConApp TyCon
tc_ [Type]
fields0_) = do
    [Maybe Type]
fields1_ <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall a b. Monoid a => Writer a (Maybe b) -> Writer a (Maybe b)
maybeRunWriter forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Writer [(Type, Type)] (Maybe Type)
cont) [Type]
fields0_
    if forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any forall a. Maybe a -> Bool
isJust [Maybe Type]
fields1_ then
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. a -> Maybe a
Just (TyCon -> [Type] -> Type
TyConApp TyCon
tc_ (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall a. a -> Maybe a -> a
fromMaybe [Type]
fields0_ [Maybe Type]
fields1_)))
    else
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
   where
    cont :: Type -> Writer [(Type, Type)] (Maybe Type)
cont = if TyCon
tc_ forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TyCon]
knownTyCons then Type -> Writer [(Type, Type)] (Maybe Type)
go else Type -> Writer [(Type, Type)] (Maybe Type)
normaliseNatEverywhere
  go Type
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing

normaliseSimplifyNat :: Type -> Writer [(Type, Type)] Type
normaliseSimplifyNat :: Type -> Writer [(Type, Type)] Type
normaliseSimplifyNat Type
ty
  | HasDebugCallStack => Type -> Type
typeKind Type
ty Type -> Type -> Bool
`eqType` Type
typeNatKind = do
      SOP TyVar CType
ty' <- Type -> Writer [(Type, Type)] (SOP TyVar CType)
normaliseNat Type
ty
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ SOP TyVar CType -> Type
reifySOP forall a b. (a -> b) -> a -> b
$ forall v c. (Ord v, Ord c) => SOP v c -> SOP v c
simplifySOP SOP TyVar CType
ty'
  | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return Type
ty

-- | Convert a 'SOP' term back to a type of /kind/ 'GHC.TypeLits.Nat'
reifySOP :: CoreSOP -> Type
reifySOP :: SOP TyVar CType -> Type
reifySOP = [Either (Product TyVar CType) (Product TyVar CType)] -> Type
combineP forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map Product TyVar CType
-> Either (Product TyVar CType) (Product TyVar CType)
negateP forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v c. SOP v c -> [Product v c]
unS
  where
    negateP :: CoreProduct -> Either CoreProduct CoreProduct
    negateP :: Product TyVar CType
-> Either (Product TyVar CType) (Product TyVar CType)
negateP (P ((I Integer
i):ps :: [Symbol TyVar CType]
ps@(Symbol TyVar CType
_:[Symbol TyVar CType]
_))) | Integer
i forall a. Eq a => a -> a -> Bool
== (-Integer
1) = forall a b. a -> Either a b
Left  (forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
ps)
    negateP (P ((I Integer
i):[Symbol TyVar CType]
ps)) | Integer
i forall a. Ord a => a -> a -> Bool
< Integer
0           = forall a b. a -> Either a b
Left  (forall v c. [Symbol v c] -> Product v c
P ((forall v c. Integer -> Symbol v c
I (forall a. Num a => a -> a
abs Integer
i))forall a. a -> [a] -> [a]
:[Symbol TyVar CType]
ps))
    negateP Product TyVar CType
ps                               = forall a b. b -> Either a b
Right Product TyVar CType
ps

    combineP :: [Either CoreProduct CoreProduct] -> Type
    combineP :: [Either (Product TyVar CType) (Product TyVar CType)] -> Type
combineP []     = Integer -> Type
mkNumLitTy Integer
0
    combineP [Either (Product TyVar CType) (Product TyVar CType)
p]    = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (\Product TyVar CType
p' -> TyCon -> [Type] -> Type
mkTyConApp TyCon
typeNatSubTyCon
                                                [Integer -> Type
mkNumLitTy Integer
0, Product TyVar CType -> Type
reifyProduct Product TyVar CType
p'])
                             Product TyVar CType -> Type
reifyProduct Either (Product TyVar CType) (Product TyVar CType)
p
    combineP [Either (Product TyVar CType) (Product TyVar CType)
p1,Either (Product TyVar CType) (Product TyVar CType)
p2] = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either
      (\Product TyVar CType
x -> forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either
               -- x neg, y neg
               (\Product TyVar CType
y -> let r :: Type
r = TyCon -> [Type] -> Type
mkTyConApp TyCon
typeNatSubTyCon [Product TyVar CType -> Type
reifyProduct Product TyVar CType
x
                                                         ,Product TyVar CType -> Type
reifyProduct Product TyVar CType
y]
                      in  TyCon -> [Type] -> Type
mkTyConApp TyCon
typeNatSubTyCon [Integer -> Type
mkNumLitTy Integer
0, Type
r])
               -- x neg, y pos
               (\Product TyVar CType
y -> TyCon -> [Type] -> Type
mkTyConApp TyCon
typeNatSubTyCon [Product TyVar CType -> Type
reifyProduct Product TyVar CType
y, Product TyVar CType -> Type
reifyProduct Product TyVar CType
x])
               Either (Product TyVar CType) (Product TyVar CType)
p2)
      (\Product TyVar CType
x -> forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either
               -- x pos, y neg
               (\Product TyVar CType
y -> TyCon -> [Type] -> Type
mkTyConApp TyCon
typeNatSubTyCon [Product TyVar CType -> Type
reifyProduct Product TyVar CType
x, Product TyVar CType -> Type
reifyProduct Product TyVar CType
y])
               -- x pos, y pos
               (\Product TyVar CType
y -> TyCon -> [Type] -> Type
mkTyConApp TyCon
typeNatAddTyCon [Product TyVar CType -> Type
reifyProduct Product TyVar CType
x, Product TyVar CType -> Type
reifyProduct Product TyVar CType
y])
               Either (Product TyVar CType) (Product TyVar CType)
p2)
      Either (Product TyVar CType) (Product TyVar CType)
p1


    combineP (Either (Product TyVar CType) (Product TyVar CType)
p:[Either (Product TyVar CType) (Product TyVar CType)]
ps)  = let es :: Type
es = [Either (Product TyVar CType) (Product TyVar CType)] -> Type
combineP [Either (Product TyVar CType) (Product TyVar CType)]
ps
                       in  forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (\Product TyVar CType
x -> TyCon -> [Type] -> Type
mkTyConApp TyCon
typeNatSubTyCon
                                                    [Type
es, Product TyVar CType -> Type
reifyProduct Product TyVar CType
x])
                                  (\Product TyVar CType
x -> TyCon -> [Type] -> Type
mkTyConApp TyCon
typeNatAddTyCon
                                                   [Product TyVar CType -> Type
reifyProduct Product TyVar CType
x, Type
es])
                                  Either (Product TyVar CType) (Product TyVar CType)
p

reifyProduct :: CoreProduct -> Type
reifyProduct :: Product TyVar CType -> Type
reifyProduct (P [Symbol TyVar CType]
ps) =
    let ps' :: [Type]
ps' = forall a b. (a -> b) -> [a] -> [b]
map Either
  (Symbol TyVar CType) (SOP TyVar CType, [Product TyVar CType])
-> Type
reifySymbol (forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Symbol TyVar CType
-> [Either
      (Symbol TyVar CType) (SOP TyVar CType, [Product TyVar CType])]
-> [Either
      (Symbol TyVar CType) (SOP TyVar CType, [Product TyVar CType])]
mergeExp [] [Symbol TyVar CType]
ps)
    in  forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 (\Type
t1 Type
t2 -> TyCon -> [Type] -> Type
mkTyConApp TyCon
typeNatMulTyCon [Type
t1,Type
t2]) [Type]
ps'
  where
    -- "2 ^ -1 * 2 ^ a" must be merged into "2 ^ (a-1)", otherwise GHC barfs
    -- at the "2 ^ -1" because of the negative exponent.
    mergeExp :: CoreSymbol -> [Either CoreSymbol (CoreSOP,[CoreProduct])]
                           -> [Either CoreSymbol (CoreSOP,[CoreProduct])]
    mergeExp :: Symbol TyVar CType
-> [Either
      (Symbol TyVar CType) (SOP TyVar CType, [Product TyVar CType])]
-> [Either
      (Symbol TyVar CType) (SOP TyVar CType, [Product TyVar CType])]
mergeExp (E SOP TyVar CType
s Product TyVar CType
p)   []     = [forall a b. b -> Either a b
Right (SOP TyVar CType
s,[Product TyVar CType
p])]
    mergeExp (E SOP TyVar CType
s1 Product TyVar CType
p1) (Either
  (Symbol TyVar CType) (SOP TyVar CType, [Product TyVar CType])
y:[Either
   (Symbol TyVar CType) (SOP TyVar CType, [Product TyVar CType])]
ys)
      | Right (SOP TyVar CType
s2,[Product TyVar CType]
p2) <- Either
  (Symbol TyVar CType) (SOP TyVar CType, [Product TyVar CType])
y
      , SOP TyVar CType
s1 forall a. Eq a => a -> a -> Bool
== SOP TyVar CType
s2
      = forall a b. b -> Either a b
Right (SOP TyVar CType
s1,(Product TyVar CType
p1forall a. a -> [a] -> [a]
:[Product TyVar CType]
p2)) forall a. a -> [a] -> [a]
: [Either
   (Symbol TyVar CType) (SOP TyVar CType, [Product TyVar CType])]
ys
      | Bool
otherwise
      = forall a b. b -> Either a b
Right (SOP TyVar CType
s1,[Product TyVar CType
p1]) forall a. a -> [a] -> [a]
: Either
  (Symbol TyVar CType) (SOP TyVar CType, [Product TyVar CType])
y forall a. a -> [a] -> [a]
: [Either
   (Symbol TyVar CType) (SOP TyVar CType, [Product TyVar CType])]
ys
    mergeExp Symbol TyVar CType
x [Either
   (Symbol TyVar CType) (SOP TyVar CType, [Product TyVar CType])]
ys = forall a b. a -> Either a b
Left Symbol TyVar CType
x forall a. a -> [a] -> [a]
: [Either
   (Symbol TyVar CType) (SOP TyVar CType, [Product TyVar CType])]
ys

reifySymbol :: Either CoreSymbol (CoreSOP,[CoreProduct]) -> Type
reifySymbol :: Either
  (Symbol TyVar CType) (SOP TyVar CType, [Product TyVar CType])
-> Type
reifySymbol (Left (I Integer
i)  )  = Integer -> Type
mkNumLitTy Integer
i
reifySymbol (Left (C CType
c)  )  = CType -> Type
unCType CType
c
reifySymbol (Left (V TyVar
v)  )  = TyVar -> Type
mkTyVarTy TyVar
v
reifySymbol (Left (E SOP TyVar CType
s Product TyVar CType
p))  = TyCon -> [Type] -> Type
mkTyConApp TyCon
typeNatExpTyCon [SOP TyVar CType -> Type
reifySOP SOP TyVar CType
s,Product TyVar CType -> Type
reifyProduct Product TyVar CType
p]
reifySymbol (Right (SOP TyVar CType
s1,[Product TyVar CType]
s2)) = TyCon -> [Type] -> Type
mkTyConApp TyCon
typeNatExpTyCon
                                         [SOP TyVar CType -> Type
reifySOP SOP TyVar CType
s1
                                         ,SOP TyVar CType -> Type
reifySOP (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
s2)
                                         ]

-- | Subtract an inequality, in order to either:
--
-- * See if the smallest solution is a natural number
-- * Cancel sums, i.e. monotonicity of addition
--
-- @
-- subtractIneq (2*y <=? 3*x ~ True)  = (-2*y + 3*x)
-- subtractIneq (2*y <=? 3*x ~ False) = (-3*x + (-1) + 2*y)
-- @
subtractIneq
  :: (CoreSOP, CoreSOP, Bool)
  -> CoreSOP
subtractIneq :: Ineq -> SOP TyVar CType
subtractIneq (SOP TyVar CType
x,SOP TyVar CType
y,Bool
isLE)
  | Bool
isLE
  = forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPAdd SOP TyVar CType
y (forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I (-Integer
1)]]) SOP TyVar CType
x)
  | Bool
otherwise
  = forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPAdd SOP TyVar CType
x (forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I (-Integer
1)]]) (forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPAdd SOP TyVar CType
y (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
1]])))

-- | Try to reverse the process of 'subtractIneq'
--
-- E.g.
--
-- @
-- subtractIneq (2*y <=? 3*x ~ True) = (-2*y + 3*x)
-- sopToIneq (-2*y+3*x) = Just (2*x <=? 3*x ~ True)
-- @
sopToIneq
  :: CoreSOP
  -> Maybe Ineq
sopToIneq :: SOP TyVar CType -> Maybe Ineq
sopToIneq (S [P ((I Integer
i):[Symbol TyVar CType]
l),Product TyVar CType
r])
  | Integer
i forall a. Ord a => a -> a -> Bool
< Integer
0
  = forall a. a -> Maybe a
Just (forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I (forall a. Num a => a -> a
negate Integer
i)]]) (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
l]),forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
r],Bool
True)
sopToIneq (S [Product TyVar CType
r,P ((I Integer
i:[Symbol TyVar CType]
l))])
  | Integer
i forall a. Ord a => a -> a -> Bool
< Integer
0
  = forall a. a -> Maybe a
Just (forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I (forall a. Num a => a -> a
negate Integer
i)]]) (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
l]),forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
r],Bool
True)
sopToIneq SOP TyVar CType
_ = forall a. Maybe a
Nothing

-- | Give the smallest solution for an inequality
ineqToSubst
  :: Ineq
  -> Maybe CoreUnify
ineqToSubst :: Ineq -> Maybe CoreUnify
ineqToSubst (SOP TyVar CType
x,S [P [V TyVar
v]],Bool
True)
  = forall a. a -> Maybe a
Just (forall v c. v -> SOP v c -> UnifyItem v c
SubstItem TyVar
v SOP TyVar CType
x)
ineqToSubst Ineq
_
  = forall a. Maybe a
Nothing

subtractionToPred
  :: TyCon
  -> (Type,Type)
  -> (PredType, Kind)
subtractionToPred :: TyCon -> (Type, Type) -> (Type, Type)
subtractionToPred TyCon
ordCond (Type
x,Type
y) =
#if MIN_VERSION_ghc(9,2,0)
  let cmpNat :: Type
cmpNat = TyCon -> [Type] -> Type
mkTyConApp TyCon
typeNatCmpTyCon [Type
y,Type
x]
      trueTc :: Type
trueTc = TyCon -> [Type] -> Type
mkTyConApp TyCon
promotedTrueDataCon []
      falseTc :: Type
falseTc = TyCon -> [Type] -> Type
mkTyConApp TyCon
promotedFalseDataCon []
      ordCmp :: Type
ordCmp = TyCon -> [Type] -> Type
mkTyConApp TyCon
ordCond
                [Type
boolTy,Type
cmpNat,Type
trueTc,Type
trueTc,Type
falseTc]
      predTy :: Type
predTy = Type -> Type -> Type
mkPrimEqPred Type
ordCmp Type
trueTc
   in (Type
predTy,Type
boolTy)
#else
  (mkPrimEqPred (mkTyConApp ordCond [y,x])
                (mkTyConApp promotedTrueDataCon [])
  ,boolTy)
#endif

-- | A substitution is essentially a list of (variable, 'SOP') pairs,
-- but we keep the original 'Ct' that lead to the substitution being
-- made, for use when turning the substitution back into constraints.
type CoreUnify = UnifyItem TyVar CType

data UnifyItem v c = SubstItem { forall v c. UnifyItem v c -> v
siVar :: v
                               , forall v c. UnifyItem v c -> SOP v c
siSOP :: SOP v c
                               }
                   | UnifyItem { forall v c. UnifyItem v c -> SOP v c
siLHS :: SOP v c
                               , forall v c. UnifyItem v c -> SOP v c
siRHS :: SOP v c
                               }
  deriving UnifyItem v c -> UnifyItem v c -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall v c. (Eq v, Eq c) => UnifyItem v c -> UnifyItem v c -> Bool
/= :: UnifyItem v c -> UnifyItem v c -> Bool
$c/= :: forall v c. (Eq v, Eq c) => UnifyItem v c -> UnifyItem v c -> Bool
== :: UnifyItem v c -> UnifyItem v c -> Bool
$c== :: forall v c. (Eq v, Eq c) => UnifyItem v c -> UnifyItem v c -> Bool
Eq

instance (Outputable v, Outputable c) => Outputable (UnifyItem v c) where
  ppr :: UnifyItem v c -> SDoc
ppr (SubstItem {v
SOP v c
siSOP :: SOP v c
siVar :: v
siSOP :: forall v c. UnifyItem v c -> SOP v c
siVar :: forall v c. UnifyItem v c -> v
..}) = forall a. Outputable a => a -> SDoc
ppr v
siVar SDoc -> SDoc -> SDoc
<+> String -> SDoc
text String
" := " SDoc -> SDoc -> SDoc
<+> forall a. Outputable a => a -> SDoc
ppr SOP v c
siSOP
  ppr (UnifyItem {SOP v c
siRHS :: SOP v c
siLHS :: SOP v c
siRHS :: forall v c. UnifyItem v c -> SOP v c
siLHS :: forall v c. UnifyItem v c -> SOP v c
..}) = forall a. Outputable a => a -> SDoc
ppr SOP v c
siLHS SDoc -> SDoc -> SDoc
<+> String -> SDoc
text String
" :~ " SDoc -> SDoc -> SDoc
<+> forall a. Outputable a => a -> SDoc
ppr SOP v c
siRHS

-- | Apply a substitution to a single normalised 'SOP' term
substsSOP :: (Ord v, Ord c) => [UnifyItem v c] -> SOP v c -> SOP v c
substsSOP :: forall v c. (Ord v, Ord c) => [UnifyItem v c] -> SOP v c -> SOP v c
substsSOP []                   SOP v c
u = SOP v c
u
substsSOP ((SubstItem {v
SOP v c
siSOP :: SOP v c
siVar :: v
siSOP :: forall v c. UnifyItem v c -> SOP v c
siVar :: forall v c. UnifyItem v c -> v
..}):[UnifyItem v c]
s) SOP v c
u = forall v c. (Ord v, Ord c) => [UnifyItem v c] -> SOP v c -> SOP v c
substsSOP [UnifyItem v c]
s (forall v c. (Ord v, Ord c) => v -> SOP v c -> SOP v c -> SOP v c
substSOP v
siVar SOP v c
siSOP SOP v c
u)
substsSOP ((UnifyItem {}):[UnifyItem v c]
s)   SOP v c
u = forall v c. (Ord v, Ord c) => [UnifyItem v c] -> SOP v c -> SOP v c
substsSOP [UnifyItem v c]
s SOP v c
u

substSOP :: (Ord v, Ord c) => v -> SOP v c -> SOP v c -> SOP v c
substSOP :: forall v c. (Ord v, Ord c) => v -> SOP v c -> SOP v c -> SOP v c
substSOP v
tv SOP v c
e = forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPAdd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall v c.
(Ord v, Ord c) =>
v -> SOP v c -> Product v c -> SOP v c
substProduct v
tv SOP v c
e) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v c. SOP v c -> [Product v c]
unS

substProduct :: (Ord v, Ord c) => v -> SOP v c -> Product v c -> SOP v c
substProduct :: forall v c.
(Ord v, Ord c) =>
v -> SOP v c -> Product v c -> SOP v c
substProduct v
tv SOP v c
e = forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall v c. (Ord v, Ord c) => v -> SOP v c -> Symbol v c -> SOP v c
substSymbol v
tv SOP v c
e) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v c. Product v c -> [Symbol v c]
unP

substSymbol :: (Ord v, Ord c) => v -> SOP v c -> Symbol v c -> SOP v c
substSymbol :: forall v c. (Ord v, Ord c) => v -> SOP v c -> Symbol v c -> SOP v c
substSymbol v
_  SOP v c
_ s :: Symbol v c
s@(I Integer
_) = forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol v c
s]]
substSymbol v
_  SOP v c
_ s :: Symbol v c
s@(C c
_) = forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol v c
s]]
substSymbol v
tv SOP v c
e (V v
tv')
  | v
tv forall a. Eq a => a -> a -> Bool
== v
tv'            = SOP v c
e
  | Bool
otherwise            = forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. v -> Symbol v c
V v
tv']]
substSymbol v
tv SOP v c
e (E SOP v c
s Product v c
p) = forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
normaliseExp (forall v c. (Ord v, Ord c) => v -> SOP v c -> SOP v c -> SOP v c
substSOP v
tv SOP v c
e SOP v c
s) (forall v c.
(Ord v, Ord c) =>
v -> SOP v c -> Product v c -> SOP v c
substProduct v
tv SOP v c
e Product v c
p)

-- | Apply a substitution to a substitution
substsSubst :: (Ord v, Ord c) => [UnifyItem v c] -> [UnifyItem v c] -> [UnifyItem v c]
substsSubst :: forall v c.
(Ord v, Ord c) =>
[UnifyItem v c] -> [UnifyItem v c] -> [UnifyItem v c]
substsSubst [UnifyItem v c]
s = forall a b. (a -> b) -> [a] -> [b]
map UnifyItem v c -> UnifyItem v c
subt
  where
    subt :: UnifyItem v c -> UnifyItem v c
subt si :: UnifyItem v c
si@(SubstItem {v
SOP v c
siSOP :: SOP v c
siVar :: v
siSOP :: forall v c. UnifyItem v c -> SOP v c
siVar :: forall v c. UnifyItem v c -> v
..}) = UnifyItem v c
si {siSOP :: SOP v c
siSOP = forall v c. (Ord v, Ord c) => [UnifyItem v c] -> SOP v c -> SOP v c
substsSOP [UnifyItem v c]
s SOP v c
siSOP}
    subt si :: UnifyItem v c
si@(UnifyItem {SOP v c
siRHS :: SOP v c
siLHS :: SOP v c
siRHS :: forall v c. UnifyItem v c -> SOP v c
siLHS :: forall v c. UnifyItem v c -> SOP v c
..}) = UnifyItem v c
si {siLHS :: SOP v c
siLHS = forall v c. (Ord v, Ord c) => [UnifyItem v c] -> SOP v c -> SOP v c
substsSOP [UnifyItem v c]
s SOP v c
siLHS, siRHS :: SOP v c
siRHS = forall v c. (Ord v, Ord c) => [UnifyItem v c] -> SOP v c -> SOP v c
substsSOP [UnifyItem v c]
s SOP v c
siRHS}
{-# INLINEABLE substsSubst #-}

-- | Result of comparing two 'SOP' terms, returning a potential substitution
-- list under which the two terms are equal.
data UnifyResult
  = Win              -- ^ Two terms are equal
  | Lose             -- ^ Two terms are /not/ equal
  | Draw [CoreUnify] -- ^ Two terms are only equal if the given substitution holds

instance Outputable UnifyResult where
  ppr :: UnifyResult -> SDoc
ppr UnifyResult
Win          = String -> SDoc
text String
"Win"
  ppr (Draw [CoreUnify]
subst) = String -> SDoc
text String
"Draw" SDoc -> SDoc -> SDoc
<+> forall a. Outputable a => a -> SDoc
ppr [CoreUnify]
subst
  ppr UnifyResult
Lose         = String -> SDoc
text String
"Lose"

-- | Given two 'SOP's @u@ and @v@, when their free variables ('fvSOP') are the
-- same, then we 'Win' if @u@ and @v@ are equal, and 'Lose' otherwise.
--
-- If @u@ and @v@ do not have the same free variables, we result in a 'Draw',
-- ware @u@ and @v@ are only equal when the returned 'CoreSubst' holds.
unifyNats :: Ct -> CoreSOP -> CoreSOP -> TcPluginM UnifyResult
unifyNats :: Ct -> SOP TyVar CType -> SOP TyVar CType -> TcPluginM UnifyResult
unifyNats Ct
ct SOP TyVar CType
u SOP TyVar CType
v = do
  String -> SDoc -> TcPluginM ()
tcPluginTrace String
"unifyNats" (forall a. Outputable a => a -> SDoc
ppr Ct
ct SDoc -> SDoc -> SDoc
$$ forall a. Outputable a => a -> SDoc
ppr SOP TyVar CType
u SDoc -> SDoc -> SDoc
$$ forall a. Outputable a => a -> SDoc
ppr SOP TyVar CType
v)
  forall (m :: * -> *) a. Monad m => a -> m a
return (Ct -> SOP TyVar CType -> SOP TyVar CType -> UnifyResult
unifyNats' Ct
ct SOP TyVar CType
u SOP TyVar CType
v)

unifyNats' :: Ct -> CoreSOP -> CoreSOP -> UnifyResult
unifyNats' :: Ct -> SOP TyVar CType -> SOP TyVar CType -> UnifyResult
unifyNats' Ct
ct SOP TyVar CType
u SOP TyVar CType
v
  = if SOP TyVar CType -> SOP TyVar CType -> Bool
eqFV SOP TyVar CType
u SOP TyVar CType
v
       then if SOP TyVar CType -> Bool
containsConstants SOP TyVar CType
u Bool -> Bool -> Bool
|| SOP TyVar CType -> Bool
containsConstants SOP TyVar CType
v
               then if SOP TyVar CType
u forall a. Eq a => a -> a -> Bool
== SOP TyVar CType
v
                       then UnifyResult
Win
                       else [CoreUnify] -> UnifyResult
Draw (forall a. (a -> Bool) -> [a] -> [a]
filter CoreUnify -> Bool
diffFromConstraint (Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers Ct
ct SOP TyVar CType
u SOP TyVar CType
v))
               else if SOP TyVar CType
u forall a. Eq a => a -> a -> Bool
== SOP TyVar CType
v
                       then UnifyResult
Win
                       else UnifyResult
Lose
       else [CoreUnify] -> UnifyResult
Draw (forall a. (a -> Bool) -> [a] -> [a]
filter CoreUnify -> Bool
diffFromConstraint (Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers Ct
ct SOP TyVar CType
u SOP TyVar CType
v))
  where
    -- A unifier is only a unifier if differs from the original constraint
    diffFromConstraint :: CoreUnify -> Bool
diffFromConstraint (UnifyItem SOP TyVar CType
x SOP TyVar CType
y) = Bool -> Bool
not (SOP TyVar CType
x forall a. Eq a => a -> a -> Bool
== SOP TyVar CType
u Bool -> Bool -> Bool
&& SOP TyVar CType
y forall a. Eq a => a -> a -> Bool
== SOP TyVar CType
v)
    diffFromConstraint CoreUnify
_               = Bool
True

-- | Find unifiers for two SOP terms
--
-- Can find the following unifiers:
--
-- @
-- t ~ a + b          ==>  [t := a + b]
-- a + b ~ t          ==>  [t := a + b]
-- (a + c) ~ (b + c)  ==>  \[a := b\]
-- (2*a) ~ (2*b)      ==>  [a := b]
-- (2 + a) ~ 5        ==>  [a := 3]
-- (i * a) ~ j        ==>  [a := div j i], when (mod j i == 0)
-- @
--
-- However, given a wanted:
--
-- @
-- [W] t ~ a + b
-- @
--
-- this function returns @[]@, or otherwise we \"solve\" the constraint by
-- finding a unifier equal to the constraint.
--
-- However, given a wanted:
--
-- @
-- [W] (a + c) ~ (b + c)
-- @
--
-- we do return the unifier:
--
-- @
-- [a := b]
-- @
unifiers :: Ct -> CoreSOP -> CoreSOP -> [CoreUnify]
unifiers :: Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers Ct
ct u :: SOP TyVar CType
u@(S [P [V TyVar
x]]) SOP TyVar CType
v
  = case Type -> Pred
classifyPredType forall a b. (a -> b) -> a -> b
$ CtEvidence -> Type
ctEvPred forall a b. (a -> b) -> a -> b
$ Ct -> CtEvidence
ctEvidence Ct
ct of
      EqPred EqRel
NomEq Type
t1 Type
_
        | Type -> CType
CType (SOP TyVar CType -> Type
reifySOP SOP TyVar CType
u) forall a. Eq a => a -> a -> Bool
/= Type -> CType
CType Type
t1 Bool -> Bool -> Bool
|| CtEvidence -> Bool
isGiven (Ct -> CtEvidence
ctEvidence Ct
ct) -> [forall v c. v -> SOP v c -> UnifyItem v c
SubstItem TyVar
x SOP TyVar CType
v]
      Pred
_ -> []
unifiers Ct
ct SOP TyVar CType
u v :: SOP TyVar CType
v@(S [P [V TyVar
x]])
  = case Type -> Pred
classifyPredType forall a b. (a -> b) -> a -> b
$ CtEvidence -> Type
ctEvPred forall a b. (a -> b) -> a -> b
$ Ct -> CtEvidence
ctEvidence Ct
ct of
      EqPred EqRel
NomEq Type
_ Type
t2
        | Type -> CType
CType (SOP TyVar CType -> Type
reifySOP SOP TyVar CType
v) forall a. Eq a => a -> a -> Bool
/= Type -> CType
CType Type
t2 Bool -> Bool -> Bool
|| CtEvidence -> Bool
isGiven (Ct -> CtEvidence
ctEvidence Ct
ct) -> [forall v c. v -> SOP v c -> UnifyItem v c
SubstItem TyVar
x SOP TyVar CType
u]
      Pred
_ -> []
unifiers Ct
ct u :: SOP TyVar CType
u@(S [P [C CType
_]]) SOP TyVar CType
v
  = case Type -> Pred
classifyPredType forall a b. (a -> b) -> a -> b
$ CtEvidence -> Type
ctEvPred forall a b. (a -> b) -> a -> b
$ Ct -> CtEvidence
ctEvidence Ct
ct of
      EqPred EqRel
NomEq Type
t1 Type
t2
        | Type -> CType
CType (SOP TyVar CType -> Type
reifySOP SOP TyVar CType
u) forall a. Eq a => a -> a -> Bool
/= Type -> CType
CType Type
t1 Bool -> Bool -> Bool
|| Type -> CType
CType (SOP TyVar CType -> Type
reifySOP SOP TyVar CType
v) forall a. Eq a => a -> a -> Bool
/= Type -> CType
CType Type
t2 -> [forall v c. SOP v c -> SOP v c -> UnifyItem v c
UnifyItem SOP TyVar CType
u SOP TyVar CType
v]
      Pred
_ -> []
unifiers Ct
ct SOP TyVar CType
u v :: SOP TyVar CType
v@(S [P [C CType
_]])
  = case Type -> Pred
classifyPredType forall a b. (a -> b) -> a -> b
$ CtEvidence -> Type
ctEvPred forall a b. (a -> b) -> a -> b
$ Ct -> CtEvidence
ctEvidence Ct
ct of
      EqPred EqRel
NomEq Type
t1 Type
t2
        | Type -> CType
CType (SOP TyVar CType -> Type
reifySOP SOP TyVar CType
u) forall a. Eq a => a -> a -> Bool
/= Type -> CType
CType Type
t1 Bool -> Bool -> Bool
|| Type -> CType
CType (SOP TyVar CType -> Type
reifySOP SOP TyVar CType
v) forall a. Eq a => a -> a -> Bool
/= Type -> CType
CType Type
t2 -> [forall v c. SOP v c -> SOP v c -> UnifyItem v c
UnifyItem SOP TyVar CType
u SOP TyVar CType
v]
      Pred
_ -> []
unifiers Ct
ct SOP TyVar CType
u SOP TyVar CType
v             = Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
ct SOP TyVar CType
u SOP TyVar CType
v

unifiers' :: Ct -> CoreSOP -> CoreSOP -> [CoreUnify]
unifiers' :: Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
_ct (S [P [V TyVar
x]]) (S [])        = [forall v c. v -> SOP v c -> UnifyItem v c
SubstItem TyVar
x (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
0]])]
unifiers' Ct
_ct (S [])        (S [P [V TyVar
x]]) = [forall v c. v -> SOP v c -> UnifyItem v c
SubstItem TyVar
x (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
0]])]

unifiers' Ct
_ct (S [P [V TyVar
x]]) SOP TyVar CType
s             = [forall v c. v -> SOP v c -> UnifyItem v c
SubstItem TyVar
x SOP TyVar CType
s]
unifiers' Ct
_ct SOP TyVar CType
s             (S [P [V TyVar
x]]) = [forall v c. v -> SOP v c -> UnifyItem v c
SubstItem TyVar
x SOP TyVar CType
s]

unifiers' Ct
_ct s1 :: SOP TyVar CType
s1@(S [P [C CType
_]]) SOP TyVar CType
s2               = [forall v c. SOP v c -> SOP v c -> UnifyItem v c
UnifyItem SOP TyVar CType
s1 SOP TyVar CType
s2]
unifiers' Ct
_ct SOP TyVar CType
s1               s2 :: SOP TyVar CType
s2@(S [P [C CType
_]]) = [forall v c. SOP v c -> SOP v c -> UnifyItem v c
UnifyItem SOP TyVar CType
s1 SOP TyVar CType
s2]


-- (z ^ a) ~ (z ^ b) ==> [a := b]
unifiers' Ct
ct (S [P [E SOP TyVar CType
s1 Product TyVar CType
p1]]) (S [P [E SOP TyVar CType
s2 Product TyVar CType
p2]])
  | SOP TyVar CType
s1 forall a. Eq a => a -> a -> Bool
== SOP TyVar CType
s2 = Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
ct (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
p1]) (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
p2])

-- (2*e ^ d) ~ (2*e*a*c) ==> [a*c := 2*e ^ (d-1)]
unifiers' Ct
ct (S [P [E (S [P [Symbol TyVar CType]
s1]) Product TyVar CType
p1]]) (S [P [Symbol TyVar CType]
p2])
  | forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Symbol TyVar CType]
p2) [Symbol TyVar CType]
s1
  = let base :: [Symbol TyVar CType]
base = forall a. Eq a => [a] -> [a] -> [a]
intersect [Symbol TyVar CType]
s1 [Symbol TyVar CType]
p2
        diff :: [Symbol TyVar CType]
diff = [Symbol TyVar CType]
p2 forall a. Eq a => [a] -> [a] -> [a]
\\ [Symbol TyVar CType]
s1
    in  Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers Ct
ct (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
diff]) (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. SOP v c -> Product v c -> Symbol v c
E (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
base]) (forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I (-Integer
1)]),forall v c. SOP v c -> Product v c -> Symbol v c
E (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
base]) Product TyVar CType
p1]])

unifiers' Ct
ct (S [P [Symbol TyVar CType]
p2]) (S [P [E (S [P [Symbol TyVar CType]
s1]) Product TyVar CType
p1]])
  | forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Symbol TyVar CType]
p2) [Symbol TyVar CType]
s1
  = let base :: [Symbol TyVar CType]
base = forall a. Eq a => [a] -> [a] -> [a]
intersect [Symbol TyVar CType]
s1 [Symbol TyVar CType]
p2
        diff :: [Symbol TyVar CType]
diff = [Symbol TyVar CType]
p2 forall a. Eq a => [a] -> [a] -> [a]
\\ [Symbol TyVar CType]
s1
    in  Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers Ct
ct (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. SOP v c -> Product v c -> Symbol v c
E (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
base]) (forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I (-Integer
1)]),forall v c. SOP v c -> Product v c -> Symbol v c
E (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
base]) Product TyVar CType
p1]]) (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
diff])

-- (i ^ a) ~ j ==> [a := round (logBase i j)], when `i` and `j` are integers,
-- and `ceiling (logBase i j) == floor (logBase i j)`
unifiers' Ct
ct (S [P [E (S [P [I Integer
i]]) Product TyVar CType
p]]) (S [P [I Integer
j]])
  = case Integer -> Integer -> Maybe Integer
integerLogBase Integer
i Integer
j of
      Just Integer
k  -> Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
ct (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
p]) (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
k]])
      Maybe Integer
Nothing -> []

unifiers' Ct
ct (S [P [I Integer
j]]) (S [P [E (S [P [I Integer
i]]) Product TyVar CType
p]])
  = case Integer -> Integer -> Maybe Integer
integerLogBase Integer
i Integer
j of
      Just Integer
k  -> Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
ct (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
p]) (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
k]])
      Maybe Integer
Nothing -> []

-- a^d * a^e ~ a^c ==> [c := d + e]
unifiers' Ct
ct (S [P [E SOP TyVar CType
s1 Product TyVar CType
p1]]) (S [Product TyVar CType
p2]) = case Product TyVar CType
-> Maybe ([SOP TyVar CType], [Product TyVar CType])
collectBases Product TyVar CType
p2 of
  Just (SOP TyVar CType
b:[SOP TyVar CType]
bs,[Product TyVar CType]
ps) | forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. Eq a => a -> a -> Bool
== SOP TyVar CType
s1) (SOP TyVar CType
bforall a. a -> [a] -> [a]
:[SOP TyVar CType]
bs) ->
    Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
ct (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
p1]) (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
ps)
  Maybe ([SOP TyVar CType], [Product TyVar CType])
_ -> []

unifiers' Ct
ct (S [Product TyVar CType
p2]) (S [P [E SOP TyVar CType
s1 Product TyVar CType
p1]]) = case Product TyVar CType
-> Maybe ([SOP TyVar CType], [Product TyVar CType])
collectBases Product TyVar CType
p2 of
  Just (SOP TyVar CType
b:[SOP TyVar CType]
bs,[Product TyVar CType]
ps) | forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. Eq a => a -> a -> Bool
== SOP TyVar CType
s1) (SOP TyVar CType
bforall a. a -> [a] -> [a]
:[SOP TyVar CType]
bs) ->
    Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
ct (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
ps) (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
p1])
  Maybe ([SOP TyVar CType], [Product TyVar CType])
_ -> []

-- (i * a) ~ j ==> [a := div j i]
-- Where 'a' is a variable, 'i' and 'j' are integer literals, and j `mod` i == 0
unifiers' Ct
ct (S [P ((I Integer
i):[Symbol TyVar CType]
ps)]) (S [P [I Integer
j]]) =
  case Integer -> Integer -> Maybe Integer
safeDiv Integer
j Integer
i of
    Just Integer
k -> Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
ct (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
ps]) (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
k]])
    Maybe Integer
_      -> []

unifiers' Ct
ct (S [P [I Integer
j]]) (S [P ((I Integer
i):[Symbol TyVar CType]
ps)]) =
  case Integer -> Integer -> Maybe Integer
safeDiv Integer
j Integer
i of
    Just Integer
k -> Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
ct (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
ps]) (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
k]])
    Maybe Integer
_      -> []

-- (2*a) ~ (2*b) ==> [a := b]
-- unifiers' ct (S [P (p:ps1)]) (S [P (p':ps2)])
--     | p == p'   = unifiers' ct (S [P ps1]) (S [P ps2])
--     | otherwise = []
unifiers' Ct
ct (S [P [Symbol TyVar CType]
ps1]) (S [P [Symbol TyVar CType]
ps2])
    | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Symbol TyVar CType]
psx  = []
    | Bool
otherwise = Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
ct (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
ps1'']) (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
ps2''])
  where
    ps1' :: [Symbol TyVar CType]
ps1'  = [Symbol TyVar CType]
ps1 forall a. Eq a => [a] -> [a] -> [a]
\\ [Symbol TyVar CType]
psx
    ps2' :: [Symbol TyVar CType]
ps2'  = [Symbol TyVar CType]
ps2 forall a. Eq a => [a] -> [a] -> [a]
\\ [Symbol TyVar CType]
psx
    ps1'' :: [Symbol TyVar CType]
ps1'' | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Symbol TyVar CType]
ps1' = [forall v c. Integer -> Symbol v c
I Integer
1]
          | Bool
otherwise = [Symbol TyVar CType]
ps1'
    ps2'' :: [Symbol TyVar CType]
ps2'' | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Symbol TyVar CType]
ps2' = [forall v c. Integer -> Symbol v c
I Integer
1]
          | Bool
otherwise = [Symbol TyVar CType]
ps2'
    psx :: [Symbol TyVar CType]
psx  = forall a. Eq a => [a] -> [a] -> [a]
intersect [Symbol TyVar CType]
ps1 [Symbol TyVar CType]
ps2

-- (2 + a) ~ 5 ==> [a := 3]
unifiers' Ct
ct (S ((P [I Integer
i]):[Product TyVar CType]
ps1)) (S ((P [I Integer
j]):[Product TyVar CType]
ps2))
    | Integer
i forall a. Ord a => a -> a -> Bool
< Integer
j     = Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
ct (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
ps1) (forall v c. [Product v c] -> SOP v c
S ((forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I (Integer
jforall a. Num a => a -> a -> a
-Integer
i)])forall a. a -> [a] -> [a]
:[Product TyVar CType]
ps2))
    | Integer
i forall a. Ord a => a -> a -> Bool
> Integer
j     = Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
ct (forall v c. [Product v c] -> SOP v c
S ((forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I (Integer
iforall a. Num a => a -> a -> a
-Integer
j)])forall a. a -> [a] -> [a]
:[Product TyVar CType]
ps1)) (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
ps2)

-- (a + c) ~ (b + c) ==> [a := b]
unifiers' Ct
ct s1 :: SOP TyVar CType
s1@(S [Product TyVar CType]
ps1) s2 :: SOP TyVar CType
s2@(S [Product TyVar CType]
ps2) = case SOP TyVar CType -> Maybe Ineq
sopToIneq SOP TyVar CType
k1 of
  Just (SOP TyVar CType
s1',SOP TyVar CType
s2',Bool
_)
    | SOP TyVar CType
s1' forall a. Eq a => a -> a -> Bool
/= SOP TyVar CType
s1 Bool -> Bool -> Bool
|| SOP TyVar CType
s2' forall a. Eq a => a -> a -> Bool
/= SOP TyVar CType
s1
    , forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Bool -> Bool -> Bool
(&&) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second forall a. Set a -> Bool
Set.null) (forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (SOP TyVar CType -> WriterT (Set CType) Maybe Bool
isNatural SOP TyVar CType
s1'))
    , forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Bool -> Bool -> Bool
(&&) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second forall a. Set a -> Bool
Set.null) (forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (SOP TyVar CType -> WriterT (Set CType) Maybe Bool
isNatural SOP TyVar CType
s2'))
    -> Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
ct SOP TyVar CType
s1' SOP TyVar CType
s2'
  Maybe Ineq
_ | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Product TyVar CType]
psx
    , forall (t :: * -> *) a. Foldable t => t a -> Int
length [Product TyVar CType]
ps1 forall a. Eq a => a -> a -> Bool
== forall (t :: * -> *) a. Foldable t => t a -> Int
length [Product TyVar CType]
ps2
    -> case forall a. Eq a => [a] -> [a]
nub (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Product TyVar CType
x Product TyVar CType
y -> Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
ct (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
x]) (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
y])) [Product TyVar CType]
ps1 [Product TyVar CType]
ps2)) of
        []                             -> Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers'' Ct
ct (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
ps1) (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
ps2)
        [CoreUnify
k] | forall (t :: * -> *) a. Foldable t => t a -> Int
length [Product TyVar CType]
ps1 forall a. Eq a => a -> a -> Bool
== forall (t :: * -> *) a. Foldable t => t a -> Int
length [Product TyVar CType]
ps2 -> [CoreUnify
k]
        [CoreUnify]
_                              -> []
    | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Product TyVar CType]
psx
    , CtEvidence -> Bool
isGiven (Ct -> CtEvidence
ctEvidence Ct
ct)
    -> Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers'' Ct
ct (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
ps1) (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
ps2)
    | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Product TyVar CType]
psx
    -> []
  Maybe Ineq
_ -> Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers' Ct
ct (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
ps1'') (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
ps2'')
  where
    k1 :: SOP TyVar CType
k1 = Ineq -> SOP TyVar CType
subtractIneq (SOP TyVar CType
s1,SOP TyVar CType
s2,Bool
True)
    ps1' :: [Product TyVar CType]
ps1'  = [Product TyVar CType]
ps1 forall a. Eq a => [a] -> [a] -> [a]
\\ [Product TyVar CType]
psx
    ps2' :: [Product TyVar CType]
ps2'  = [Product TyVar CType]
ps2 forall a. Eq a => [a] -> [a] -> [a]
\\ [Product TyVar CType]
psx
    ps1'' :: [Product TyVar CType]
ps1'' | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Product TyVar CType]
ps1' = [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
0]]
          | Bool
otherwise = [Product TyVar CType]
ps1'
    ps2'' :: [Product TyVar CType]
ps2'' | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Product TyVar CType]
ps2' = [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
0]]
          | Bool
otherwise = [Product TyVar CType]
ps2'
    psx :: [Product TyVar CType]
psx = forall a. Eq a => [a] -> [a] -> [a]
intersect [Product TyVar CType]
ps1 [Product TyVar CType]
ps2

unifiers'' :: Ct -> CoreSOP -> CoreSOP -> [CoreUnify]
unifiers'' :: Ct -> SOP TyVar CType -> SOP TyVar CType -> [CoreUnify]
unifiers'' Ct
ct (S [P [I Integer
i],P [V TyVar
v]]) SOP TyVar CType
s2
  | CtEvidence -> Bool
isGiven (Ct -> CtEvidence
ctEvidence Ct
ct) = [forall v c. v -> SOP v c -> UnifyItem v c
SubstItem TyVar
v (forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPAdd SOP TyVar CType
s2 (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I (forall a. Num a => a -> a
negate Integer
i)]]))]
unifiers'' Ct
ct SOP TyVar CType
s1 (S [P [I Integer
i],P [V TyVar
v]])
  | CtEvidence -> Bool
isGiven (Ct -> CtEvidence
ctEvidence Ct
ct) = [forall v c. v -> SOP v c -> UnifyItem v c
SubstItem TyVar
v (forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPAdd SOP TyVar CType
s1 (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I (forall a. Num a => a -> a
negate Integer
i)]]))]
unifiers'' Ct
_ SOP TyVar CType
_ SOP TyVar CType
_ = []

collectBases :: CoreProduct -> Maybe ([CoreSOP],[CoreProduct])
collectBases :: Product TyVar CType
-> Maybe ([SOP TyVar CType], [Product TyVar CType])
collectBases = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. [(a, b)] -> ([a], [b])
unzip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall {v} {c}. Symbol v c -> Maybe (SOP v c, Product v c)
go forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v c. Product v c -> [Symbol v c]
unP
  where
    go :: Symbol v c -> Maybe (SOP v c, Product v c)
go (E SOP v c
s1 Product v c
p1) = forall a. a -> Maybe a
Just (SOP v c
s1,Product v c
p1)
    go Symbol v c
_         = forall a. Maybe a
Nothing

-- | Find the 'TyVar' in a 'CoreSOP'
fvSOP :: CoreSOP -> UniqSet TyVar
fvSOP :: SOP TyVar CType -> UniqSet TyVar
fvSOP = forall a. [UniqSet a] -> UniqSet a
unionManyUniqSets forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map Product TyVar CType -> UniqSet TyVar
fvProduct forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v c. SOP v c -> [Product v c]
unS

fvProduct :: CoreProduct -> UniqSet TyVar
fvProduct :: Product TyVar CType -> UniqSet TyVar
fvProduct = forall a. [UniqSet a] -> UniqSet a
unionManyUniqSets forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map Symbol TyVar CType -> UniqSet TyVar
fvSymbol forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v c. Product v c -> [Symbol v c]
unP

fvSymbol :: CoreSymbol -> UniqSet TyVar
fvSymbol :: Symbol TyVar CType -> UniqSet TyVar
fvSymbol (I Integer
_)   = forall a. UniqSet a
emptyUniqSet
fvSymbol (C CType
_)   = forall a. UniqSet a
emptyUniqSet
fvSymbol (V TyVar
v)   = forall a. Uniquable a => a -> UniqSet a
unitUniqSet TyVar
v
fvSymbol (E SOP TyVar CType
s Product TyVar CType
p) = SOP TyVar CType -> UniqSet TyVar
fvSOP SOP TyVar CType
s forall a. UniqSet a -> UniqSet a -> UniqSet a
`unionUniqSets` Product TyVar CType -> UniqSet TyVar
fvProduct Product TyVar CType
p

eqFV :: CoreSOP -> CoreSOP -> Bool
eqFV :: SOP TyVar CType -> SOP TyVar CType -> Bool
eqFV = forall a. Eq a => a -> a -> Bool
(==) forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` SOP TyVar CType -> UniqSet TyVar
fvSOP

containsConstants :: CoreSOP -> Bool
containsConstants :: SOP TyVar CType -> Bool
containsConstants =
  forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Symbol TyVar CType -> Bool
symbolContainsConstant forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v c. Product v c -> [Symbol v c]
unP) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v c. SOP v c -> [Product v c]
unS
  where
    symbolContainsConstant :: Symbol TyVar CType -> Bool
symbolContainsConstant Symbol TyVar CType
c = case Symbol TyVar CType
c of
      C {} -> Bool
True
      E SOP TyVar CType
s Product TyVar CType
p -> SOP TyVar CType -> Bool
containsConstants SOP TyVar CType
s Bool -> Bool -> Bool
|| SOP TyVar CType -> Bool
containsConstants (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
p])
      Symbol TyVar CType
_ -> Bool
False

safeDiv :: Integer -> Integer -> Maybe Integer
safeDiv :: Integer -> Integer -> Maybe Integer
safeDiv Integer
i Integer
j
  | Integer
j forall a. Eq a => a -> a -> Bool
== Integer
0    = forall a. a -> Maybe a
Just Integer
0
  | Bool
otherwise = case forall a. Integral a => a -> a -> (a, a)
divMod Integer
i Integer
j of
                  (Integer
k,Integer
0) -> forall a. a -> Maybe a
Just Integer
k
                  (Integer, Integer)
_     -> forall a. Maybe a
Nothing

-- | Given `x` and `y`, return `Just n` when
--
-- `ceiling (logBase x y) == floor (logBase x y)`
integerLogBase :: Integer -> Integer -> Maybe Integer
integerLogBase :: Integer -> Integer -> Maybe Integer
integerLogBase Integer
x Integer
y | Integer
x forall a. Ord a => a -> a -> Bool
> Integer
1 Bool -> Bool -> Bool
&& Integer
y forall a. Ord a => a -> a -> Bool
> Integer
0 =
  let z1 :: Int#
z1 = Integer -> Integer -> Int#
integerLogBase# Integer
x Integer
y
      z2 :: Int#
z2 = Integer -> Integer -> Int#
integerLogBase# Integer
x (Integer
yforall a. Num a => a -> a -> a
-Integer
1)
  in  if Int# -> Bool
isTrue# (Int#
z1 Int# -> Int# -> Int#
==# Int#
z2)
         then forall a. Maybe a
Nothing
         else forall a. a -> Maybe a
Just (Int# -> Integer
smallInteger Int#
z1)
integerLogBase Integer
_ Integer
_ = forall a. Maybe a
Nothing

isNatural :: CoreSOP -> WriterT (Set CType) Maybe Bool
isNatural :: SOP TyVar CType -> WriterT (Set CType) Maybe Bool
isNatural (S [])           = forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
isNatural (S [P []])       = forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
isNatural (S [P (I Integer
i:[Symbol TyVar CType]
ps)])
  | Integer
i forall a. Ord a => a -> a -> Bool
>= Integer
0    = SOP TyVar CType -> WriterT (Set CType) Maybe Bool
isNatural (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
ps])
  | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
isNatural (S [P (V TyVar
_:[Symbol TyVar CType]
ps)]) = SOP TyVar CType -> WriterT (Set CType) Maybe Bool
isNatural (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
ps])
isNatural (S [P (E SOP TyVar CType
s Product TyVar CType
p:[Symbol TyVar CType]
ps)]) = do
  Bool
sN <- SOP TyVar CType -> WriterT (Set CType) Maybe Bool
isNatural SOP TyVar CType
s
  Bool
pN <- SOP TyVar CType -> WriterT (Set CType) Maybe Bool
isNatural (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
p])
  if Bool
sN Bool -> Bool -> Bool
&& Bool
pN
     then SOP TyVar CType -> WriterT (Set CType) Maybe Bool
isNatural (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
ps])
     else forall w (m :: * -> *) a. m (a, w) -> WriterT w m a
WriterT forall a. Maybe a
Nothing
-- We give up for all other products for now
isNatural (S [P (C CType
c:[Symbol TyVar CType]
ps)]) = do
  forall (m :: * -> *) w. Monad m => w -> WriterT w m ()
tell (forall a. a -> Set a
Set.singleton CType
c)
  SOP TyVar CType -> WriterT (Set CType) Maybe Bool
isNatural (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
ps])
-- Adding two natural numbers is also a natural number
isNatural (S (Product TyVar CType
p:[Product TyVar CType]
ps)) = do
  Bool
pN <- SOP TyVar CType -> WriterT (Set CType) Maybe Bool
isNatural (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
p])
  Bool
pK <- SOP TyVar CType -> WriterT (Set CType) Maybe Bool
isNatural (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
ps)
  case (Bool
pN,Bool
pK) of
    (Bool
True,Bool
True)   -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True  -- both are natural
    (Bool
False,Bool
False) -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False -- both are non-natural
    (Bool, Bool)
_             -> forall w (m :: * -> *) a. m (a, w) -> WriterT w m a
WriterT forall a. Maybe a
Nothing
    -- if one is natural and the other isn't, then their sum *might* be natural,
    -- but we simply cant be sure.

-- | Try to solve inequalities
solveIneq
  :: Word
  -- ^ Solving depth
  -> Ineq
  -- ^ Inequality we want to solve
  -> Ineq
  -- ^ Given/proven inequality
  -> WriterT (Set CType) Maybe Bool
  -- ^ Solver result
  --
  -- * /Nothing/: exhausted solver steps
  --
  -- * /Just True/: inequality is solved
  --
  -- * /Just False/: solver is unable to solve inequality, note that this does
  -- __not__ mean the wanted inequality does not hold.
solveIneq :: Word -> Ineq -> Ineq -> WriterT (Set CType) Maybe Bool
solveIneq Word
0 Ineq
_ Ineq
_ = forall a. WriterT (Set CType) Maybe a
noRewrite
solveIneq Word
k want :: Ineq
want@(SOP TyVar CType
_,SOP TyVar CType
_,Bool
True) have :: Ineq
have@(SOP TyVar CType
_,SOP TyVar CType
_,Bool
True)
  | Ineq
want forall a. Eq a => a -> a -> Bool
== Ineq
have
  = forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
  | Bool
otherwise
  = do
    let -- Apply all the rules, and get all the successful ones
        new :: [([(Ineq, Ineq)], Set CType)]
new     = forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
f -> forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
f Ineq
want Ineq
have)) [Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]]
ineqRules
        -- Recurse down with all the transformed equations
        solved :: [([(Bool, Set CType)], Set CType)]
solved  = forall a b. (a -> b) -> [a] -> [b]
map (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (Word -> Ineq -> Ineq -> WriterT (Set CType) Maybe Bool
solveIneq (Word
kforall a. Num a => a -> a -> a
-Word
1))))) [([(Ineq, Ineq)], Set CType)]
new
        -- For the results of every recursive call, find the one that yields
        -- 'True' and has the smallest set of constraints.
        solved1 :: [((Bool, Set CType), Set CType)]
solved1 = forall a b. (a -> b) -> [a] -> [b]
map (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first forall a. [(Bool, Set a)] -> (Bool, Set a)
solvedInEqSmallestConstraint) [([(Bool, Set CType)], Set CType)]
solved
        -- Union the constraints from the corresponding rewrites with the
        -- constraints from the recursive results
        solved2 :: [(Bool, Set CType)]
solved2 = forall a b. (a -> b) -> [a] -> [b]
map (\((Bool
b,Set CType
s1),Set CType
s2) -> (Bool
b,forall a. Ord a => Set a -> Set a -> Set a
Set.union Set CType
s1 Set CType
s2)) [((Bool, Set CType), Set CType)]
solved1
        -- From these results, again find the single result that yields 'True'
        -- and has the smallest set of constraints.
        solved3 :: (Bool, Set CType)
solved3 = forall a. [(Bool, Set a)] -> (Bool, Set a)
solvedInEqSmallestConstraint [(Bool, Set CType)]
solved2
    if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [([(Bool, Set CType)], Set CType)]
solved then
      forall a. WriterT (Set CType) Maybe a
noRewrite
    else do
      forall w (m :: * -> *) a. m (a, w) -> WriterT w m a
WriterT (forall a. a -> Maybe a
Just (Bool, Set CType)
solved3)

solveIneq Word
_ Ineq
_ Ineq
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False

-- Find the solved inequality with the fewest number of constraints
solvedInEqSmallestConstraint :: [(Bool,Set a)] -> (Bool, Set a)
solvedInEqSmallestConstraint :: forall a. [(Bool, Set a)] -> (Bool, Set a)
solvedInEqSmallestConstraint = forall {a}. (Bool, Set a) -> [(Bool, Set a)] -> (Bool, Set a)
go (Bool
False, forall a. Set a
Set.empty)
 where
  go :: (Bool, Set a) -> [(Bool, Set a)] -> (Bool, Set a)
go (Bool, Set a)
bs [] = (Bool, Set a)
bs
  go (Bool
b,Set a
s) ((Bool
b1,Set a
s1):[(Bool, Set a)]
solved)
    | Bool -> Bool
not Bool
b Bool -> Bool -> Bool
&& Bool
b1
    = (Bool, Set a) -> [(Bool, Set a)] -> (Bool, Set a)
go (Bool
b1,Set a
s1) [(Bool, Set a)]
solved
    | Bool
b Bool -> Bool -> Bool
&& Bool
b1
    , forall a. Set a -> Int
Set.size Set a
s forall a. Ord a => a -> a -> Bool
>  forall a. Set a -> Int
Set.size Set a
s1
    = (Bool, Set a) -> [(Bool, Set a)] -> (Bool, Set a)
go (Bool
b1,Set a
s1) [(Bool, Set a)]
solved
    | Bool
otherwise
    = (Bool, Set a) -> [(Bool, Set a)] -> (Bool, Set a)
go (Bool
b,Set a
s) [(Bool, Set a)]
solved

-- | Try to instantly solve an inequality by using the inequality solver using
-- @1 <=? 1 ~ True@ as the given constraint.
instantSolveIneq
  :: Word
  -- ^ Solving depth
  -> Ineq
  -- ^ Inequality we want to solve
  -> WriterT (Set CType) Maybe Bool
instantSolveIneq :: Word -> Ineq -> WriterT (Set CType) Maybe Bool
instantSolveIneq Word
k Ineq
u = Word -> Ineq -> Ineq -> WriterT (Set CType) Maybe Bool
solveIneq Word
k Ineq
u (forall {v} {c}. SOP v c
one,forall {v} {c}. SOP v c
one,Bool
True)
 where
  one :: SOP v c
one = forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
1]]

type Ineq = (CoreSOP, CoreSOP, Bool)
type IneqRule = Ineq -> Ineq  -> WriterT (Set CType) Maybe [(Ineq,Ineq)]

noRewrite :: WriterT (Set CType) Maybe a
noRewrite :: forall a. WriterT (Set CType) Maybe a
noRewrite = forall w (m :: * -> *) a. m (a, w) -> WriterT w m a
WriterT forall a. Maybe a
Nothing

ineqRules
  :: [IneqRule]
ineqRules :: [Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]]
ineqRules =
  [ Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
leTrans
  , Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
plusMonotone
  , Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
timesMonotone
  , Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
powMonotone
  , Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
pow2MonotoneSpecial
  , Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
haveSmaller
  , Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
haveBigger
  ]

-- | Transitivity of inequality
leTrans :: IneqRule
leTrans :: Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
leTrans want :: Ineq
want@(SOP TyVar CType
a,SOP TyVar CType
b,Bool
le) (SOP TyVar CType
x,SOP TyVar CType
y,Bool
_)
  -- want: 1 <=? y ~ True
  -- have: 2 <=? y ~ True
  --
  -- new want: want
  -- new have: 1 <=? y ~ True
  | S [P [I Integer
a']] <- SOP TyVar CType
a
  , S [P [I Integer
x']] <- SOP TyVar CType
x
  , Integer
x' forall a. Ord a => a -> a -> Bool
>= Integer
a'
  = forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Ineq
want,(SOP TyVar CType
a,SOP TyVar CType
y,Bool
le))]
  -- want: y <=? 10 ~ True
  -- have: y <=? 9 ~ True
  --
  -- new want: want
  -- new have: y <=? 10 ~ True
  | S [P [I Integer
b']] <- SOP TyVar CType
b
  , S [P [I Integer
y']] <- SOP TyVar CType
y
  , Integer
y' forall a. Ord a => a -> a -> Bool
< Integer
b'
  = forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Ineq
want,(SOP TyVar CType
x,SOP TyVar CType
b,Bool
le))]
leTrans Ineq
_ Ineq
_ = forall a. WriterT (Set CType) Maybe a
noRewrite

-- | Monotonicity of addition
--
-- We use SOP normalization to apply this rule by e.g.:
--
-- * Given: (2*x+1) <= (3*x-1)
-- * Turn to: (3*x-1) - (2*x+1)
-- * SOP version: -2 + x
-- * Convert back to inequality: 2 <= x
plusMonotone :: IneqRule
plusMonotone :: Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
plusMonotone Ineq
want Ineq
have
  | Just Ineq
want' <- SOP TyVar CType -> Maybe Ineq
sopToIneq (Ineq -> SOP TyVar CType
subtractIneq Ineq
want)
  , Ineq
want' forall a. Eq a => a -> a -> Bool
/= Ineq
want
  = forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Ineq
want',Ineq
have)]
  | Just Ineq
have' <- SOP TyVar CType -> Maybe Ineq
sopToIneq (Ineq -> SOP TyVar CType
subtractIneq Ineq
have)
  , Ineq
have' forall a. Eq a => a -> a -> Bool
/= Ineq
have
  = forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Ineq
want,Ineq
have')]
plusMonotone Ineq
_ Ineq
_ = forall a. WriterT (Set CType) Maybe a
noRewrite

-- | Make the `a` of a given `a <= b` smaller
haveSmaller :: IneqRule
haveSmaller :: Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
haveSmaller Ineq
want Ineq
have
  | (S (Product TyVar CType
x:Product TyVar CType
y:[Product TyVar CType]
ys),SOP TyVar CType
us,Bool
True) <- Ineq
have
  = forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Ineq
want,(forall v c. [Product v c] -> SOP v c
S (Product TyVar CType
xforall a. a -> [a] -> [a]
:[Product TyVar CType]
ys),SOP TyVar CType
us,Bool
True))
    ,(Ineq
want,(forall v c. [Product v c] -> SOP v c
S (Product TyVar CType
yforall a. a -> [a] -> [a]
:[Product TyVar CType]
ys),SOP TyVar CType
us,Bool
True))
    ]
  | (S [P [I Integer
1]], S [P (I Integer
_:p :: [Symbol TyVar CType]
p@(Symbol TyVar CType
_:[Symbol TyVar CType]
_))],Bool
True) <- Ineq
have
  = forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Ineq
want,(forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
1]],forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
p],Bool
True))]
haveSmaller Ineq
_ Ineq
_ = forall a. WriterT (Set CType) Maybe a
noRewrite

-- | Make the `b` of a given `a <= b` bigger
haveBigger :: IneqRule
haveBigger :: Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
haveBigger Ineq
want Ineq
have
  | (SOP TyVar CType
_ ,S [Product TyVar CType]
vs,Bool
True) <- Ineq
want
  , (SOP TyVar CType
as,S [Product TyVar CType]
bs,Bool
True) <- Ineq
have
  , let vs' :: [Product TyVar CType]
vs' = [Product TyVar CType]
vs forall a. Eq a => [a] -> [a] -> [a]
\\ [Product TyVar CType]
bs
  , Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Product TyVar CType]
vs')
  -- want : a <= x + 1
  -- have : y <= x
  --
  -- new want: want
  -- new have: y <= x + 1
  = do
    -- Ensure that we're actually making the RHS larger
    Bool
b <- SOP TyVar CType -> WriterT (Set CType) Maybe Bool
isNatural (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
vs')
    if Bool
b then
      forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Ineq
want,(SOP TyVar CType
as,forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPAdd (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
bs) (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType]
vs'),Bool
True))]
    else
      forall a. WriterT (Set CType) Maybe a
noRewrite
haveBigger Ineq
_ Ineq
_ = forall a. WriterT (Set CType) Maybe a
noRewrite

-- | Monotonicity of multiplication
timesMonotone :: IneqRule
timesMonotone :: Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
timesMonotone want :: Ineq
want@(SOP TyVar CType
a,SOP TyVar CType
b,Bool
le) have :: Ineq
have@(SOP TyVar CType
x,SOP TyVar CType
y,Bool
_)
  -- want: C*a <=? b ~ True
  -- have: x <=? y ~ True
  --
  -- new want: want
  -- new have: C*a <=? C*y ~ True
  | S [P a' :: [Symbol TyVar CType]
a'@(Symbol TyVar CType
_:Symbol TyVar CType
_:[Symbol TyVar CType]
_)] <- SOP TyVar CType
a
  , S [P [Symbol TyVar CType]
x'] <- SOP TyVar CType
x
  , S [P [Symbol TyVar CType]
y'] <- SOP TyVar CType
y
  , let ax :: [Symbol TyVar CType]
ax = [Symbol TyVar CType]
a' forall a. Eq a => [a] -> [a] -> [a]
\\ [Symbol TyVar CType]
x'
  , let ay :: [Symbol TyVar CType]
ay = [Symbol TyVar CType]
a' forall a. Eq a => [a] -> [a] -> [a]
\\ [Symbol TyVar CType]
y'
  -- Ensure we don't repeat this rule over and over
  , Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Symbol TyVar CType]
ax)
  , Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Symbol TyVar CType]
ay)
  -- Pick the smallest product
  , let az :: SOP TyVar CType
az = if forall (t :: * -> *) a. Foldable t => t a -> Int
length [Symbol TyVar CType]
ax forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length [Symbol TyVar CType]
ay then forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
ax] else forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
ay]
  = forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Ineq
want,(forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul SOP TyVar CType
az SOP TyVar CType
x, forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul SOP TyVar CType
az SOP TyVar CType
y,Bool
le))]

  -- want: a <=? C*b ~ True
  -- have: x <=? y ~ True
  --
  -- new want: want
  -- new have: C*a <=? C*y ~ True
  | S [P b' :: [Symbol TyVar CType]
b'@(Symbol TyVar CType
_:Symbol TyVar CType
_:[Symbol TyVar CType]
_)] <- SOP TyVar CType
b
  , S [P [Symbol TyVar CType]
x'] <- SOP TyVar CType
x
  , S [P [Symbol TyVar CType]
y'] <- SOP TyVar CType
y
  , let bx :: [Symbol TyVar CType]
bx = [Symbol TyVar CType]
b' forall a. Eq a => [a] -> [a] -> [a]
\\ [Symbol TyVar CType]
x'
  , let by :: [Symbol TyVar CType]
by = [Symbol TyVar CType]
b' forall a. Eq a => [a] -> [a] -> [a]
\\ [Symbol TyVar CType]
y'
  -- Ensure we don't repeat this rule over and over
  , Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Symbol TyVar CType]
bx)
  , Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Symbol TyVar CType]
by)
  -- Pick the smallest product
  , let bz :: SOP TyVar CType
bz = if forall (t :: * -> *) a. Foldable t => t a -> Int
length [Symbol TyVar CType]
bx forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length [Symbol TyVar CType]
by then forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
bx] else forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
by]
  = forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Ineq
want,(forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul SOP TyVar CType
bz SOP TyVar CType
x, forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul SOP TyVar CType
bz SOP TyVar CType
y,Bool
le))]

  -- want: a <=? b ~ True
  -- have: C*x <=? y ~ True
  --
  -- new want: C*a <=? C*b ~ True
  -- new have: have
  | S [P x' :: [Symbol TyVar CType]
x'@(Symbol TyVar CType
_:Symbol TyVar CType
_:[Symbol TyVar CType]
_)] <- SOP TyVar CType
x
  , S [P [Symbol TyVar CType]
a'] <- SOP TyVar CType
a
  , S [P [Symbol TyVar CType]
b'] <- SOP TyVar CType
b
  , let xa :: [Symbol TyVar CType]
xa = [Symbol TyVar CType]
x' forall a. Eq a => [a] -> [a] -> [a]
\\ [Symbol TyVar CType]
a'
  , let xb :: [Symbol TyVar CType]
xb = [Symbol TyVar CType]
x' forall a. Eq a => [a] -> [a] -> [a]
\\ [Symbol TyVar CType]
b'
  -- Ensure we don't repeat this rule over and over
  , Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Symbol TyVar CType]
xa)
  , Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Symbol TyVar CType]
xb)
  -- Pick the smallest product
  , let xz :: SOP TyVar CType
xz = if forall (t :: * -> *) a. Foldable t => t a -> Int
length [Symbol TyVar CType]
xa forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length [Symbol TyVar CType]
xb then forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
xa] else forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
xb]
  = forall (f :: * -> *) a. Applicative f => a -> f a
pure [((forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul SOP TyVar CType
xz SOP TyVar CType
a, forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul SOP TyVar CType
xz SOP TyVar CType
b,Bool
le),Ineq
have)]

  -- want: a <=? b ~ True
  -- have: x <=? C*y ~ True
  --
  -- new want: C*a <=? C*b ~ True
  -- new have: have
  | S [P y' :: [Symbol TyVar CType]
y'@(Symbol TyVar CType
_:Symbol TyVar CType
_:[Symbol TyVar CType]
_)] <- SOP TyVar CType
y
  , S [P [Symbol TyVar CType]
a'] <- SOP TyVar CType
a
  , S [P [Symbol TyVar CType]
b'] <- SOP TyVar CType
b
  , let ya :: [Symbol TyVar CType]
ya = [Symbol TyVar CType]
y' forall a. Eq a => [a] -> [a] -> [a]
\\ [Symbol TyVar CType]
a'
  , let yb :: [Symbol TyVar CType]
yb = [Symbol TyVar CType]
y' forall a. Eq a => [a] -> [a] -> [a]
\\ [Symbol TyVar CType]
b'
  -- Ensure we don't repeat this rule over and over
  , Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Symbol TyVar CType]
ya)
  , Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Symbol TyVar CType]
yb)
  -- Pick the smallest product
  , let yz :: SOP TyVar CType
yz = if forall (t :: * -> *) a. Foldable t => t a -> Int
length [Symbol TyVar CType]
ya forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length [Symbol TyVar CType]
yb then forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
ya] else forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [Symbol TyVar CType]
yb]
  = forall (f :: * -> *) a. Applicative f => a -> f a
pure [((forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul SOP TyVar CType
yz SOP TyVar CType
a, forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul SOP TyVar CType
yz SOP TyVar CType
b,Bool
le),Ineq
have)]

timesMonotone Ineq
_ Ineq
_ = forall a. WriterT (Set CType) Maybe a
noRewrite

-- | Monotonicity of exponentiation
powMonotone :: IneqRule
powMonotone :: Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
powMonotone Ineq
want (SOP TyVar CType
x, S [P [E SOP TyVar CType
yS Product TyVar CType
yP]],Bool
le)
  = case SOP TyVar CType
x of
      S [P [E SOP TyVar CType
xS Product TyVar CType
xP]]
        -- want: XXX
        -- have: 2^x <=? 2^y ~ True
        --
        -- new want: want
        -- new have: x <=? y ~ True
        | SOP TyVar CType
xS forall a. Eq a => a -> a -> Bool
== SOP TyVar CType
yS
        -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Ineq
want,(forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
xP],forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
yP],Bool
le))]
        -- want: XXX
        -- have: x^2 <=? y^2 ~ True
        --
        -- new want: want
        -- new have: x <=? y ~ True
        | Product TyVar CType
xP forall a. Eq a => a -> a -> Bool
== Product TyVar CType
yP
        -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Ineq
want,(SOP TyVar CType
xS,SOP TyVar CType
yS,Bool
le))]
        -- want: XXX
        -- have: 2 <=? 2 ^ x ~ True
        --
        -- new want: want
        -- new have: 1 <=? x ~ True
      SOP TyVar CType
_ | SOP TyVar CType
x forall a. Eq a => a -> a -> Bool
== SOP TyVar CType
yS
        -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Ineq
want,(forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
1]],forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
yP],Bool
le))]
      SOP TyVar CType
_ -> forall a. WriterT (Set CType) Maybe a
noRewrite

powMonotone (SOP TyVar CType
a,S [P [E SOP TyVar CType
bS Product TyVar CType
bP]],Bool
le) Ineq
have
  = case SOP TyVar CType
a of
      S [P [E SOP TyVar CType
aS Product TyVar CType
aP]]
        -- want: 2^x <=? 2^y ~ True
        -- have: XXX
        --
        -- new want: x <=? y ~ True
        -- new have: have
        | SOP TyVar CType
aS forall a. Eq a => a -> a -> Bool
== SOP TyVar CType
bS
        -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [((forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
aP],forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
bP],Bool
le),Ineq
have)]
        -- want: x^2 <=? y^2 ~ True
        -- have: XXX
        --
        -- new want: x <=? y ~ True
        -- new have: have
        | Product TyVar CType
aP forall a. Eq a => a -> a -> Bool
== Product TyVar CType
bP
        -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [((SOP TyVar CType
aS,SOP TyVar CType
bS,Bool
le),Ineq
have)]
        -- want: 2 <=? 2 ^ x ~ True
        -- have: XXX
        --
        -- new want: 1 <=? x ~ True
        -- new have: XXX
      SOP TyVar CType
_ | SOP TyVar CType
a forall a. Eq a => a -> a -> Bool
== SOP TyVar CType
bS
        -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [((forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
1]],forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
bP],Bool
le),Ineq
have)]
      SOP TyVar CType
_ -> forall a. WriterT (Set CType) Maybe a
noRewrite

powMonotone Ineq
_ Ineq
_ = forall a. WriterT (Set CType) Maybe a
noRewrite

-- | Try to get the power-of-2 factors, and apply the monotonicity of
-- exponentiation rule.
--
-- TODO: I wish we could generalize to find arbitrary factors, but currently
-- I don't know how.
pow2MonotoneSpecial :: IneqRule
pow2MonotoneSpecial :: Ineq -> Ineq -> WriterT (Set CType) Maybe [(Ineq, Ineq)]
pow2MonotoneSpecial (SOP TyVar CType
a,SOP TyVar CType
b,Bool
le) Ineq
have
  -- want: 4 * 4^x <=? 8^x ~ True
  -- have: XXX
  --
  -- want as pow 2 factors: 2^(2+2*x) <=? 2^(3*x) ~ True
  --
  -- new want: 2+2*x <=? 3*x ~ True
  -- new have: have
  | Just SOP TyVar CType
a' <- Integer -> SOP TyVar CType -> Maybe (SOP TyVar CType)
facSOP Integer
2 SOP TyVar CType
a
  , Just SOP TyVar CType
b' <- Integer -> SOP TyVar CType -> Maybe (SOP TyVar CType)
facSOP Integer
2 SOP TyVar CType
b
  = forall (f :: * -> *) a. Applicative f => a -> f a
pure [((SOP TyVar CType
a',SOP TyVar CType
b',Bool
le),Ineq
have)]
pow2MonotoneSpecial Ineq
want (SOP TyVar CType
x,SOP TyVar CType
y,Bool
le)
  -- want: XXX
  -- have:4 * 4^x <=? 8^x ~ True
  --
  -- have as pow 2 factors: 2^(2+2*x) <=? 2^(3*x) ~ True
  --
  -- new want: want
  -- new have: 2+2*x <=? 3*x ~ True
  | Just SOP TyVar CType
x' <- Integer -> SOP TyVar CType -> Maybe (SOP TyVar CType)
facSOP Integer
2 SOP TyVar CType
x
  , Just SOP TyVar CType
y' <- Integer -> SOP TyVar CType -> Maybe (SOP TyVar CType)
facSOP Integer
2 SOP TyVar CType
y
  = forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Ineq
want,(SOP TyVar CType
x',SOP TyVar CType
y',Bool
le))]
pow2MonotoneSpecial Ineq
_ Ineq
_ = forall a. WriterT (Set CType) Maybe a
noRewrite

-- | Get the power of /N/ factors of a SOP term
facSOP
  :: Integer
  -- ^ The power /N/
  -> CoreSOP
  -> Maybe CoreSOP
facSOP :: Integer -> SOP TyVar CType -> Maybe (SOP TyVar CType)
facSOP Integer
n (S [P [Symbol TyVar CType]
ps]) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall v c. [Product v c] -> SOP v c
S forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall v c. SOP v c -> [Product v c]
unS) (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Integer -> Symbol TyVar CType -> Maybe (SOP TyVar CType)
facSymbol Integer
n) [Symbol TyVar CType]
ps)
facSOP Integer
_ SOP TyVar CType
_          = forall a. Maybe a
Nothing

-- | Get the power of /N/ factors of a Symbol
facSymbol
  :: Integer
  -- ^ The power
  -> CoreSymbol
  -> Maybe CoreSOP
facSymbol :: Integer -> Symbol TyVar CType -> Maybe (SOP TyVar CType)
facSymbol Integer
n (I Integer
i)
  | Just Integer
j <- Integer -> Integer -> Maybe Integer
integerLogBase Integer
n Integer
i
  = forall a. a -> Maybe a
Just (forall v c. [Product v c] -> SOP v c
S [forall v c. [Symbol v c] -> Product v c
P [forall v c. Integer -> Symbol v c
I Integer
j]])
facSymbol Integer
n (E SOP TyVar CType
s Product TyVar CType
p)
  | Just SOP TyVar CType
s' <- Integer -> SOP TyVar CType -> Maybe (SOP TyVar CType)
facSOP Integer
n SOP TyVar CType
s
  = forall a. a -> Maybe a
Just (forall v c. (Ord v, Ord c) => SOP v c -> SOP v c -> SOP v c
mergeSOPMul SOP TyVar CType
s' (forall v c. [Product v c] -> SOP v c
S [Product TyVar CType
p]))
facSymbol Integer
_ Symbol TyVar CType
_ = forall a. Maybe a
Nothing