-- | Look for opportunity to solve goals by instantiating variables.
module Cryptol.TypeCheck.Solver.Improve where

import qualified Data.Set as Set
import Control.Applicative
import Control.Monad

import Cryptol.Utils.Patterns

import Cryptol.TypeCheck.Type
import Cryptol.TypeCheck.SimpType as Mk
import Cryptol.TypeCheck.Solver.Types
import Cryptol.TypeCheck.Solver.Numeric.Interval
import Cryptol.TypeCheck.TypePat
import Cryptol.TypeCheck.Subst



-- | Improvements from a bunch of propositions.
-- Invariant:
-- the substitions should be already applied to the new sub-goals, if any.
improveProps :: Bool -> Ctxt -> [Prop] -> Match (Subst,[Prop])
improveProps :: Bool -> Ctxt -> [Prop] -> Match (Subst, [Prop])
improveProps Bool
impSkol Ctxt
ctxt [Prop]
ps0 = Subst -> [Prop] -> Match (Subst, [Prop])
forall (m :: * -> *).
MonadPlus m =>
Subst -> [Prop] -> m (Subst, [Prop])
loop Subst
emptySubst [Prop]
ps0
  where
  loop :: Subst -> [Prop] -> m (Subst, [Prop])
loop Subst
su [Prop]
props = case Subst -> [Prop] -> [Prop] -> (Subst, [Prop])
go Subst
emptySubst [] [Prop]
props of
                    (Subst
newSu,[Prop]
newProps)
                      | Subst -> Bool
isEmptySubst Subst
newSu ->
                        if Subst -> Bool
isEmptySubst Subst
su then m (Subst, [Prop])
forall (m :: * -> *) a. MonadPlus m => m a
mzero else (Subst, [Prop]) -> m (Subst, [Prop])
forall (m :: * -> *) a. Monad m => a -> m a
return (Subst
su,[Prop]
props)
                      | Bool
otherwise -> Subst -> [Prop] -> m (Subst, [Prop])
loop (Subst
newSu Subst -> Subst -> Subst
@@ Subst
su) [Prop]
newProps

  go :: Subst -> [Prop] -> [Prop] -> (Subst, [Prop])
go Subst
su [Prop]
subs [] = (Subst
su,[Prop]
subs)
  go Subst
su [Prop]
subs (Prop
p : [Prop]
ps) =
    case Match (Subst, [Prop]) -> Maybe (Subst, [Prop])
forall a. Match a -> Maybe a
matchMaybe (Bool -> Ctxt -> Prop -> Match (Subst, [Prop])
improveProp Bool
impSkol Ctxt
ctxt Prop
p) of
      Maybe (Subst, [Prop])
Nothing            -> Subst -> [Prop] -> [Prop] -> (Subst, [Prop])
go Subst
su (Prop
pProp -> [Prop] -> [Prop]
forall a. a -> [a] -> [a]
:[Prop]
subs) [Prop]
ps
      Just (Subst
suNew,[Prop]
psNew) -> Subst -> [Prop] -> [Prop] -> (Subst, [Prop])
go (Subst
suNew Subst -> Subst -> Subst
@@ Subst
su) ([Prop]
psNew [Prop] -> [Prop] -> [Prop]
forall a. [a] -> [a] -> [a]
++ Subst -> [Prop] -> [Prop]
forall t. TVars t => Subst -> t -> t
apSubst Subst
suNew [Prop]
subs)
                                             (Subst -> [Prop] -> [Prop]
forall t. TVars t => Subst -> t -> t
apSubst Subst
su [Prop]
ps)


-- | Improvements from a proposition.
-- Invariant:
-- the substitions should be already applied to the new sub-goals, if any.
improveProp :: Bool -> Ctxt -> Prop -> Match (Subst,[Prop])
improveProp :: Bool -> Ctxt -> Prop -> Match (Subst, [Prop])
improveProp Bool
impSkol Ctxt
ctxt Prop
prop =
  Bool -> Ctxt -> Prop -> Match (Subst, [Prop])
improveEq Bool
impSkol Ctxt
ctxt Prop
prop Match (Subst, [Prop])
-> Match (Subst, [Prop]) -> Match (Subst, [Prop])
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|>
  Bool -> Prop -> Match (Subst, [Prop])
improveLit Bool
impSkol Prop
prop
  -- XXX: others

-- Whenever we have `Literal n [m]a`,
-- we can learn that `a = Bit`
improveLit :: Bool -> Prop -> Match (Subst, [Prop])
improveLit :: Bool -> Prop -> Match (Subst, [Prop])
improveLit Bool
impSkol Prop
prop =
  do (Prop
_,Prop
t) <- Pat Prop (Prop, Prop)
aLiteral Prop
prop
     (Prop
_,Prop
b) <- Pat Prop (Prop, Prop)
aSeq Prop
t
     TVar
a     <- Pat Prop TVar
aTVar Prop
b
     Bool -> Match () -> Match ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
