{-# LANGUAGE OverloadedStrings #-}
module Wingman.Simplify
( simplify
) where
import Data.Generics (GenericT, everywhere, mkT)
import Data.List.Extra (unsnoc)
import Data.Monoid (Endo (..))
import Development.IDE.GHC.Compat
import GHC.SourceGen (var)
import GHC.SourceGen.Expr (lambda)
import Wingman.CodeGen.Utils
import Wingman.GHC (containsHsVar, fromPatCompat)
------------------------------------------------------------------------------
-- | A pattern over the otherwise (extremely) messy AST for lambdas.
pattern Lambda :: [Pat GhcPs] -> HsExpr GhcPs -> HsExpr GhcPs
pattern Lambda pats body <-
HsLam _
(MG {mg_alts = L _ [L _
(Match { m_pats = fmap fromPatCompat -> pats
, m_grhss = GRHSs {grhssGRHSs = [L _ (
GRHS _ [] (L _ body))]}
})]})
where
-- If there are no patterns to bind, just stick in the body
Lambda [] body = body
Lambda pats body = lambda pats body
------------------------------------------------------------------------------
-- | Simlify an expression.
simplify :: LHsExpr GhcPs -> LHsExpr GhcPs
simplify
= head
. drop 3 -- Do three passes; this should be good enough for the limited
-- amount of gas we give to auto
. iterate (everywhere $ foldEndo
[ simplifyEtaReduce
, simplifyRemoveParens
, simplifyCompose
])
------------------------------------------------------------------------------
-- | Like 'foldMap' but for endomorphisms.
foldEndo :: Foldable t => t (a -> a) -> a -> a
foldEndo = appEndo . foldMap Endo
------------------------------------------------------------------------------
-- | Perform an eta reduction. For example, transforms @\x -> (f g) x@ into
-- @f g@.
simplifyEtaReduce :: GenericT
simplifyEtaReduce = mkT $ \case
Lambda
[VarPat _ (L _ pat)]
(HsVar _ (L _ a)) | pat == a ->
var "id"
Lambda
(unsnoc -> Just (pats, (VarPat _ (L _ pat))))
(HsApp _ (L _ f) (L _ (HsVar _ (L _ a))))
| pat == a
-- We can only perform this simplifiation if @pat@ is otherwise unused.
, not (containsHsVar pat f) ->
Lambda pats f
x -> x
------------------------------------------------------------------------------
-- | Perform an eta-reducing function composition. For example, transforms
-- @\x -> f (g (h x))@ into @f . g . h@.
simplifyCompose :: GenericT
simplifyCompose = mkT $ \case
Lambda
(unsnoc -> Just (pats, (VarPat _ (L _ pat))))
(unroll -> (fs@(_:_), (HsVar _ (L _ a))))
| pat == a
-- We can only perform this simplifiation if @pat@ is otherwise unused.
, not (containsHsVar pat fs) ->
Lambda pats (foldr1 (infixCall ".") fs)
x -> x
------------------------------------------------------------------------------
-- | Removes unnecessary parentheses on any token that doesn't need them.
simplifyRemoveParens :: GenericT
simplifyRemoveParens = mkT $ \case
HsPar _ (L _ x) | isAtomicHsExpr x -> x
(x :: HsExpr GhcPs) -> x
------------------------------------------------------------------------------
-- | Unrolls a right-associative function application of the form
-- @HsApp f (HsApp g (HsApp h x))@ into @([f, g, h], x)@.
unroll :: HsExpr GhcPs -> ([HsExpr GhcPs], HsExpr GhcPs)
unroll (HsPar _ (L _ x)) = unroll x
unroll (HsApp _ (L _ f) (L _ a)) =
let (fs, r) = unroll a
in (f : fs, r)
unroll x = ([], x)