{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE TupleSections             #-}
{-# LANGUAGE UndecidableInstances      #-}
{-# LANGUAGE FlexibleContexts          #-}

{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

-- | This module contains functions for recursively "rewriting"
--   GHC core using "rules".

module Language.Haskell.Liquid.Transforms.Rewrite
  ( -- * Top level rewrite function
    rewriteBinds

  -- * Low-level Rewriting Function
  -- , rewriteWith

  -- * Rewrite Rule
  -- ,  RewriteRule

  ) where

import           Liquid.GHC.API as Ghc hiding (get, showPpr, substExpr)
import           Language.Haskell.Liquid.GHC.TypeRep ()
import           Data.Maybe     (fromMaybe, isJust, mapMaybe)
import           Control.Monad.State hiding (lift)
import           Language.Haskell.Liquid.Misc (Nat)
import           Language.Haskell.Liquid.GHC.Play (sub, substExpr)
import           Language.Haskell.Liquid.GHC.Misc (unTickExpr, isTupleId, mkAlive)
import           Language.Haskell.Liquid.Types.Errors (impossible)
import           Language.Haskell.Liquid.UX.Config  (Config, noSimplifyCore)
import qualified Data.List as L
import qualified Data.HashMap.Strict as M

--------------------------------------------------------------------------------
-- | Top-level rewriter --------------------------------------------------------
--------------------------------------------------------------------------------
rewriteBinds :: Config -> [CoreBind] -> [CoreBind]
rewriteBinds :: Config -> [CoreBind] -> [CoreBind]
rewriteBinds Config
cfg
  | Config -> Bool
simplifyCore Config
cfg
  = (CoreBind -> CoreBind) -> [CoreBind] -> [CoreBind]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (CoreBind -> CoreBind
normalizeTuples
       (CoreBind -> CoreBind)
-> (CoreBind -> CoreBind) -> CoreBind -> CoreBind
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RewriteRule -> CoreBind -> CoreBind
rewriteBindWith RewriteRule
undollar
       (CoreBind -> CoreBind)
-> (CoreBind -> CoreBind) -> CoreBind -> CoreBind
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoreBind -> CoreBind
tidyTuples
       (CoreBind -> CoreBind)
-> (CoreBind -> CoreBind) -> CoreBind -> CoreBind
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RewriteRule -> CoreBind -> CoreBind
rewriteBindWith RewriteRule
inlineLoopBreakerTx
       (CoreBind -> CoreBind)
-> (CoreBind -> CoreBind) -> CoreBind -> CoreBind
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoreBind -> CoreBind
inlineLoopBreaker
       (CoreBind -> CoreBind)
-> (CoreBind -> CoreBind) -> CoreBind -> CoreBind
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RewriteRule -> CoreBind -> CoreBind
rewriteBindWith RewriteRule
strictifyLazyLets
       (CoreBind -> CoreBind)
-> (CoreBind -> CoreBind) -> CoreBind -> CoreBind
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoreBind -> CoreBind
inlineFailCases)
  | Bool
otherwise
  = [CoreBind] -> [CoreBind]
forall a. a -> a
id

simplifyCore :: Config -> Bool
simplifyCore :: Config -> Bool
simplifyCore = Bool -> Bool
not (Bool -> Bool) -> (Config -> Bool) -> Config -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Config -> Bool
noSimplifyCore

undollar :: RewriteRule
undollar :: RewriteRule
undollar Expr Var
e
    | Just (Expr Var
f, Expr Var
a) <- Expr Var -> Maybe (Expr Var, Expr Var)
splitDollarApp Expr Var
e =
      RewriteRule
forall a. a -> Maybe a
Just RewriteRule -> RewriteRule
forall a b. (a -> b) -> a -> b
$ Expr Var -> Expr Var -> Expr Var
forall b. Expr b -> Expr b -> Expr b
App Expr Var
f Expr Var
a
    | Bool
otherwise = Maybe (Expr Var)
forall a. Maybe a
Nothing

tidyTuples :: CoreBind -> CoreBind
tidyTuples :: CoreBind -> CoreBind
tidyTuples CoreBind
ce = case CoreBind
ce of
   NonRec Var
x Expr Var
e -> Var -> Expr Var -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Var
x (State [((AltCon, Var), [Var])] (Expr Var)
-> [((AltCon, Var), [Var])] -> Expr Var
forall s a. State s a -> s -> a
evalState (Expr Var -> State [((AltCon, Var), [Var])] (Expr Var)
forall {f :: * -> *}.
MonadState [((AltCon, Var), [Var])] f =>
Expr Var -> f (Expr Var)
go Expr Var
e) [])
   Rec [(Var, Expr Var)]
xs -> [(Var, Expr Var)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec ([(Var, Expr Var)] -> CoreBind) -> [(Var, Expr Var)] -> CoreBind
forall a b. (a -> b) -> a -> b
$ ((Var, Expr Var) -> (Var, Expr Var))
-> [(Var, Expr Var)] -> [(Var, Expr Var)]
forall a b. (a -> b) -> [a] -> [b]
map ((Expr Var -> Expr Var) -> (Var, Expr Var) -> (Var, Expr Var)
forall a b. (a -> b) -> (Var, a) -> (Var, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Expr Var
e -> State [((AltCon, Var), [Var])] (Expr Var)
-> [((AltCon, Var), [Var])] -> Expr Var
forall s a. State s a -> s -> a
evalState (Expr Var -> State [((AltCon, Var), [Var])] (Expr Var)
forall {f :: * -> *}.
MonadState [((AltCon, Var), [Var])] f =>
Expr Var -> f (Expr Var)
go Expr Var
e) [])) [(Var, Expr Var)]
xs
  where
    go :: Expr Var -> f (Expr Var)
go (Tick CoreTickish
t Expr Var
e)
      = CoreTickish -> Expr Var -> Expr Var
forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
t (Expr Var -> Expr Var) -> f (Expr Var) -> f (Expr Var)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr Var -> f (Expr Var)
go Expr Var
e
    go (Let (NonRec Var
x Expr Var
ex) Expr Var
e)
      = do ex' <- Expr Var -> f (Expr Var)
go Expr Var
ex
           e'  <- go e
           return $ Let (NonRec x ex') e'
    go (Let (Rec [(Var, Expr Var)]
bes) Expr Var
e)
      = CoreBind -> Expr Var -> Expr Var
forall b. Bind b -> Expr b -> Expr b
Let (CoreBind -> Expr Var -> Expr Var)
-> f CoreBind -> f (Expr Var -> Expr Var)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([(Var, Expr Var)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec ([(Var, Expr Var)] -> CoreBind)
-> f [(Var, Expr Var)] -> f CoreBind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Var, Expr Var) -> f (Var, Expr Var))
-> [(Var, Expr Var)] -> f [(Var, Expr Var)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Var, Expr Var) -> f (Var, Expr Var)
goRec [(Var, Expr Var)]
bes) f (Expr Var -> Expr Var) -> f (Expr Var) -> f (Expr Var)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr Var -> f (Expr Var)
go Expr Var
e
    go (Case (Var Var
v) Var
x Type
t [Alt Var]
alts)
      = Expr Var -> Var -> Type -> [Alt Var] -> Expr Var
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (Var -> Expr Var
forall b. Var -> Expr b
Var Var
v) Var
x Type
t ([Alt Var] -> Expr Var) -> f [Alt Var] -> f (Expr Var)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Alt Var -> f (Alt Var)) -> [Alt Var] -> f [Alt Var]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Var -> Alt Var -> f (Alt Var)
forall {m :: * -> *} {b}.
(MonadState [((AltCon, b), [Var])] m, Eq b) =>
b -> Alt Var -> m (Alt Var)
goAltR Var
v) [Alt Var]
alts
    go (Case Expr Var
e Var
x Type
t [Alt Var]
alts)
      = Expr Var -> Var -> Type -> [Alt Var] -> Expr Var
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case Expr Var
e Var
x Type
t ([Alt Var] -> Expr Var) -> f [Alt Var] -> f (Expr Var)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Alt Var -> f (Alt Var)) -> [Alt Var] -> f [Alt Var]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Alt Var -> f (Alt Var)
goAlt [Alt Var]
alts
    go (App Expr Var
e1 Expr Var
e2)
      = Expr Var -> Expr Var -> Expr Var
forall b. Expr b -> Expr b -> Expr b
App (Expr Var -> Expr Var -> Expr Var)
-> f (Expr Var) -> f (Expr Var -> Expr Var)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr Var -> f (Expr Var)
go Expr Var
e1 f (Expr Var -> Expr Var) -> f (Expr Var) -> f (Expr Var)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr Var -> f (Expr Var)
go Expr Var
e2
    go (Lam Var
x Expr Var
e)
      = Var -> Expr Var -> Expr Var
forall b. b -> Expr b -> Expr b
Lam Var
x (Expr Var -> Expr Var) -> f (Expr Var) -> f (Expr Var)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr Var -> f (Expr Var)
go Expr Var
e
    go (Cast Expr Var
e CoercionR
c)
      = (Expr Var -> CoercionR -> Expr Var
forall b. Expr b -> CoercionR -> Expr b
`Cast` CoercionR
c) (Expr Var -> Expr Var) -> f (Expr Var) -> f (Expr Var)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr Var -> f (Expr Var)
go Expr Var
e
    go Expr Var
e
      = Expr Var -> f (Expr Var)
forall a. a -> f a
forall (m :: * -> *) a. Monad m => a -> m a
return Expr Var
e

    goRec :: (Var, Expr Var) -> f (Var, Expr Var)
goRec (Var
x, Expr Var
e)
      = (Var
x,) (Expr Var -> (Var, Expr Var)) -> f (Expr Var) -> f (Var, Expr Var)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr Var -> f (Expr Var)
go Expr Var
e

    goAlt :: Alt Var -> f (Alt Var)
goAlt (Alt AltCon
c [Var]
bs Expr Var
e)
      = AltCon -> [Var] -> Expr Var -> Alt Var
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
c [Var]
bs (Expr Var -> Alt Var) -> f (Expr Var) -> f (Alt Var)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr Var -> f (Expr Var)
go Expr Var
e

    goAltR :: b -> Alt Var -> m (Alt Var)
goAltR b
v (Alt AltCon
c [Var]
bs Expr Var
e)
      = do m <- m [((AltCon, b), [Var])]
forall s (m :: * -> *). MonadState s m => m s
get
           case L.lookup (c,v) m of
            Just [Var]
bs' -> Alt Var -> m (Alt Var)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (AltCon -> [Var] -> Expr Var -> Alt Var
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
c [Var]
bs' ([Var] -> [Var] -> Expr Var -> Expr Var
substTuple [Var]
bs' [Var]
bs Expr Var
e))
            Maybe [Var]