impSkol (Match () -> Match ()) -> Match () -> Match ()
forall a b. (a -> b) -> a -> b
$ Bool -> Match ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (TVar -> Bool
isFreeTV TVar
a)
     let su :: Subst
su = TVar -> Prop -> Subst
uncheckedSingleSubst TVar
a Prop
tBit
     (Subst, [Prop]) -> Match (Subst, [Prop])
forall (m :: * -> *) a. Monad m => a -> m a
return (Subst
su, [])


-- | Improvements from equality constraints.
-- Invariant:
-- the substitions should be already applied to the new sub-goals, if any.
improveEq :: Bool -> Ctxt -> Prop -> Match (Subst,[Prop])
improveEq :: Bool -> Ctxt -> Prop -> Match (Subst, [Prop])
improveEq Bool
impSkol Ctxt
fins Prop
prop =
  do (Prop
lhs,Prop
rhs) <- Pat Prop (Prop, Prop)
(|=|) Prop
prop
     Prop -> Prop -> Match (Subst, [Prop])
rewrite Prop
lhs Prop
rhs Match (Subst, [Prop])
-> Match (Subst, [Prop]) -> Match (Subst, [Prop])
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Prop -> Prop -> Match (Subst, [Prop])
rewrite Prop
rhs Prop
lhs
  where
  rewrite :: Prop -> Prop -> Match (Subst, [Prop])
rewrite Prop
this Prop
other =
    do TVar
x <- Pat Prop TVar
aTVar Prop
this
       Bool -> Match ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (TVar -> Bool
considerVar TVar
x)
       case TVar -> Prop -> Either SubstError Subst
singleSubst TVar
x Prop
other of
         Left SubstError
_ -> Match (Subst, [Prop])
forall (m :: * -> *) a. MonadPlus m => m a
mzero
         Right Subst
su -> (Subst, [Prop]) -> Match (Subst, [Prop])
forall (m :: * -> *) a. Monad m => a -> m a
return (Subst
su, [])
    Match (Subst, [Prop])
-> Match (Subst, [Prop]) -> Match (Subst, [Prop])
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|>
    do (TVar
v,Prop
s) <- Prop -> Match (TVar, Prop)
isSum Prop
this
       case TVar -> Prop -> Either SubstError Subst
singleSubst TVar
v (Prop -> Prop -> Prop
Mk.tSub Prop
other Prop
s) of
         Left SubstError
_ -> Match (Subst, [Prop])
forall (m :: * -> *) a. MonadPlus m => m a
mzero
         Right Subst
su -> (Subst, [Prop]) -> Match (Subst, [Prop])
forall (m :: * -> *) a. Monad m => a -> m a
return (Subst
su, [ Prop
other Prop -> Prop -> Prop
>== Prop
s ])

  isSum :: Prop -> Match (TVar, Prop)
isSum Prop
t = do (TVar
v,Prop
s) <- Prop
-> (Pat Prop (Prop, Prop), Pat Prop TVar, Pat Prop Prop)
-> Match (TVar, Prop)
forall thing pats res.
Matches thing pats res =>
thing -> pats -> Match res
matches Prop
t (Pat Prop (Prop, Prop)
anAdd, Pat Prop TVar
aTVar, Pat Prop Prop
forall a. Pat a a
__)
               TVar -> Prop -> Match (TVar, Prop)
forall (m :: * -> *).
(Monad m, Alternative m) =>
TVar -> Prop -> m (TVar, Prop)
valid TVar
v Prop
s
        Match (TVar, Prop) -> Match (TVar, Prop) -> Match (TVar, Prop)
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> do (Prop
s,TVar
v) <- Prop
-> (Pat Prop (Prop, Prop), Pat Prop Prop, Pat Prop TVar)
-> Match (Prop, TVar)
forall thing pats res.
Matches thing pats res =>
thing -> pats -> Match res
matches Prop
t (Pat Prop (Prop, Prop)
anAdd, Pat Prop Prop
forall a. Pat a a
__, Pat Prop TVar
aTVar)
               TVar -> Prop -> Match (TVar, Prop)
forall (m :: * -> *).
(Monad m, Alternative m) =>
TVar -> Prop -> m (TVar, Prop)
valid TVar
v Prop
s

  valid :: TVar -> Prop -> m (TVar, Prop)
