{-|
  Copyright   :  (C) 2012-2016, University of Twente,
                          2017, Google Inc.
                          2021, QBayLogic B.V.
                          2026, Martijn Bastiaan
  License     :  BSD2 (see the file LICENSE)
  Maintainer  :  QBayLogic B.V. <devops@qbaylogic.com>

  Capture-free substitution function for CoreHW
-}

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StandaloneDeriving #-}

{-# 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
  , substGlobalsInExistentials
  , substInExistentials
  , substInExistentialsList
    -- * Substitution into terms
    -- ** Substitution environments
  , Subst (..)
  , mkSubst
  , mkTvSubst
  , extendInScopeId
  , extendInScopeIdList
  , extendIdSubst
  , extendIdSubstList
  , extendGblSubstList
    -- ** Applying substitutions
  , substTm
  , maybeSubstTm
  , unsafeSubstTm
  , substAlt
  , substId
    -- * Variable renaming
  , deShadowTerm
  , deShadowAlt
  , freshenTm
  , deshadowLetExpr
    -- * Alpha equivalence
  , aeqType
  , aeqTerm
  , acmpTerm
  , aeqTickInfo
  , acmpTickInfo
  , aeqTickInfoLevels
  , acmpTickInfoLevels
    -- * Structural equivalence
  , eqTerm
  , eqType
  )
where

import           Data.Coerce               (coerce)

#if MIN_VERSION_prettyprinter(1,7,0)
import           Prettyprinter
#else
import           Data.Text.Prettyprint.Doc
#endif

import           Data.Hashable             (Hashable (hashWithSalt))
import qualified Data.List                 as List
import qualified Data.List.Extra           as List
import           Data.Maybe                (fromMaybe)
import           Data.Ord                  (comparing)
import           GHC.Stack                 (HasCallStack)
import           GHC.SrcLoc.Extra          () -- Hashable RealSrcSpan
import           GHC.Types.SrcLoc
  (SrcSpan (RealSrcSpan, UnhelpfulSpan), leftmost_smallest)

import           Clash.Core.HasFreeVars
import           Clash.Core.Pretty         (ppr, fromPpr)
import           Clash.Core.Term
  (Alt, Bind(..), Pat (..), Term (..), TickInfo (..), PrimInfo(primName))
import           Clash.Core.Type           (Type (..))
import           Clash.Core.VarEnv
import           Clash.Core.Var
  (Id, Var (..), TyVar, isGlobalId, varKey)
import qualified Clash.Data.UniqMap as UniqMap
import           Clash.Debug               (debugIsOn)
import           Clash.Pretty
import           Clash.Util

