{-|
  Copyright   :  (C) 2012-2016, University of Twente,
                          2017, Google Inc.
  License     :  BSD2 (see the file LICENSE)
  Maintainer  :  Christiaan Baaij <christiaan.baaij@gmail.com>

  Capture-free substitution function for CoreHW
-}

{-# LANGUAGE CPP               #-}
{-# LANGUAGE LambdaCase        #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns      #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}

#include "../../ClashDebug.h"

module Clash.Core.Subst
  ( -- * Substitution into types
    -- ** Substitution environments
    TvSubst (..)
  , TvSubstEnv
  -- , mkTvSubst
  , extendTvSubst
  , extendTvSubstList
    -- ** Applying substitutions
  , substTy
  , substTyWith
  , substTyInVar
    -- * Substitution into terms
    -- ** Substitution environments
  , Subst (..)
  , mkSubst
  , mkTvSubst
  , extendInScopeId
  , extendInScopeIdList
  , extendIdSubst
  , extendIdSubstList
  , extendGblSubstList
    -- ** Applying substitutions
  , substTm
    -- * Variable renaming
  , deShadowTerm
  , freshenTm
    -- * Alpha equivalence
  , aeqType
  , aeqTerm
  )
where

import           Data.Coerce               (coerce)
import           Data.Text.Prettyprint.Doc
import qualified Data.List                 as List

import           Clash.Core.FreeVars
  (noFreeVarsOfType, localFVsOfTerms, tyFVsOfTypes)
import           Clash.Core.Pretty         (ppr, fromPpr)
import           Clash.Core.Term
  (LetBinding, Pat (..), Term (..), TickInfo (..))
import           Clash.Core.Type           (Type (..))
import           Clash.Core.VarEnv
import           Clash.Core.Var            (Id, Var (..), TyVar, isGlobalId)
import           Clash.Unique
import           Clash.Util
import           Clash.Pretty

-- * Subst

-- | A substitution of 'Type's for 'TyVar's
--
-- Note [Extending the TvSubstEnv]
-- See 'TvSubst' for the invariants that must hold
--
-- This invariant allows a short-cut when the subst env is empty: if the
-- TvSubstEnv is empty, i.e. @nullVarEnv TvSubstEnv@ holds, then
-- (substTy subst ty) does nothing.
--
-- For example, consider:
--
--    (/\a -> /\b(a ~ Int) -> ... b ...) Int
--
-- We substitute Int for 'a'. The Unique of 'b' does not change, but
-- nevertheless we add 'b' to the 'TvSubstEnv' because b's kind does change
--
-- This invariant has several consequences:
--
--   * In 'substTyVarBndr', we extend TvSubstEnv if the unique has changed, or
--     if the kind has changed
--
--   * In 'substTyVar', we do not need to consult the 'InScopeSet'; the
--     TvSubstEnv is enough
--
--   * In 'substTy', we can short-circuit when TvSubstEnv is empty
type TvSubstEnv = VarEnv Type

-- | Type substitution
--
-- The following invariants must hold:
--
--   1. The 'InScopeSet' is needed only to guide the generation of fresh uniques
--
--   2. In particular, the kind of the type variables in the 'InScopeSet' is not
--      relevant.
--
--   3. The substitution is only applied once
--
-- Note [Apply Once]
--
-- We might instantiate @forall a b. ty@ with the types @[a, b]@ or @[b, a]@.
-- So the substitution might go like @[a -> b, b -> a]@. A similar situation
-- arises in terms when we find a redex like @(/\a -> /\b -> e) b a@. Then we
-- also end up with a substitution that permutes variables. Other variations
-- happen to; for example @[a -> (a,b)]@.
--
-- SO A TvSubst MUST BE APPLIED PRECISELY ONCE, OR THINGS MIGHT LOOP
--
-- Note [The substitution invariant]
--
-- When calling (substTy subst ty) it should be the case that the 'InScopeSet'
-- is a superset of both:
--
--   * The free variables of the range of the substitution
--
--   * The free variables of /ty/ minus the domain of the substitution
data TvSubst
  = TvSubst InScopeSet -- Variable in scope /after/ substitution
            TvSubstEnv -- Substitution for types

instance ClashPretty TvSubst where
  clashPretty :: TvSubst -> Doc ()
clashPretty (TvSubst ins :: InScopeSet
ins tenv :: TvSubstEnv
tenv) =
    Doc () -> Doc ()
forall ann. Doc ann -> Doc ann
brackets (Doc () -> Doc ()) -> Doc () -> Doc ()
forall a b. (a -> b) -> a -> b
$ [Doc ()] -> Doc ()
forall ann. [Doc ann] -> Doc ann
sep [ "TvSubst"
                   , Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
nest 2 ("In scope:" Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> InScopeSet -> Doc ()
forall a. ClashPretty a => a -> Doc ()
clashPretty InScopeSet
ins)
                   , Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
nest 2 ("Type env:" Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> TvSubstEnv -> Doc ()
forall a. ClashPretty a => a -> Doc ()
clashPretty TvSubstEnv
tenv)]

-- | A substitution  of 'Term's for 'Id's
--
-- Note [Extending the Subst]
--
-- For a term 'Subst', which binds 'Id's as well, we make a different choice for
-- Ids than we do for TyVars.
--
-- For TyVars see 'TvSubstEnv's Note [Extending the TvSubstEnv]
--
-- For Ids, we have a different invariant:
--
--   The IdSubstEnv is extended only when the Unique on an Id changes.
--   Otherwise, we just extend the InScopeSet
--
-- In consequence:
--
--   * If all subst envs are empty, substsTm would be a no-op
--
--     However, substTm still goes ahead and substitutes. Reason: we may want
--     to replace existing Ids with new ones from the in-scope set, to avoid
--     space leaks.
--
--   * In substIdBndr, we extend the 'IdSubstEnv' only when the unique changes
--
--   * If TvSubstEnv and IdSubstEnv are all empty, substExpr does nothing
--     (Note that the above rule for 'substIdBndr' maintains this property.)
--
--   * In 'lookupIdSubst', we must look up the Id in the in-scope set, because
--     it may contain non-trivial changes. Exmaple:
--
--     (/\a -> \x:a. ... x ...) Int
--
--     We extend the 'TvSubstEnv' with a @[a |-> Int]@; but x's unique does not
--     change so we only extend the in-scope set. Then we must look up in the
--     in-scope set when we find the occurrence of x.
--
--   * The requirement to look  up the Id in the in-scope set means that we
--     must not take no-op short cut when the 'IdSubstEnv' is empty. We must
--     still look up ever Id in the in-scope set.
--
--   * (However, we don't need to do so for the expression found in the
--     IdSubstEnv, whose range is assumed to be correct wrt the in-scope set)
type IdSubstEnv = VarEnv Term

-- | A substitution environment containing containing both 'Id' and 'TyVar'
-- substitutions.
--
-- Some invariants apply to how you use the substitution:
--
--   1. The 'InScopeSet' contains at least those 'Id's and 'TyVar's that will
--      be in scope /after/ applying the substitution  to a term. Precisely,
--      the in-scope set must be a superset of the free variables of the
--      substitution range that might possibly clash with locally-bound
--      variables in the thing being substituted in.
--
--   2. You may only apply the substitution once. See 'TvSubst'
--
-- There are various ways of setting up the in-scope set such that the first of
-- of these invariants holds:
--
--   * Arrange that the in-scope set really is all the things in scope
--
--   * Arrange that it's the  free vars of the range of the substitution
--
--   * Make it empty, if you know that all the free variables of the
--     substitution are fresh, and hence can´t possibly clash
data Subst
  = Subst
  { Subst -> InScopeSet
substInScope :: InScopeSet -- Variables in scope /after/ substitution
  , Subst -> IdSubstEnv
substTmEnv   :: IdSubstEnv -- Substitution for terms
  , Subst -> TvSubstEnv
substTyEnv   :: TvSubstEnv -- Substitution for types
  , Subst -> IdSubstEnv
substGblEnv  :: IdSubstEnv -- Substitution of globals (in terms)
  }

emptySubst
  :: Subst
emptySubst :: Subst
emptySubst = InScopeSet -> IdSubstEnv -> TvSubstEnv -> IdSubstEnv -> Subst
Subst InScopeSet
emptyInScopeSet IdSubstEnv
forall a. VarEnv a
emptyVarEnv TvSubstEnv
forall a. VarEnv a
emptyVarEnv IdSubstEnv
forall a. VarEnv a
emptyVarEnv

-- | An empty substitution, starting the variables currently in scope
mkSubst
  :: InScopeSet
  -> Subst
mkSubst :: InScopeSet -> Subst
mkSubst is :: InScopeSet
is = InScopeSet -> IdSubstEnv -> TvSubstEnv -> IdSubstEnv -> Subst
Subst InScopeSet
is IdSubstEnv
forall a. VarEnv a
emptyVarEnv TvSubstEnv
forall a. VarEnv a
emptyVarEnv IdSubstEnv
forall a. VarEnv a
emptyVarEnv

-- | Create a type substitution
mkTvSubst
  :: InScopeSet
  -> VarEnv Type
  -> Subst
mkTvSubst :: InScopeSet -> TvSubstEnv -> Subst
mkTvSubst is :: InScopeSet
is env :: TvSubstEnv
env = InScopeSet -> IdSubstEnv -> TvSubstEnv -> IdSubstEnv -> Subst
Subst InScopeSet
is IdSubstEnv
forall a. VarEnv a
emptyVarEnv TvSubstEnv
env IdSubstEnv
forall a. VarEnv a
emptyVarEnv

-- | Generates the in-scope set for the 'Subst' from the types in the incoming
-- environment.
--
-- Should only be used the type we're substituting into has no free variables
-- outside of the domain of substitution
zipTvSubst
  :: [TyVar]
  -> [Type]
  -> Subst
zipTvSubst :: [TyVar] -> [Type] -> Subst
zipTvSubst tvs :: [TyVar]
tvs tys :: [Type]
tys
  | Bool
debugIsOn
  , [TyVar] -> [Type] -> Bool
forall a b. [a] -> [b] -> Bool
neLength [TyVar]
tvs [Type]
tys
  = String -> Doc ClashAnnotation -> Subst -> Subst
forall ann a. String -> Doc ann -> a -> a
pprTrace "zipTvSubst" ([TyVar] -> Doc ClashAnnotation
forall p. PrettyPrec p => p -> Doc ClashAnnotation
ppr [TyVar]
tvs Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
forall a. Semigroup a => a -> a -> a
<> Doc ClashAnnotation
forall ann. Doc ann
line Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
forall a. Semigroup a => a -> a -> a
<> [Type] -> Doc ClashAnnotation
forall p. PrettyPrec p => p -> Doc ClashAnnotation
ppr [Type]
tys) Subst
emptySubst
  | Bool
otherwise
  = InScopeSet -> IdSubstEnv -> TvSubstEnv -> IdSubstEnv -> Subst
Subst (VarSet -> InScopeSet
mkInScopeSet ([Type] -> VarSet
forall (f :: * -> *). Foldable f => f Type -> VarSet
tyFVsOfTypes [Type]
tys)) IdSubstEnv
forall a. VarEnv a
emptyVarEnv TvSubstEnv
tenv IdSubstEnv
forall a. VarEnv a
emptyVarEnv
 where
  tenv :: TvSubstEnv
tenv = [TyVar] -> [Type] -> TvSubstEnv
zipTyEnv [TyVar]
tvs [Type]
tys

zipTyEnv
  :: [TyVar]
  -> [Type]
  -> VarEnv Type
zipTyEnv :: [TyVar] -> [Type] -> TvSubstEnv
zipTyEnv tvs :: [TyVar]
tvs tys :: [Type]
tys = [(TyVar, Type)] -> TvSubstEnv
forall a b. [(Var a, b)] -> VarEnv b
mkVarEnv ([TyVar] -> [Type] -> [(TyVar, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zipEqual [TyVar]
tvs [Type]
tys)

-- | Extend the substitution environment with a new 'Id' substitution
extendIdSubst
  :: Subst
  -> Id
  -> Term
  -> Subst
extendIdSubst :: Subst -> Id -> Term -> Subst
extendIdSubst (Subst is :: InScopeSet
is env :: IdSubstEnv
env tenv :: TvSubstEnv
tenv genv :: IdSubstEnv
genv) i :: Id
i e :: Term
e =
  InScopeSet -> IdSubstEnv -> TvSubstEnv -> IdSubstEnv -> Subst
Subst InScopeSet
is (Id -> Term -> IdSubstEnv -> IdSubstEnv
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv Id
i Term
e IdSubstEnv
env) TvSubstEnv
tenv IdSubstEnv
genv

-- | Extend the substitution environment with a list of 'Id' substitutions
extendIdSubstList
  :: Subst
  -> [(Id,Term)]
  -> Subst
extendIdSubstList :: Subst -> [(Id, Term)] -> Subst
extendIdSubstList (Subst is :: InScopeSet
is env :: IdSubstEnv
env tenv :: TvSubstEnv
tenv genv :: IdSubstEnv
genv) es :: [(Id, Term)]
es =
  InScopeSet -> IdSubstEnv -> TvSubstEnv -> IdSubstEnv -> Subst
Subst InScopeSet
is (IdSubstEnv -> [(Id, Term)] -> IdSubstEnv
forall a b. VarEnv a -> [(Var b, a)] -> VarEnv a
extendVarEnvList IdSubstEnv
env [(Id, Term)]
es) TvSubstEnv
tenv IdSubstEnv
genv

-- | Extend the substitution environment with a list of global 'Id' substitutions
extendGblSubstList
  :: Subst
  -> [(Id,Term)]
  -> Subst
extendGblSubstList :: Subst -> [(Id, Term)] -> Subst
extendGblSubstList (Subst is :: InScopeSet
is env :: IdSubstEnv
env tenv :: TvSubstEnv
tenv genv :: IdSubstEnv
genv) es :: [(Id, Term)]
es =
  InScopeSet -> IdSubstEnv -> TvSubstEnv -> IdSubstEnv -> Subst
Subst InScopeSet
is IdSubstEnv
env TvSubstEnv
tenv (IdSubstEnv -> [(Id, Term)] -> IdSubstEnv
forall a b. VarEnv a -> [(Var b, a)] -> VarEnv a
extendVarEnvList IdSubstEnv
genv [(Id, Term)]
es)

-- | Extend the substitution environment with a new 'TyVar' substitution
extendTvSubst
  :: Subst
  -> TyVar
  -> Type
  -> Subst
extendTvSubst :: Subst -> TyVar -> Type -> Subst
extendTvSubst (Subst is :: InScopeSet
is env :: IdSubstEnv
env tenv :: TvSubstEnv
tenv genv :: IdSubstEnv
genv) tv :: TyVar
tv t :: Type
t =
  InScopeSet -> IdSubstEnv -> TvSubstEnv -> IdSubstEnv -> Subst
Subst InScopeSet
is IdSubstEnv
env (TyVar -> Type -> TvSubstEnv -> TvSubstEnv
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv TyVar
tv Type
t TvSubstEnv
tenv) IdSubstEnv
genv

-- | Extend the substitution environment with a list of 'TyVar' substitutions
extendTvSubstList
  :: Subst
  -> [(TyVar, Type)]
  -> Subst
extendTvSubstList :: Subst -> [(TyVar, Type)] -> Subst
extendTvSubstList (Subst is :: InScopeSet
is env :: IdSubstEnv
env tenv :: TvSubstEnv
tenv genv :: IdSubstEnv
genv) ts :: [(TyVar, Type)]
ts =
  InScopeSet -> IdSubstEnv -> TvSubstEnv -> IdSubstEnv -> Subst
Subst InScopeSet
is IdSubstEnv
env (TvSubstEnv -> [(TyVar, Type)] -> TvSubstEnv
forall a b. VarEnv a -> [(Var b, a)] -> VarEnv a
extendVarEnvList TvSubstEnv
tenv [(TyVar, Type)]
ts) IdSubstEnv
genv

-- | Add an 'Id' to the in-scope set: as a side effect, remove any existing
-- substitutions for it.
extendInScopeId
  :: Subst
  -> Id
  -> Subst
extendInScopeId :: Subst -> Id -> Subst
extendInScopeId (Subst inScope :: InScopeSet
inScope env :: IdSubstEnv
env tenv :: TvSubstEnv
tenv genv :: IdSubstEnv
genv) id' :: Id
id' =
  InScopeSet -> IdSubstEnv -> TvSubstEnv -> IdSubstEnv -> Subst
Subst InScopeSet
inScope' IdSubstEnv
env' TvSubstEnv
tenv IdSubstEnv
genv
 where
  inScope' :: InScopeSet
inScope' = InScopeSet -> Id -> InScopeSet
forall a. InScopeSet -> Var a -> InScopeSet
extendInScopeSet InScopeSet
inScope Id
id'
  env' :: IdSubstEnv
env'     = IdSubstEnv -> Id -> IdSubstEnv
forall a b. VarEnv a -> Var b -> VarEnv a
delVarEnv IdSubstEnv
env Id
id'

-- | Add 'Id's to the in-scope set. See also 'extendInScopeId'
extendInScopeIdList
  :: Subst
  -> [Id]
  -> Subst
extendInScopeIdList :: Subst -> [Id] -> Subst
extendInScopeIdList (Subst inScope :: InScopeSet
inScope env :: IdSubstEnv
env tenv :: TvSubstEnv
tenv genv :: IdSubstEnv
genv) ids :: [Id]
ids =
  InScopeSet -> IdSubstEnv -> TvSubstEnv -> IdSubstEnv -> Subst
Subst InScopeSet
inScope' IdSubstEnv
env' TvSubstEnv
tenv IdSubstEnv
genv
 where
  inScope' :: InScopeSet
inScope' = InScopeSet -> [Id] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
inScope [Id]
ids
  env' :: IdSubstEnv
env'     = IdSubstEnv -> [Id] -> IdSubstEnv
forall a b. VarEnv a -> [Var b] -> VarEnv a
delVarEnvList IdSubstEnv
env [Id]
ids

-- | Substitute within a 'Type'
--
-- The substitution has to satisfy the invariant described in
-- 'TvSubst's Note [The substitution environment]
substTy
  :: HasCallStack
  => Subst
  -> Type
  -> Type
substTy :: Subst -> Type -> Type
substTy (Subst inScope :: InScopeSet
inScope _ tvS :: TvSubstEnv
tvS _) ty :: Type
ty
  | TvSubstEnv -> Bool
forall a. VarEnv a -> Bool
nullVarEnv TvSubstEnv
tvS
  = Type
ty
  | Bool
otherwise
  = TvSubst -> [Type] -> Type -> Type
forall a. HasCallStack => TvSubst -> [Type] -> a -> a
checkValidSubst TvSubst
s' [Type
ty] (HasCallStack => TvSubst -> Type -> Type
TvSubst -> Type -> Type
substTy' TvSubst
s' Type
ty)
 where
  s' :: TvSubst
s' = InScopeSet -> TvSubstEnv -> TvSubst
TvSubst InScopeSet
inScope TvSubstEnv
tvS

-- | Substitute within a 'TyVar'. See 'substTy'.
substTyInVar
  :: HasCallStack
  => Subst
  -> Var a
  -> Var a
substTyInVar :: Subst -> Var a -> Var a
substTyInVar subst :: Subst
subst tyVar :: Var a
tyVar =
  Var a
tyVar { varType :: Type
varType = (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst (Var a -> Type
forall a. Var a -> Type
varType Var a
tyVar)) }

-- | Like 'substTy', but skips the checks for the invariants described in
-- 'TvSubts' Note [The substitution environment]. Should be used inside this
-- module only.
substTyUnchecked
  :: HasCallStack
  => TvSubst
  -> Type
  -> Type
substTyUnchecked :: TvSubst -> Type -> Type
substTyUnchecked subst :: TvSubst
subst@(TvSubst _ tvS :: TvSubstEnv
tvS) ty :: Type
ty
  | TvSubstEnv -> Bool
forall a. VarEnv a -> Bool
nullVarEnv TvSubstEnv
tvS
  = Type
ty
  | Bool
otherwise
  = HasCallStack => TvSubst -> Type -> Type
TvSubst -> Type -> Type
substTy' TvSubst
subst Type
ty

-- | This checks if the substitution satisfies the invariant from 'TvSbust's
-- Note [The substitution invariant].
checkValidSubst
  :: HasCallStack
  => TvSubst
  -> [Type]
  -> a
  -> a
checkValidSubst :: TvSubst -> [Type] -> a -> a
checkValidSubst subst :: TvSubst
subst@(TvSubst inScope :: InScopeSet
inScope tenv :: TvSubstEnv
tenv) tys :: [Type]
tys a :: a
a =
  WARN( not (isValidSubst subst),
        "inScope" <+> clashPretty inScope <> line <>
        "tenv" <+> clashPretty tenv <> line <>
        "tenvFVs" <+> clashPretty (tyFVsOfTypes tenv) <> line <>
        "tys" <+> fromPpr tys)
  WARN( not tysFVsInSope,
       "inScope" <+> clashPretty inScope <> line <>
       "tenv" <+> clashPretty tenv <> line <>
       "tys" <+> fromPpr tys <> line <>
       "needsInScope" <+> clashPretty needsInScope)
  a
a
 where
  needsInScope :: VarSet
needsInScope = (Int -> Type -> VarSet -> VarSet) -> VarSet -> TvSubstEnv -> VarSet
forall a b. (Int -> a -> b -> b) -> b -> UniqMap a -> b
foldrWithUnique (\k :: Int
k _ s :: VarSet
s -> Int -> VarSet -> VarSet
delVarSetByKey Int
k VarSet
s)
                   ([Type] -> VarSet
forall (f :: * -> *). Foldable f => f Type -> VarSet
tyFVsOfTypes [Type]
tys)
                   TvSubstEnv
tenv
  tysFVsInSope :: Bool
tysFVsInSope = VarSet
needsInScope VarSet -> InScopeSet -> Bool
`varSetInScope` InScopeSet
inScope

-- | When calling 'substTy' it should be the case that the in-scope set in the
-- substitution is a superset of the free variables of the range of the
-- substitution.
--
-- See also 'TvSubst's Note [The substitution invariant].
isValidSubst
  :: TvSubst
  -> Bool
isValidSubst :: TvSubst -> Bool
isValidSubst (TvSubst inScope :: InScopeSet
inScope tenv :: TvSubstEnv
tenv) = VarSet
tenvFVs VarSet -> InScopeSet -> Bool
`varSetInScope` InScopeSet
inScope
 where
  tenvFVs :: VarSet
tenvFVs = TvSubstEnv -> VarSet
forall (f :: * -> *). Foldable f => f Type -> VarSet
tyFVsOfTypes TvSubstEnv
tenv

-- | The work-horse of 'substTy'
substTy'
  :: HasCallStack
  => TvSubst
  -> Type
  -> Type
substTy' :: TvSubst -> Type -> Type
substTy' subst :: TvSubst
subst = Type -> Type
go where
  go :: Type -> Type
go = \case
    VarTy tv :: TyVar
tv -> TvSubst -> TyVar -> Type
substTyVar TvSubst
subst TyVar
tv
    ForAllTy tv :: TyVar
tv ty :: Type
ty -> case TvSubst -> TyVar -> (TvSubst, TyVar)
substTyVarBndr TvSubst
subst TyVar
tv of
      (subst' :: TvSubst
subst', tv' :: TyVar
tv') -> TyVar -> Type -> Type
ForAllTy TyVar
tv' (HasCallStack => TvSubst -> Type -> Type
TvSubst -> Type -> Type
substTy' TvSubst
subst' Type
ty)
    AppTy fun :: Type
fun arg :: Type
arg -> Type -> Type -> Type
AppTy (Type -> Type
go Type
fun) (Type -> Type
go Type
arg)
    ty :: Type
ty -> Type
ty

-- | Substitute a variable with a type if it's within the substitution's domain.
--
-- Does not substitute within the kind of free variables.
substTyVar
  :: TvSubst
  -> TyVar
  -> Type
substTyVar :: TvSubst -> TyVar -> Type
substTyVar (TvSubst _ tenv :: TvSubstEnv
tenv) tv :: TyVar
tv = case TyVar -> TvSubstEnv -> Maybe Type
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv TyVar
tv TvSubstEnv
tenv of
  Just ty :: Type
ty -> Type
ty
  _       -> TyVar -> Type
VarTy TyVar
tv

-- | Substitute a type variable in a binding position, returning an extended
-- substitution environment and a new type variable.
--
-- Substitutes within the kind of the type variable
substTyVarBndr
  :: TvSubst
  -> TyVar
  -> (TvSubst, TyVar)
substTyVarBndr :: TvSubst -> TyVar -> (TvSubst, TyVar)
substTyVarBndr subst :: TvSubst
subst@(TvSubst inScope :: InScopeSet
inScope tenv :: TvSubstEnv
tenv) oldVar :: TyVar
oldVar =
  ASSERT2( no_capture, clashPretty oldVar <> line
                    <> clashPretty newVar <> line
                    <> clashPretty subst )
  (InScopeSet -> TvSubstEnv -> TvSubst
TvSubst (InScopeSet
inScope InScopeSet -> TyVar -> InScopeSet
forall a. InScopeSet -> Var a -> InScopeSet
`extendInScopeSet` TyVar
newVar) TvSubstEnv
newEnv, TyVar
newVar)
 where
  newEnv :: TvSubstEnv
newEnv | Bool
noChange  = TvSubstEnv -> TyVar -> TvSubstEnv
forall a b. VarEnv a -> Var b -> VarEnv a
delVarEnv TvSubstEnv
tenv TyVar
oldVar
         | Bool
otherwise = TyVar -> Type -> TvSubstEnv -> TvSubstEnv
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv TyVar
oldVar (TyVar -> Type
VarTy TyVar
newVar) TvSubstEnv
tenv

  -- Assertion that we're not capturing something in the substitution
  no_capture :: Bool
no_capture = Bool -> Bool
not (TyVar
newVar TyVar -> VarSet -> Bool
forall a. Var a -> VarSet -> Bool
`elemVarSet` TvSubstEnv -> VarSet
forall (f :: * -> *). Foldable f => f Type -> VarSet
tyFVsOfTypes TvSubstEnv
tenv)

  oldKi :: Type
oldKi        = TyVar -> Type
forall a. Var a -> Type
varType TyVar
oldVar
  -- verify that the kind is closed
  noKindChange :: Bool
noKindChange = Type -> Bool
noFreeVarsOfType Type
oldKi
  -- noChange means that the new type variable is identical in all respects to
  -- the old type variable (same unique, same kind)
  -- See 'TvSubstEnv's Note [Extending the TvSubstEnv]
  --
  -- In that case we don't need to extend the substitution to map old to new.
  -- But instead we must zap any current substitution for the variable. For
  -- example
  --
  --   (\x.e) with subst = [x | -> e']
  --
  -- Here we must simply zap the substitution for x
  noChange :: Bool
noChange     = Bool
noKindChange Bool -> Bool -> Bool
&& (TyVar
newVar TyVar -> TyVar -> Bool
forall a. Eq a => a -> a -> Bool
== TyVar
oldVar)

  -- uniqAway ensures that the new variable is not already in scope
  newVar :: TyVar
newVar | Bool
noKindChange = InScopeSet -> TyVar -> TyVar
forall a. InScopeSet -> Var a -> Var a
uniqAway InScopeSet
inScope TyVar
oldVar
         | Bool
otherwise    = InScopeSet -> TyVar -> TyVar
forall a. InScopeSet -> Var a -> Var a
uniqAway InScopeSet
inScope
                            (TyVar
oldVar {varType :: Type
varType = HasCallStack => TvSubst -> Type -> Type
TvSubst -> Type -> Type
substTyUnchecked TvSubst
subst Type
oldKi})

-- | Substitute within a 'Type'
substTm
  :: HasCallStack
  => Doc ()
  -> Subst
  -> Term
  -> Term
substTm :: Doc () -> Subst -> Term -> Term
substTm doc :: Doc ()
doc subst :: Subst
subst = Term -> Term
go where
  go :: Term -> Term
go = \case
    Var v :: Id
v -> HasCallStack => Doc () -> Subst -> Id -> Term
Doc () -> Subst -> Id -> Term
lookupIdSubst (Doc ()
doc Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> Doc ()
forall ann. Doc ann
line Doc () -> Doc () -> Doc ()
forall a. Semigroup a => a -> a -> a
<> "subsTm") Subst
subst Id
v
    Lam v :: Id
v e :: Term
e -> case HasCallStack => Subst -> Id -> (Subst, Id)
Subst -> Id -> (Subst, Id)
substIdBndr Subst
subst Id
v of
      (subst' :: Subst
subst',v' :: Id
v') -> Id -> Term -> Term
Lam Id
v' (HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm Doc ()
doc Subst
subst' Term
e)
    TyLam v :: TyVar
v e :: Term
e -> case HasCallStack => Subst -> TyVar -> (Subst, TyVar)
Subst -> TyVar -> (Subst, TyVar)
substTyVarBndr' Subst
subst TyVar
v of
      (subst' :: Subst
subst',v' :: TyVar
v') -> TyVar -> Term -> Term
TyLam TyVar
v' (HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm Doc ()
doc Subst
subst' Term
e)
    App l :: Term
l r :: Term
r -> Term -> Term -> Term
App (Term -> Term
go Term
l) (Term -> Term
go Term
r)
    TyApp l :: Term
l r :: Type
r -> Term -> Type -> Term
TyApp (Term -> Term
go Term
l) (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst Type
r)
    Letrec bs :: [(Id, Term)]
bs e :: Term
e -> case HasCallStack =>
Doc () -> Subst -> [(Id, Term)] -> (Subst, [(Id, Term)])
Doc () -> Subst -> [(Id, Term)] -> (Subst, [(Id, Term)])
substBind Doc ()
doc Subst
subst [(Id, Term)]
bs of
      (subst' :: Subst
subst',bs' :: [(Id, Term)]
bs') -> [(Id, Term)] -> Term -> Term
Letrec [(Id, Term)]
bs' (HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm Doc ()
doc Subst
subst' Term
e)
    Case subj :: Term
subj ty :: Type
ty alts :: [Alt]
alts -> Term -> Type -> [Alt] -> Term
Case (Term -> Term
go Term
subj) (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst Type
ty) ((Alt -> Alt) -> [Alt] -> [Alt]
forall a b. (a -> b) -> [a] -> [b]
map Alt -> Alt
goAlt [Alt]
alts)
    Cast e :: Term
e t1 :: Type
t1 t2 :: Type
t2 -> Term -> Type -> Type -> Term
Cast (Term -> Term
go Term
e) (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst Type
t1) (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst Type
t2)
    Tick tick :: TickInfo
tick e :: Term
e -> TickInfo -> Term -> Term
Tick (TickInfo -> TickInfo
goTick TickInfo
tick) (Term -> Term
go Term
e)
    tm :: Term
tm -> Term
tm

  goAlt :: Alt -> Alt
goAlt (pat :: Pat
pat,alt :: Term
alt) = case Pat
pat of
    DataPat dc :: DataCon
dc tvs :: [TyVar]
tvs ids :: [Id]
ids -> case (Subst -> TyVar -> (Subst, TyVar))
-> Subst -> [TyVar] -> (Subst, [TyVar])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
List.mapAccumL HasCallStack => Subst -> TyVar -> (Subst, TyVar)
Subst -> TyVar -> (Subst, TyVar)
substTyVarBndr' Subst
subst [TyVar]
tvs of
      (subst1 :: Subst
subst1,tvs' :: [TyVar]
tvs') -> case (Subst -> Id -> (Subst, Id)) -> Subst -> [Id] -> (Subst, [Id])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
List.mapAccumL HasCallStack => Subst -> Id -> (Subst, Id)
Subst -> Id -> (Subst, Id)
substIdBndr Subst
subst1 [Id]
ids of
        (subst2 :: Subst
subst2,ids' :: [Id]
ids') -> (DataCon -> [TyVar] -> [Id] -> Pat
DataPat DataCon
dc [TyVar]
tvs' [Id]
ids',HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm Doc ()
doc Subst
subst2 Term
alt)
    _ -> (Pat
pat,Term -> Term
go Term
alt)

  goTick :: TickInfo -> TickInfo
goTick t :: TickInfo
t@(SrcSpan _)  = TickInfo
t
  goTick (NameMod m :: NameMod
m ty :: Type
ty) = NameMod -> Type -> TickInfo
NameMod NameMod
m (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst Type
ty)

-- | Find the substitution for an 'Id' in the 'Subst'
lookupIdSubst
  :: HasCallStack
  => Doc ()
  -> Subst
  -> Id
  -> Term
lookupIdSubst :: Doc () -> Subst -> Id -> Term
lookupIdSubst doc :: Doc ()
doc (Subst inScope :: InScopeSet
inScope tmS :: IdSubstEnv
tmS _ genv :: IdSubstEnv
genv) v :: Id
v
  | Id -> Bool
forall a. Var a -> Bool
isGlobalId Id
v = case Id -> IdSubstEnv -> Maybe Term
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv Id
v IdSubstEnv
genv of
                     Just e :: Term
e -> Term
e
                     _      -> Id -> Term
Var Id
v
  | Just e :: Term
e <- Id -> IdSubstEnv -> Maybe Term
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv Id
v IdSubstEnv
tmS = Term
e
  -- Vital! See 'IdSubstEnv' Note [Extending the Subst]
  | Just v' :: Var Any
v' <- InScopeSet -> Id -> Maybe (Var Any)
forall a. InScopeSet -> Var a -> Maybe (Var Any)
lookupInScope InScopeSet
inScope Id
v = Id -> Term
Var (Var Any -> Id
forall a b. Coercible a b => a -> b
coerce Var Any
v')
  | Bool
otherwise = WARN(True, "Subst.lookupIdSubst" <+> doc <+> fromPpr v)
                Id -> Term
Var Id
v

-- | Substitute an 'Id' for another one according to the 'Subst' given,
-- returning the result and an update 'Subst' that should be used in subsequent
-- substitutions.
substIdBndr
  :: HasCallStack
  => Subst
  -> Id
  -> (Subst,Id)
substIdBndr :: Subst -> Id -> (Subst, Id)
substIdBndr subst :: Subst
subst@(Subst inScope :: InScopeSet
inScope env :: IdSubstEnv
env tenv :: TvSubstEnv
tenv genv :: IdSubstEnv
genv) oldId :: Id
oldId =
  (InScopeSet -> IdSubstEnv -> TvSubstEnv -> IdSubstEnv -> Subst
Subst (InScopeSet
inScope InScopeSet -> Id -> InScopeSet
forall a. InScopeSet -> Var a -> InScopeSet
`extendInScopeSet` Id
newId) IdSubstEnv
newEnv TvSubstEnv
tenv IdSubstEnv
genv, Id
newId)
 where
  id1 :: Id
id1 = InScopeSet -> Id -> Id
forall a. InScopeSet -> Var a -> Var a
uniqAway InScopeSet
inScope Id
oldId
  newId :: Id
newId | Bool
noTypeChange = Id
id1
        | Bool
otherwise    = Id
id1 {varType :: Type
varType = HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst (Id -> Type
forall a. Var a -> Type
varType Id
id1)}

  oldTy :: Type
oldTy = Id -> Type
forall a. Var a -> Type
varType Id
oldId
  noTypeChange :: Bool
noTypeChange = TvSubstEnv -> Bool
forall a. VarEnv a -> Bool
nullVarEnv TvSubstEnv
tenv Bool -> Bool -> Bool
|| Type -> Bool
noFreeVarsOfType Type
oldTy

  -- Extend the substitution if the unique has changed.
  --
  -- In case it hasn't changed we don't need to extend the substitution to map
  -- old to new. But instead we must zap any current substitution for the
  -- variable. For example
  --
  --   (\x.e) with subst = [x | -> e']
  --
  -- Here we must simply zap the substitution for x
  newEnv :: IdSubstEnv
newEnv | Bool
noChange  = IdSubstEnv -> Id -> IdSubstEnv
forall a b. VarEnv a -> Var b -> VarEnv a
delVarEnv IdSubstEnv
env Id
oldId
         | Bool
otherwise = Id -> Term -> IdSubstEnv -> IdSubstEnv
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv Id
oldId (Id -> Term
Var Id
newId) IdSubstEnv
env

  -- See Note [Extending the Subst] why it's not necessary to check noTypeChange
  noChange :: Bool
noChange = Id
id1 Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
oldId

-- | Like 'substTyVarBndr' but takes a 'Subst' instead of a 'TvSubst'
substTyVarBndr'
  :: HasCallStack
  => Subst
  -> TyVar
  -> (Subst,TyVar)
substTyVarBndr' :: Subst -> TyVar -> (Subst, TyVar)
substTyVarBndr' (Subst inScope :: InScopeSet
inScope tmS :: IdSubstEnv
tmS tyS :: TvSubstEnv
tyS tgS :: IdSubstEnv
tgS) tv :: TyVar
tv =
  case TvSubst -> TyVar -> (TvSubst, TyVar)
substTyVarBndr (InScopeSet -> TvSubstEnv -> TvSubst
TvSubst InScopeSet
inScope TvSubstEnv
tyS) TyVar
tv of
    (TvSubst inScope' :: InScopeSet
inScope' tyS' :: TvSubstEnv
tyS',tv' :: TyVar
tv') -> (InScopeSet -> IdSubstEnv -> TvSubstEnv -> IdSubstEnv -> Subst
Subst InScopeSet
inScope' IdSubstEnv
tmS TvSubstEnv
tyS' IdSubstEnv
tgS, TyVar
tv')

-- | Apply a substitution to an entire set of let-bindings, additionally
-- returning an updated 'Subst' that should be used by subsequent substitutions.
substBind
  :: HasCallStack
  => Doc ()
  -> Subst
  -> [LetBinding]
  -> (Subst,[LetBinding])
substBind :: Doc () -> Subst -> [(Id, Term)] -> (Subst, [(Id, Term)])
substBind doc :: Doc ()
doc subst :: Subst
subst xs :: [(Id, Term)]
xs =
  (Subst
subst',[Id] -> [Term] -> [(Id, Term)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
bndrs' [Term]
rhss')
 where
  (bndrs :: [Id]
bndrs,rhss :: [Term]
rhss)    = [(Id, Term)] -> ([Id], [Term])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Id, Term)]
xs
  (subst' :: Subst
subst',bndrs' :: [Id]
bndrs') = (Subst -> Id -> (Subst, Id)) -> Subst -> [Id] -> (Subst, [Id])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
List.mapAccumL HasCallStack => Subst -> Id -> (Subst, Id)
Subst -> Id -> (Subst, Id)
substIdBndr Subst
subst [Id]
bndrs
  rhss' :: [Term]
rhss'           = (Term -> Term) -> [Term] -> [Term]
forall a b. (a -> b) -> [a] -> [b]
map (HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm ("substBind" Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ()
doc) Subst
subst') [Term]
rhss

-- | Type substitution, see 'zipTvSubst'
--
-- Works only if the domain of the substitution is superset of the type being
-- substituted into
substTyWith
  :: HasCallStack
  => [TyVar]
  -> [Type]
  -> Type
  -> Type
substTyWith :: [TyVar] -> [Type] -> Type -> Type
substTyWith tvs :: [TyVar]
tvs tys :: [Type]
tys =
  ASSERT( tvs `equalLength` tys )
  HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy ([TyVar] -> [Type] -> Subst
zipTvSubst [TyVar]
tvs [Type]
tys)

-- | Ensure that non of the binders in an expression shadow each-other, nor
-- conflict with he in-scope set
deShadowTerm
  :: HasCallStack
  => InScopeSet
  -> Term
  -> Term
deShadowTerm :: InScopeSet -> Term -> Term
deShadowTerm is :: InScopeSet
is e :: Term
e = HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm "deShadowTerm" (InScopeSet -> Subst
mkSubst InScopeSet
is) Term
e

-- | A much stronger variant of `deShadowTerm` that ensures that all bound
-- variables are unique.
--
-- Also returns an extended 'InScopeSet' additionally containing the (renamed)
-- unique bound variables of the term.
freshenTm
  :: InScopeSet
  -- ^ Current set of variables in scope
  -> Term
  -> (InScopeSet, Term)
freshenTm :: InScopeSet -> Term -> (InScopeSet, Term)
freshenTm is0 :: InScopeSet
is0 = Subst -> Term -> (InScopeSet, Term)
go (InScopeSet -> Subst
mkSubst InScopeSet
is0) where
  go :: Subst -> Term -> (InScopeSet, Term)
go subst0 :: Subst
subst0 = \case
    Var v :: Id
v -> (Subst -> InScopeSet
substInScope Subst
subst0, HasCallStack => Doc () -> Subst -> Id -> Term
Doc () -> Subst -> Id -> Term
lookupIdSubst "freshenTm" Subst
subst0 Id
v)
    Lam v :: Id
v e :: Term
e -> case HasCallStack => Subst -> Id -> (Subst, Id)
Subst -> Id -> (Subst, Id)
substIdBndr Subst
subst0 Id
v of
      (subst1 :: Subst
subst1,v' :: Id
v') -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst1 Term
e of
        (is2 :: InScopeSet
is2,e' :: Term
e') -> (InScopeSet
is2, Id -> Term -> Term
Lam Id
v' Term
e')
    TyLam v :: TyVar
v e :: Term
e -> case HasCallStack => Subst -> TyVar -> (Subst, TyVar)
Subst -> TyVar -> (Subst, TyVar)
substTyVarBndr' Subst
subst0 TyVar
v of
      (subst1 :: Subst
subst1,v' :: TyVar
v') -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst1 Term
e of
        (is2 :: InScopeSet
is2,e' :: Term
e') -> (InScopeSet
is2,TyVar -> Term -> Term
TyLam TyVar
v' Term
e')
    App l :: Term
l r :: Term
r -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 Term
l of
      (is1 :: InScopeSet
is1,l' :: Term
l') -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 {substInScope :: InScopeSet
substInScope = InScopeSet
is1} Term
r of
        (is2 :: InScopeSet
is2,r' :: Term
r') -> (InScopeSet
is2, Term -> Term -> Term
App Term
l' Term
r')
    TyApp l :: Term
l r :: Type
r -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 Term
l of
      (is1 :: InScopeSet
is1,l' :: Term
l') -> (InScopeSet
is1, Term -> Type -> Term
TyApp Term
l' (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst0 Type
r))
    Letrec bs :: [(Id, Term)]
bs e :: Term
e -> case Subst -> [(Id, Term)] -> (Subst, [(Id, Term)])
goBind Subst
subst0 [(Id, Term)]
bs of
      (subst1 :: Subst
subst1,bs' :: [(Id, Term)]
bs') -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst1 Term
e of
        (is2 :: InScopeSet
is2,e' :: Term
e') -> (InScopeSet
is2,[(Id, Term)] -> Term -> Term
Letrec [(Id, Term)]
bs' Term
e')
    Case subj :: Term
subj ty :: Type
ty alts :: [Alt]
alts -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 Term
subj of
      (is1 :: InScopeSet
is1,subj' :: Term
subj') -> case (InScopeSet -> Alt -> (InScopeSet, Alt))
-> InScopeSet -> [Alt] -> (InScopeSet, [Alt])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
List.mapAccumL (\isN :: InScopeSet
isN -> Subst -> Alt -> (InScopeSet, Alt)
goAlt Subst
subst0 {substInScope :: InScopeSet
substInScope = InScopeSet
isN}) InScopeSet
is1 [Alt]
alts of
        (is2 :: InScopeSet
is2,alts' :: [Alt]
alts') -> (InScopeSet
is2, Term -> Type -> [Alt] -> Term
Case Term
subj' (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst0 Type
ty) [Alt]
alts')
    Cast e :: Term
e t1 :: Type
t1 t2 :: Type
t2 -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 Term
e of
      (is1 :: InScopeSet
is1, e' :: Term
e') -> (InScopeSet
is1, Term -> Type -> Type -> Term
Cast Term
e' (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst0 Type
t1) (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst0 Type
t2))
    Tick tick :: TickInfo
tick e :: Term
e -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 Term
e of
       (is1 :: InScopeSet
is1, e' :: Term
e') -> (InScopeSet
is1, TickInfo -> Term -> Term
Tick (Subst -> TickInfo -> TickInfo
goTick Subst
subst0 TickInfo
tick) Term
e')
    tm :: Term
tm -> (Subst -> InScopeSet
substInScope Subst
subst0, Term
tm)

  goBind :: Subst -> [(Id, Term)] -> (Subst, [(Id, Term)])
goBind subst0 :: Subst
subst0 xs :: [(Id, Term)]
xs =
    let (bndrs :: [Id]
bndrs,rhss :: [Term]
rhss)    = [(Id, Term)] -> ([Id], [Term])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Id, Term)]
xs
        (subst1 :: Subst
subst1,bndrs' :: [Id]
bndrs') = (Subst -> Id -> (Subst, Id)) -> Subst -> [Id] -> (Subst, [Id])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
List.mapAccumL HasCallStack => Subst -> Id -> (Subst, Id)
Subst -> Id -> (Subst, Id)
substIdBndr Subst
subst0 [Id]
bndrs
        (is2 :: InScopeSet
is2,rhss' :: [Term]
rhss')     = (InScopeSet -> Term -> (InScopeSet, Term))
-> InScopeSet -> [Term] -> (InScopeSet, [Term])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
List.mapAccumL (\isN :: InScopeSet
isN -> Subst -> Term -> (InScopeSet, Term)
go Subst
subst1 {substInScope :: InScopeSet
substInScope = InScopeSet
isN})
                                         (Subst -> InScopeSet
substInScope Subst
subst1)
                                         [Term]
rhss
    in  (Subst
subst1 {substInScope :: InScopeSet
substInScope = InScopeSet
is2},[Id] -> [Term] -> [(Id, Term)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
bndrs' [Term]
rhss')

  goAlt :: Subst -> Alt -> (InScopeSet, Alt)
goAlt subst0 :: Subst
subst0 (pat :: Pat
pat,alt :: Term
alt) = case Pat
pat of
    DataPat dc :: DataCon
dc tvs :: [TyVar]
tvs ids :: [Id]
ids -> case (Subst -> TyVar -> (Subst, TyVar))
-> Subst -> [TyVar] -> (Subst, [TyVar])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
List.mapAccumL HasCallStack => Subst -> TyVar -> (Subst, TyVar)
Subst -> TyVar -> (Subst, TyVar)
substTyVarBndr' Subst
subst0 [TyVar]
tvs of
      (subst1 :: Subst
subst1,tvs' :: [TyVar]
tvs') -> case (Subst -> Id -> (Subst, Id)) -> Subst -> [Id] -> (Subst, [Id])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
List.mapAccumL HasCallStack => Subst -> Id -> (Subst, Id)
Subst -> Id -> (Subst, Id)
substIdBndr Subst
subst1 [Id]
ids of
        (subst2 :: Subst
subst2,ids' :: [Id]
ids') -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst2 Term
alt of
          (is3 :: InScopeSet
is3,alt' :: Term
alt') -> (InScopeSet
is3,(DataCon -> [TyVar] -> [Id] -> Pat
DataPat DataCon
dc [TyVar]
tvs' [Id]
ids',Term
alt'))
    _ -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 Term
alt of
      (is1 :: InScopeSet
is1,alt' :: Term
alt') -> (InScopeSet
is1,(Pat
pat,Term
alt'))

  goTick :: Subst -> TickInfo -> TickInfo
goTick subst0 :: Subst
subst0 (NameMod m :: NameMod
m ty :: Type
ty) = NameMod -> Type -> TickInfo
NameMod NameMod
m (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst0 Type
ty)
  goTick _      tick :: TickInfo
tick           = TickInfo
tick

-- * AEQ

-- | Alpha equality for types
aeqType
  :: Type
  -> Type
  -> Bool
aeqType :: Type -> Type -> Bool
aeqType t1 :: Type
t1 t2 :: Type
t2 = RnEnv -> Type -> Type -> Ordering
acmpType' RnEnv
rnEnv Type
t1 Type
t2 Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ
 where
  rnEnv :: RnEnv
rnEnv = InScopeSet -> RnEnv
mkRnEnv (VarSet -> InScopeSet
mkInScopeSet ([Type] -> VarSet
forall (f :: * -> *). Foldable f => f Type -> VarSet
tyFVsOfTypes [Type
t1,Type
t2]))

-- | Alpha comparison for types
acmpType
  :: Type
  -> Type
  -> Ordering
acmpType :: Type -> Type -> Ordering
acmpType t1 :: Type
t1 t2 :: Type
t2 = RnEnv -> Type -> Type -> Ordering
acmpType' (InScopeSet -> RnEnv
mkRnEnv InScopeSet
inScope) Type
t1 Type
t2
 where
  inScope :: InScopeSet
inScope = VarSet -> InScopeSet
mkInScopeSet ([Type] -> VarSet
forall (f :: * -> *). Foldable f => f Type -> VarSet
tyFVsOfTypes [Type
t1,Type
t2])

-- | Alpha comparison for types. Faster than 'acmpType' as it doesn't need to
-- calculate the free variables to create the 'InScopeSet'
acmpType'
  :: RnEnv
  -> Type
  -> Type
  -> Ordering
acmpType' :: RnEnv -> Type -> Type -> Ordering
acmpType' = RnEnv -> Type -> Type -> Ordering
go
 where
  go :: RnEnv -> Type -> Type -> Ordering
go env :: RnEnv
env (VarTy tv1 :: TyVar
tv1) (VarTy tv2 :: TyVar
tv2) = TyVar -> TyVar -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (RnEnv -> TyVar -> TyVar
rnOccLTy RnEnv
env TyVar
tv1) (RnEnv -> TyVar -> TyVar
rnOccRTy RnEnv
env TyVar
tv2)
  go _   (ConstTy c1 :: ConstTy
c1) (ConstTy c2 :: ConstTy
c2) = ConstTy -> ConstTy -> Ordering
forall a. Ord a => a -> a -> Ordering
compare ConstTy
c1 ConstTy
c2
  go env :: RnEnv
env (ForAllTy tv1 :: TyVar
tv1 t1 :: Type
t1) (ForAllTy tv2 :: TyVar
tv2 t2 :: Type
t2) =
    RnEnv -> Type -> Type -> Ordering
go RnEnv
env (TyVar -> Type
forall a. Var a -> Type
varType TyVar
tv1) (TyVar -> Type
forall a. Var a -> Type
varType TyVar
tv2) Ordering -> Ordering -> Ordering
`thenCompare` RnEnv -> Type -> Type -> Ordering
go (RnEnv -> TyVar -> TyVar -> RnEnv
rnTyBndr RnEnv
env TyVar
tv1 TyVar
tv2) Type
t1 Type
t2
  go env :: RnEnv
env (AppTy s1 :: Type
s1 t1 :: Type
t1) (AppTy s2 :: Type
s2 t2 :: Type
t2) =
    RnEnv -> Type -> Type -> Ordering
go RnEnv
env Type
s1 Type
s2 Ordering -> Ordering -> Ordering
`thenCompare` RnEnv -> Type -> Type -> Ordering
go RnEnv
env Type
t1 Type
t2
  go _ (LitTy l1 :: LitTy
l1) (LitTy l2 :: LitTy
l2) = LitTy -> LitTy -> Ordering
forall a. Ord a => a -> a -> Ordering
compare LitTy
l1 LitTy
l2
  go _ t1 :: Type
t1 t2 :: Type
t2 = Word -> Word -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Type -> Word
getRank Type
t1) (Type -> Word
getRank Type
t2)

  getRank :: Type -> Word
  getRank :: Type -> Word
getRank (VarTy {})    = 0
  getRank (LitTy {})    = 1
  getRank (ConstTy {})  = 2
  getRank (AnnType {})  = 3
  getRank (AppTy {})    = 4
  getRank (ForAllTy {}) = 5

-- | Alpha equality for terms
aeqTerm
  :: Term
  -> Term
  -> Bool
aeqTerm :: Term -> Term -> Bool
aeqTerm t1 :: Term
t1 t2 :: Term
t2 = InScopeSet -> Term -> Term -> Bool
aeqTerm' InScopeSet
inScope Term
t1 Term
t2
 where
  inScope :: InScopeSet
inScope = VarSet -> InScopeSet
mkInScopeSet ([Term] -> VarSet
forall (f :: * -> *). Foldable f => f Term -> VarSet
localFVsOfTerms [Term
t1,Term
t2])

-- | Alpha equality for terms. Faster than 'aeqTerm' as it doesn't need to
-- calculate the free variables to create the 'InScopeSet'
aeqTerm'
  :: InScopeSet
  -- ^ Superset of variables in scope of the left and right term
  -> Term
  -> Term
  -> Bool
aeqTerm' :: InScopeSet -> Term -> Term -> Bool
aeqTerm' inScope :: InScopeSet
inScope t1 :: Term
t1 t2 :: Term
t2 = InScopeSet -> Term -> Term -> Ordering
acmpTerm' InScopeSet
inScope Term
t1 Term
t2 Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ

-- | Alpha comparison for types
acmpTerm
  :: Term
  -> Term
  -> Ordering
acmpTerm :: Term -> Term -> Ordering
acmpTerm t1 :: Term
t1 t2 :: Term
t2 = InScopeSet -> Term -> Term -> Ordering
acmpTerm' InScopeSet
inScope Term
t1 Term
t2
 where
  inScope :: InScopeSet
inScope = VarSet -> InScopeSet
mkInScopeSet ([Term] -> VarSet
forall (f :: * -> *). Foldable f => f Term -> VarSet
localFVsOfTerms [Term
t1,Term
t2])

-- | Alpha comparison for types. Faster than 'acmpTerm' as it doesn't need to
-- calculate the free variables to create the 'InScopeSet'
acmpTerm'
  :: InScopeSet
  -- ^ Superset of variables in scope of the left and right term
  -> Term
  -> Term
  -> Ordering
acmpTerm' :: InScopeSet -> Term -> Term -> Ordering
acmpTerm' inScope :: InScopeSet
inScope = RnEnv -> Term -> Term -> Ordering
go (InScopeSet -> RnEnv
mkRnEnv InScopeSet
inScope)
 where
  thenCmpTm :: Ordering -> Ordering -> Ordering
thenCmpTm EQ  rel :: Ordering
rel = Ordering
rel
  thenCmpTm rel :: Ordering
rel _   = Ordering
rel

  go :: RnEnv -> Term -> Term -> Ordering
go env :: RnEnv
env (Var id1 :: Id
id1) (Var id2 :: Id
id2)   = Id -> Id -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (RnEnv -> Id -> Id
rnOccLId RnEnv
env Id
id1) (RnEnv -> Id -> Id
rnOccRId RnEnv
env Id
id2)
  go _   (Data dc1 :: DataCon
dc1) (Data dc2 :: DataCon
dc2) = DataCon -> DataCon -> Ordering
forall a. Ord a => a -> a -> Ordering
compare DataCon
dc1 DataCon
dc2
  go _   (Literal l1 :: Literal
l1) (Literal l2 :: Literal
l2) = Literal -> Literal -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Literal
l1 Literal
l2
  go _   (Prim p1 :: Text
p1 _) (Prim p2 :: Text
p2 _) = Text -> Text -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Text
p1 Text
p2
  go env :: RnEnv
env (Lam b1 :: Id
b1 e1 :: Term
e1) (Lam b2 :: Id
b2 e2 :: Term
e2) =
    RnEnv -> Type -> Type -> Ordering
acmpType' RnEnv
env (Id -> Type
forall a. Var a -> Type
varType Id
b1) (Id -> Type
forall a. Var a -> Type
varType Id
b2) Ordering -> Ordering -> Ordering
`thenCompare`
    RnEnv -> Term -> Term -> Ordering
go (RnEnv -> Id -> Id -> RnEnv
rnTmBndr RnEnv
env Id
b1 Id
b2) Term
e1 Term
e2
  go env :: RnEnv
env (TyLam b1 :: TyVar
b1 e1 :: Term
e1) (TyLam b2 :: TyVar
b2 e2 :: Term
e2) =
    RnEnv -> Type -> Type -> Ordering
acmpType' RnEnv
env (TyVar -> Type
forall a. Var a -> Type
varType TyVar
b1) (TyVar -> Type
forall a. Var a -> Type
varType TyVar
b2) Ordering -> Ordering -> Ordering
`thenCompare`
    RnEnv -> Term -> Term -> Ordering
go (RnEnv -> TyVar -> TyVar -> RnEnv
rnTyBndr RnEnv
env TyVar
b1 TyVar
b2) Term
e1 Term
e2
  go env :: RnEnv
env (App l1 :: Term
l1 r1 :: Term
r1) (App l2 :: Term
l2 r2 :: Term
r2) =
    RnEnv -> Term -> Term -> Ordering
go RnEnv
env Term
l1 Term
l2 Ordering -> Ordering -> Ordering
`thenCompare` RnEnv -> Term -> Term -> Ordering
go RnEnv
env Term
r1 Term
r2
  go env :: RnEnv
env (TyApp l1 :: Term
l1 r1 :: Type
r1) (TyApp l2 :: Term
l2 r2 :: Type
r2) =
    RnEnv -> Term -> Term -> Ordering
go RnEnv
env Term
l1 Term
l2 Ordering -> Ordering -> Ordering
`thenCompare` RnEnv -> Type -> Type -> Ordering
acmpType' RnEnv
env Type
r1 Type
r2
  go env :: RnEnv
env (Letrec bs1 :: [(Id, Term)]
bs1 e1 :: Term
e1) (Letrec bs2 :: [(Id, Term)]
bs2 e2 :: Term
e2) =
    Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare ([(Id, Term)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Id, Term)]
bs1) ([(Id, Term)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Id, Term)]
bs2) Ordering -> Ordering -> Ordering
`thenCompare`
    (Ordering -> Ordering -> Ordering)
-> Ordering -> [Ordering] -> Ordering
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Ordering -> Ordering -> Ordering
thenCmpTm Ordering
EQ ((Term -> Term -> Ordering) -> [Term] -> [Term] -> [Ordering]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (RnEnv -> Term -> Term -> Ordering
go RnEnv
env') [Term]
rhs1 [Term]
rhs2) Ordering -> Ordering -> Ordering
`thenCompare`
    RnEnv -> Term -> Term -> Ordering
go RnEnv
env' Term
e1 Term
e2
   where
    (ids1 :: [Id]
ids1,rhs1 :: [Term]
rhs1) = [(Id, Term)] -> ([Id], [Term])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Id, Term)]
bs1
    (ids2 :: [Id]
ids2,rhs2 :: [Term]
rhs2) = [(Id, Term)] -> ([Id], [Term])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Id, Term)]
bs2
    env' :: RnEnv
env' = RnEnv -> [Id] -> [Id] -> RnEnv
rnTmBndrs RnEnv
env [Id]
ids1 [Id]
ids2
  go env :: RnEnv
env (Case e1 :: Term
e1 _ a1 :: [Alt]
a1) (Case e2 :: Term
e2 _ a2 :: [Alt]
a2) =
    Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare ([Alt] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Alt]
a1) ([Alt] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Alt]
a2) Ordering -> Ordering -> Ordering
`thenCompare`
    RnEnv -> Term -> Term -> Ordering
go RnEnv
env Term
e1 Term
e2 Ordering -> Ordering -> Ordering
`thenCompare`
    (Ordering -> Ordering -> Ordering)
-> Ordering -> [Ordering] -> Ordering
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Ordering -> Ordering -> Ordering
thenCmpTm Ordering
EQ ((Alt -> Alt -> Ordering) -> [Alt] -> [Alt] -> [Ordering]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (RnEnv -> Alt -> Alt -> Ordering
goAlt RnEnv
env) [Alt]
a1 [Alt]
a2)
  go env :: RnEnv
env (Cast e1 :: Term
e1 l1 :: Type
l1 r1 :: Type
r1) (Cast e2 :: Term
e2 l2 :: Type
l2 r2 :: Type
r2) =
    RnEnv -> Term -> Term -> Ordering
go RnEnv
env Term
e1 Term
e2 Ordering -> Ordering -> Ordering
`thenCompare`
    RnEnv -> Type -> Type -> Ordering
acmpType' RnEnv
env Type
l1 Type
l2 Ordering -> Ordering -> Ordering
`thenCompare`
    RnEnv -> Type -> Type -> Ordering
acmpType' RnEnv
env Type
r1 Type
r2
  -- Look through ticks for aeq
  go env :: RnEnv
env (Tick _ e1 :: Term
e1) e2 :: Term
e2 = RnEnv -> Term -> Term -> Ordering
go RnEnv
env Term
e1 Term
e2
  go env :: RnEnv
env e1 :: Term
e1 (Tick _ e2 :: Term
e2) = RnEnv -> Term -> Term -> Ordering
go RnEnv
env Term
e1 Term
e2

  go _ e1 :: Term
e1 e2 :: Term
e2 = Word -> Word -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Term -> Word
getRank Term
e1) (Term -> Word
getRank Term
e2)

  goAlt :: RnEnv -> Alt -> Alt -> Ordering
goAlt env :: RnEnv
env (DataPat c1 :: DataCon
c1 tvs1 :: [TyVar]
tvs1 ids1 :: [Id]
ids1,e1 :: Term
e1) (DataPat c2 :: DataCon
c2 tvs2 :: [TyVar]
tvs2 ids2 :: [Id]
ids2,e2 :: Term
e2) =
    DataCon -> DataCon -> Ordering
forall a. Ord a => a -> a -> Ordering
compare DataCon
c1 DataCon
c2 Ordering -> Ordering -> Ordering
`thenCompare` RnEnv -> Term -> Term -> Ordering
go RnEnv
env' Term
e1 Term
e2
   where
    env' :: RnEnv
env' = RnEnv -> [Id] -> [Id] -> RnEnv
rnTmBndrs (RnEnv -> [TyVar] -> [TyVar] -> RnEnv
rnTyBndrs RnEnv
env [TyVar]
tvs1 [TyVar]
tvs2) [Id]
ids1 [Id]
ids2
  goAlt env :: RnEnv
env (c1 :: Pat
c1,e1 :: Term
e1) (c2 :: Pat
c2,e2 :: Term
e2) =
    Pat -> Pat -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Pat
c1 Pat
c2 Ordering -> Ordering -> Ordering
`thenCompare` RnEnv -> Term -> Term -> Ordering
go RnEnv
env Term
e1 Term
e2

  getRank :: Term -> Word
  getRank :: Term -> Word
getRank = \case
    Var {}     -> 0
    Data {}    -> 1
    Literal {} -> 2
    Prim {}    -> 3
    Cast {}    -> 4
    App {}     -> 5
    TyApp {}   -> 6
    Lam {}     -> 7
    TyLam {}   -> 8
    Letrec {}  -> 9
    Case {}    -> 10
    Tick {}    -> 11

thenCompare :: Ordering -> Ordering -> Ordering
thenCompare :: Ordering -> Ordering -> Ordering
thenCompare EQ rel :: Ordering
rel = Ordering
rel
thenCompare rel :: Ordering
rel _  = Ordering
rel

instance Eq Type where
  == :: Type -> Type -> Bool
(==) = Type -> Type -> Bool
aeqType

instance Ord Type where
  compare :: Type -> Type -> Ordering
compare = Type -> Type -> Ordering
acmpType

instance Eq Term where
  == :: Term -> Term -> Bool
(==) = Term -> Term -> Bool
aeqTerm

instance Ord Term where
  compare :: Term -> Term -> Ordering
compare = Term -> Term -> Ordering
acmpTerm