-- Copyright (c) Facebook, Inc. and its affiliates.
--
-- This source code is licensed under the MIT license found in the
-- LICENSE file in the root directory of this source tree.
--
{-# LANGUAGE CPP #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Retrie.Context
  ( ContextUpdater
  , updateContext
  , emptyContext
  ) where

import Control.Monad.IO.Class
import Data.Char (isDigit)
import Data.Either (partitionEithers)
import Data.Generics hiding (Fixity)
import Data.List
import Data.Maybe

import Retrie.AlphaEnv
import Retrie.ExactPrint
import Retrie.Fixity
import Retrie.FreeVars
import Retrie.GHC
import Retrie.Substitution
import Retrie.SYB
import Retrie.Types
import Retrie.Universe

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

-- | Type of context update functions for 'apply'.
-- When defining your own 'ContextUpdater', you probably want to extend
-- 'updateContext' using SYB combinators such as 'mkQ' and 'extQ'.
type ContextUpdater = forall m. MonadIO m => GenericCU (TransformT m) Context

-- | Default context update function.
updateContext :: forall m. MonadIO m => GenericCU (TransformT m) Context
updateContext :: forall (m :: * -> *). MonadIO m => GenericCU (TransformT m) Context
updateContext Context
c Int
i =
  TransformT m Context -> a -> TransformT m Context
forall a b. a -> b -> a
const (Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return Context
c)
    (a -> TransformT m Context)
-> (HsExpr (GhcPass 'Parsed) -> TransformT m Context)
-> a
-> TransformT m Context
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` (Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Context -> TransformT m Context)
-> (HsExpr (GhcPass 'Parsed) -> Context)
-> HsExpr (GhcPass 'Parsed)
-> TransformT m Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HsExpr (GhcPass 'Parsed) -> Context
updExp)
    (a -> TransformT m Context)
-> (HsType (GhcPass 'Parsed) -> TransformT m Context)
-> a
-> TransformT m Context
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` (Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Context -> TransformT m Context)
-> (HsType (GhcPass 'Parsed) -> Context)
-> HsType (GhcPass 'Parsed)
-> TransformT m Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HsType (GhcPass 'Parsed) -> Context
updType)
    (a -> TransformT m Context)
-> (Match
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
    -> TransformT m Context)
-> a
-> TransformT m Context
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` (Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Context -> TransformT m Context)
-> (Match
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
    -> Context)
-> Match
     (GhcPass 'Parsed)
     (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> TransformT m Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Match (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed)) -> Context
Match
  (GhcPass 'Parsed)
  (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> Context
updMatch)
    (a -> TransformT m Context)
-> (GRHSs
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
    -> TransformT m Context)
-> a
-> TransformT m Context
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` (Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Context -> TransformT m Context)
-> (GRHSs
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
    -> Context)
-> GRHSs
     (GhcPass 'Parsed)
     (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> TransformT m Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GRHSs (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed)) -> Context
GRHSs
  (GhcPass 'Parsed)
  (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> Context
updGRHSs)
    (a -> TransformT m Context)
-> (GRHS
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
    -> TransformT m Context)
-> a
-> TransformT m Context
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` (Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Context -> TransformT m Context)
-> (GRHS
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
    -> Context)
-> GRHS
     (GhcPass 'Parsed)
     (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> TransformT m Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GRHS (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed)) -> Context
GRHS
  (GhcPass 'Parsed)
  (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> Context
updGRHS)
    (a -> TransformT m Context)
-> (StmtLR
      (GhcPass 'Parsed)
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
    -> TransformT m Context)
-> a
-> TransformT m Context
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` (Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Context -> TransformT m Context)
-> (StmtLR
      (GhcPass 'Parsed)
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
    -> Context)
-> StmtLR
     (GhcPass 'Parsed)
     (GhcPass 'Parsed)
     (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> TransformT m Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stmt (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed)) -> Context
StmtLR
  (GhcPass 'Parsed)
  (GhcPass 'Parsed)
  (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> Context
updStmt)
    (a -> TransformT m Context)
-> (Pat (GhcPass 'Parsed) -> TransformT m Context)
-> a
-> TransformT m Context
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` (Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Context -> TransformT m Context)
-> (Pat (GhcPass 'Parsed) -> Context)
-> Pat (GhcPass 'Parsed)
-> TransformT m Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (GhcPass 'Parsed) -> Context
updPat)
    (a -> TransformT m Context)
-> ([GenLocated
       SrcSpanAnnA
       (StmtLR
          (GhcPass 'Parsed)
          (GhcPass 'Parsed)
          (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed))))]
    -> TransformT m Context)
-> a
-> TransformT m Context
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` [LStmt (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed))]
-> TransformT m Context
[GenLocated
   SrcSpanAnnA
   (StmtLR
      (GhcPass 'Parsed)
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed))))]
-> TransformT m Context
updStmtList
    (a -> TransformT m Context)
-> (HsBind (GhcPass 'Parsed) -> TransformT m Context)
-> a
-> TransformT m Context
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` (Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Context -> TransformT m Context)
-> (HsBind (GhcPass 'Parsed) -> Context)
-> HsBind (GhcPass 'Parsed)
-> TransformT m Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HsBind (GhcPass 'Parsed) -> Context
updHsBind)
    (a -> TransformT m Context)
-> (TyClDecl (GhcPass 'Parsed) -> TransformT m Context)
-> a
-> TransformT m Context
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` (Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Context -> TransformT m Context)
-> (TyClDecl (GhcPass 'Parsed) -> Context)
-> TyClDecl (GhcPass 'Parsed)
-> TransformT m Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyClDecl (GhcPass 'Parsed) -> Context
updTyClDecl)
  where
    neverParen :: Context
neverParen = Context
c { ctxtParentPrec = NeverParen }

    updExp :: HsExpr GhcPs -> Context
    updExp :: HsExpr (GhcPass 'Parsed) -> Context
updExp HsApp{} =
#if __GLASGOW_HASKELL__ < 908
      Context
c { ctxtParentPrec = HasPrec $ Fixity (SourceText "HsApp") (10 + i - firstChild) InfixL }
#else
      c { ctxtParentPrec = HasPrec $ Fixity (SourceText (fsLit "HsApp")) (10 + i - firstChild) InfixL }
#endif
    -- Reason for 10 + i: (i is index of child, 0 = left, 1 = right)
    -- In left child, prec is 10, so HsApp child will NOT get paren'd
    -- In right child, prec is 11, so every child gets paren'd (unless atomic)
    updExp (OpApp XOpApp (GhcPass 'Parsed)
_ LHsExpr (GhcPass 'Parsed)
_ LHsExpr (GhcPass 'Parsed)
op LHsExpr (GhcPass 'Parsed)
_) = Context
c { ctxtParentPrec = HasPrec $ lookupOp op (ctxtFixityEnv c) }
#if __GLASGOW_HASKELL__ < 904
    updExp (HsLet _ lbs _) = addInScope neverParen $ collectLocalBinders CollNoDictBinders lbs
#else
    updExp (HsLet XLet (GhcPass 'Parsed)
_ LHsToken "let" (GhcPass 'Parsed)
_ HsLocalBinds (GhcPass 'Parsed)
lbs LHsToken "in" (GhcPass 'Parsed)
_ LHsExpr (GhcPass 'Parsed)
_) = Context -> [RdrName] -> Context
addInScope Context
neverParen ([RdrName] -> Context) -> [RdrName] -> Context
forall a b. (a -> b) -> a -> b
$ CollectFlag (GhcPass 'Parsed)
-> HsLocalBinds (GhcPass 'Parsed) -> [IdP (GhcPass 'Parsed)]
forall (idL :: Pass) (idR :: Pass).
CollectPass (GhcPass idL) =>
CollectFlag (GhcPass idL)
-> HsLocalBindsLR (GhcPass idL) (GhcPass idR)
-> [IdP (GhcPass idL)]
collectLocalBinders CollectFlag (GhcPass 'Parsed)
forall p. CollectFlag p
CollNoDictBinders HsLocalBinds (GhcPass 'Parsed)
lbs
#endif
    updExp HsExpr (GhcPass 'Parsed)
_ = Context
neverParen

    updType :: HsType GhcPs -> Context
    updType :: HsType (GhcPass 'Parsed) -> Context
updType HsAppTy{}
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
firstChild = Context
c { ctxtParentPrec = IsHsAppsTy }
    updType HsType (GhcPass 'Parsed)
_ = Context
neverParen

    updMatch :: Match GhcPs (LHsExpr GhcPs) -> Context
    updMatch :: Match (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed)) -> Context
updMatch
      | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2  -- m_pats field
      = Context -> [RdrName] -> Context
addInScope Context
c{ctxtParentPrec = IsLhs} ([RdrName] -> Context)
-> (Match
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
    -> [RdrName])
-> Match
     (GhcPass 'Parsed)
     (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CollectFlag (GhcPass 'Parsed)
-> [XRec (GhcPass 'Parsed) (Pat (GhcPass 'Parsed))]
-> [IdP (GhcPass 'Parsed)]
forall p. CollectPass p => CollectFlag p -> [LPat p] -> [IdP p]
collectPatsBinders CollectFlag (GhcPass 'Parsed)
forall p. CollectFlag p
CollNoDictBinders ([XRec (GhcPass 'Parsed) (Pat (GhcPass 'Parsed))] -> [RdrName])
-> (Match
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
    -> [XRec (GhcPass 'Parsed) (Pat (GhcPass 'Parsed))])
-> Match
     (GhcPass 'Parsed)
     (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> [RdrName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Match
  (GhcPass 'Parsed)
  (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> [XRec (GhcPass 'Parsed) (Pat (GhcPass 'Parsed))]
forall p body. Match p body -> [LPat p]
m_pats
      | Bool
otherwise = Context -> [RdrName] -> Context
addInScope Context
neverParen ([RdrName] -> Context)
-> (Match
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
    -> [RdrName])
-> Match
     (GhcPass 'Parsed)
     (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CollectFlag (GhcPass 'Parsed)
-> [XRec (GhcPass 'Parsed) (Pat (GhcPass 'Parsed))]
-> [IdP (GhcPass 'Parsed)]
forall p. CollectPass p => CollectFlag p -> [LPat p] -> [IdP p]
collectPatsBinders CollectFlag (GhcPass 'Parsed)
forall p. CollectFlag p
CollNoDictBinders ([XRec (GhcPass 'Parsed) (Pat (GhcPass 'Parsed))] -> [RdrName])
-> (Match
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
    -> [XRec (GhcPass 'Parsed) (Pat (GhcPass 'Parsed))])
-> Match
     (GhcPass 'Parsed)
     (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> [RdrName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Match
  (GhcPass 'Parsed)
  (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> [XRec (GhcPass 'Parsed) (Pat (GhcPass 'Parsed))]
forall p body. Match p body -> [LPat p]
m_pats
      where

    updGRHSs :: GRHSs GhcPs (LHsExpr GhcPs) -> Context
    updGRHSs :: GRHSs (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed)) -> Context
updGRHSs = Context -> [RdrName] -> Context
addInScope Context
neverParen ([RdrName] -> Context)
-> (GRHSs
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
    -> [RdrName])
-> GRHSs
     (GhcPass 'Parsed)
     (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CollectFlag (GhcPass 'Parsed)
-> HsLocalBinds (GhcPass 'Parsed) -> [IdP (GhcPass 'Parsed)]
forall (idL :: Pass) (idR :: Pass).
CollectPass (GhcPass idL) =>
CollectFlag (GhcPass idL)
-> HsLocalBindsLR (GhcPass idL) (GhcPass idR)
-> [IdP (GhcPass idL)]
collectLocalBinders CollectFlag (GhcPass 'Parsed)
forall p. CollectFlag p
CollNoDictBinders (HsLocalBinds (GhcPass 'Parsed) -> [RdrName])
-> (GRHSs
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
    -> HsLocalBinds (GhcPass 'Parsed))
-> GRHSs
     (GhcPass 'Parsed)
     (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> [RdrName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GRHSs
  (GhcPass 'Parsed)
  (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> HsLocalBinds (GhcPass 'Parsed)
forall p body. GRHSs p body -> HsLocalBinds p
grhssLocalBinds

    updGRHS :: GRHS GhcPs (LHsExpr GhcPs) -> Context
#if __GLASGOW_HASKELL__ < 900
    updGRHS XGRHS{} = neverParen
#endif
    updGRHS :: GRHS (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed)) -> Context
updGRHS (GRHS XCGRHS (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed))
_ [LStmt (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed))]
gs LHsExpr (GhcPass 'Parsed)
_)
        -- binders are in scope over the body (right child) only
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
firstChild = Context -> [RdrName] -> Context
addInScope Context
neverParen [RdrName]
bs
      | Bool
otherwise = (Context, [RdrName]) -> Context
forall a b. (a, b) -> a
fst ((Context, [RdrName]) -> Context)
-> (Context, [RdrName]) -> Context
forall a b. (a -> b) -> a -> b
$ Context -> [RdrName] -> (Context, [RdrName])
updateSubstitution Context
neverParen [RdrName]
bs
      where
        bs :: [IdP (GhcPass 'Parsed)]
bs = CollectFlag (GhcPass 'Parsed)
-> [LStmtLR
      (GhcPass 'Parsed)
      (GhcPass 'Parsed)
      (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))]
-> [IdP (GhcPass 'Parsed)]
forall (idL :: Pass) (idR :: Pass) body.
CollectPass (GhcPass idL) =>
CollectFlag (GhcPass idL)
-> [LStmtLR (GhcPass idL) (GhcPass idR) body]
-> [IdP (GhcPass idL)]
collectLStmtsBinders CollectFlag (GhcPass 'Parsed)
forall p. CollectFlag p
CollNoDictBinders [LStmt (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed))]
[LStmtLR
   (GhcPass 'Parsed)
   (GhcPass 'Parsed)
   (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))]
gs

    updStmt :: Stmt GhcPs (LHsExpr GhcPs) -> Context
    updStmt :: Stmt (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed)) -> Context
updStmt Stmt (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed))
_ = Context
neverParen

    updStmtList :: [LStmt GhcPs (LHsExpr GhcPs)] -> TransformT m Context
    updStmtList :: [LStmt (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed))]
-> TransformT m Context
updStmtList [] = Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return Context
neverParen
    updStmtList (LStmt (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed))
ls:[LStmt (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed))]
_)
        -- binders are in scope over tail of list (right child)
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = Context
-> [RdrName]
-> GenLocated
     SrcSpanAnnA
     (StmtLR
        (GhcPass 'Parsed)
        (GhcPass 'Parsed)
        (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed))))
-> TransformT m Context
forall k (m :: * -> *).
(Matchable k, MonadIO m) =>
Context -> [RdrName] -> k -> TransformT m Context
insertDependentRewrites Context
neverParen [RdrName]
bs LStmt (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed))
GenLocated
  SrcSpanAnnA
  (StmtLR
     (GhcPass 'Parsed)
     (GhcPass 'Parsed)
     (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed))))
ls
        -- lets are recursive in do-blocks
      | L SrcSpanAnnA
_ (LetStmt XLetStmt
  (GhcPass 'Parsed)
  (GhcPass 'Parsed)
  (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
_ HsLocalBinds (GhcPass 'Parsed)
bnds) <- LStmt (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed))
ls =
          Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Context -> TransformT m Context)
-> Context -> TransformT m Context
forall a b. (a -> b) -> a -> b
$ Context -> [RdrName] -> Context
addInScope Context
neverParen ([RdrName] -> Context) -> [RdrName] -> Context
forall a b. (a -> b) -> a -> b
$ CollectFlag (GhcPass 'Parsed)
-> HsLocalBinds (GhcPass 'Parsed) -> [IdP (GhcPass 'Parsed)]
forall (idL :: Pass) (idR :: Pass).
CollectPass (GhcPass idL) =>
CollectFlag (GhcPass idL)
-> HsLocalBindsLR (GhcPass idL) (GhcPass idR)
-> [IdP (GhcPass idL)]
collectLocalBinders CollectFlag (GhcPass 'Parsed)
forall p. CollectFlag p
CollNoDictBinders HsLocalBinds (GhcPass 'Parsed)
bnds
      | Bool
otherwise = Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Context -> TransformT m Context)
-> Context -> TransformT m Context
forall a b. (a -> b) -> a -> b
$ (Context, [RdrName]) -> Context
forall a b. (a, b) -> a
fst ((Context, [RdrName]) -> Context)
-> (Context, [RdrName]) -> Context
forall a b. (a -> b) -> a -> b
$ Context -> [RdrName] -> (Context, [RdrName])
updateSubstitution Context
neverParen [RdrName]
bs
      where
        bs :: [IdP (GhcPass 'Parsed)]
bs = CollectFlag (GhcPass 'Parsed)
-> LStmtLR
     (GhcPass 'Parsed)
     (GhcPass 'Parsed)
     (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
-> [IdP (GhcPass 'Parsed)]
forall (idL :: Pass) (idR :: Pass) body.
CollectPass (GhcPass idL) =>
CollectFlag (GhcPass idL)
-> LStmtLR (GhcPass idL) (GhcPass idR) body -> [IdP (GhcPass idL)]
collectLStmtBinders CollectFlag (GhcPass 'Parsed)
forall p. CollectFlag p
CollNoDictBinders LStmt (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed))
LStmtLR
  (GhcPass 'Parsed)
  (GhcPass 'Parsed)
  (GenLocated SrcSpanAnnA (HsExpr (GhcPass 'Parsed)))
ls

    updHsBind :: HsBind GhcPs -> Context
    updHsBind :: HsBind (GhcPass 'Parsed) -> Context
updHsBind FunBind{XFunBind (GhcPass 'Parsed) (GhcPass 'Parsed)
LIdP (GhcPass 'Parsed)
MatchGroup (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed))
fun_ext :: XFunBind (GhcPass 'Parsed) (GhcPass 'Parsed)
fun_id :: LIdP (GhcPass 'Parsed)
fun_matches :: MatchGroup (GhcPass 'Parsed) (LHsExpr (GhcPass 'Parsed))
fun_ext :: forall idL idR. HsBindLR idL idR -> XFunBind idL idR
fun_id :: forall idL idR. HsBindLR idL idR -> LIdP idL
fun_matches :: forall idL idR. HsBindLR idL idR -> MatchGroup idR (LHsExpr idR)
..} =
      let rdr :: RdrName
rdr = GenLocated SrcSpanAnnN RdrName -> RdrName
forall l e. GenLocated l e -> e
unLoc LIdP (GhcPass 'Parsed)
GenLocated SrcSpanAnnN RdrName
fun_id
      in Context -> [RdrName] -> Context
addBinders (Context -> [RdrName] -> Context
addInScope Context
neverParen [RdrName
rdr]) [RdrName
rdr]
    updHsBind HsBind (GhcPass 'Parsed)
_ = Context
neverParen

    updTyClDecl :: TyClDecl GhcPs -> Context
    updTyClDecl :: TyClDecl (GhcPass 'Parsed) -> Context
updTyClDecl SynDecl{XSynDecl (GhcPass 'Parsed)
LIdP (GhcPass 'Parsed)
LHsType (GhcPass 'Parsed)
LexicalFixity
LHsQTyVars (GhcPass 'Parsed)
tcdSExt :: XSynDecl (GhcPass 'Parsed)
tcdLName :: LIdP (GhcPass 'Parsed)
tcdTyVars :: LHsQTyVars (GhcPass 'Parsed)
tcdFixity :: LexicalFixity
tcdRhs :: LHsType (GhcPass 'Parsed)
tcdSExt :: forall pass. TyClDecl pass -> XSynDecl pass
tcdLName :: forall pass. TyClDecl pass -> LIdP pass
tcdTyVars :: forall pass. TyClDecl pass -> LHsQTyVars pass
tcdFixity :: forall pass. TyClDecl pass -> LexicalFixity
tcdRhs :: forall pass. TyClDecl pass -> LHsType pass
..} = Context -> [RdrName] -> Context
addInScope Context
neverParen [GenLocated SrcSpanAnnN RdrName -> RdrName
forall l e. GenLocated l e -> e
unLoc LIdP (GhcPass 'Parsed)
GenLocated SrcSpanAnnN RdrName
tcdLName]
    updTyClDecl DataDecl{XDataDecl (GhcPass 'Parsed)
LIdP (GhcPass 'Parsed)
LexicalFixity
LHsQTyVars (GhcPass 'Parsed)
HsDataDefn (GhcPass 'Parsed)
tcdLName :: forall pass. TyClDecl pass -> LIdP pass
tcdTyVars :: forall pass. TyClDecl pass -> LHsQTyVars pass
tcdFixity :: forall pass. TyClDecl pass -> LexicalFixity
tcdDExt :: XDataDecl (GhcPass 'Parsed)
tcdLName :: LIdP (GhcPass 'Parsed)
tcdTyVars :: LHsQTyVars (GhcPass 'Parsed)
tcdFixity :: LexicalFixity
tcdDataDefn :: HsDataDefn (GhcPass 'Parsed)
tcdDExt :: forall pass. TyClDecl pass -> XDataDecl pass
tcdDataDefn :: forall pass. TyClDecl pass -> HsDataDefn pass
..} = Context -> [RdrName] -> Context
addInScope Context
neverParen [GenLocated SrcSpanAnnN RdrName -> RdrName
forall l e. GenLocated l e -> e
unLoc LIdP (GhcPass 'Parsed)
GenLocated SrcSpanAnnN RdrName
tcdLName]
    updTyClDecl ClassDecl{[LSig (GhcPass 'Parsed)]
[LDocDecl (GhcPass 'Parsed)]
[LTyFamDefltDecl (GhcPass 'Parsed)]
[LFamilyDecl (GhcPass 'Parsed)]
[LHsFunDep (GhcPass 'Parsed)]
Maybe (LHsContext (GhcPass 'Parsed))
XClassDecl (GhcPass 'Parsed)
LIdP (GhcPass 'Parsed)
LayoutInfo (GhcPass 'Parsed)
LHsBinds (GhcPass 'Parsed)
LexicalFixity
LHsQTyVars (GhcPass 'Parsed)
tcdLName :: forall pass. TyClDecl pass -> LIdP pass
tcdTyVars :: forall pass. TyClDecl pass -> LHsQTyVars pass
tcdFixity :: forall pass. TyClDecl pass -> LexicalFixity
tcdCExt :: XClassDecl (GhcPass 'Parsed)
tcdLayout :: LayoutInfo (GhcPass 'Parsed)
tcdCtxt :: Maybe (LHsContext (GhcPass 'Parsed))
tcdLName :: LIdP (GhcPass 'Parsed)
tcdTyVars :: LHsQTyVars (GhcPass 'Parsed)
tcdFixity :: LexicalFixity
tcdFDs :: [LHsFunDep (GhcPass 'Parsed)]
tcdSigs :: [LSig (GhcPass 'Parsed)]
tcdMeths :: LHsBinds (GhcPass 'Parsed)
tcdATs :: [LFamilyDecl (GhcPass 'Parsed)]
tcdATDefs :: [LTyFamDefltDecl (GhcPass 'Parsed)]
tcdDocs :: [LDocDecl (GhcPass 'Parsed)]
tcdCExt :: forall pass. TyClDecl pass -> XClassDecl pass
tcdLayout :: forall pass. TyClDecl pass -> LayoutInfo pass
tcdCtxt :: forall pass. TyClDecl pass -> Maybe (LHsContext pass)
tcdFDs :: forall pass. TyClDecl pass -> [LHsFunDep pass]
tcdSigs :: forall pass. TyClDecl pass -> [LSig pass]
tcdMeths :: forall pass. TyClDecl pass -> LHsBinds pass
tcdATs :: forall pass. TyClDecl pass -> [LFamilyDecl pass]
tcdATDefs :: forall pass. TyClDecl pass -> [LTyFamDefltDecl pass]
tcdDocs :: forall pass. TyClDecl pass -> [LDocDecl pass]
..} = Context -> [RdrName] -> Context
addInScope Context
neverParen [GenLocated SrcSpanAnnN RdrName -> RdrName
forall l e. GenLocated l e -> e
unLoc LIdP (GhcPass 'Parsed)
GenLocated SrcSpanAnnN RdrName
tcdLName]
    updTyClDecl TyClDecl (GhcPass 'Parsed)
_ = Context
neverParen

    updPat :: Pat GhcPs -> Context
    updPat :: Pat (GhcPass 'Parsed) -> Context
updPat Pat (GhcPass 'Parsed)
_ = Context
neverParen

-- | Create an empty 'Context' with given 'FixityEnv', rewriter, and dependent
-- rewrite generator.
emptyContext :: FixityEnv -> Rewriter -> Rewriter -> Context
emptyContext :: FixityEnv -> Rewriter -> Rewriter -> Context
emptyContext FixityEnv
ctxtFixityEnv Rewriter
ctxtRewriter Rewriter
ctxtDependents = Context{[RdrName]
Maybe Substitution
FixityEnv
AlphaEnv
Rewriter
ParentPrec
forall {a}. [a]
forall {a}. Maybe a
ctxtParentPrec :: ParentPrec
ctxtFixityEnv :: FixityEnv
ctxtFixityEnv :: FixityEnv
ctxtRewriter :: Rewriter
ctxtDependents :: Rewriter
ctxtBinders :: forall {a}. [a]
ctxtInScope :: AlphaEnv
ctxtParentPrec :: ParentPrec
ctxtSubst :: forall {a}. Maybe a
ctxtBinders :: [RdrName]
ctxtDependents :: Rewriter
ctxtInScope :: AlphaEnv
ctxtRewriter :: Rewriter
ctxtSubst :: Maybe Substitution
..}
  where
    ctxtBinders :: [a]
ctxtBinders = []
    ctxtInScope :: AlphaEnv
ctxtInScope = AlphaEnv
emptyAlphaEnv
    ctxtParentPrec :: ParentPrec
ctxtParentPrec = ParentPrec
NeverParen
    ctxtSubst :: Maybe a
ctxtSubst = Maybe a
forall {a}. Maybe a
Nothing

-- Deal with Trees-That-Grow adding extension points
-- as the first child everywhere.
firstChild :: Int
firstChild :: Int
firstChild = Int
1

-- | Add dependent rewrites to 'ctxtRewriter' if necessary.
insertDependentRewrites
  :: (Matchable k, MonadIO m) => Context -> [RdrName] -> k -> TransformT m Context
insertDependentRewrites :: forall k (m :: * -> *).
(Matchable k, MonadIO m) =>
Context -> [RdrName] -> k -> TransformT m Context
insertDependentRewrites Context
c [RdrName]
bs k
x = do
  MatchResult k
r <- (RewriterResult Universe -> RewriterResult Universe)
-> Context -> Rewriter -> k -> TransformT m (MatchResult k)
forall ast (m :: * -> *).
(Matchable ast, MonadIO m) =>
(RewriterResult Universe -> RewriterResult Universe)
-> Context -> Rewriter -> ast -> TransformT m (MatchResult ast)
runRewriter RewriterResult Universe -> RewriterResult Universe
forall a. a -> a
id Context
c (Context -> Rewriter
ctxtDependents Context
c) k
x
  let
    c' :: Context
c' = Context -> [RdrName] -> Context
addInScope Context
c [RdrName]
bs
  case MatchResult k
r of
    MatchResult k
NoMatch -> Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return Context
c'
    MatchResult Substitution
_ Template{Maybe [Rewrite Universe]
Annotated k
AnnotatedImports
tTemplate :: Annotated k
tImports :: AnnotatedImports
tDependents :: Maybe [Rewrite Universe]
tTemplate :: forall ast. Template ast -> Annotated ast
tImports :: forall ast. Template ast -> AnnotatedImports
tDependents :: forall ast. Template ast -> Maybe [Rewrite Universe]
..} -> do
      let
        rrs :: [Rewrite Universe]
rrs = [Rewrite Universe]
-> Maybe [Rewrite Universe] -> [Rewrite Universe]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [Rewrite Universe]
tDependents
        ds :: [Rewrite Universe]
ds = [Rewrite Universe] -> [Rewrite Universe]
forall ast. [Rewrite ast] -> [Rewrite ast]
rewritesWithDependents [Rewrite Universe]
rrs
        f :: [Rewrite Universe] -> Rewriter
f = (Rewrite Universe -> Rewriter) -> [Rewrite Universe] -> Rewriter
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (AlphaEnv -> Rewrite Universe -> Rewriter
forall ast. Matchable ast => AlphaEnv -> Rewrite ast -> Rewriter
mkLocalRewriter (AlphaEnv -> Rewrite Universe -> Rewriter)
-> AlphaEnv -> Rewrite Universe -> Rewriter
forall a b. (a -> b) -> a -> b
$ Context -> AlphaEnv
ctxtInScope Context
c')
      Context -> TransformT m Context
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return Context
c'
        { ctxtRewriter = f rrs <> ctxtRewriter c'
        , ctxtDependents = f ds <> ctxtDependents c'
        }

-- | Add set of binders to 'ctxtInScope'.
addInScope :: Context -> [RdrName] -> Context
addInScope :: Context -> [RdrName] -> Context
addInScope Context
c [RdrName]
bs =
  Context
c' { ctxtInScope = foldr extendAlphaEnv (ctxtInScope c') bs' }
  where
    (Context
c', [RdrName]
bs') = Context -> [RdrName] -> (Context, [RdrName])
updateSubstitution Context
c [RdrName]
bs

-- | Add set of binders to 'ctxtBinders'.
addBinders :: Context -> [RdrName] -> Context
addBinders :: Context -> [RdrName] -> Context
addBinders Context
c [RdrName]
bs = Context
c { ctxtBinders = bs ++ ctxtBinders c }

-- Capture-avoiding substitution
--------------------------------------------------------------------------------

-- | Update the Context's substitution appropriately for a set of binders.
-- Returns a new Context and a potentially alpha-renamed set of binders.
updateSubstitution :: Context -> [RdrName] -> (Context, [RdrName])
updateSubstitution :: Context -> [RdrName] -> (Context, [RdrName])
updateSubstitution Context
c [RdrName]
rdrs =
  case Context -> Maybe Substitution
ctxtSubst Context
c of
    Maybe Substitution
Nothing -> (Context
c, [RdrName]
rdrs)
    Just Substitution
sub ->
      let
        -- This prevents substituting for 'x' under a binding for 'x'.
        sub' :: Substitution
sub' = Substitution -> [FastString] -> Substitution
deleteSubst Substitution
sub ([FastString] -> Substitution) -> [FastString] -> Substitution
forall a b. (a -> b) -> a -> b
$ (RdrName -> FastString) -> [RdrName] -> [FastString]
forall a b. (a -> b) -> [a] -> [b]
map RdrName -> FastString
rdrFS [RdrName]
rdrs
        -- Compute free vars of substitution that could possibly be captured.
        fvs :: FreeVars
fvs = Substitution -> FreeVars
substFVs Substitution
sub'
        -- Partition binders into noncapturing and capturing.
        ([RdrName]
noncapturing, [(RdrName, RdrName)]
capturing) =
          [Either RdrName (RdrName, RdrName)]
-> ([RdrName], [(RdrName, RdrName)])
forall a b. [Either a b] -> ([a], [b])
partitionEithers ([Either RdrName (RdrName, RdrName)]
 -> ([RdrName], [(RdrName, RdrName)]))
-> [Either RdrName (RdrName, RdrName)]
-> ([RdrName], [(RdrName, RdrName)])
forall a b. (a -> b) -> a -> b
$ (RdrName -> Either RdrName (RdrName, RdrName))
-> [RdrName] -> [Either RdrName (RdrName, RdrName)]
forall a b. (a -> b) -> [a] -> [b]
map (FreeVars -> RdrName -> Either RdrName (RdrName, RdrName)
updateBinder FreeVars
fvs) [RdrName]
rdrs
        -- Extend substitution with alpha-renamings.
        alphaSub :: Substitution
alphaSub = (Substitution -> (FastString, HoleVal) -> Substitution)
-> Substitution -> [(FastString, HoleVal)] -> Substitution
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((FastString -> HoleVal -> Substitution)
-> (FastString, HoleVal) -> Substitution
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((FastString -> HoleVal -> Substitution)
 -> (FastString, HoleVal) -> Substitution)
-> (Substitution -> FastString -> HoleVal -> Substitution)
-> Substitution
-> (FastString, HoleVal)
-> Substitution
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Substitution -> FastString -> HoleVal -> Substitution
extendSubst) Substitution
sub'
          [ (RdrName -> FastString
rdrFS RdrName
rdr, RdrName -> HoleVal
HoleRdr RdrName
rdr') | (RdrName
rdr, RdrName
rdr') <- [(RdrName, RdrName)]
capturing ]
        -- There are no telescopes in source Haskell, so order doesn't matter.
        -- Capturing should be rare, so put it first to avoid quadratic append.
        rdrs' :: [RdrName]
rdrs' = ((RdrName, RdrName) -> RdrName)
-> [(RdrName, RdrName)] -> [RdrName]
forall a b. (a -> b) -> [a] -> [b]
map (RdrName, RdrName) -> RdrName
forall a b. (a, b) -> b
snd [(RdrName, RdrName)]
capturing [RdrName] -> [RdrName] -> [RdrName]
forall a. [a] -> [a] -> [a]
++ [RdrName]
noncapturing
      in (Context
c { ctxtSubst = Just alphaSub }, [RdrName]
rdrs')

-- | Check if RdrName is in FreeVars.
--
-- If so, return a pair of it and its new name (Right).
-- If not, return it unchanged (Left).
updateBinder :: FreeVars -> RdrName -> Either RdrName (RdrName, RdrName)
updateBinder :: FreeVars -> RdrName -> Either RdrName (RdrName, RdrName)
updateBinder FreeVars
fvs RdrName
rdr
  | RdrName -> FreeVars -> Bool
elemFVs RdrName
rdr FreeVars
fvs = (RdrName, RdrName) -> Either RdrName (RdrName, RdrName)
forall a b. b -> Either a b
Right (RdrName
rdr, RdrName -> FreeVars -> RdrName
renameBinder RdrName
rdr FreeVars
fvs)
  | Bool
otherwise = RdrName -> Either RdrName (RdrName, RdrName)
forall a b. a -> Either a b
Left RdrName
rdr

-- | Given a RdrName, rename it to something not in given FreeVars.
--
--   x => x1
--   x1 => x2
--   x9 => x10
--
-- etc.
--
-- Only works on unqualified RdrNames. This is fine, as we only use this to
-- rename local binders.
renameBinder :: RdrName -> FreeVars -> RdrName
renameBinder :: RdrName -> FreeVars -> RdrName
renameBinder RdrName
rdr FreeVars
fvs = [RdrName] -> RdrName
forall a. HasCallStack => [a] -> a
head
  [ RdrName
rdr'
  | Int
i <- [Int
n..]
  , let rdr' :: RdrName
rdr' = FastString -> RdrName
mkVarUnqual (FastString -> RdrName) -> FastString -> RdrName
forall a b. (a -> b) -> a -> b
$ String -> FastString
mkFastString (String -> FastString) -> String -> FastString
forall a b. (a -> b) -> a -> b
$ String
baseName String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
i
  , Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ RdrName
rdr' RdrName -> FreeVars -> Bool
`elemFVs` FreeVars
fvs
  ]
  where
    (String
ds, String
rest) = (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
span Char -> Bool
isDigit (String -> (String, String)) -> String -> (String, String)
forall a b. (a -> b) -> a -> b
$ String -> String
forall a. [a] -> [a]
reverse (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ OccName -> String
occNameString (OccName -> String) -> OccName -> String
forall a b. (a -> b) -> a -> b
$ RdrName -> OccName
forall name. HasOccName name => name -> OccName
occName RdrName
rdr

    baseName :: String
baseName = String -> String
forall a. [a] -> [a]
reverse String
rest

    n :: Int
    n :: Int
n | String -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
ds = Int
1
      | Bool
otherwise = String -> Int
forall a. Read a => String -> a
read (String -> String
forall a. [a] -> [a]
reverse String
ds) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1