    {-# LANGUAGE Safe #-}

    -- | 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 a. m a
    forall (m :: * -> *) a. MonadPlus m => m a
    mzero else (Subst, [Prop]) -> m (Subst, [Prop])
    forall a. a -> m a
    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 a. Match a -> Match a -> Match a
    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 a. a -> Match a
    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 a. Match a -> Match a -> Match a
    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 a. Match a
    forall (m :: * -> *) a. MonadPlus m => m a
    mzero
             Right Subst
    su -> (Subst, [Prop]) -> Match (Subst, [Prop])
    forall a. a -> Match a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Subst
    su, [])
        Match (Subst, [Prop])
    -> Match (Subst, [Prop]) -> Match (Subst, [Prop])
    forall a. Match a -> Match a -> Match a
    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 a. Match a
    forall (m :: * -> *) a. MonadPlus m => m a
    mzero
             Right Subst
    su -> (Subst, [Prop]) -> Match (Subst, [Prop])
    forall a. a -> Match a
    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. a -> Match 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 a. Match a -> Match a -> Match a
    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. a -> Match 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 a. a -> m a
    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)
    -}