Nothing  -> do let bs' :: [Var]
bs' = Var -> Var
mkAlive (Var -> Var) -> [Var] -> [Var]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Var]
bs
                           ([((AltCon, b), [Var])] -> [((AltCon, b), [Var])]) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((AltCon
c,b
v),[Var]
bs')((AltCon, b), [Var])
-> [((AltCon, b), [Var])] -> [((AltCon, b), [Var])]
forall a. a -> [a] -> [a]
:)
                           Alt Var -> m (Alt Var)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (AltCon -> [Var] -> Expr Var -> Alt Var
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
c [Var]
bs' Expr Var
e)



normalizeTuples :: CoreBind -> CoreBind
normalizeTuples :: CoreBind -> CoreBind
normalizeTuples CoreBind
cb
  | NonRec Var
x Expr Var
e <- CoreBind
cb
  = Var -> Expr Var -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Var
x (Expr Var -> CoreBind) -> Expr Var -> CoreBind
forall a b. (a -> b) -> a -> b
$ Expr Var -> Expr Var
go Expr Var
e
  | Rec [(Var, Expr Var)]
xes <- CoreBind
cb
  = let ([Var]
xs,[Expr Var]
es) = [(Var, Expr Var)] -> ([Var], [Expr Var])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Var, Expr Var)]
xes in
    [(Var, Expr Var)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec ([(Var, Expr Var)] -> CoreBind) -> [(Var, Expr Var)] -> CoreBind
forall a b. (a -> b) -> a -> b
$ [Var] -> [Expr Var] -> [(Var, Expr Var)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Var]
xs (Expr Var -> Expr Var
go (Expr Var -> Expr Var) -> [Expr Var] -> [Expr Var]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Expr Var]
es)
  where
    go :: Expr Var -> Expr Var
go (Let (NonRec Var
x Expr Var
ex) Expr Var
e)
      | Case Expr Var