valid TVar
v Prop
s = do let i :: Interval
i = Map TVar Interval -> Prop -> Interval
typeInterval (Ctxt -> Map TVar Interval
intervals Ctxt
fins) Prop
s
                 Bool -> m ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (TVar -> Bool
considerVar TVar
v Bool -> Bool -> Bool
&& TVar
v TVar -> Set TVar -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` Prop -> Set TVar
forall t. FVS t => t -> Set TVar
fvs Prop
s Bool -> Bool -> Bool
&& Interval -> Bool
iIsFin Interval
i)
                 (TVar, Prop) -> m (TVar, Prop)
forall (m :: * -> *) a. Monad m => a -> m a
return (TVar
v,Prop
s)

  considerVar :: TVar -> Bool
considerVar TVar
x = Bool
impSkol Bool -> Bool -> Bool
|| TVar -> Bool
isFreeTV TVar
x


--------------------------------------------------------------------------------

-- XXX

{-
-- | When given an equality constraint, attempt to rewrite it to the form `?x =
-- ...`, by moving all occurrences of `?x` to the LHS, and any other variables
-- to the RHS.  This will only work when there's only one unification variable
-- present in the prop.

tryRewrteEqAsSubst :: Ctxt -> Type -> Type -> Maybe (TVar,Type)
tryRewrteEqAsSubst fins t1 t2 =
  do let vars = Set.toList (Set.filter isFreeTV (fvs (t1,t2)))
     listToMaybe $ sortBy (flip compare `on` rank)
                 $ catMaybes [ tryRewriteEq fins var t1 t2 | var <- vars ]


-- | Rank a rewrite, favoring expressions that have fewer subtractions than
-- additions.
rank :: (TVar,Type) -> Int
rank (_,ty) = go ty
  where

  go (TCon (TF TCAdd) ts) = sum (map go ts) + 1
  go (TCon (TF TCSub) ts) = sum (map go ts) - 1
  go (TCon (TF TCMul) ts) = sum (map go ts) + 1
  go (TCon (TF TCDiv) ts) = sum (map go ts) - 1
  go (TCon _          ts) = sum (map go ts)
  go _                    = 0


-- | Rewrite an equation with respect to a unification variable ?x, into the
-- form `?x = t`.  There are two interesting cases to consider (four with
-- symmetry):
--
--  * ?x = ty
--  * expr containing ?x = expr
--
-- In the first case, we just return the type variable and the type, but in the
-- second we try to rewrite the equation until it's in the form of the first
-- case.
tryRewriteEq :: Map TVar Interval -> TVar -> Type -> Type -> Maybe (TVar,Type)
tryRewriteEq fins uvar l r =
  msum [ do guard (uvarTy == l && uvar `Set.notMember` rfvs)
            return (uvar, r)

       , do guard (uvarTy == r && uvar `Set.notMember` lfvs)
            return (uvar, l)

       , do guard (uvar `Set.notMember` rfvs)
            ty <- rewriteLHS fins uvar l r
            return (uvar,ty)

       , do guard (uvar `Set.notMember` lfvs)
            ty <- rewriteLHS fins uvar r l
            return (uvar,ty)
       ]

  where

  uvarTy = TVar uvar

  lfvs   = fvs l
  rfvs   = fvs r


-- | Check that a type contains only finite type variables.
allFin :: Map TVar Interval -> Type -> Bool
allFin ints ty = iIsFin (typeInterval ints ty)


-- | Rewrite an equality until the LHS is just `uvar`. Return the rewritten RHS.
--
-- There are a few interesting cases when rewriting the equality:
--
--  A o B = R  when `uvar` is only present in A
--  A o B = R  when `uvar` is only present in B
--
-- In the first case, as we only consider addition and subtraction, the
-- rewriting will continue on the left, after moving the `B` side to the RHS of
-- the equation.  In the second case, if the operation is addition, the `A` side
-- will be moved to the RHS, with rewriting continuing in `B`. However, in the
-- case of subtraction, the `B` side is moved to the RHS, and rewriting
-- continues on the RHS instead.
--
-- In both cases, if the operation is addition, rewriting will only continue if
-- the operand being moved to the RHS is known to be finite. If this check was
-- not done, we would end up violating the well-definedness condition for
-- subtraction (for a, b: well defined (a - b) iff fin b).
rewriteLHS :: Map TVar Interval -> TVar -> Type -> Type -> Maybe Type
rewriteLHS fins uvar = go
  where

  go (TVar tv) rhs | tv == uvar = return rhs

  go (TCon (TF tf) [x,y]) rhs =
    do let xfvs = fvs x
           yfvs = fvs y

           inX  = Set.member uvar xfvs
           inY  = Set.member uvar yfvs

       if | inX && inY -> mzero
          | inX        -> balanceR x tf y rhs
          | inY        -> balanceL x tf y rhs
          | otherwise  -> mzero


  -- discard type synonyms, the rewriting will make them no longer apply
  go (TUser _ _ l) rhs =
       go l rhs

  -- records won't work here.
  go _ _ =
       mzero


  -- invert the type function to balance the equation, when the variable occurs
  -- on the LHS of the expression `x tf y`
  balanceR x TCAdd y rhs = do guardFin y
                              go x (tSub rhs y)
  balanceR x TCSub y rhs = go x (tAdd rhs y)
  balanceR _ _     _ _   = mzero


  -- invert the type function to balance the equation, when the variable occurs
  -- on the RHS of the expression `x tf y`
  balanceL x TCAdd y rhs = do guardFin y
                              go y (tSub rhs x)
  balanceL x TCSub y rhs = go (tAdd rhs y) x
  balanceL _ _     _ _   = mzero


  -- guard that the type is finite
  --
  -- XXX this ignores things like `min x inf` where x is finite, and just
  -- assumes that it won't work.
  guardFin ty = guard (allFin fins ty)
-}