-- * 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 InScopeSet
ins 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 [ Doc ()
"TvSubst"
                   , Int -> Doc () -> Doc ()
forall ann. Int -> Doc ann -> Doc ann
nest Int
2 (Doc ()
"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 Int
2 (Doc ()
"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 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 InScopeSet
is 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 [TyVar]
tvs [Type]
tys
  | Bool
debugIsOn
  , Bool -> Bool
not ([TyVar] -> [Type] -> Bool
forall a b. [a] -> [b] -> Bool
List.equalLength [TyVar]
tvs [Type]
tys)
  = String -> Doc ClashAnnotation -> Subst -> Subst
forall ann a. String -> Doc ann -> a -> a
pprTrace String
"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 a. HasFreeVars a => a -> VarSet
freeVarsOf [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 [TyVar]
tvs [Type]
tys = [(TyVar, Type)] -> TvSubstEnv
forall a b. [(Var a, b)] -> VarEnv b
mkVarEnv ([TyVar] -> [Type] -> [(TyVar, Type)]
forall a b. HasCallStack => [a] -> [b] -> [(a, b)]
List.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 InScopeSet
is IdSubstEnv
env TvSubstEnv
tenv IdSubstEnv
genv) Id
i 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 InScopeSet
is IdSubstEnv
env TvSubstEnv
tenv IdSubstEnv
genv) [(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 InScopeSet
is IdSubstEnv
env TvSubstEnv
tenv IdSubstEnv
genv) [(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 InScopeSet
is IdSubstEnv
env TvSubstEnv
tenv IdSubstEnv
genv) TyVar
tv 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 InScopeSet
is IdSubstEnv
env TvSubstEnv
tenv IdSubstEnv
genv) [(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 InScopeSet
inScope IdSubstEnv
env TvSubstEnv
tenv IdSubstEnv
genv) 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 InScopeSet
inScope IdSubstEnv
env TvSubstEnv
tenv IdSubstEnv
genv) [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 :: HasCallStack => Subst -> Type -> Type
substTy (Subst InScopeSet
inScope IdSubstEnv
_ TvSubstEnv
tvS IdSubstEnv
_) 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 :: forall a. HasCallStack => Subst -> Var a -> Var a
substTyInVar Subst
subst Var a
tyVar =
  Var a
tyVar { varType = (substTy subst (varType 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 :: HasCallStack => TvSubst -> Type -> Type
substTyUnchecked subst :: TvSubst
subst@(TvSubst InScopeSet
_ TvSubstEnv
tvS) 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

-- Safely substitute global type variables in a list of potentially
-- shadowing type variables.
substGlobalsInExistentials
  :: HasCallStack
  => InScopeSet
  -- ^ Variables in scope
  -> [TyVar]
  -- ^ List of existentials to apply the substitution for
  -> [(TyVar, Type)]
  -- ^ Substitutions
  -> [TyVar]
substGlobalsInExistentials :: HasCallStack => InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
substGlobalsInExistentials InScopeSet
is [TyVar]
exts [(TyVar, Type)]
substs0 = [TyVar]
result
  -- TODO: Is is actually possible that existentials shadow each other? If they
  -- TODO: can't, we can remove this function
  where
    iss :: [InScopeSet]
iss     = (InScopeSet -> TyVar -> InScopeSet)
-> InScopeSet -> [TyVar] -> [InScopeSet]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl InScopeSet -> TyVar -> InScopeSet
forall a. InScopeSet -> Var a -> InScopeSet
extendInScopeSet InScopeSet
is [TyVar]
exts
    substs1 :: [Subst]
substs1 = (InScopeSet -> Subst) -> [InScopeSet] -> [Subst]
forall a b. (a -> b) -> [a] -> [b]
map (\InScopeSet
is_ -> Subst -> [(TyVar, Type)] -> Subst
extendTvSubstList (InScopeSet -> Subst
mkSubst InScopeSet
is_) [(TyVar, Type)]
substs0) [InScopeSet]
iss
    result :: [TyVar]
result  = (Subst -> TyVar -> TyVar) -> [Subst] -> [TyVar] -> [TyVar]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Subst -> TyVar -> TyVar
forall a. HasCallStack => Subst -> Var a -> Var a
substTyInVar [Subst]
substs1 [TyVar]
exts

-- | Safely substitute type variables in a list of existentials. This function
-- will account for cases where existentials shadow each other.
substInExistentialsList
  :: HasCallStack
  => InScopeSet
  -- ^ Variables in scope
  -> [TyVar]
  -- ^ List of existentials to apply the substitution for
  -> [(TyVar, Type)]
  -- ^ Substitutions
  -> [TyVar]
substInExistentialsList :: HasCallStack => InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
substInExistentialsList InScopeSet
is [TyVar]
exts [(TyVar, Type)]
substs =
  ([TyVar] -> (TyVar, Type) -> [TyVar])
-> [TyVar] -> [(TyVar, Type)] -> [TyVar]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (HasCallStack => InScopeSet -> [TyVar] -> (TyVar, Type) -> [TyVar]
InScopeSet -> [TyVar] -> (TyVar, Type) -> [TyVar]
substInExistentials InScopeSet
is) [TyVar]
exts [(TyVar, Type)]
substs

-- | Safely substitute a type variable in a list of existentials. This function
-- will account for cases where existentials shadow each other.
substInExistentials
  :: HasCallStack
  => InScopeSet
  -- ^ Variables in scope
  -> [TyVar]
  -- ^ List of existentials to apply the substitution for
  -> (TyVar, Type)
  -- ^ Substitution
  -> [TyVar]
substInExistentials :: HasCallStack => InScopeSet -> [TyVar] -> (TyVar, Type) -> [TyVar]
substInExistentials InScopeSet
is [TyVar]
exts subst :: (TyVar, Type)
subst@(TyVar
typeVar, Type
_type) =
  -- TODO: Is is actually possible that existentials shadow each other? If they
  -- TODO: can't, we can remove this function
  case TyVar -> [TyVar] -> [Int]
forall a. Eq a => a -> [a] -> [Int]
List.elemIndices TyVar
typeVar [TyVar]
exts of
    [] ->
      -- We're not replacing any of the existentials, but a global variable
      HasCallStack => InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
substGlobalsInExistentials InScopeSet
is [TyVar]
exts [(TyVar, Type)
subst]
    ([Int] -> Int
forall a. HasCallStack => [a] -> a
last -> Int
i) ->
      -- We're replacing an existential. That means we're not touching any
      -- variables that were introduced before it. For all variables after it,
      -- it is as we would replace global variables in them.
      Int -> [TyVar] -> [TyVar]
forall a. Int -> [a] -> [a]
take (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [TyVar]
exts [TyVar] -> [TyVar] -> [TyVar]
forall a. [a] -> [a] -> [a]
++ HasCallStack => InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
substGlobalsInExistentials InScopeSet
is (Int -> [TyVar] -> [TyVar]
forall a. Int -> [a] -> [a]
drop (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [TyVar]
exts) [(TyVar, Type)
subst]

-- | This checks if the substitution satisfies the invariant from 'TvSubst's
-- Note [The substitution invariant].
checkValidSubst
  :: HasCallStack
  => TvSubst
  -> [Type]
  -> a
  -> a
checkValidSubst :: forall a. HasCallStack => TvSubst -> [Type] -> a -> a
checkValidSubst subst :: TvSubst
subst@(TvSubst InScopeSet
inScope TvSubstEnv
tenv) [Type]
tys a
a =
  WARN( not (isValidSubst subst),
        "inScope" <+> clashPretty inScope <> line <>
        "tenv" <+> clashPretty tenv <> line <>
        "tenvFVs" <+> clashPretty (freeVarsOf 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 = (Unique -> Type -> VarSet -> VarSet)
-> VarSet -> TvSubstEnv -> VarSet
forall a b. (Unique -> a -> b -> b) -> b -> UniqMap a -> b
UniqMap.foldrWithUnique (\Unique
k Type
_ VarSet
s -> Unique -> VarSet -> VarSet
delVarSetByKey Unique
k VarSet
s)
                   ([Type] -> VarSet
forall a. HasFreeVars a => a -> VarSet
freeVarsOf [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 InScopeSet
inScope TvSubstEnv
tenv) = VarSet
tenvFVs VarSet -> InScopeSet -> Bool
`varSetInScope` InScopeSet
inScope
 where
  tenvFVs :: VarSet
tenvFVs = TvSubstEnv -> VarSet
forall a. HasFreeVars a => a -> VarSet
freeVarsOf TvSubstEnv
tenv

-- | The work-horse of 'substTy'
substTy'
  :: HasCallStack
  => TvSubst
  -> Type
  -> Type
substTy' :: HasCallStack => TvSubst -> Type -> Type
substTy' TvSubst
subst = Type -> Type
go where
  go :: Type -> Type
go = \case
    VarTy TyVar
tv -> TvSubst -> TyVar -> Type
substTyVar TvSubst
subst TyVar
tv
    ForAllTy TyVar
tv Type
ty -> case TvSubst -> TyVar -> (TvSubst, TyVar)
substTyVarBndr TvSubst
subst TyVar
tv of
      (TvSubst
subst', TyVar
tv') -> TyVar -> Type -> Type
ForAllTy TyVar
tv' (HasCallStack => TvSubst -> Type -> Type
TvSubst -> Type -> Type
substTy' TvSubst
subst' Type
ty)
    AppTy Type
fun Type
arg -> Type -> Type -> Type
AppTy (Type -> Type
go Type
fun) (Type -> Type
go Type
arg)
    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 InScopeSet
_ TvSubstEnv
tenv) TyVar
tv = case TyVar -> TvSubstEnv -> Maybe Type
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv TyVar
tv TvSubstEnv
tenv of
  Just Type
ty -> Type
ty
  Maybe Type
_       -> 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 InScopeSet
inScope TvSubstEnv
tenv) 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 a. HasFreeVars a => a -> VarSet
freeVarsOf 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
forall a. HasFreeVars a => a -> Bool
isClosed 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. (Uniquable a, ClashPretty a) => InScopeSet -> a -> a
uniqAway InScopeSet
inScope TyVar
oldVar
         | Bool
otherwise    = InScopeSet -> TyVar -> TyVar
forall a. (Uniquable a, ClashPretty a) => InScopeSet -> a -> a
uniqAway InScopeSet
inScope
                            (TyVar
oldVar {varType = substTyUnchecked subst oldKi})

-- | Substitute within a 'Term'. Just return original term if given
-- substitution is "Nothing".
maybeSubstTm
  :: HasCallStack
  => Doc ()
  -> Maybe Subst
  -> Term
  -> Term
maybeSubstTm :: HasCallStack => Doc () -> Maybe Subst -> Term -> Term
maybeSubstTm Doc ()
_doc Maybe Subst
Nothing = Term -> Term
forall a. a -> a
id
maybeSubstTm Doc ()
doc (Just Subst
s) = HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm Doc ()
doc Subst
s

-- | Substitute within a 'Term'
substTm
  :: HasCallStack
  => Doc ()
  -> Subst
  -> Term
  -> Term
substTm :: HasCallStack => Doc () -> Subst -> Term -> Term
substTm Doc ()
doc Subst
subst = Term -> Term
go where
  go :: Term -> Term
go = \case
    Var 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
<> Doc ()
"subsTm") Subst
subst Id
v
    Lam Id
v Term
e -> case HasCallStack => Subst -> Id -> (Subst, Id)
Subst -> Id -> (Subst, Id)
substIdBndr Subst
subst Id
v of
      (Subst
subst',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 TyVar
v Term
e -> case HasCallStack => Subst -> TyVar -> (Subst, TyVar)
Subst -> TyVar -> (Subst, TyVar)
substTyVarBndr' Subst
subst TyVar
v of
      (Subst
subst',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 Term
l Term
r -> Term -> Term -> Term
App (Term -> Term
go Term
l) (Term -> Term
go Term
r)
    TyApp Term
l Type
r -> Term -> Type -> Term
TyApp (Term -> Term
go Term
l) (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst Type
r)
    Let Bind Term
bs Term
e -> case HasCallStack => Doc () -> Subst -> Bind Term -> (Subst, Bind Term)
Doc () -> Subst -> Bind Term -> (Subst, Bind Term)
substBind Doc ()
doc Subst
subst Bind Term
bs of
      (Subst
subst',Bind Term
bs') -> Bind Term -> Term -> Term
Let Bind Term
bs' (HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm Doc ()
doc Subst
subst' Term
e)
    Case Term
subj Type
ty [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 Term
e Type
t1 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 TickInfo
tick Term
e -> TickInfo -> Term -> Term
Tick (TickInfo -> TickInfo
goTick TickInfo
tick) (Term -> Term
go Term
e)
    Term
tm -> Term
tm

  goAlt :: Alt -> Alt
goAlt (Pat
pat,Term
alt) = case Pat
pat of
    DataPat DataCon
dc [TyVar]
tvs [Id]
ids -> case (Subst -> TyVar -> (Subst, TyVar))
-> Subst -> [TyVar] -> (Subst, [TyVar])
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
List.mapAccumL HasCallStack => Subst -> TyVar -> (Subst, TyVar)
Subst -> TyVar -> (Subst, TyVar)
substTyVarBndr' Subst
subst [TyVar]
tvs of
      (Subst
subst1,[TyVar]
tvs') -> case (Subst -> Id -> (Subst, Id)) -> Subst -> [Id] -> (Subst, [Id])
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
List.mapAccumL HasCallStack => Subst -> Id -> (Subst, Id)
Subst -> Id -> (Subst, Id)
substIdBndr Subst
subst1 [Id]
ids of
        (Subst
subst2,[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
pat,Term -> Term
go Term
alt)

  goTick :: TickInfo -> TickInfo
goTick t :: TickInfo
t@(SrcSpan SrcSpan
_)  = TickInfo
t
  goTick (NameMod NameMod
m Type
ty) = NameMod -> Type -> TickInfo
NameMod NameMod
m (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst Type
ty)
  goTick t :: TickInfo
t@TickInfo
DeDup        = TickInfo
t
  goTick t :: TickInfo
t@TickInfo
NoDeDup      = TickInfo
t
  goTick (Attributes Type
ty Term
tm) = Type -> Term -> TickInfo
Attributes (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst Type
ty) (Term -> Term
go Term
tm)

-- | Like 'substTm', but doesn't account for shadowing or free variable capture.
--
-- An example of shadowing: in @(x, \x -> x)@ with a substitution @x |-> 5@
-- 'unsafeSubstTm' will yield @(5, \x -> 5)@, whereas a safe substitution would
-- yield @(5, \x -> x)@.
--
-- An example of free variable capture: in @\x -> y@ and a substitution
-- @y |-> f x@, 'unsafeSubstTm' will yield @\x -> f x@, whereas a safe substitution
-- would yield @\x' -> f x@.
--
-- You should therefore only use this function if:
--
--   1. No binder in the term binds variables in the domain of the local
--      substitution.
--
--   2. No binder in the term has the same unique as a local, free variable of any
--      replacement term in the range of either substitution.
--
-- Both conditions hold if the term is deshadowed (see 'deShadowTerm') with
-- respect to an in-scope set that contains the domains of the substitutions
-- and the free variables of the replacement terms.
--
-- Note that global substituations are always safe. They never have local free
-- variables and cannot be introduced through a binder. I.e., neither rule 1
-- nor 2 applies.
unsafeSubstTm
  :: VarEnv Term
  -- ^ Substitution: global variable to replacement term
  -> VarEnv Term
  -- ^ Substitution: local variable to replacement term
  -> Term
  -- ^ Term to substitute in
  -> Term
unsafeSubstTm :: IdSubstEnv -> IdSubstEnv -> Term -> Term
unsafeSubstTm IdSubstEnv
globals IdSubstEnv
locals = \Term
term -> Term -> Maybe Term -> Term
forall a. a -> Maybe a -> a
fromMaybe Term
term (Term -> Maybe Term
go Term
term)
 where
  go :: Term -> Maybe Term
  go :: Term -> Maybe Term
go = \case
    Var Id
v
      | Id -> Bool
forall a. Var a -> Bool
isGlobalId Id
v -> Id -> IdSubstEnv -> Maybe Term
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv Id
v IdSubstEnv
globals
      | Bool
otherwise -> Id -> IdSubstEnv -> Maybe Term
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv Id
v IdSubstEnv
locals
    Lam Id
v Term
e -> Id -> Term -> Term
Lam Id
v (Term -> Term) -> Maybe Term -> Maybe Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Term -> Maybe Term
go Term
e
    TyLam TyVar
tv Term
e -> TyVar -> Term -> Term
TyLam TyVar
tv (Term -> Term) -> Maybe Term -> Maybe Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Term -> Maybe Term
go Term
e
    App Term
l Term
r -> case (Term -> Maybe Term
go Term
l, Term -> Maybe Term
go Term
r) of
      (Maybe Term
Nothing, Maybe Term
Nothing) -> Maybe Term
forall a. Maybe a
Nothing
      (Maybe Term
l1, Maybe Term
r1) -> Term -> Maybe Term
forall a. a -> Maybe a
Just (Term -> Term -> Term
App (Term -> Maybe Term -> Term
forall a. a -> Maybe a -> a
fromMaybe Term
l Maybe Term
l1) (Term -> Maybe Term -> Term
forall a. a -> Maybe a -> a
fromMaybe Term
r Maybe Term
r1))
    TyApp Term
e Type
ty -> (Term -> Type -> Term
`TyApp` Type
ty) (Term -> Term) -> Maybe Term -> Maybe Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Term -> Maybe Term
go Term
e
    Let Bind Term
bs Term
body -> case (Bind Term -> Maybe (Bind Term)
goBind Bind Term
bs, Term -> Maybe Term
go Term
body) of
      (Maybe (Bind Term)
Nothing, Maybe Term
Nothing) -> Maybe Term
forall a. Maybe a
Nothing
      (Maybe (Bind Term)
bs1, Maybe Term
body1) -> Term -> Maybe Term
forall a. a -> Maybe a
Just (Bind Term -> Term -> Term
Let (Bind Term -> Maybe (Bind Term) -> Bind Term
forall a. a -> Maybe a -> a
fromMaybe Bind Term
bs Maybe (Bind Term)
bs1) (Term -> Maybe Term -> Term
forall a. a -> Maybe a -> a
fromMaybe Term
body Maybe Term
body1))
    Case Term
subject Type
ty [Alt]
alternatives ->
      case (Term -> Maybe Term
go Term
subject, (Alt -> Maybe Alt) -> [Alt] -> Maybe [Alt]
forall a. (a -> Maybe a) -> [a] -> Maybe [a]
goList Alt -> Maybe Alt
forall {a}. (a, Term) -> Maybe (a, Term)
goAlternative [Alt]
alternatives) of
        (Maybe Term
Nothing, Maybe [Alt]
Nothing) -> Maybe Term
forall a. Maybe a
Nothing
        (Maybe Term
subject1, Maybe [Alt]
alternatives1) ->
          Term -> Maybe Term
forall a. a -> Maybe a
Just (Term -> Type -> [Alt] -> Term
Case (Term -> Maybe Term -> Term
forall a. a -> Maybe a -> a
fromMaybe Term
subject Maybe Term
subject1) Type
ty
                     ([Alt] -> Maybe [Alt] -> [Alt]
forall a. a -> Maybe a -> a
fromMaybe [Alt]
alternatives Maybe [Alt]
alternatives1))
    Cast Term
e Type
t1 Type
t2 -> (\Term
e1 -> Term -> Type -> Type -> Term
Cast Term
e1 Type
t1 Type
t2) (Term -> Term) -> Maybe Term -> Maybe Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Term -> Maybe Term
go Term
e
    Tick TickInfo
tickInfo Term
e -> case (TickInfo -> Maybe TickInfo
goTickInfo TickInfo
tickInfo, Term -> Maybe Term
go Term
e) of
      (Maybe TickInfo
Nothing, Maybe Term
Nothing) -> Maybe Term
forall a. Maybe a
Nothing
      (Maybe TickInfo
tick1, Maybe Term
e1) -> Term -> Maybe Term
forall a. a -> Maybe a
Just (TickInfo -> Term -> Term
Tick (TickInfo -> Maybe TickInfo -> TickInfo
forall a. a -> Maybe a -> a
fromMaybe TickInfo
tickInfo Maybe TickInfo
tick1) (Term -> Maybe Term -> Term
forall a. a -> Maybe a -> a
fromMaybe Term
e Maybe Term
e1))
    Data{} -> Maybe Term
forall a. Maybe a
Nothing
    Literal{} -> Maybe Term
forall a. Maybe a
Nothing
    Prim{} -> Maybe Term
forall a. Maybe a
Nothing

  goBind :: Bind Term -> Maybe (Bind Term)
goBind (NonRec Id
v Term
rhs) = Id -> Term -> Bind Term
forall a. Id -> a -> Bind a
NonRec Id
v (Term -> Bind Term) -> Maybe Term -> Maybe (Bind Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Term -> Maybe Term
go Term
rhs
  goBind (Rec [(Id, Term)]
bindings0) = [(Id, Term)] -> Bind Term
forall a. [(Id, a)] -> Bind a
Rec ([(Id, Term)] -> Bind Term)
-> Maybe [(Id, Term)] -> Maybe (Bind Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Id, Term) -> Maybe (Id, Term))
-> [(Id, Term)] -> Maybe [(Id, Term)]
forall a. (a -> Maybe a) -> [a] -> Maybe [a]
goList (Id, Term) -> Maybe (Id, Term)
forall {a}. (a, Term) -> Maybe (a, Term)
goBinding [(Id, Term)]
bindings0

  goBinding :: (a, Term) -> Maybe (a, Term)
goBinding (a
v, Term
rhs) = (,) a
v (Term -> (a, Term)) -> Maybe Term -> Maybe (a, Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Term -> Maybe Term
go Term
rhs

  goAlternative :: (a, Term) -> Maybe (a, Term)
goAlternative (a
pat, Term
alternative) = (,) a
pat (Term -> (a, Term)) -> Maybe Term -> Maybe (a, Term)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Term -> Maybe Term
go Term
alternative

  -- Types contain no term variables, so of the tick constructors only
  -- 'Attributes', which carries a term, needs a traversal.
  goTickInfo :: TickInfo -> Maybe TickInfo
goTickInfo (Attributes Type
ty Term
e) = Type -> Term -> TickInfo
Attributes Type
ty (Term -> TickInfo) -> Maybe Term -> Maybe TickInfo
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Term -> Maybe Term
go Term
e
  goTickInfo SrcSpan{} = Maybe TickInfo
forall a. Maybe a
Nothing
  goTickInfo NameMod{} = Maybe TickInfo
forall a. Maybe a
Nothing
  goTickInfo TickInfo
DeDup = Maybe TickInfo
forall a. Maybe a
Nothing
  goTickInfo TickInfo
NoDeDup = Maybe TickInfo
forall a. Maybe a
Nothing

  goList :: (a -> Maybe a) -> [a] -> Maybe [a]
  goList :: forall a. (a -> Maybe a) -> [a] -> Maybe [a]
goList a -> Maybe a
f = [a] -> Maybe [a]
goElements
   where
    goElements :: [a] -> Maybe [a]
goElements [] = Maybe [a]
forall a. Maybe a
Nothing
    goElements (a
x:[a]
xs) = case (a -> Maybe a
f a
x, [a] -> Maybe [a]
goElements [a]
xs) of
      (Maybe a
Nothing, Maybe [a]
Nothing) -> Maybe [a]
forall a. Maybe a
Nothing
      (Maybe a
x1, Maybe [a]
xs1) -> [a] -> Maybe [a]
forall a. a -> Maybe a
Just (a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe a
x Maybe a
x1 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> Maybe [a] -> [a]
forall a. a -> Maybe a -> a
fromMaybe [a]
xs Maybe [a]
xs1)

-- | Substitute within a case-alternative
substAlt
  :: HasCallStack
  => Doc ()
  -> Subst
  -- ^ The substitution
  -> (Pat, Term)
  -- ^ The alternative in which to apply the substitution
  -> (Pat, Term)
substAlt :: HasCallStack => Doc () -> Subst -> Alt -> Alt
substAlt Doc ()
doc Subst
subst (Pat
pat,Term
alt) = case Pat
pat of
  DataPat DataCon
dc [TyVar]
tvs [Id]
ids -> case (Subst -> TyVar -> (Subst, TyVar))
-> Subst -> [TyVar] -> (Subst, [TyVar])
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
List.mapAccumL HasCallStack => Subst -> TyVar -> (Subst, TyVar)
Subst -> TyVar -> (Subst, TyVar)
substTyVarBndr' Subst
subst [TyVar]
tvs of
    (Subst
subst1,[TyVar]
tvs1) -> case (Subst -> Id -> (Subst, Id)) -> Subst -> [Id] -> (Subst, [Id])
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
List.mapAccumL HasCallStack => Subst -> Id -> (Subst, Id)
Subst -> Id -> (Subst, Id)
substIdBndr Subst
subst1 [Id]
ids of
      (Subst
subst2,[Id]
ids1) -> (DataCon -> [TyVar] -> [Id] -> Pat
DataPat DataCon
dc [TyVar]
tvs1 [Id]
ids1,HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm Doc ()
doc Subst
subst2 Term
alt)
  Pat
_ -> (Pat
pat, HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm Doc ()
doc Subst
subst Term
alt)

substId
  :: HasCallStack
  => Subst
  -> Id
  -> Id
substId :: HasCallStack => Subst -> Id -> Id
substId Subst
subst Id
oldId = (Subst, Id) -> Id
forall a b. (a, b) -> b
snd ((Subst, Id) -> Id) -> (Subst, Id) -> Id
forall a b. (a -> b) -> a -> b
$ HasCallStack => Subst -> Id -> (Subst, Id)
Subst -> Id -> (Subst, Id)
substIdBndr Subst
subst Id
oldId

-- | Find the substitution for an 'Id' in the 'Subst'
lookupIdSubst
  :: HasCallStack
  => Doc ()
  -> Subst
  -> Id
  -> Term
lookupIdSubst :: HasCallStack => Doc () -> Subst -> Id -> Term
lookupIdSubst Doc ()
doc (Subst InScopeSet
inScope IdSubstEnv
tmS TvSubstEnv
_ IdSubstEnv
genv) 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 Term
e -> Term
e
                     Maybe Term
_      -> Id -> Term
Var Id
v
  | Just 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]
  --
  -- TODO: We match on Id here to workaround an issue where type variables
  -- TODO: "shadow" term variables. Omitting the check would make 'lookupIdSubst'
  -- TODO: potentially replace an "Id" with a TyVar. For more information:
  -- TODO:
  -- TODO:   https://github.com/clash-lang/clash-compiler/issues/1046
  -- TODO:
  | Just v' :: Var Any
v'@(Id {}) <- 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 :: HasCallStack => Subst -> Id -> (Subst, Id)
substIdBndr subst :: Subst
subst@(Subst InScopeSet
inScope IdSubstEnv
env TvSubstEnv
tenv IdSubstEnv
genv) 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. (Uniquable a, ClashPretty a) => InScopeSet -> a -> a
uniqAway InScopeSet
inScope Id
oldId
  newId :: Id
newId | Bool
noTypeChange = Id
id1
        | Bool
otherwise    = Id
id1 {varType = substTy subst (varType 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
forall a. HasFreeVars a => a -> Bool
isClosed 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' :: HasCallStack => Subst -> TyVar -> (Subst, TyVar)
substTyVarBndr' (Subst InScopeSet
inScope IdSubstEnv
tmS TvSubstEnv
tyS IdSubstEnv
tgS) TyVar
tv =
  case TvSubst -> TyVar -> (TvSubst, TyVar)
substTyVarBndr (InScopeSet -> TvSubstEnv -> TvSubst
TvSubst InScopeSet
inScope TvSubstEnv
tyS) TyVar
tv of
    (TvSubst InScopeSet
inScope' TvSubstEnv
tyS',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
  -> Bind Term
  -> (Subst, Bind Term)
substBind :: HasCallStack => Doc () -> Subst -> Bind Term -> (Subst, Bind Term)
substBind Doc ()
doc Subst
subst (NonRec Id
i Term
x) =
  (Subst
subst', Id -> Term -> Bind Term
forall a. Id -> a -> Bind a
NonRec Id
i' Term
x')
 where
  (Subst
subst', Id
i') = HasCallStack => Subst -> Id -> (Subst, Id)
Subst -> Id -> (Subst, Id)
substIdBndr Subst
subst Id
i
  x' :: Term
x' = HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm (Doc ()
"substBind" Doc () -> Doc () -> Doc ()
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ()
doc) Subst
subst Term
x

substBind Doc ()
doc Subst
subst (Rec [(Id, Term)]
xs) =
  (Subst
subst', [(Id, Term)] -> Bind Term
forall a. [(Id, a)] -> Bind a
Rec ([Id] -> [Term] -> [(Id, Term)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
bndrs' [Term]
rhss'))
 where
  ([Id]
bndrs,[Term]
rhss)    = [(Id, Term)] -> ([Id], [Term])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Id, Term)]
xs
  (Subst
subst',[Id]
bndrs') = (Subst -> Id -> (Subst, Id)) -> Subst -> [Id] -> (Subst, [Id])
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
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 (Doc ()
"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 :: HasCallStack => [TyVar] -> [Type] -> Type -> Type
substTyWith [TyVar]
tvs [Type]
tys =
  ASSERT( List.equalLength tvs 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 :: HasCallStack => InScopeSet -> Term -> Term
deShadowTerm InScopeSet
is Term
e = HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm Doc ()
"deShadowTerm" (InScopeSet -> Subst
mkSubst InScopeSet
is) Term
e

-- | Ensure that non of the binders in an alternative shadow each-other, nor
-- conflict with the in-scope set
deShadowAlt ::
  HasCallStack =>
  InScopeSet ->
  (Pat, Term) ->
  (Pat, Term)
deShadowAlt :: HasCallStack => InScopeSet -> Alt -> Alt
deShadowAlt InScopeSet
is = HasCallStack => Doc () -> Subst -> Alt -> Alt
Doc () -> Subst -> Alt -> Alt
substAlt Doc ()
"deShadowAlt" (InScopeSet -> Subst
mkSubst InScopeSet
is)

-- | Ensure that non of the let-bindings of a let-expression shadow w.r.t the
-- in-scope set
deshadowLetExpr
  :: HasCallStack
  => InScopeSet
  -- ^ Current InScopeSet
  -> Bind Term
  -- ^ Bindings of the let-expression
  -> Term
  -- ^ The body of the let-expression
  -> (Bind Term, Term)
  -- ^ Deshadowed let-bindings, where let-bound expressions and the let-body
  -- properly reference the renamed variables
deshadowLetExpr :: HasCallStack =>
InScopeSet -> Bind Term -> Term -> (Bind Term, Term)
deshadowLetExpr InScopeSet
is Bind Term
bs Term
e =
  case HasCallStack => Doc () -> Subst -> Bind Term -> (Subst, Bind Term)
Doc () -> Subst -> Bind Term -> (Subst, Bind Term)
substBind Doc ()
"deshadowLetBindings" (InScopeSet -> Subst
mkSubst InScopeSet
is) Bind Term
bs of
    (Subst
s1,Bind Term
bs1) -> (Bind Term
bs1, HasCallStack => Doc () -> Subst -> Term -> Term
Doc () -> Subst -> Term -> Term
substTm Doc ()
"deShadowLetBody" Subst
s1 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 InScopeSet
is0 = Subst -> Term -> (InScopeSet, Term)
go (InScopeSet -> Subst
mkSubst InScopeSet
is0) where
  go :: Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 = \case
    Var Id
v -> (Subst -> InScopeSet
substInScope Subst
subst0, HasCallStack => Doc () -> Subst -> Id -> Term
Doc () -> Subst -> Id -> Term
lookupIdSubst Doc ()
"freshenTm" Subst
subst0 Id
v)
    Lam Id
v Term
e -> case HasCallStack => Subst -> Id -> (Subst, Id)
Subst -> Id -> (Subst, Id)
substIdBndr Subst
subst0 Id
v of
      (Subst
subst1,Id
v') -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst1 Term
e of
        (InScopeSet
is2,Term
e') -> (InScopeSet
is2, Id -> Term -> Term
Lam Id
v' Term
e')
    TyLam TyVar
v Term
e -> case HasCallStack => Subst -> TyVar -> (Subst, TyVar)
Subst -> TyVar -> (Subst, TyVar)
substTyVarBndr' Subst
subst0 TyVar
v of
      (Subst
subst1,TyVar
v') -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst1 Term
e of
        (InScopeSet
is2,Term
e') -> (InScopeSet
is2,TyVar -> Term -> Term
TyLam TyVar
v' Term
e')
    App Term
l Term
r -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 Term
l of
      (InScopeSet
is1,Term
l') -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 {substInScope = is1} Term
r of
        (InScopeSet
is2,Term
r') -> (InScopeSet
is2, Term -> Term -> Term
App Term
l' Term
r')
    TyApp Term
l Type
r -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 Term
l of
      (InScopeSet
is1,Term
l') -> (InScopeSet
is1, Term -> Type -> Term
TyApp Term
l' (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst0 Type
r))
    Let Bind Term
bs Term
e -> case Subst -> Bind Term -> (Subst, Bind Term)
goBind Subst
subst0 Bind Term
bs of
      (Subst
subst1,Bind Term
bs') -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst1 Term
e of
        (InScopeSet
is2,Term
e') -> (InScopeSet
is2,Bind Term -> Term -> Term
Let Bind Term
bs' Term
e')
    Case Term
subj Type
ty [Alt]
alts -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 Term
subj of
      (InScopeSet
is1,Term
subj') -> case (InScopeSet -> Alt -> (InScopeSet, Alt))
-> InScopeSet -> [Alt] -> (InScopeSet, [Alt])
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
List.mapAccumL (\InScopeSet
isN -> Subst -> Alt -> (InScopeSet, Alt)
goAlt Subst
subst0 {substInScope = isN}) InScopeSet
is1 [Alt]
alts of
        (InScopeSet
is2,[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 Term
e Type
t1 Type
t2 -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 Term
e of
      (InScopeSet
is1, 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 TickInfo
tick Term
e -> case Subst -> TickInfo -> (InScopeSet, TickInfo)
goTick Subst
subst0 TickInfo
tick of
      (InScopeSet
is1, TickInfo
tick') -> case Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 {substInScope = is1} Term
e of
        (InScopeSet
is2, Term
e') -> (InScopeSet
is2, TickInfo -> Term -> Term
Tick TickInfo
tick' Term
e')
    Term
tm -> (Subst -> InScopeSet
substInScope Subst
subst0, Term
tm)

  goBind :: Subst -> Bind Term -> (Subst, Bind Term)
goBind Subst
subst0 (NonRec Id
i Term
x) =
    let (Subst
subst1, Id
i') = HasCallStack => Subst -> Id -> (Subst, Id)
Subst -> Id -> (Subst, Id)
substIdBndr Subst
subst0 Id
i
        (InScopeSet
is2, Term
x') = Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 Term
x
     in (Subst
subst1 { substInScope = extendInScopeSet is2 i' }, Id -> Term -> Bind Term
forall a. Id -> a -> Bind a
NonRec Id
i' Term
x')

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

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

  goTick :: Subst -> TickInfo -> (InScopeSet, TickInfo)
goTick Subst
subst0 t :: TickInfo
t@(SrcSpan SrcSpan
_) = (Subst -> InScopeSet
substInScope Subst
subst0, TickInfo
t)
  goTick Subst
subst0 (NameMod NameMod
m Type
ty) =
    (Subst -> InScopeSet
substInScope Subst
subst0, NameMod -> Type -> TickInfo
NameMod NameMod
m (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst0 Type
ty))
  goTick Subst
subst0 t :: TickInfo
t@TickInfo
DeDup = (Subst -> InScopeSet
substInScope Subst
subst0, TickInfo
t)
  goTick Subst
subst0 t :: TickInfo
t@TickInfo
NoDeDup = (Subst -> InScopeSet
substInScope Subst
subst0, TickInfo
t)
  goTick Subst
subst0 (Attributes Type
ty Term
tm) = case Subst -> Term -> (InScopeSet, Term)
go Subst
subst0 Term
tm of
    (InScopeSet
is1, Term
tm') -> (InScopeSet
is1, Type -> Term -> TickInfo
Attributes (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst0 Type
ty) Term
tm')

-- * AEQ

-- | Alpha equality for types
aeqType
  :: Type
  -> Type
  -> Bool
aeqType :: Type -> Type -> Bool
aeqType Type
t1 Type
t2 = Type -> Type -> Ordering
acmpType Type
t1 Type
t2 Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ
{-# INLINE aeqType #-}

-- | Alpha comparison for types
acmpType
  :: Type
  -> Type
  -> Ordering
acmpType :: Type -> Type -> Ordering
acmpType = Int -> VarEnv Int -> VarEnv Int -> Type -> Type -> Ordering
acmpTypeLevels Int
0 VarEnv Int
forall a. VarEnv a
emptyVarEnv VarEnv Int
forall a. VarEnv a
emptyVarEnv

-- | Structural equality on 'Type'
eqType
  :: Type
  -> Type
  -> Bool
eqType :: Type -> Type -> Bool
eqType = Type -> Type -> Bool
go
 where
  go :: Type -> Type -> Bool
go (VarTy TyVar
tv1) (VarTy TyVar
tv2) = TyVar
tv1 TyVar -> TyVar -> Bool
forall a. Eq a => a -> a -> Bool
== TyVar
tv2
  go (ConstTy ConstTy
c1) (ConstTy ConstTy
c2) = ConstTy
c1 ConstTy -> ConstTy -> Bool
forall a. Eq a => a -> a -> Bool
== ConstTy
c2
  go (ForAllTy TyVar
tv1 Type
t1) (ForAllTy TyVar
tv2 Type
t2) =
    TyVar
tv1 TyVar -> TyVar -> Bool
forall a. Eq a => a -> a -> Bool
== TyVar
tv2 Bool -> Bool -> Bool
&& Type -> Type -> Bool
go (TyVar -> Type
forall a. Var a -> Type
varType TyVar
tv1) (TyVar -> Type
forall a. Var a -> Type
varType TyVar
tv2) Bool -> Bool -> Bool
&& Type -> Type -> Bool
go Type
t1 Type
t2
  go (AppTy Type
s1 Type
t1) (AppTy Type
s2 Type
t2) = Type -> Type -> Bool
go Type
s1 Type
s2 Bool -> Bool -> Bool
&& Type -> Type -> Bool
go Type
t1 Type
t2
  go (LitTy LitTy
l1) (LitTy LitTy
l2) = LitTy
l1 LitTy -> LitTy -> Bool
forall a. Eq a => a -> a -> Bool
== LitTy
l2
  go (AnnType [Attr Text]
a1 Type
t1) (AnnType [Attr Text]
a2 Type
t2) = [Attr Text]
a1 [Attr Text] -> [Attr Text] -> Bool
forall a. Eq a => a -> a -> Bool
== [Attr Text]
a2 Bool -> Bool -> Bool
&& Type -> Type -> Bool
go Type
t1 Type
t2
  go Type
_ Type
_ = Bool
False

-- | Alpha equality for terms
aeqTerm
  :: Term
  -> Term
  -> Bool
aeqTerm :: Term -> Term -> Bool
aeqTerm Term
t1 Term
t2 = Term -> Term -> Ordering
acmpTerm Term
t1 Term
t2 Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ
{-# INLINE aeqTerm #-}

-- | Alpha comparison for terms
acmpTerm
  :: Term
  -> Term
  -> Ordering
acmpTerm :: Term -> Term -> Ordering
acmpTerm =
  Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> Term
-> Term
-> Ordering
acmpTermLevels Int
0 VarEnv Int
forall a. VarEnv a
emptyVarEnv VarEnv Int
forall a. VarEnv a
emptyVarEnv VarEnv Int
forall a. VarEnv a
emptyVarEnv VarEnv Int
forall a. VarEnv a
emptyVarEnv

-- | Alpha equality for ticks
aeqTickInfo
  :: TickInfo
  -> TickInfo
  -> Bool
aeqTickInfo :: TickInfo -> TickInfo -> Bool
aeqTickInfo TickInfo
t1 TickInfo
t2 = TickInfo -> TickInfo -> Ordering
acmpTickInfo TickInfo
t1 TickInfo
t2 Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ
{-# INLINE aeqTickInfo #-}

-- | Alpha comparison for ticks
acmpTickInfo
  :: TickInfo
  -> TickInfo
  -> Ordering
acmpTickInfo :: TickInfo -> TickInfo -> Ordering
acmpTickInfo =
  Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> TickInfo
-> TickInfo
-> Ordering
acmpTickInfoLevels Int
0 VarEnv Int
forall a. VarEnv a
emptyVarEnv VarEnv Int
forall a. VarEnv a
emptyVarEnv VarEnv Int
forall a. VarEnv a
emptyVarEnv VarEnv Int
forall a. VarEnv a
emptyVarEnv

-- | Structural equality on 'Term'
eqTerm :: Term -> Term -> Bool
eqTerm :: Term -> Term -> Bool
eqTerm = Term -> Term -> Bool
go
 where
  go :: Term -> Term -> Bool
go (Var Id
id1) (Var Id
id2) = Id
id1 Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
id2
  go (Data DataCon
dc1) (Data DataCon
dc2) = DataCon
dc1 DataCon -> DataCon -> Bool
forall a. Eq a => a -> a -> Bool
== DataCon
dc2
  go (Literal Literal
l1) (Literal Literal
l2) = Literal
l1 Literal -> Literal -> Bool
forall a. Eq a => a -> a -> Bool
== Literal
l2
  go (Prim PrimInfo
p1) (Prim PrimInfo
p2) = PrimInfo -> Text
primName PrimInfo
p1 Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== PrimInfo -> Text
primName PrimInfo
p2
  go (Lam Id
b1 Term
e1) (Lam Id
b2 Term
e2) =
    Id
b1 Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
b2 Bool -> Bool -> Bool
&& Type -> Type -> Bool
eqType (Id -> Type
forall a. Var a -> Type
varType Id
b1) (Id -> Type
forall a. Var a -> Type
varType Id
b2) Bool -> Bool -> Bool
&& Term -> Term -> Bool
go Term
e1 Term
e2
  go (TyLam TyVar
b1 Term
e1) (TyLam TyVar
b2 Term
e2) =
    TyVar
b1 TyVar -> TyVar -> Bool
forall a. Eq a => a -> a -> Bool
== TyVar
b2 Bool -> Bool -> Bool
&& Type -> Type -> Bool
eqType (TyVar -> Type
forall a. Var a -> Type
varType TyVar
b1) (TyVar -> Type
forall a. Var a -> Type
varType TyVar
b2) Bool -> Bool -> Bool
&& Term -> Term -> Bool
go Term
e1 Term
e2
  go (App Term
l1 Term
r1) (App Term
l2 Term
r2) = Term -> Term -> Bool
go Term
l1 Term
l2 Bool -> Bool -> Bool
&& Term -> Term -> Bool
go Term
r1 Term
r2
  go (TyApp Term
l1 Type
r1) (TyApp Term
l2 Type
r2) = Term -> Term -> Bool
go Term
l1 Term
l2 Bool -> Bool -> Bool
&& Type -> Type -> Bool
eqType Type
r1 Type
r2
  go (Let Bind Term
bs1 Term
e1) (Let Bind Term
bs2 Term
e2) =
    Term -> Term -> Bool
go Term
e1 Term
e2 Bool -> Bool -> Bool
&&
    Bind Term -> Bind Term -> Bool
goBind Bind Term
bs1 Bind Term
bs2
   where
    goBind :: Bind Term -> Bind Term -> Bool
goBind (NonRec Id
b1 Term
r1) (NonRec Id
b2 Term
r2) =
      -- No need to check types of NonRec bindings, when the RHSs match the
      -- types must be the same.
      Id
b1 Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
b2 Bool -> Bool -> Bool
&& Term -> Term -> Bool
go Term
r1 Term
r2
    goBind (Rec [(Id, Term)]
brs1) (Rec [(Id, Term)]
brs2) =
      ((Id, Term) -> (Id, Term) -> Bool)
-> [(Id, Term)] -> [(Id, Term)] -> Bool
forall a b. (a -> b -> Bool) -> [a] -> [b] -> Bool
List.all2
        (\(Id
b1,Term
r1) (Id
b2,Term
r2) ->
          Id
b1 Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
b2 Bool -> Bool -> Bool
&&
          -- We need to check the types of Rec bindings, because:
          --
          -- letrec (x : Bool) = x in X
          --
          -- is not structurally equivalent to
          --
          -- letrec (x : Int) = x in x
          Type -> Type -> Bool
eqType (Id -> Type
forall a. Var a -> Type
varType Id
b1) (Id -> Type
forall a. Var a -> Type
varType Id
b2) Bool -> Bool -> Bool
&&
          Term -> Term -> Bool
go Term
r1 Term
r2)
        [(Id, Term)]
brs1 [(Id, Term)]
brs2
    goBind Bind Term
_ Bind Term
_ = Bool
False
  -- Note [Case result types and alpha-equivalence]
  go (Case Term
e1 Type
_ [Alt]
a1) (Case Term
e2 Type
_ [Alt]
a2) =
    Term -> Term -> Bool
go Term
e1 Term
e2 Bool -> Bool -> Bool
&&
    (Alt -> Alt -> Bool) -> [Alt] -> [Alt] -> Bool
forall a b. (a -> b -> Bool) -> [a] -> [b] -> Bool
List.all2 Alt -> Alt -> Bool
forall {a}. Eq a => (a, Term) -> (a, Term) -> Bool
goAlt [Alt]
a1 [Alt]
a2
   where
    goAlt :: (a, Term) -> (a, Term) -> Bool
goAlt (a
p1,Term
r1) (a
p2,Term
r2) = a
p1 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
p2 Bool -> Bool -> Bool
&& Term -> Term -> Bool
go Term
r1 Term
r2
  go (Cast Term
e1 Type
l1 Type
r1) (Cast Term
e2 Type
l2 Type
r2) =
    Term -> Term -> Bool
go Term
e1 Term
e2 Bool -> Bool -> Bool
&&
    Type -> Type -> Bool
eqType Type
l1 Type
l2 Bool -> Bool -> Bool
&&
    Type -> Type -> Bool
eqType Type
r1 Type
r2
  go (Tick TickInfo
t1 Term
e1) (Tick TickInfo
t2 Term
e2) = TickInfo -> TickInfo -> Bool
goTick TickInfo
t1 TickInfo
t2 Bool -> Bool -> Bool
&& Term -> Term -> Bool
go Term
e1 Term
e2
  go Term
_ Term
_ = Bool
False

  -- @Eq TickInfo@ compares 'NameMod' and 'Attributes' with alpha-equivalence,
  -- which is too coarse here: this is structural equality.
  goTick :: TickInfo -> TickInfo -> Bool
goTick (SrcSpan SrcSpan
s1) (SrcSpan SrcSpan
s2) = SrcSpan
s1 SrcSpan -> SrcSpan -> Bool
forall a. Eq a => a -> a -> Bool
== SrcSpan
s2
  goTick (NameMod NameMod
m1 Type
t1) (NameMod NameMod
m2 Type
t2) = NameMod
m1 NameMod -> NameMod -> Bool
forall a. Eq a => a -> a -> Bool
== NameMod
m2 Bool -> Bool -> Bool
&& Type -> Type -> Bool
eqType Type
t1 Type
t2
  goTick TickInfo
DeDup TickInfo
DeDup = Bool
True
  goTick TickInfo
NoDeDup TickInfo
NoDeDup = Bool
True
  goTick (Attributes Type
t1 Term
a1) (Attributes Type
t2 Term
a2) = Type -> Type -> Bool
eqType Type
t1 Type
t2 Bool -> Bool -> Bool
&& Term -> Term -> Bool
go Term
a1 Term
a2
  goTick TickInfo
_ TickInfo
_ = Bool
False

{- Note [Case result types and alpha-equivalence]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
'acmpTermLevels' (and hence 'aeqTerm', 'acmpTerm', @Eq Term@ and @Ord Term@)
does not compare the result type stored in a 'Case' constructor, and 'eqTerm'
follows it. This is OK for case expressions with at least one alternative: the
result type is determined by the alternatives. Every alternative's right-hand
side has exactly the result type, and (alpha-)equal right-hand sides have
(alpha-)equal types, so comparing the result type as well would be redundant
work.

For a case expression with no alternatives the result type is /not/ determined
by the subterms, so @(case x of {}) :: A@ and @(case x of {}) :: B@ compare
equal even though their types differ. Such case expressions don't occur in
CoreHW: @Clash.GHC.GHC2Core@ turns them into an @undefined@ or @undefinedX@.
-}

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

{- Note [Numbering binders by De Bruijn level]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Alpha comparison compares two variable occurrences by their De Bruijn level:
how many binders enclose the binder an occurrence resolves to. Both sides are
walked in lockstep, so at any point they are under equally many binders, and a
level therefore identifies a binder position rather than a name. Two
occurrences are alpha-equal exactly when they resolve to the same level.

Binders that come into scope together, such as those of a 'Rec' or a 'DataPat',
get consecutive levels: they are numbered as if they were nested.
-}

-- | Compare a pair of variable occurrences, under the levels of the binders
-- enclosing them. See Note [Numbering binders by De Bruijn level].
acmpOccLevels ::
  -- | Levels of the binders enclosing the left occurrence
  VarEnv Int ->
  -- | Levels of the binders enclosing the right occurrence
  VarEnv Int ->
  Var a ->
  Var a ->
  Ordering
acmpOccLevels :: forall a. VarEnv Int -> VarEnv Int -> Var a -> Var a -> Ordering
acmpOccLevels VarEnv Int
envL VarEnv Int
envR Var a
v1 Var a
v2 =
  case (Var a -> VarEnv Int -> Maybe Int
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv Var a
v1 VarEnv Int
envL, Var a -> VarEnv Int -> Maybe Int
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv Var a
v2 VarEnv Int
envR) of
    -- Both bound: equal exactly when bound at the same level
    (Just Int
lvlL, Just Int
lvlR) -> Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
lvlL Int
lvlR
    -- Neither bound: compare the variables themselves
    (Maybe Int
Nothing, Maybe Int
Nothing) -> (Unique, Maybe IdScope) -> (Unique, Maybe IdScope) -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Var a -> (Unique, Maybe IdScope)
forall a. Var a -> (Unique, Maybe IdScope)
varKey Var a
v1) (Var a -> (Unique, Maybe IdScope)
forall a. Var a -> (Unique, Maybe IdScope)
varKey Var a
v2)
    -- A bound variable is never equal to a free one. Which way round is
    -- arbitrary, it only has to be consistent to keep the order total.
    (Just Int
_, Maybe Int
Nothing) -> Ordering
LT
    (Maybe Int
Nothing, Just Int
_) -> Ordering
GT

-- | Give a group of binders that come into scope together, such as a 'Rec' or
-- a 'DataPat', consecutive levels, and return the level following the group.
extendLevels ::
  Int ->
  [Var a] ->
  [Var a] ->
  VarEnv Int ->
  VarEnv Int ->
  (Int, VarEnv Int, VarEnv Int)
extendLevels :: forall a.
Int
-> [Var a]
-> [Var a]
-> VarEnv Int
-> VarEnv Int
-> (Int, VarEnv Int, VarEnv Int)
extendLevels Int
lvl [Var a]
vs1 [Var a]
vs2 VarEnv Int
envL VarEnv Int
envR =
  ((Int, VarEnv Int, VarEnv Int)
 -> (Var a, Var a) -> (Int, VarEnv Int, VarEnv Int))
-> (Int, VarEnv Int, VarEnv Int)
-> [(Var a, Var a)]
-> (Int, VarEnv Int, VarEnv Int)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (Int, VarEnv Int, VarEnv Int)
-> (Var a, Var a) -> (Int, VarEnv Int, VarEnv Int)
forall {a} {b} {b}.
Num a =>
(a, VarEnv a, VarEnv a)
-> (Var b, Var b) -> (a, VarEnv a, VarEnv a)
one (Int
lvl, VarEnv Int
envL, VarEnv Int
envR) ([Var a] -> [Var a] -> [(Var a, Var a)]
forall a b. HasCallStack => [a] -> [b] -> [(a, b)]
List.zipEqual [Var a]
vs1 [Var a]
vs2)
 where
  one :: (a, VarEnv a, VarEnv a)
-> (Var b, Var b) -> (a, VarEnv a, VarEnv a)
one (!a
l, VarEnv a
eL, VarEnv a
eR) (Var b
v1, Var b
v2) =
    (a
l a -> a -> a
forall a. Num a => a -> a -> a
+ a
1, Var b -> a -> VarEnv a -> VarEnv a
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv Var b
v1 a
l VarEnv a
eL, Var b -> a -> VarEnv a -> VarEnv a
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv Var b
v2 a
l VarEnv a
eR)

-- | Alpha comparison for 'Type's, under the binders enclosing them.
-- See Note [Numbering binders by De Bruijn level].
acmpTypeLevels ::
  -- | Number of enclosing binders
  Int ->
  -- | Levels of the type binders enclosing the left type
  VarEnv Int ->
  -- | Levels of the type binders enclosing the right type
  VarEnv Int ->
  Type ->
  Type ->
  Ordering
acmpTypeLevels :: Int -> VarEnv Int -> VarEnv Int -> Type -> Type -> Ordering
acmpTypeLevels !Int
lvl VarEnv Int
tyL VarEnv Int
tyR = Type -> Type -> Ordering
go
 where
  go :: Type -> Type -> Ordering
  go :: Type -> Type -> Ordering
go (VarTy TyVar
tv1) (VarTy TyVar
tv2) = VarEnv Int -> VarEnv Int -> TyVar -> TyVar -> Ordering
forall a. VarEnv Int -> VarEnv Int -> Var a -> Var a -> Ordering
acmpOccLevels VarEnv Int
tyL VarEnv Int
tyR TyVar
tv1 TyVar
tv2
  go (ConstTy ConstTy
c1) (ConstTy ConstTy
c2) = ConstTy -> ConstTy -> Ordering
forall a. Ord a => a -> a -> Ordering
compare ConstTy
c1 ConstTy
c2
  go (ForAllTy TyVar
tv1 Type
t1) (ForAllTy TyVar
tv2 Type
t2) =
    Type -> Type -> Ordering
go (TyVar -> Type
forall a. Var a -> Type
varType TyVar
tv1) (TyVar -> Type
forall a. Var a -> Type
varType TyVar
tv2) Ordering -> Ordering -> Ordering
`thenCompare`
      Int -> VarEnv Int -> VarEnv Int -> Type -> Type -> Ordering
acmpTypeLevels
        (Int
lvl Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (TyVar -> Int -> VarEnv Int -> VarEnv Int
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv TyVar
tv1 Int
lvl VarEnv Int
tyL) (TyVar -> Int -> VarEnv Int -> VarEnv Int
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv TyVar
tv2 Int
lvl VarEnv Int
tyR) Type
t1 Type
t2
  go (AppTy Type
s1 Type
t1) (AppTy Type
s2 Type
t2) = Type -> Type -> Ordering
go Type
s1 Type
s2 Ordering -> Ordering -> Ordering
`thenCompare` Type -> Type -> Ordering
go Type
t1 Type
t2
  go (LitTy LitTy
l1) (LitTy LitTy
l2) = LitTy -> LitTy -> Ordering
forall a. Ord a => a -> a -> Ordering
compare LitTy
l1 LitTy
l2
  go (AnnType [Attr Text]
a1 Type
t1) (AnnType [Attr Text]
a2 Type
t2) = [Attr Text] -> [Attr Text] -> Ordering
forall a. Ord a => a -> a -> Ordering
compare [Attr Text]
a1 [Attr Text]
a2 Ordering -> Ordering -> Ordering
`thenCompare` Type -> Type -> Ordering
go Type
t1 Type
t2
  go Type
t1 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 {})    = Word
0
  getRank (LitTy {})    = Word
1
  getRank (ConstTy {})  = Word
2
  getRank (AnnType {})  = Word
3
  getRank (AppTy {})    = Word
4
  getRank (ForAllTy {}) = Word
5

-- | Alpha comparison for 'Term's. See 'acmpTypeLevels'.
acmpTermLevels ::
  -- | Number of enclosing binders
  Int ->
  -- | Levels of the term binders enclosing the left term
  VarEnv Int ->
  -- | Levels of the term binders enclosing the right term
  VarEnv Int ->
  -- | Levels of the type binders enclosing the left term
  VarEnv Int ->
  -- | Levels of the type binders enclosing the right term
  VarEnv Int ->
  Term ->
  Term ->
  Ordering
acmpTermLevels :: Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> Term
-> Term
-> Ordering
acmpTermLevels !Int
lvl VarEnv Int
tmL VarEnv Int
tmR VarEnv Int
tyL VarEnv Int
tyR = Term -> Term -> Ordering
go
 where
  goType :: Type -> Type -> Ordering
goType = Int -> VarEnv Int -> VarEnv Int -> Type -> Type -> Ordering
acmpTypeLevels Int
lvl VarEnv Int
tyL VarEnv Int
tyR
  goTick :: TickInfo -> TickInfo -> Ordering
goTick = Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> TickInfo
-> TickInfo
-> Ordering
acmpTickInfoLevels Int
lvl VarEnv Int
tmL VarEnv Int
tmR VarEnv Int
tyL VarEnv Int
tyR

  -- Compare underneath one more term or type binder
  underTmBndr :: Var b -> Var b -> Term -> Term -> Ordering
underTmBndr Var b
b1 Var b
b2 =
    Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> Term
-> Term
-> Ordering
acmpTermLevels
      (Int
lvl Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Var b -> Int -> VarEnv Int -> VarEnv Int
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv Var b
b1 Int
lvl VarEnv Int
tmL) (Var b -> Int -> VarEnv Int -> VarEnv Int
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv Var b
b2 Int
lvl VarEnv Int
tmR) VarEnv Int
tyL VarEnv Int
tyR
  underTyBndr :: Var b -> Var b -> Term -> Term -> Ordering
underTyBndr Var b
b1 Var b
b2 =
    Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> Term
-> Term
-> Ordering
acmpTermLevels
      (Int
lvl Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) VarEnv Int
tmL VarEnv Int
tmR (Var b -> Int -> VarEnv Int -> VarEnv Int
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv Var b
b1 Int
lvl VarEnv Int
tyL) (Var b -> Int -> VarEnv Int -> VarEnv Int
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv Var b
b2 Int
lvl VarEnv Int
tyR)

  go :: Term -> Term -> Ordering
  go :: Term -> Term -> Ordering
go (Var Id
id1) (Var Id
id2) = Id -> Id -> Ordering
goVar Id
id1 Id
id2
  go (Data DataCon
dc1) (Data DataCon
dc2) = DataCon -> DataCon -> Ordering
forall a. Ord a => a -> a -> Ordering
compare DataCon
dc1 DataCon
dc2
  go (Literal Literal
l1) (Literal Literal
l2) = Literal -> Literal -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Literal
l1 Literal
l2
  go (Prim PrimInfo
p1) (Prim PrimInfo
p2) = (PrimInfo -> Text) -> PrimInfo -> PrimInfo -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing PrimInfo -> Text
primName PrimInfo
p1 PrimInfo
p2
  go (Lam Id
b1 Term
e1) (Lam Id
b2 Term
e2) =
    Type -> Type -> Ordering
goType (Id -> Type
forall a. Var a -> Type
varType Id
b1) (Id -> Type
forall a. Var a -> Type
varType Id
b2) Ordering -> Ordering -> Ordering
`thenCompare` Id -> Id -> Term -> Term -> Ordering
forall {b} {b}. Var b -> Var b -> Term -> Term -> Ordering
underTmBndr Id
b1 Id
b2 Term
e1 Term
e2
  go (TyLam TyVar
b1 Term
e1) (TyLam TyVar
b2 Term
e2) =
    Type -> Type -> Ordering
goType (TyVar -> Type
forall a. Var a -> Type
varType TyVar
b1) (TyVar -> Type
forall a. Var a -> Type
varType TyVar
b2) Ordering -> Ordering -> Ordering
`thenCompare` TyVar -> TyVar -> Term -> Term -> Ordering
forall {b} {b}. Var b -> Var b -> Term -> Term -> Ordering
underTyBndr TyVar
b1 TyVar
b2 Term
e1 Term
e2
  go (App Term
l1 Term
r1) (App Term
l2 Term
r2) = Term -> Term -> Ordering
go Term
l1 Term
l2 Ordering -> Ordering -> Ordering
`thenCompare` Term -> Term -> Ordering
go Term
r1 Term
r2
  go (TyApp Term
l1 Type
r1) (TyApp Term
l2 Type
r2) = Term -> Term -> Ordering
go Term
l1 Term
l2 Ordering -> Ordering -> Ordering
`thenCompare` Type -> Type -> Ordering
goType Type
r1 Type
r2
  go (Let (NonRec Id
i1 Term
x1) Term
e1) (Let (NonRec Id
i2 Term
x2) Term
e2) =
    Term -> Term -> Ordering
go Term
x1 Term
x2 Ordering -> Ordering -> Ordering
`thenCompare` Id -> Id -> Term -> Term -> Ordering
forall {b} {b}. Var b -> Var b -> Term -> Term -> Ordering
underTmBndr Id
i1 Id
i2 Term
e1 Term
e2
  go (Let (Rec [(Id, Term)]
bs1) Term
e1) (Let (Rec [(Id, Term)]
bs2) Term
e2) =
    -- The lengths are compared first: the binder lists are only numbered
    -- against each other when they match
    Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare ([(Id, Term)] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [(Id, Term)]
bs1) ([(Id, Term)] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [(Id, Term)]
bs2) Ordering -> Ordering -> Ordering
`thenCompare`
      let ([Id]
ids1, [Term]
rhs1) = [(Id, Term)] -> ([Id], [Term])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Id, Term)]
bs1
          ([Id]
ids2, [Term]
rhs2) = [(Id, Term)] -> ([Id], [Term])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Id, Term)]
bs2
          (Int
lvl', VarEnv Int
tmL', VarEnv Int
tmR') = Int
-> [Id]
-> [Id]
-> VarEnv Int
-> VarEnv Int
-> (Int, VarEnv Int, VarEnv Int)
forall a.
Int
-> [Var a]
-> [Var a]
-> VarEnv Int
-> VarEnv Int
-> (Int, VarEnv Int, VarEnv Int)
extendLevels Int
lvl [Id]
ids1 [Id]
ids2 VarEnv Int
tmL VarEnv Int
tmR
          under :: Term -> Term -> Ordering
under = Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> Term
-> Term
-> Ordering
acmpTermLevels Int
lvl' VarEnv Int
tmL' VarEnv Int
tmR' VarEnv Int
tyL VarEnv Int
tyR
      -- Note that we compare types, because:
      --
      --   let (x :: Int) = x in x
      --
      -- is not alpha equivalent to:
      --
      --   let (x :: Word) = x in x
      --
      in (Type -> Type -> Ordering) -> [Type] -> [Type] -> Ordering
forall a. (a -> a -> Ordering) -> [a] -> [a] -> Ordering
goList Type -> Type -> Ordering
goType ((Id -> Type) -> [Id] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Id -> Type
forall a. Var a -> Type
varType [Id]
ids1) ((Id -> Type) -> [Id] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Id -> Type
forall a. Var a -> Type
varType [Id]
ids2) Ordering -> Ordering -> Ordering
`thenCompare`
           ((Term -> Term -> Ordering) -> [Term] -> [Term] -> Ordering
forall a. (a -> a -> Ordering) -> [a] -> [a] -> Ordering
goList Term -> Term -> Ordering
under [Term]
rhs1 [Term]
rhs2 Ordering -> Ordering -> Ordering
`thenCompare` Term -> Term -> Ordering
under Term
e1 Term
e2)
  -- Note [Case result types and alpha-equivalence]
  go (Case Term
e1 Type
_ [Alt]
a1) (Case Term
e2 Type
_ [Alt]
a2) =
    Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare ([Alt] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Alt]
a1) ([Alt] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Alt]
a2) Ordering -> Ordering -> Ordering
`thenCompare`
      (Term -> Term -> Ordering
go Term
e1 Term
e2 Ordering -> Ordering -> Ordering
`thenCompare` [Alt] -> [Alt] -> Ordering
goAlts [Alt]
a1 [Alt]
a2)
  go (Cast Term
e1 Type
l1 Type
r1) (Cast Term
e2 Type
l2 Type
r2) =
    Term -> Term -> Ordering
go Term
e1 Term
e2 Ordering -> Ordering -> Ordering
`thenCompare` (Type -> Type -> Ordering
goType Type
l1 Type
l2 Ordering -> Ordering -> Ordering
`thenCompare` Type -> Type -> Ordering
goType Type
r1 Type
r2)
  go (Tick TickInfo
t1 Term
e1) (Tick TickInfo
t2 Term
e2) = TickInfo -> TickInfo -> Ordering
goTick TickInfo
t1 TickInfo
t2 Ordering -> Ordering -> Ordering
`thenCompare` Term -> Term -> Ordering
go Term
e1 Term
e2
  go Term
e1 Term
e2 = Word -> Word -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Term -> Word
getRank Term
e1) (Term -> Word
getRank Term
e2)

  goList :: (a -> a -> Ordering) -> [a] -> [a] -> Ordering
  goList :: forall a. (a -> a -> Ordering) -> [a] -> [a] -> Ordering
goList a -> a -> Ordering
cmp (a
x : [a]
xs) (a
y : [a]
ys) = a -> a -> Ordering
cmp a
x a
y Ordering -> Ordering -> Ordering
`thenCompare` (a -> a -> Ordering) -> [a] -> [a] -> Ordering
forall a. (a -> a -> Ordering) -> [a] -> [a] -> Ordering
goList a -> a -> Ordering
cmp [a]
xs [a]
ys
  goList a -> a -> Ordering
_ [a]
_ [a]
_ = Ordering
EQ

  goAlts :: [Alt] -> [Alt] -> Ordering
  goAlts :: [Alt] -> [Alt] -> Ordering
goAlts (Alt
x : [Alt]
xs) (Alt
y : [Alt]
ys) = Alt -> Alt -> Ordering
goAlt Alt
x Alt
y Ordering -> Ordering -> Ordering
`thenCompare` [Alt] -> [Alt] -> Ordering
goAlts [Alt]
xs [Alt]
ys
  goAlts [Alt]
_ [Alt]
_ = Ordering
EQ

  goAlt :: Alt -> Alt -> Ordering
  goAlt :: Alt -> Alt -> Ordering
goAlt (DataPat DataCon
c1 [TyVar]
tvs1 [Id]
ids1, Term
e1) (DataPat DataCon
c2 [TyVar]
tvs2 [Id]
ids2, Term
e2) =
    -- Two 'DataPat's for the same 'DataCon' necessarily bind equally many
    -- variables, so 'extendLevels' erroring on lists of unequal length is the
    -- right response to Core that is already ill-formed
    DataCon -> DataCon -> Ordering
forall a. Ord a => a -> a -> Ordering
compare DataCon
c1 DataCon
c2 Ordering -> Ordering -> Ordering
`thenCompare`
      let (Int
lvlTy, VarEnv Int
tyL', VarEnv Int
tyR') = Int
-> [TyVar]
-> [TyVar]
-> VarEnv Int
-> VarEnv Int
-> (Int, VarEnv Int, VarEnv Int)
forall a.
Int
-> [Var a]
-> [Var a]
-> VarEnv Int
-> VarEnv Int
-> (Int, VarEnv Int, VarEnv Int)
extendLevels Int
lvl [TyVar]
tvs1 [TyVar]
tvs2 VarEnv Int
tyL VarEnv Int
tyR
          (Int
lvl', VarEnv Int
tmL', VarEnv Int
tmR') = Int
-> [Id]
-> [Id]
-> VarEnv Int
-> VarEnv Int
-> (Int, VarEnv Int, VarEnv Int)
forall a.
Int
-> [Var a]
-> [Var a]
-> VarEnv Int
-> VarEnv Int
-> (Int, VarEnv Int, VarEnv Int)
extendLevels Int
lvlTy [Id]
ids1 [Id]
ids2 VarEnv Int
tmL VarEnv Int
tmR
      in Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> Term
-> Term
-> Ordering
acmpTermLevels Int
lvl' VarEnv Int
tmL' VarEnv Int
tmR' VarEnv Int
tyL' VarEnv Int
tyR' Term
e1 Term
e2
  goAlt (Pat
c1, Term
e1) (Pat
c2, Term
e2) = Pat -> Pat -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Pat
c1 Pat
c2 Ordering -> Ordering -> Ordering
`thenCompare` Term -> Term -> Ordering
go Term
e1 Term
e2

  goVar :: Id -> Id -> Ordering
  goVar :: Id -> Id -> Ordering
goVar Id
id1 Id
id2
    -- A global is never bound by an enclosing binder, so it never resolves to
    -- a level. Checked before the environments, because those are keyed on
    -- unique alone and a global may share a unique with a bound local.
    | Id -> Bool
forall a. Var a -> Bool
isGlobalId Id
id1 Bool -> Bool -> Bool
|| Id -> Bool
forall a. Var a -> Bool
isGlobalId Id
id2 = (Unique, Maybe IdScope) -> (Unique, Maybe IdScope) -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Id -> (Unique, Maybe IdScope)
forall a. Var a -> (Unique, Maybe IdScope)
varKey Id
id1) (Id -> (Unique, Maybe IdScope)
forall a. Var a -> (Unique, Maybe IdScope)
varKey Id
id2)
    | Bool
otherwise = VarEnv Int -> VarEnv Int -> Id -> Id -> Ordering
forall a. VarEnv Int -> VarEnv Int -> Var a -> Var a -> Ordering
acmpOccLevels VarEnv Int
tmL VarEnv Int
tmR Id
id1 Id
id2

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

-- | Alpha equality for ticks, under the binders enclosing them.
-- See 'acmpTickInfoLevels'.
aeqTickInfoLevels ::
  -- | Number of enclosing binders
  Int ->
  -- | Levels of the term binders enclosing the left tick
  VarEnv Int ->
  -- | Levels of the term binders enclosing the right tick
  VarEnv Int ->
  -- | Levels of the type binders enclosing the left tick
  VarEnv Int ->
  -- | Levels of the type binders enclosing the right tick
  VarEnv Int ->
  TickInfo ->
  TickInfo ->
  Bool
aeqTickInfoLevels :: Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> TickInfo
-> TickInfo
-> Bool
aeqTickInfoLevels Int
lvl VarEnv Int
tmL VarEnv Int
tmR VarEnv Int
tyL VarEnv Int
tyR TickInfo
t1 TickInfo
t2 =
  Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> TickInfo
-> TickInfo
-> Ordering
acmpTickInfoLevels Int
lvl VarEnv Int
tmL VarEnv Int
tmR VarEnv Int
tyL VarEnv Int
tyR TickInfo
t1 TickInfo
t2 Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ
{-# INLINE aeqTickInfoLevels #-}

-- | Alpha comparison for ticks, under the binders enclosing them.
--
-- The 'Type' in 'NameMod' and the 'Term' in 'Attributes' live in the scope
-- enclosing the tick, so they are compared under the same binders as the term
-- the tick is attached to. 'SrcSpan's are compared with 'leftmost_smallest',
-- which treats all unhelpful spans alike, so it is coarser than @Eq TickInfo@.
acmpTickInfoLevels ::
  -- | Number of enclosing binders
  Int ->
  -- | Levels of the term binders enclosing the left tick
  VarEnv Int ->
  -- | Levels of the term binders enclosing the right tick
  VarEnv Int ->
  -- | Levels of the type binders enclosing the left tick
  VarEnv Int ->
  -- | Levels of the type binders enclosing the right tick
  VarEnv Int ->
  TickInfo ->
  TickInfo ->
  Ordering
acmpTickInfoLevels :: Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> TickInfo
-> TickInfo
-> Ordering
acmpTickInfoLevels !Int
lvl VarEnv Int
tmL VarEnv Int
tmR VarEnv Int
tyL VarEnv Int
tyR = TickInfo -> TickInfo -> Ordering
go
 where
  goType :: Type -> Type -> Ordering
goType = Int -> VarEnv Int -> VarEnv Int -> Type -> Type -> Ordering
acmpTypeLevels Int
lvl VarEnv Int
tyL VarEnv Int
tyR
  goTerm :: Term -> Term -> Ordering
goTerm = Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> VarEnv Int
-> Term
-> Term
-> Ordering
acmpTermLevels Int
lvl VarEnv Int
tmL VarEnv Int
tmR VarEnv Int
tyL VarEnv Int
tyR

  go :: TickInfo -> TickInfo -> Ordering
  go :: TickInfo -> TickInfo -> Ordering
go (SrcSpan SrcSpan
s1) (SrcSpan SrcSpan
s2) = SrcSpan -> SrcSpan -> Ordering
leftmost_smallest SrcSpan
s1 SrcSpan
s2
  go (NameMod NameMod
m1 Type
t1) (NameMod NameMod
m2 Type
t2) = NameMod -> NameMod -> Ordering
forall a. Ord a => a -> a -> Ordering
compare NameMod
m1 NameMod
m2 Ordering -> Ordering -> Ordering
`thenCompare` Type -> Type -> Ordering
goType Type
t1 Type
t2
  go (Attributes Type
t1 Term
a1) (Attributes Type
t2 Term
a2) =
    Type -> Type -> Ordering
goType Type
t1 Type
t2 Ordering -> Ordering -> Ordering
`thenCompare` Term -> Term -> Ordering
goTerm Term
a1 Term
a2
  go TickInfo
t1 TickInfo
t2 = Word -> Word -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (TickInfo -> Word
getRank TickInfo
t1) (TickInfo -> Word
getRank TickInfo
t2)

  getRank :: TickInfo -> Word
  getRank :: TickInfo -> Word
getRank = \case
    SrcSpan {}    -> Word
0
    NameMod {}    -> Word
1
    TickInfo
DeDup         -> Word
2
    TickInfo
NoDeDup       -> Word
3
    Attributes {} -> Word
4

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

-- * Alpha hashing
--
-- Hashing is modulo alpha equivalence, so that it agrees with @Eq Term@ and
-- @Eq Type@: alpha-equivalent terms hash alike. Only that direction holds,
-- terms that are not alpha-equivalent may collide as with any hash. A bound
-- variable is hashed by its De Bruijn level, see
-- Note [Numbering binders by De Bruijn level].

-- | Mix the tag of a constructor into a salt, so that terms differing only in
-- which constructor they use do not hash alike.
hashTag :: Int -> Int -> Int
hashTag :: Int -> Int -> Int
hashTag = Int -> Int -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt

-- | Give a group of binders that come into scope together, such as a 'Rec' or a
-- 'DataPat', consecutive levels, and return the level following the group.
extendLevelsOf :: Int -> [Var a] -> VarEnv Int -> (Int, VarEnv Int)
extendLevelsOf :: forall a. Int -> [Var a] -> VarEnv Int -> (Int, VarEnv Int)
extendLevelsOf Int
lvl [Var a]
vs VarEnv Int
env = ((Int, VarEnv Int) -> Var a -> (Int, VarEnv Int))
-> (Int, VarEnv Int) -> [Var a] -> (Int, VarEnv Int)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (Int, VarEnv Int) -> Var a -> (Int, VarEnv Int)
forall {a} {b}. Num a => (a, VarEnv a) -> Var b -> (a, VarEnv a)
one (Int
lvl, VarEnv Int
env) [Var a]
vs
 where
  one :: (a, VarEnv a) -> Var b -> (a, VarEnv a)
one (!a
l, VarEnv a
e) Var b
v = (a
l a -> a -> a
forall a. Num a => a -> a -> a
+ a
1, Var b -> a -> VarEnv a -> VarEnv a
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv Var b
v a
l VarEnv a
e)

-- | Hash a 'SrcSpan' only as finely as @Eq Term@ tells one apart: all
-- 'UnhelpfulSpan's alike, and a 'RealSrcSpan' by its file name and its start
-- and end positions. The structural @Hashable SrcSpan@ orphan from
-- "GHC.SrcLoc.Extra" is finer than that, as it takes in the buffer span and the
-- reason of an unhelpful span as well, and would tell alpha-equivalent terms
-- apart.
hashSrcSpan :: Int -> SrcSpan -> Int
hashSrcSpan :: Int -> SrcSpan -> Int
hashSrcSpan Int
salt = \case
  RealSrcSpan RealSrcSpan
realSrcSpan Maybe BufSpan
_bufSpan -> Int -> (Int, RealSrcSpan) -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
salt (Int
0 :: Int, RealSrcSpan
realSrcSpan)
  UnhelpfulSpan UnhelpfulSpanReason
_reason -> Int -> Int -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
salt (Int
1 :: Int)

-- | Hash a 'Type' modulo alpha, under the binders enclosing it.
-- See Note [Numbering binders by De Bruijn level].
aTypeHashLevels ::
  -- | Number of enclosing binders
  Int ->
  -- | Levels of the type binders enclosing the type
  VarEnv Int ->
  -- | Salt
  Int ->
  Type ->
  Int
aTypeHashLevels :: Int -> VarEnv Int -> Int -> Type -> Int
aTypeHashLevels !Int
lvl VarEnv Int
tyEnv = Int -> Type -> Int
go
 where
  go :: Int -> Type -> Int
  go :: Int -> Type -> Int
go Int
salt = \case
    VarTy TyVar
tv ->
      -- N.B.: Variables are hashed with a "tag" to differentiate between, e.g.,
      --       a bound local variable bound at level "2" and a free local variable
      --       with unique "2".
      case TyVar -> VarEnv Int -> Maybe Int
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv TyVar
tv VarEnv Int
tyEnv of
        Just Int
boundLvl -> Int -> (Int, Int) -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
salt (Int
0 :: Int, Int
boundLvl)
        Maybe Int
Nothing -> Int -> (Int, Unique) -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
salt (Int
1 :: Int, TyVar -> Unique
forall a. Var a -> Unique
varUniq TyVar
tv)
    LitTy LitTy
l -> Int -> LitTy -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt (Int -> Int -> Int
hashTag Int
salt Int
1) LitTy
l
    ConstTy ConstTy
c -> Int -> ConstTy -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt (Int -> Int -> Int
hashTag Int
salt Int
2) ConstTy
c
    AnnType [Attr Text]
attrs Type
t -> Int -> Type -> Int
go (Int -> [Attr Text] -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt (Int -> Int -> Int
hashTag Int
salt Int
3) [Attr Text]
attrs) Type
t
    AppTy Type
t1 Type
t2 -> Int -> Type -> Int
go (Int -> Type -> Int
go (Int -> Int -> Int
hashTag Int
salt Int
4) Type
t1) Type
t2
    ForAllTy TyVar
tv Type
t ->
      Int -> VarEnv Int -> Int -> Type -> Int
aTypeHashLevels (Int
lvl Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (TyVar -> Int -> VarEnv Int -> VarEnv Int
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv TyVar
tv Int
lvl VarEnv Int
tyEnv)
        (Int -> Type -> Int
go (Int -> Int -> Int
hashTag Int
salt Int
5) (TyVar -> Type
forall a. Var a -> Type
varType TyVar
tv)) Type
t

-- | Hash a 'Term' modulo alpha, under the binders enclosing it.
-- See Note [Numbering binders by De Bruijn level].
aTermHashLevels ::
  -- | Number of enclosing binders
  Int ->
  -- | Levels of the term binders enclosing the term
  VarEnv Int ->
  -- | Levels of the type binders enclosing the term
  VarEnv Int ->
  -- | Salt
  Int ->
  Term ->
  Int
aTermHashLevels :: Int -> VarEnv Int -> VarEnv Int -> Int -> Term -> Int
aTermHashLevels !Int
lvl VarEnv Int
tmEnv VarEnv Int
tyEnv = Int -> Term -> Int
go
 where
  goType :: Int -> Type -> Int
goType = Int -> VarEnv Int -> Int -> Type -> Int
aTypeHashLevels Int
lvl VarEnv Int
tyEnv

  underTmBndr :: Var b -> Int -> Term -> Int
underTmBndr Var b
b = Int -> VarEnv Int -> VarEnv Int -> Int -> Term -> Int
aTermHashLevels (Int
lvl Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Var b -> Int -> VarEnv Int -> VarEnv Int
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv Var b
b Int
lvl VarEnv Int
tmEnv) VarEnv Int
tyEnv
  underTyBndr :: Var b -> Int -> Term -> Int
underTyBndr Var b
b = Int -> VarEnv Int -> VarEnv Int -> Int -> Term -> Int
aTermHashLevels (Int
lvl Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) VarEnv Int
tmEnv (Var b -> Int -> VarEnv Int -> VarEnv Int
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv Var b
b Int
lvl VarEnv Int
tyEnv)

  go :: Int -> Term -> Int
  go :: Int -> Term -> Int
go Int
salt = \case
    Var Id
i -> Int -> Id -> Int
goVar (Int -> Int -> Int
hashTag Int
salt Int
0) Id
i
    Data DataCon
dc -> Int -> DataCon -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt (Int -> Int -> Int
hashTag Int
salt Int
1) DataCon
dc
    Literal Literal
l -> Int -> Literal -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt (Int -> Int -> Int
hashTag Int
salt Int
2) Literal
l
    -- A primitive is identified by its name
    Prim PrimInfo
p -> Int -> Text -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt (Int -> Int -> Int
hashTag Int
salt Int
3) (PrimInfo -> Text
primName PrimInfo
p)
    Cast Term
e Type
t1 Type
t2 -> Int -> Type -> Int
goType (Int -> Type -> Int
goType (Int -> Term -> Int
go (Int -> Int -> Int
hashTag Int
salt Int
4) Term
e) Type
t1) Type
t2
    App Term
e1 Term
e2 -> Int -> Term -> Int
go (Int -> Term -> Int
go (Int -> Int -> Int
hashTag Int
salt Int
5) Term
e1) Term
e2
    TyApp Term
e Type
t -> Int -> Type -> Int
goType (Int -> Term -> Int
go (Int -> Int -> Int
hashTag Int
salt Int
6) Term
e) Type
t
    Lam Id
b Term
e -> Id -> Int -> Term -> Int
forall {b}. Var b -> Int -> Term -> Int
underTmBndr Id
b (Int -> Type -> Int
goType (Int -> Int -> Int
hashTag Int
salt Int
7) (Id -> Type
forall a. Var a -> Type
varType Id
b)) Term
e
    TyLam TyVar
b Term
e -> TyVar -> Int -> Term -> Int
forall {b}. Var b -> Int -> Term -> Int
underTyBndr TyVar
b (Int -> Type -> Int
goType (Int -> Int -> Int
hashTag Int
salt Int
8) (TyVar -> Type
forall a. Var a -> Type
varType TyVar
b)) Term
e
    -- A 'NonRec' binder's type is pinned down by its right-hand side, so it is
    -- left out. A 'Rec' binder may occur in its own right-hand side, and then
    -- it is not: @let x = x in x@ is the same term whether @x@ is an Int or a
    -- Bool.
    Let (NonRec Id
i Term
x) Term
e -> Id -> Int -> Term -> Int
forall {b}. Var b -> Int -> Term -> Int
underTmBndr Id
i (Int -> Term -> Int
go (Int -> Int -> Int
hashTag Int
salt Int
9) Term
x) Term
e
    Let (Rec [(Id, Term)]
bs) Term
e ->
      let ([Id]
ids, [Term]
rhss) = [(Id, Term)] -> ([Id], [Term])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Id, Term)]
bs
          (Int
lvl', VarEnv Int
tmEnv') = Int -> [Id] -> VarEnv Int -> (Int, VarEnv Int)
forall a. Int -> [Var a] -> VarEnv Int -> (Int, VarEnv Int)
extendLevelsOf Int
lvl [Id]
ids VarEnv Int
tmEnv
          under :: Int -> Term -> Int
under = Int -> VarEnv Int -> VarEnv Int -> Int -> Term -> Int
aTermHashLevels Int
lvl' VarEnv Int
tmEnv' VarEnv Int
tyEnv
          types :: Int
types = (Int -> Type -> Int) -> Int -> [Type] -> Int
forall a. (Int -> a -> Int) -> Int -> [a] -> Int
goList Int -> Type -> Int
goType (Int -> Int -> Int
hashTag Int
salt Int
10) ((Id -> Type) -> [Id] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Id -> Type
forall a. Var a -> Type
varType [Id]
ids)
      in Int -> Term -> Int
under ((Int -> Term -> Int) -> Int -> [Term] -> Int
forall a. (Int -> a -> Int) -> Int -> [a] -> Int
goList Int -> Term -> Int
under Int
types [Term]
rhss) Term
e
    -- Note [Case result types and alpha-equivalence]
    Case Term
subj Type
_ty [Alt]
alts -> (Int -> Alt -> Int) -> Int -> [Alt] -> Int
forall a. (Int -> a -> Int) -> Int -> [a] -> Int
goList Int -> Alt -> Int
goAlt (Int -> Term -> Int
go (Int -> Int -> Int
hashTag Int
salt Int
11) Term
subj) [Alt]
alts
    Tick TickInfo
tick Term
e -> Int -> Term -> Int
go (Int -> TickInfo -> Int
goTick (Int -> Int -> Int
hashTag Int
salt Int
12) TickInfo
tick) Term
e

  -- Hash a list, mixing in its length, so that a list does not hash like one of
  -- its prefixes
  goList :: (Int -> a -> Int) -> Int -> [a] -> Int
  goList :: forall a. (Int -> a -> Int) -> Int -> [a] -> Int
goList Int -> a -> Int
hashElement Int
salt [a]
xs =
    Int -> Int -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt ((Int -> a -> Int) -> Int -> [a] -> Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' Int -> a -> Int
hashElement Int
salt [a]
xs) ([a] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [a]
xs)

  goAlt :: Int -> Alt -> Int
  goAlt :: Int -> Alt -> Int
goAlt Int
salt = \case
    -- A 'DataCon' fixes how many variables its 'DataPat' binds, so the binder
    -- counts need no hashing of their own
    (DataPat DataCon
dc [TyVar]
tvs [Id]
ids, Term
e) ->
      let (Int
lvlTy, VarEnv Int
tyEnv') = Int -> [TyVar] -> VarEnv Int -> (Int, VarEnv Int)
forall a. Int -> [Var a] -> VarEnv Int -> (Int, VarEnv Int)
extendLevelsOf Int
lvl [TyVar]
tvs VarEnv Int
tyEnv
          (Int
lvl', VarEnv Int
tmEnv') = Int -> [Id] -> VarEnv Int -> (Int, VarEnv Int)
forall a. Int -> [Var a] -> VarEnv Int -> (Int, VarEnv Int)
extendLevelsOf Int
lvlTy [Id]
ids VarEnv Int
tmEnv
      in Int -> VarEnv Int -> VarEnv Int -> Int -> Term -> Int
aTermHashLevels Int
lvl' VarEnv Int
tmEnv' VarEnv Int
tyEnv'
           (Int -> DataCon -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt (Int -> Int -> Int
hashTag Int
salt Int
0) DataCon
dc) Term
e
    (LitPat Literal
l, Term
e) -> Int -> Term -> Int
go (Int -> Literal -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt (Int -> Int -> Int
hashTag Int
salt Int
1) Literal
l) Term
e
    (Pat
DefaultPat, Term
e) -> Int -> Term -> Int
go (Int -> Int -> Int
hashTag Int
salt Int
2) Term
e

  -- N.B.: Variables are hashed with a "tag" to differentiate between, e.g., a
  --       local free variable with unique "2" and a bound free variable which
  --       happens to be bound at level 2.
  goVar :: Int -> Id -> Int
  goVar :: Int -> Id -> Int
goVar Int
salt Id
i
    -- Global, never bound with respect to 'tmEnv'
    | Id -> Bool
forall a. Var a -> Bool
isGlobalId Id
i = Int -> (Int, Unique) -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
salt (Int
0 :: Int, Id -> Unique
forall a. Var a -> Unique
varUniq Id
i)
    -- Local, bound variable
    | Just Int
boundLvl <- Id -> VarEnv Int -> Maybe Int
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv Id
i VarEnv Int
tmEnv = Int -> (Int, Int) -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
salt (Int
1 :: Int, Int
boundLvl)
    -- Free, local variable
    | Bool
otherwise = Int -> (Int, Unique) -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
salt (Int
2 :: Int, Id -> Unique
forall a. Var a -> Unique
varUniq Id
i)

  -- A tick's payload lives in the scope enclosing the tick, so it is hashed
  -- under the enclosing binders rather than in isolation
  goTick :: Int -> TickInfo -> Int
  goTick :: Int -> TickInfo -> Int
goTick Int
salt = \case
    SrcSpan SrcSpan
s -> Int -> SrcSpan -> Int
hashSrcSpan (Int -> Int -> Int
hashTag Int
salt Int
0) SrcSpan
s
    NameMod NameMod
m Type
t -> Int -> Type -> Int
goType (Int -> NameMod -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt (Int -> Int -> Int
hashTag Int
salt Int
1) NameMod
m) Type
t
    TickInfo
DeDup -> Int -> Int -> Int
hashTag Int
salt Int
2
    TickInfo
NoDeDup -> Int -> Int -> Int
hashTag Int
salt Int
3
    Attributes Type
t Term
e -> Int -> Term -> Int
go (Int -> Type -> Int
goType (Int -> Int -> Int
hashTag Int
salt Int
4) Type
t) Term
e

-- | Hash a 'Type' modulo alpha.
-- See Note [Numbering binders by De Bruijn level].
aTypeHashWithSalt :: Int -> Type -> Int
aTypeHashWithSalt :: Int -> Type -> Int
aTypeHashWithSalt = Int -> VarEnv Int -> Int -> Type -> Int
aTypeHashLevels Int
0 VarEnv Int
forall a. VarEnv a
emptyVarEnv

-- | Hash a 'Term' modulo alpha.
-- See Note [Numbering binders by De Bruijn level].
aTermHashWithSalt :: Int -> Term -> Int
aTermHashWithSalt :: Int -> Term -> Int
aTermHashWithSalt = Int -> VarEnv Int -> VarEnv Int -> Int -> Term -> Int
aTermHashLevels Int
0 VarEnv Int
forall a. VarEnv a
emptyVarEnv VarEnv Int
forall a. VarEnv a
emptyVarEnv

instance Hashable Type where
  hashWithSalt :: Int -> Type -> Int
hashWithSalt = Int -> Type -> Int
aTypeHashWithSalt

instance Hashable Term where
  hashWithSalt :: Int -> Term -> Int
hashWithSalt = Int -> Term -> Int
aTermHashWithSalt