_ Var
_ Type
_ [Alt Var]
alts  <- Expr Var -> Expr Var
unTickExpr Expr Var
ex
      , [Alt AltCon
_ [Var]
vs (Var Var
z)] <- [Alt Var]
alts
      , Var
z Var -> [Var] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Var]
vs
      = CoreBind -> Expr Var -> Expr Var
forall b. Bind b -> Expr b -> Expr b
Let (Var -> Expr Var -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Var
z (Expr Var -> Expr Var
go Expr Var
ex)) ([Var] -> [Var] -> Expr Var -> Expr Var
substTuple [Var
z] [Var
x] (Expr Var -> Expr Var
go Expr Var
e))
    go (Let (NonRec Var
x Expr Var
ex) Expr Var
e)
      = CoreBind -> Expr Var -> Expr Var
forall b. Bind b -> Expr b -> Expr b
Let (Var -> Expr Var -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Var
x (Expr Var -> Expr Var
go Expr Var
ex)) (Expr Var -> Expr Var
go Expr Var
e)
    go (Let (Rec [(Var, Expr Var)]
xes) Expr Var
e)
      = CoreBind -> Expr Var -> Expr Var
forall b. Bind b -> Expr b -> Expr b
Let ([(Var, Expr Var)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec ((Expr Var -> Expr Var) -> (Var, Expr Var) -> (Var, Expr Var)
forall a b. (a -> b) -> (Var, a) -> (Var, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Expr Var -> Expr Var
go ((Var, Expr Var) -> (Var, Expr Var))
-> [(Var, Expr Var)] -> [(Var, Expr Var)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Var, Expr Var)]
xes)) (Expr Var -> Expr Var
go Expr Var
e)
    go (App Expr Var
e1 Expr Var
e2)
      = Expr Var -> Expr Var -> Expr Var
forall b. Expr b -> Expr b -> Expr b
App (Expr Var -> Expr Var
go Expr Var
e1) (Expr Var -> Expr Var
go Expr Var
e2)
    go (Lam Var
x Expr Var
e)
      = Var -> Expr Var -> Expr Var
forall b. b -> Expr b -> Expr b
Lam Var
x (Expr Var -> Expr Var
go Expr Var
e)
    go (Case Expr Var
e Var
b Type
t [Alt Var]
alt)
      = Expr Var -> Var -> Type -> [Alt Var] -> Expr Var
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (Expr Var -> Expr Var
go Expr Var
e) Var
b Type
t ((\(Alt AltCon
c [Var]
bs Expr Var
e') -> AltCon -> [Var] -> Expr Var -> Alt Var
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
c [Var]
bs (Expr Var -> Expr Var
go Expr Var
e')) (Alt Var -> Alt Var) -> [Alt Var] -> [Alt Var]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Alt Var]
alt)
    go (Cast Expr Var
e CoercionR
c)
      = Expr Var -> CoercionR -> Expr Var
forall b. Expr b -> CoercionR -> Expr b
Cast (Expr Var -> Expr Var
go Expr Var
e) CoercionR
c
    go (Tick CoreTickish
t Expr Var
e)
      = CoreTickish -> Expr Var -> Expr Var
forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
t (Expr Var -> Expr Var
go Expr Var
e)
    go (Type Type
t)
      = Type -> Expr Var
forall b. Type -> Expr b
Type Type
t
    go (Coercion CoercionR
c)
      = CoercionR -> Expr Var
forall b. CoercionR -> Expr b
Coercion CoercionR
c
    go (Lit Literal
l)
      = Literal -> Expr Var
forall b. Literal -> Expr b
Lit Literal
l
    go (Var Var
x)
      = Var -> Expr Var
forall b. Var -> Expr b
Var Var
x


--------------------------------------------------------------------------------
-- | A @RewriteRule@ is a function that maps a CoreExpr to another
--------------------------------------------------------------------------------
type RewriteRule = CoreExpr -> Maybe CoreExpr
--------------------------------------------------------------------------------

--------------------------------------------------------------------------------
rewriteBindWith :: RewriteRule -> CoreBind -> CoreBind
--------------------------------------------------------------------------------
rewriteBindWith :: RewriteRule -> CoreBind -> CoreBind
rewriteBindWith RewriteRule
r (NonRec Var
x Expr Var
e) = Var -> Expr Var -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Var
x (RewriteRule -> Expr Var -> Expr Var
rewriteWith RewriteRule
r Expr Var
e)
rewriteBindWith RewriteRule
r (Rec [(Var, Expr Var)]
xes)    = [(Var, Expr Var)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec    ((Expr Var -> Expr Var) -> (Var, Expr Var) -> (Var, Expr Var)
forall a b. (a -> b) -> (Var, a) -> (Var, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (RewriteRule -> Expr Var -> Expr Var
rewriteWith RewriteRule
r) ((Var, Expr Var) -> (Var, Expr Var))
-> [(Var, Expr Var)] -> [(Var, Expr Var)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Var, Expr Var)]
xes)

--------------------------------------------------------------------------------
rewriteWith :: RewriteRule -> CoreExpr -> CoreExpr
--------------------------------------------------------------------------------
rewriteWith :: RewriteRule -> Expr Var -> Expr Var
rewriteWith RewriteRule
tx           = Expr Var -> Expr Var
go
  where
    go :: Expr Var -> Expr Var
go                   = Expr Var -> Expr Var
step (Expr Var -> Expr Var)
-> (Expr Var -> Expr Var) -> Expr Var -> Expr Var
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr Var -> Expr Var
txTop
    txTop :: Expr Var -> Expr Var
txTop Expr Var
e              = Expr Var -> Maybe (Expr Var) -> Expr Var
forall a. a -> Maybe a -> a
fromMaybe Expr Var
e (RewriteRule
tx Expr Var
e)
    goB :: CoreBind -> CoreBind
goB (Rec [(Var, Expr Var)]
xes)        = [(Var, Expr Var)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec         ((Expr Var -> Expr Var) -> (Var, Expr Var) -> (Var, Expr Var)
forall a b. (a -> b) -> (Var, a) -> (Var, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Expr Var -> Expr Var
go ((Var, Expr Var) -> (Var, Expr Var))
-> [(Var, Expr Var)] -> [(Var, Expr Var)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Var, Expr Var)]
xes)
    goB (NonRec Var
x Expr Var
e)     = Var -> Expr Var -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Var
x    (Expr Var -> Expr Var
go Expr Var
e)
    step :: Expr Var -> Expr Var
step (Let CoreBind
b Expr Var
e)       = CoreBind -> Expr Var -> Expr Var
forall b. Bind b -> Expr b -> Expr b
Let (CoreBind -> CoreBind
goB CoreBind
b) (Expr Var -> Expr Var
go Expr Var
e)
    step (App Expr Var
e Expr Var
e')      = Expr Var -> Expr Var -> Expr Var
forall b. Expr b -> Expr b -> Expr b
App (Expr Var -> Expr Var
go Expr Var
e)  (Expr Var -> Expr Var
go Expr Var
e')
    step (Lam Var
x Expr Var
e)       = Var -> Expr Var -> Expr Var
forall b. b -> Expr b -> Expr b
Lam Var
x       (Expr Var -> Expr Var
go Expr Var
e)
    step (Cast Expr Var
e CoercionR
c)      = Expr Var -> CoercionR -> Expr Var
forall b. Expr b -> CoercionR -> Expr b
Cast (Expr Var -> Expr Var
go Expr Var
e) CoercionR
c
    step (Tick CoreTickish
t Expr Var
e)      = CoreTickish -> Expr Var -> Expr Var
forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
t      (Expr Var -> Expr Var
go Expr Var
e)
    step (Case Expr Var
e Var
x Type
t [Alt Var]
cs) = Expr Var -> Var -> Type -> [Alt Var] -> Expr Var
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (Expr Var -> Expr Var
go Expr Var
e) Var
x Type
t ((\(Alt AltCon
c [Var]
bs Expr Var
e') -> AltCon -> [Var] -> Expr Var -> Alt Var
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
c [Var]
bs (Expr Var -> Expr Var
go Expr Var
e')) (Alt Var -> Alt Var) -> [Alt Var] -> [Alt Var]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Alt Var]
cs)
    step e :: Expr Var
e@(Type Type
_)      = Expr Var
e
    step e :: Expr Var
e@(Lit Literal
_)       = Expr Var
e
    step e :: Expr Var
e@(Var Var
_)       = Expr Var
e
    step e :: Expr Var
e@(Coercion CoercionR
_)  = Expr Var
e


--------------------------------------------------------------------------------
-- | Rewriting Pattern-Match-Tuples --------------------------------------------
--------------------------------------------------------------------------------

-- | Transforms
--
-- > let ds = case e0 of
-- >            pat -> (x1,...,xn)
-- >     y1 = proj1 ds
-- >     ...
-- >     yn = projn ds
-- >  in e1
--
--  to
--
-- > case e0 of
-- >   pat -> e1[y1 := x1,..., yn := xn]
--
-- Note that the transformation changes the meaning of the expression if
-- evaluation order matters. But it changes it in a way that LH cannot
-- distinguish.
--
-- Also transforms a variant of the above
--
-- > let y1 = case v of
-- >            C x1 ... xn -> xi
-- >     y2 = proj2 v
-- >     ...
-- >     yn = projn v
-- >  in e1
--
--  to
--
-- > case v of
-- >   C x1 ... xn -> e1[y1 := x1,..., yn := xn]
--
-- The purpose of the transformations is to unpack all of the variables in
-- @pat@ at once in a single scope when verifying @e1@, which allows LH to
-- see the dependencies between the fields of @pat@.
--
strictifyLazyLets :: RewriteRule
strictifyLazyLets :: RewriteRule
strictifyLazyLets (Let (NonRec Var
x e :: Expr Var
e@(Case Expr Var
_ Var
_ Type
_ [Alt (DataAlt DataCon
_) [Var]
_ Expr Var
_])) Expr Var
rest)
  | Just ([Var]
bs, [Var]
bs') <- Expr Var -> Maybe ([Var], [Var])
onlyHasATupleInNestedCases Expr Var
e
  , [Var] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([Var]
bs' [Var] -> [Var] -> [Var]
forall a. Eq a => [a] -> [a] -> [a]
L.\\ [Var]
bs) -- All variables are from the pattern and occur only once
  , let n :: Int
n = [Var] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Var]
bs'
  , Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1
  =
    let ([(Var, Expr Var)]
nrbinds, Expr Var
e') = Int -> Expr Var -> ([(Var, Expr Var)], Expr Var)
takeBinds Int
n Expr Var
rest
        fields :: [(Maybe Int, (Var, Expr Var))]
fields = [ (Var -> Expr Var -> Maybe Int
isProjectionOf Var
x Expr Var
ce, (Var, Expr Var)
b) | b :: (Var, Expr Var)
b@(Var
_, Expr Var
ce) <- [(Var, Expr Var)]
nrbinds ]
        ([(Maybe Int, (Var, Expr Var))]
projs, [(Maybe Int, (Var, Expr Var))]
otherBinds) = ((Maybe Int, (Var, Expr Var)) -> Bool)
-> [(Maybe Int, (Var, Expr Var))]
-> ([(Maybe Int, (Var, Expr Var))], [(Maybe Int, (Var, Expr Var))])
forall a. (a -> Bool) -> [a] -> ([a], [a])
L.partition (Maybe Int -> Bool
forall a. Maybe a -> Bool
isJust (Maybe Int -> Bool)
-> ((Maybe Int, (Var, Expr Var)) -> Maybe Int)
-> (Maybe Int, (Var, Expr Var))
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe Int, (Var, Expr Var)) -> Maybe Int
forall a b. (a, b) -> a
fst) [(Maybe Int, (Var, Expr Var))]
fields
        ss :: [(Var, Var)]
ss = [ ([Var]
bs' [Var] -> Int -> Var
forall a. HasCallStack => [a] -> Int -> a
!! Int
i, Var
v) | (Just Int
i, (Var
v, Expr Var
_)) <- [(Maybe Int, (Var, Expr Var))]
projs ]
        e'' :: Expr Var
e'' = ((Maybe Int, (Var, Expr Var)) -> Expr Var -> Expr Var)
-> Expr Var -> [(Maybe Int, (Var, Expr Var))] -> Expr Var
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\(Maybe Int
_, (Var
v, Expr Var
ce)) -> CoreBind -> Expr Var -> Expr Var
forall b. Bind b -> Expr b -> Expr b
Let (Var -> Expr Var -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Var
v Expr Var
ce)) Expr Var
e' [(Maybe Int, (Var, Expr Var))]
otherBinds
     in RewriteRule
forall a. a -> Maybe a
Just RewriteRule -> RewriteRule
forall a b. (a -> b) -> a -> b
$ CoreBind -> Expr Var -> Expr Var
forall b. Bind b -> Expr b -> Expr b
Let (Var -> Expr Var -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Var
x Expr Var
e) (Expr Var -> Expr Var) -> Expr Var -> Expr Var
forall a b. (a -> b) -> a -> b
$
        Type -> [(Var, Var)] -> Expr Var -> Expr Var -> Expr Var
replaceAltInNestedCases (HasDebugCallStack => Expr Var -> Type
Expr Var -> Type
Ghc.exprType Expr Var
e') [(Var, Var)]
ss Expr Var
e'' Expr Var
e

strictifyLazyLets (Let (NonRec Var
x e :: Expr Var
e@(Case Expr Var
e0 Var
_ Type
_ [Alt (DataAlt DataCon
_) [Var]
bs Expr Var
_])) Expr Var
rest)
  | Just Var
v0 <- Expr Var -> Maybe Var
isVar Expr Var
e0
  , Just Int
i0 <- Var -> Expr Var -> Maybe Int
isProjectionOf Var
v0 Expr Var
e
  , let n :: Int
n = [Var] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Var]
bs
  , Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1
  =
    let ([(Var, Expr Var)]
nrbinds, Expr Var
e') = Int -> Expr Var -> ([(Var, Expr Var)], Expr Var)
takeBinds (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Expr Var
rest
        fields :: [(Maybe Int, (Var, Expr Var))]
fields = [ (Var -> Expr Var -> Maybe Int
isProjectionOf Var
v0 Expr Var
ce, (Var, Expr Var)
b) | b :: (Var, Expr Var)
b@(Var
_, Expr Var
ce) <- [(Var, Expr Var)]
nrbinds ]
        ([(Maybe Int, (Var, Expr Var))]
projs, [(Maybe Int, (Var, Expr Var))]
otherBinds) = ((Maybe Int, (Var, Expr Var)) -> Bool)
-> [(Maybe Int, (Var, Expr Var))]
-> ([(Maybe Int, (Var, Expr Var))], [(Maybe Int, (Var, Expr Var))])
forall a. (a -> Bool) -> [a] -> ([a], [a])
L.partition (Maybe Int -> Bool
forall a. Maybe a -> Bool
isJust (Maybe Int -> Bool)
-> ((Maybe Int, (Var, Expr Var)) -> Maybe Int)
-> (Maybe Int, (Var, Expr Var))
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe Int, (Var, Expr Var)) -> Maybe Int
forall a b. (a, b) -> a
fst) [(Maybe Int, (Var, Expr Var))]
fields
        ss :: [(Var, Var)]
ss = [ ([Var]
bs [Var] -> Int -> Var
forall a. HasCallStack => [a] -> Int -> a
!! Int
i, Var
v) | (Just Int
i, (Var
v, Expr Var
_)) <- (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i0, (Var
x, Expr Var
e)) (Maybe Int, (Var, Expr Var))
-> [(Maybe Int, (Var, Expr Var))] -> [(Maybe Int, (Var, Expr Var))]
forall a. a -> [a] -> [a]
: [(Maybe Int, (Var, Expr Var))]
projs ]
        e'' :: Expr Var
e'' = ((Maybe Int, (Var, Expr Var)) -> Expr Var -> Expr Var)
-> Expr Var -> [(Maybe Int, (Var, Expr Var))] -> Expr Var
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\(Maybe Int
_, (Var
v, Expr Var
ce)) -> CoreBind -> Expr Var -> Expr Var
forall b. Bind b -> Expr b -> Expr b
Let (Var -> Expr Var -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Var
v Expr Var
ce)) Expr Var
e' [(Maybe Int, (Var, Expr Var))]
otherBinds
     in RewriteRule
forall a. a -> Maybe a
Just RewriteRule -> RewriteRule
forall a b. (a -> b) -> a -> b
$ Type -> [(Var, Var)] -> Expr Var -> Expr Var -> Expr Var
replaceAltInNestedCases (HasDebugCallStack => Expr Var -> Type
Expr Var -> Type
Ghc.exprType Expr Var
e') [(Var, Var)]
ss Expr Var
e'' Expr Var
e

strictifyLazyLets Expr Var
_
  = Maybe (Expr Var)
forall a. Maybe a
Nothing

-- | Replaces an expression at the end of a sequence of nested cases with a
-- single alternative.
replaceAltInNestedCases
  :: Type
  -> [(Var, Var)]
  -> CoreExpr -- ^ The expression to place at the end of the nested cases
  -> CoreExpr -- ^ The expression with the nested cases
  -> CoreExpr
replaceAltInNestedCases :: Type -> [(Var, Var)] -> Expr Var -> Expr Var -> Expr Var
replaceAltInNestedCases Type
t [(Var, Var)]
ss Expr Var
ef = Expr Var -> Expr Var
go
  where
    go :: Expr Var -> Expr Var
go (Case Expr Var
e0 Var
v Type
_ [Alt AltCon
c [Var]
bs Expr Var
e1]) =
      let bs' :: [Var]
bs' = [ Var -> Maybe Var -> Var
forall a. a -> Maybe a -> a
fromMaybe Var
b (Var -> [(Var, Var)] -> Maybe Var
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Var
b [(Var, Var)]
ss) | Var
b <- [Var]
bs ]
       in Expr Var -> Var -> Type -> [Alt Var] -> Expr Var
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case Expr Var
e0 Var
v Type
t [AltCon -> [Var] -> Expr Var -> Alt Var
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
c [Var]
bs' (Expr Var -> Expr Var
go Expr Var
e1)]
    go Expr Var
_ = Expr Var
ef


-- | Takes at most n binds from an expression that starts with n non-recursive
-- lets.
takeBinds  :: Nat -> CoreExpr -> ([(Var, CoreExpr)], CoreExpr)
takeBinds :: Int -> Expr Var -> ([(Var, Expr Var)], Expr Var)
takeBinds Int
nat Expr Var
ce = Int -> Expr Var -> ([(Var, Expr Var)], Expr Var)
forall {t} {a}.
(Eq t, Num t) =>
t -> Expr a -> ([(a, Expr a)], Expr a)
go Int
nat Expr Var
ce
    where
      go :: t -> Expr a -> ([(a, Expr a)], Expr a)
go t
0 Expr a
e = ([], Expr a
e)
      go t
n (Let (NonRec a
x Expr a
e) Expr a
e') =
        let ([(a, Expr a)]
xes, Expr a
e'') = t -> Expr a -> ([(a, Expr a)], Expr a)
go (t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
1) Expr a
e'
         in ((a
x,Expr a
e) (a, Expr a) -> [(a, Expr a)] -> [(a, Expr a)]
forall a. a -> [a] -> [a]
: [(a, Expr a)]
xes, Expr a
e'')
      go t
_ Expr a
e = ([], Expr a
e)

-- | Checks that the binding is a projections of some data constructor.
-- | Yields the index of the field being projected
isProjectionOf :: Var -> CoreExpr -> Maybe Int
isProjectionOf :: Var -> Expr Var -> Maybe Int
isProjectionOf Var
v (Case Expr Var
xe Var
_ Type
_ [Alt (DataAlt DataCon
_) [Var]
ys (Var Var
y)])
  | Just Var
xv <- Expr Var -> Maybe Var
isVar Expr Var
xe
  , Var
v Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
xv = Var
y Var -> [Var] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
`L.elemIndex` [Var]
ys
isProjectionOf Var
_ Expr Var
_ = Maybe Int
forall a. Maybe a
Nothing

--------------------------------------------------------------------------------
-- | `substTuple xs ys e'` returns e' [y1 := x1,...,yn := xn]
--------------------------------------------------------------------------------
substTuple :: [Var] -> [Var] -> CoreExpr -> CoreExpr
substTuple :: [Var] -> [Var] -> Expr Var -> Expr Var
substTuple [Var]
xs [Var]
ys = HashMap Var Var -> Expr Var -> Expr Var
substExpr ([(Var, Var)] -> HashMap Var Var
forall k v. Hashable k => [(k, v)] -> HashMap k v
M.fromList ([(Var, Var)] -> HashMap Var Var)
-> [(Var, Var)] -> HashMap Var Var
forall a b. (a -> b) -> a -> b
$ [Var] -> [Var] -> [(Var, Var)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Var]
ys [Var]
xs)

-- | Yields the tuple of variables at the end of nested cases with
-- a single alternative each.
--
-- > case e0 of
-- >   pat0 -> case e1 of
-- >     pat1 -> (x1,...,xn)
--
-- Yields both the bound variables of the patterns, and the
-- variables @x1,...,xn@
onlyHasATupleInNestedCases :: CoreExpr -> Maybe ([Var], [Var])
onlyHasATupleInNestedCases :: Expr Var -> Maybe ([Var], [Var])
onlyHasATupleInNestedCases = [[Var]] -> Expr Var -> Maybe ([Var], [Var])
go []
  where
    go :: [[Var]] -> Expr Var -> Maybe ([Var], [Var])
go [[Var]]
bss (Case Expr Var
_ Var
_ Type
_ [Alt (DataAlt DataCon
_) [Var]
bs Expr Var
e]) = [[Var]] -> Expr Var -> Maybe ([Var], [Var])
go ([Var]
bs[Var] -> [[Var]] -> [[Var]]
forall a. a -> [a] -> [a]
:[[Var]]
bss) Expr Var
e
    go [[Var]]
bss Expr Var
e = ([[Var]] -> [Var]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Var]]
bss,) ([Var] -> ([Var], [Var])) -> Maybe [Var] -> Maybe ([Var], [Var])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr Var -> Maybe [Var]
isTuple Expr Var
e

isTuple :: CoreExpr -> Maybe [Var]
isTuple :: Expr Var -> Maybe [Var]
isTuple Expr Var
e
  | (Var Var
t, [Expr Var]
es) <- Expr Var -> (Expr Var, [Expr Var])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs Expr Var
e
  , Var -> Bool
isTupleId Var
t
  , Just [Var]
xs     <- (Expr Var -> Maybe Var) -> [Expr Var] -> Maybe [Var]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Expr Var -> Maybe Var
isVar ([Expr Var] -> [Expr Var]
forall a. [a] -> [a]
secondHalf [Expr Var]
es)
  = [Var] -> Maybe [Var]
forall a. a -> Maybe a
Just [Var]
xs
  | Bool
otherwise
  = Maybe [Var]
forall a. Maybe a
Nothing

isVar :: CoreExpr -> Maybe Var
isVar :: Expr Var -> Maybe Var
isVar (Var Var
x) = Var -> Maybe Var
forall a. a -> Maybe a
Just Var
x
isVar (Tick CoreTickish
_ Expr Var
e) = Expr Var -> Maybe Var
isVar Expr Var
e
isVar Expr Var
_       = Maybe Var
forall a. Maybe a
Nothing

secondHalf :: [a] -> [a]
secondHalf :: forall a. [a] -> [a]
secondHalf [a]
xs = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) [a]
xs
  where
    n :: Int
n         = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs


inlineLoopBreakerTx :: RewriteRule
inlineLoopBreakerTx :: RewriteRule
inlineLoopBreakerTx (Let CoreBind
b Expr Var
e) = RewriteRule
forall a. a -> Maybe a
Just RewriteRule -> RewriteRule
forall a b. (a -> b) -> a -> b
$ CoreBind -> Expr Var -> Expr Var
forall b. Bind b -> Expr b -> Expr b
Let (CoreBind -> CoreBind
inlineLoopBreaker CoreBind
b) Expr Var
e
inlineLoopBreakerTx Expr Var
_ = Maybe (Expr Var)
forall a. Maybe a
Nothing

-- | Changes top level bindings of the form
--
-- > v = \x1...xn ->
-- >   letrec v0 = \y0...ym -> e0
-- >       in v0 xj..xn
--
-- to
--
-- > v = \x1...xj y0...ym ->
-- >   e0 [ v0 := v x1...xj y0...ym ]
--
inlineLoopBreaker :: Bind Id -> Bind Id
inlineLoopBreaker :: CoreBind -> CoreBind
inlineLoopBreaker (NonRec Var
x Expr Var
e)
    | Just (Var
lbx, Expr Var
lbe, [Expr Var]
lbargs) <- Expr Var -> Maybe (Var, Expr Var, [Expr Var])
hasLoopBreaker Expr Var
be =
       let asPrefix :: [Var]
asPrefix = Int -> [Var] -> [Var]
forall a. Int -> [a] -> [a]
take ([Var] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Var]
as Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Expr Var] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr Var]
lbargs) [Var]
as
           lbe' :: Expr Var
lbe' = HashMap Var (Expr Var) -> Expr Var -> Expr Var
forall a. Subable a => HashMap Var (Expr Var) -> a -> a
sub (Var -> Expr Var -> HashMap Var (Expr Var)
forall k v. Hashable k => k -> v -> HashMap k v
M.singleton Var
lbx ([Var] -> Expr Var
ecall [Var]
asPrefix)) Expr Var
lbe
        in [(Var, Expr Var)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(Var
x, [Var] -> Expr Var -> Expr Var
forall b. [b] -> Expr b -> Expr b
mkLams ([Var]
αs [Var] -> [Var] -> [Var]
forall a. [a] -> [a] -> [a]
++ [Var]
asPrefix) ([CoreBind] -> Expr Var -> Expr Var
forall b. [Bind b] -> Expr b -> Expr b
mkLets [CoreBind]
nrbinds Expr Var
lbe'))]
  where
    ([Var]
αs, [Var]
as, Expr Var
e') = Expr Var -> ([Var], [Var], Expr Var)
collectTyAndValBinders Expr Var
e
    ([CoreBind]
nrbinds, Expr Var
be) = Expr Var -> ([CoreBind], Expr Var)
forall t. Expr t -> ([Bind t], Expr t)
collectNonRecLets Expr Var
e'

    ecall :: [Var] -> Expr Var
ecall [Var]
xs = (Expr Var -> Expr Var -> Expr Var)
-> Expr Var -> [Expr Var] -> Expr Var
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' Expr Var -> Expr Var -> Expr Var
forall b. Expr b -> Expr b -> Expr b
App ((Expr Var -> Expr Var -> Expr Var)
-> Expr Var -> [Expr Var] -> Expr Var
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' Expr Var -> Expr Var -> Expr Var
forall b. Expr b -> Expr b -> Expr b
App (Var -> Expr Var
forall b. Var -> Expr b
Var Var
x) (Type -> Expr Var
forall b. Type -> Expr b
Type (Type -> Expr Var) -> (Var -> Type) -> Var -> Expr Var
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Var -> Type
TyVarTy (Var -> Expr Var) -> [Var] -> [Expr Var]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Var]
αs)) (Var -> Expr Var
forall b. Var -> Expr b
Var (Var -> Expr Var) -> [Var] -> [Expr Var]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Var]
xs)

    hasLoopBreaker :: CoreExpr -> Maybe (Var, CoreExpr, [CoreExpr])
    hasLoopBreaker :: Expr Var -> Maybe (Var, Expr Var, [Expr Var])
hasLoopBreaker (Let (Rec [(Var
x1, Expr Var
e1)]) Expr Var
e2)
      | Bool -> Bool
not (InlinePragma -> Bool
isNoInlinePragma (Var -> InlinePragma
idInlinePragma Var
x1))
      , (Var Var
x2, [Expr Var]
args) <- Expr Var -> (Expr Var, [Expr Var])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs Expr Var
e2
      , Var -> Bool
isLoopBreaker Var
x1
      , Var
x1 Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
x2
      , (Expr Var -> Bool) -> [Expr Var] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Maybe Var -> Bool
forall a. Maybe a -> Bool
isJust (Maybe Var -> Bool) -> (Expr Var -> Maybe Var) -> Expr Var -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr Var -> Maybe Var
isVar) [Expr Var]
args
      , [Var] -> [Var] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
L.isSuffixOf ((Expr Var -> Maybe Var) -> [Expr Var] -> [Var]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Expr Var -> Maybe Var
forall {b}. Expr b -> Maybe Var
getVar [Expr Var]
args) [Var]
as
      = (Var, Expr Var, [Expr Var]) -> Maybe (Var, Expr Var, [Expr Var])
forall a. a -> Maybe a
Just (Var
x1, Expr Var
e1, [Expr Var]
args)
    hasLoopBreaker Expr Var
_ = Maybe (Var, Expr Var, [Expr Var])
forall a. Maybe a
Nothing

    isLoopBreaker :: Var -> Bool
isLoopBreaker =  OccInfo -> Bool
isStrongLoopBreaker (OccInfo -> Bool) -> (Var -> OccInfo) -> Var -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IdInfo -> OccInfo
occInfo (IdInfo -> OccInfo) -> (Var -> IdInfo) -> Var -> OccInfo
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HasDebugCallStack => Var -> IdInfo
Var -> IdInfo
idInfo

    getVar :: Expr b -> Maybe Var
getVar (Var Var
v) = Var -> Maybe Var
forall a. a -> Maybe a
Just Var
v
    getVar Expr b
_ = Maybe Var
forall a. Maybe a
Nothing

inlineLoopBreaker CoreBind
bs
  = CoreBind
bs

collectNonRecLets :: Expr t -> ([Bind t], Expr t)
collectNonRecLets :: forall t. Expr t -> ([Bind t], Expr t)
collectNonRecLets = [Bind t] -> Expr t -> ([Bind t], Expr t)
forall {b}. [Bind b] -> Expr b -> ([Bind b], Expr b)
go []
  where go :: [Bind b] -> Expr b -> ([Bind b], Expr b)
go [Bind b]
bs (Let b :: Bind b
b@(NonRec b
_ Expr b
_) Expr b
e') = [Bind b] -> Expr b -> ([Bind b], Expr b)
go (Bind b
bBind b -> [Bind b] -> [Bind b]
forall a. a -> [a] -> [a]
:[Bind b]
bs) Expr b
e'
        go [Bind b]
bs Expr b
e'                      = ([Bind b] -> [Bind b]
forall a. [a] -> [a]
reverse [Bind b]
bs, Expr b
e')

-- | Inlines bindings of the form
--
-- > let v = \x -> e0
-- >  in e1
--
-- whenever all of the following hold:
--  * "fail" is a prefix of variable @v@,
--  * @x@ is not free in @e0@, and
--  * v is applied to some value in @e1@.
--
-- In addition to inlining, this function also beta reduces
-- the resulting expressions @(\x -> e0) a@ by replacing them
-- with @e0@.
--
inlineFailCases :: CoreBind -> CoreBind
inlineFailCases :: CoreBind -> CoreBind
inlineFailCases = [(Var, Expr Var)] -> CoreBind -> CoreBind
go []
  where
    go :: [(Var, Expr Var)] -> CoreBind -> CoreBind
go [(Var, Expr Var)]
su (Rec [(Var, Expr Var)]
xes)    = [(Var, Expr Var)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec ((Expr Var -> Expr Var) -> (Var, Expr Var) -> (Var, Expr Var)
forall a b. (a -> b) -> (Var, a) -> (Var, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([(Var, Expr Var)] -> Expr Var -> Expr Var
go' [(Var, Expr Var)]
su) ((Var, Expr Var) -> (Var, Expr Var))
-> [(Var, Expr Var)] -> [(Var, Expr Var)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Var, Expr Var)]
xes)
    go [(Var, Expr Var)]
su (NonRec Var
x Expr Var
e) = Var -> Expr Var -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Var
x ([(Var, Expr Var)] -> Expr Var -> Expr Var
go' [(Var, Expr Var)]
su Expr Var
e)

    go' :: [(Var, Expr Var)] -> Expr Var -> Expr Var
go' [(Var, Expr Var)]
su (App (Var Var
x) Expr Var
_)       | Var -> Bool
isFailId Var
x, Just Expr Var
e <- Var -> [(Var, Expr Var)] -> Maybe (Expr Var)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
getFailExpr Var
x [(Var, Expr Var)]
su = Expr Var
e
    go' [(Var, Expr Var)]
su (Let (NonRec Var
x Expr Var
ex) Expr Var
e) | Var -> Bool
isFailId Var
x   = [(Var, Expr Var)] -> Expr Var -> Expr Var
go' (Var -> Expr Var -> [(Var, Expr Var)] -> [(Var, Expr Var)]
forall {a}. a -> Expr Var -> [(a, Expr Var)] -> [(a, Expr Var)]
addFailExpr Var
x ([(Var, Expr Var)] -> Expr Var -> Expr Var
go' [(Var, Expr Var)]
su Expr Var
ex) [(Var, Expr Var)]
su) Expr Var
e

    go' [(Var, Expr Var)]
su (App Expr Var
e1 Expr Var
e2)      = Expr Var -> Expr Var -> Expr Var
forall b. Expr b -> Expr b -> Expr b
App ([(Var, Expr Var)] -> Expr Var -> Expr Var
go' [(Var, Expr Var)]
su Expr Var
e1) ([(Var, Expr Var)] -> Expr Var -> Expr Var
go' [(Var, Expr Var)]
su Expr Var
e2)
    go' [(Var, Expr Var)]
su (Lam Var
x Expr Var
e)        = Var -> Expr Var -> Expr Var
forall b. b -> Expr b -> Expr b
Lam Var
x ([(Var, Expr Var)] -> Expr Var -> Expr Var
go' [(Var, Expr Var)]
su Expr Var
e)
    go' [(Var, Expr Var)]
su (Let CoreBind
xs Expr Var
e)       = CoreBind -> Expr Var -> Expr Var
forall b. Bind b -> Expr b -> Expr b
Let ([(Var, Expr Var)] -> CoreBind -> CoreBind
go [(Var, Expr Var)]
su CoreBind
xs) ([(Var, Expr Var)] -> Expr Var -> Expr Var
go' [(Var, Expr Var)]
su Expr Var
e)
    go' [(Var, Expr Var)]
su (Case Expr Var
e Var
x Type
t [Alt Var]
alt) = Expr Var -> Var -> Type -> [Alt Var] -> Expr Var
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case ([(Var, Expr Var)] -> Expr Var -> Expr Var
go' [(Var, Expr Var)]
su Expr Var
e) Var
x Type
t ([(Var, Expr Var)] -> Alt Var -> Alt Var
goalt [(Var, Expr Var)]
su (Alt Var -> Alt Var) -> [Alt Var] -> [Alt Var]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Alt Var]
alt)
    go' [(Var, Expr Var)]
su (Cast Expr Var
e CoercionR
c)       = Expr Var -> CoercionR -> Expr Var
forall b. Expr b -> CoercionR -> Expr b
Cast ([(Var, Expr Var)] -> Expr Var -> Expr Var
go' [(Var, Expr Var)]
su Expr Var
e) CoercionR
c
    go' [(Var, Expr Var)]
su (Tick CoreTickish
t Expr Var
e)       = CoreTickish -> Expr Var -> Expr Var
forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
t ([(Var, Expr Var)] -> Expr Var -> Expr Var
go' [(Var, Expr Var)]
su Expr Var
e)
    go' [(Var, Expr Var)]
_  Expr Var
e                = Expr Var
e

    goalt :: [(Var, Expr Var)] -> Alt Var -> Alt Var
goalt [(Var, Expr Var)]
su (Alt AltCon
c [Var]
xs Expr Var
e)   = AltCon -> [Var] -> Expr Var -> Alt Var
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
c [Var]
xs ([(Var, Expr Var)] -> Expr Var -> Expr Var
go' [(Var, Expr Var)]
su Expr Var
e)

    isFailId :: Var -> Bool
isFailId Var
x  = Var -> Bool
isLocalId Var
x Bool -> Bool -> Bool
&& Name -> Bool
isSystemName (Var -> Name
varName Var
x) Bool -> Bool -> Bool
&& [Char] -> [Char] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
L.isPrefixOf [Char]
"fail" (Var -> [Char]
forall a. NamedThing a => a -> [Char]
getOccString Var
x)
    getFailExpr :: a -> [(a, b)] -> Maybe b
getFailExpr = a -> [(a, b)] -> Maybe b
forall a b. Eq a => a -> [(a, b)] -> Maybe b
L.lookup

    addFailExpr :: a -> Expr Var -> [(a, Expr Var)] -> [(a, Expr Var)]
addFailExpr a
x (Lam Var
v Expr Var
e) [(a, Expr Var)]
su
      | Bool -> Bool
not (Var -> VarSet -> Bool
elemVarSet Var
v (VarSet -> Bool) -> VarSet -> Bool
forall a b. (a -> b) -> a -> b
$ Expr Var -> VarSet
exprFreeVars Expr Var
e)  = (a
x, Expr Var
e)(a, Expr Var) -> [(a, Expr Var)] -> [(a, Expr Var)]
forall a. a -> [a] -> [a]
:[(a, Expr Var)]
su
    addFailExpr a
_ Expr Var
_         [(a, Expr Var)]
_  = Maybe SrcSpan -> [Char] -> [(a, Expr Var)]
forall a. Maybe SrcSpan -> [Char] -> a
impossible Maybe SrcSpan
forall a. Maybe a
Nothing [Char]
"internal error" -- this cannot happen