module Wingman.CaseSplit
( mkFirstAgda
, iterateSplit
, splitToDecl
) where
import Data.Bool (bool)
import Data.Data
import Data.Generics
import Data.Set (Set)
import qualified Data.Set as S
import Development.IDE.GHC.Compat
import GHC.Exts (IsString (fromString))
import GHC.SourceGen (funBindsWithFixity, match, wildP)
import Wingman.GHC
import Wingman.Types
------------------------------------------------------------------------------
-- | Construct an 'AgdaMatch' from patterns in scope (should be the LHS of the
-- match) and a body.
mkFirstAgda :: [Pat GhcPs] -> HsExpr GhcPs -> AgdaMatch
mkFirstAgda pats (Lambda pats' body) = mkFirstAgda (pats <> pats') body
mkFirstAgda pats body = AgdaMatch pats body
------------------------------------------------------------------------------
-- | Transform an 'AgdaMatch' whose body is a case over a bound pattern, by
-- splitting it into multiple matches: one for each alternative of the case.
agdaSplit :: AgdaMatch -> [AgdaMatch]
agdaSplit (AgdaMatch pats (Case (HsVar _ (L _ var)) matches))
-- Ensure the thing we're destructing is actually a pattern that's been
-- bound.
| containsVar var pats
= do
(pat, body) <- matches
-- TODO(sandy): use an at pattern if necessary
pure $ AgdaMatch (rewriteVarPat var pat pats) $ unLoc body
agdaSplit x = [x]
------------------------------------------------------------------------------
-- | Replace unused bound patterns with wild patterns.
wildify :: AgdaMatch -> AgdaMatch
wildify (AgdaMatch pats body) =
let make_wild = bool id (wildifyT (allOccNames body)) $ not $ containsHole body
in AgdaMatch (make_wild pats) body
------------------------------------------------------------------------------
-- | Helper function for 'wildify'.
wildifyT :: Data a => Set OccName -> a -> a
wildifyT (S.map occNameString -> used) = everywhere $ mkT $ \case
VarPat _ (L _ var) | S.notMember (occNameString $ occName var) used -> wildP
(x :: Pat GhcPs) -> x
------------------------------------------------------------------------------
-- | Determine whether the given 'RdrName' exists as a 'VarPat' inside of @a@.
containsVar :: Data a => RdrName -> a -> Bool
containsVar name = everything (||) $
mkQ False (\case
VarPat _ (L _ var) -> eqRdrName name var
(_ :: Pat GhcPs) -> False
)
`extQ` \case
HsRecField lbl _ True -> eqRdrName name $ unLoc $ rdrNameFieldOcc $ unLoc lbl
(_ :: HsRecField' (FieldOcc GhcPs) (PatCompat GhcPs)) -> False
------------------------------------------------------------------------------
-- | Replace a 'VarPat' with the given @'Pat' GhcPs@.
rewriteVarPat :: Data a => RdrName -> Pat GhcPs -> a -> a
rewriteVarPat name rep = everywhere $
mkT (\case
VarPat _ (L _ var) | eqRdrName name var -> rep
(x :: Pat GhcPs) -> x
)
`extT` \case
HsRecField lbl _ True
| eqRdrName name $ unLoc $ rdrNameFieldOcc $ unLoc lbl
-> HsRecField lbl (toPatCompat rep) False
(x :: HsRecField' (FieldOcc GhcPs) (PatCompat GhcPs)) -> x
------------------------------------------------------------------------------
-- | Construct an 'HsDecl' from a set of 'AgdaMatch'es.
splitToDecl
:: Maybe LexicalFixity
-> OccName -- ^ The name of the function
-> [AgdaMatch]
-> LHsDecl GhcPs
splitToDecl fixity name ams = do
traceX "fixity" fixity $
noLoc $
funBindsWithFixity fixity (fromString . occNameString . occName $ name) $ do
AgdaMatch pats body <- ams
pure $ match pats body
------------------------------------------------------------------------------
-- | Sometimes 'agdaSplit' exposes another opportunity to do 'agdaSplit'. This
-- function runs it a few times, hoping it will find a fixpoint.
iterateSplit :: AgdaMatch -> [AgdaMatch]
iterateSplit am =
let iterated = iterate (agdaSplit =<<) $ pure am
in fmap wildify . (!! 5) $ iterated