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

  Term representation in the CoreHW language: System F + LetRec + Case
-}

{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE UndecidableInstances #-}

module Clash.Core.Term
  ( Term (.., Letrec)
  , mkAbstraction
  , mkTyLams
  , mkLams
  , mkApps
  , mkTyApps
  , mkTmApps
  , mkTicks
  , TmName
  , varToId
  , Bind(..)
  , LetBinding
  , Pat (..)
  , patIds
  , patVars
  , Alt
  , TickInfo (..)
  , stripTicks
  , stripAllTicks
  , partitionTicks
  , NameMod (..)
  , PrimInfo (..)
  , PrimUnfolding (..)
  , IsMultiPrim (..)
  , MultiPrimInfo (..)
  , WorkInfo (..)
  , CoreContext (..)
  , Context
  , isLambdaBodyCtx
  , isTickCtx
  , walkTerm
  , collectArgs
  , collectArgsTicks
  , collectTicks
  , collectTermIds
  , collectBndrs
  , primArg
  ) where

-- External Modules
import Control.DeepSeq
import Data.Binary                             (Binary)
import Data.Coerce                             (coerce)
import qualified Data.DList                    as DList
import Data.Either                             (lefts, rights)
#if !MIN_VERSION_base(4,20,0)
import Data.Foldable                           (foldl')
#endif
import Data.Hashable                           (Hashable)
import Data.Maybe                              (catMaybes)
import Data.List                               (nub, partition)
import Data.Text                               (Text)
import GHC.Generics
#if MIN_VERSION_ghc(9,0,0)
import GHC.Types.SrcLoc                        (SrcSpan, leftmost_smallest)
#else
import SrcLoc                                  (SrcSpan, leftmost_smallest)
#endif

-- Internal Modules
import Clash.Core.DataCon                      (DataCon)
import Clash.Core.Literal                      (Literal)
import Clash.Core.Name                         (Name (..))
import {-# SOURCE #-} Clash.Core.Subst         () -- instance Eq/Ord Type
import {-# SOURCE #-} Clash.Core.Type          (Type)
import Clash.Core.Var                          (Var, Id, TyVar)
import Clash.Util                              (curLoc, thenCompare)

-- | Term representation in the CoreHW language: System F + LetRec + Case
data Term
  = Var     !Id                             -- ^ Variable reference
  | Data    !DataCon                        -- ^ Datatype constructor
  | Literal !Literal                        -- ^ Literal
  | Prim    !PrimInfo                       -- ^ Primitive
  | Lam     !Id Term                        -- ^ Term-abstraction
  | TyLam   !TyVar Term                     -- ^ Type-abstraction
  | App     !Term !Term                     -- ^ Application
  | TyApp   !Term !Type                     -- ^ Type-application
  | Let     !(Bind Term) Term               -- ^ Recursive let-binding
  | Case    !Term !Type [Alt]               -- ^ Case-expression: subject, type of
                                            -- alternatives, list of alternatives
  | Cast    !Term !Type !Type               -- ^ Cast a term from one type to another
  | Tick    !TickInfo !Term                 -- ^ Annotated term
  deriving (Int -> Term -> ShowS
[Term] -> ShowS
Term -> String
(Int -> Term -> ShowS)
-> (Term -> String) -> ([Term] -> ShowS) -> Show Term
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Term] -> ShowS
$cshowList :: [Term] -> ShowS
show :: Term -> String
$cshow :: Term -> String
showsPrec :: Int -> Term -> ShowS
$cshowsPrec :: Int -> Term -> ShowS
Show, (forall x. Term -> Rep Term x)
-> (forall x. Rep Term x -> Term) -> Generic Term
forall x. Rep Term x -> Term
forall x. Term -> Rep Term x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Term x -> Term
$cfrom :: forall x. Term -> Rep Term x
Generic, Term -> ()
(Term -> ()) -> NFData Term
forall a. (a -> ()) -> NFData a
rnf :: Term -> ()
$crnf :: Term -> ()
NFData, Get Term
[Term] -> Put
Term -> Put
(Term -> Put) -> Get Term -> ([Term] -> Put) -> Binary Term
forall t. (t -> Put) -> Get t -> ([t] -> Put) -> Binary t
putList :: [Term] -> Put
$cputList :: [Term] -> Put
get :: Get Term
$cget :: Get Term
put :: Term -> Put
$cput :: Term -> Put
Binary)

-- TODO When it is possible, remove this pattern.
pattern Letrec :: [LetBinding] -> Term -> Term
pattern $bLetrec :: [LetBinding] -> Term -> Term
$mLetrec :: forall r. Term -> ([LetBinding] -> Term -> r) -> (Void# -> r) -> r
Letrec bs x <- Let (bindToList -> bs) x
 where
  Letrec [LetBinding]
bs Term
x = Bind Term -> Term -> Term
Let ([LetBinding] -> Bind Term
forall a. [(Id, a)] -> Bind a
Rec [LetBinding]
bs) Term
x

bindToList :: Bind a -> [(Id, a)]
bindToList :: Bind a -> [(Id, a)]
bindToList (NonRec Id
i a
x) = [(Id
i, a
x)]
bindToList (Rec [(Id, a)]
xs) = [(Id, a)]
xs

data TickInfo
  = SrcSpan !SrcSpan
  -- ^ Source tick, will get added by GHC by running clash with `-g`
  | NameMod !NameMod !Type
  -- ^ Modifier for naming module instantiations and registers, are added by
  -- the user by using the functions @Clash.Magic.[prefixName,suffixName,setName]@
  | DeDup
  -- ^ Deduplicate, i.e. try to share expressions between multiple branches.
  | NoDeDup
  -- ^ Do not deduplicate, i.e. /keep/, an expression inside a case-alternative;
  -- do not try to share expressions between multiple branches.
  deriving (TickInfo -> TickInfo -> Bool
(TickInfo -> TickInfo -> Bool)
-> (TickInfo -> TickInfo -> Bool) -> Eq TickInfo
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: TickInfo -> TickInfo -> Bool
$c/= :: TickInfo -> TickInfo -> Bool
== :: TickInfo -> TickInfo -> Bool
$c== :: TickInfo -> TickInfo -> Bool
Eq, Int -> TickInfo -> ShowS
[TickInfo] -> ShowS
TickInfo -> String
(Int -> TickInfo -> ShowS)
-> (TickInfo -> String) -> ([TickInfo] -> ShowS) -> Show TickInfo
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TickInfo] -> ShowS
$cshowList :: [TickInfo] -> ShowS
show :: TickInfo -> String
$cshow :: TickInfo -> String
showsPrec :: Int -> TickInfo -> ShowS
$cshowsPrec :: Int -> TickInfo -> ShowS
Show, (forall x. TickInfo -> Rep TickInfo x)
-> (forall x. Rep TickInfo x -> TickInfo) -> Generic TickInfo
forall x. Rep TickInfo x -> TickInfo
forall x. TickInfo -> Rep TickInfo x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep TickInfo x -> TickInfo
$cfrom :: forall x. TickInfo -> Rep TickInfo x
Generic, TickInfo -> ()
(TickInfo -> ()) -> NFData TickInfo
forall a. (a -> ()) -> NFData a
rnf :: TickInfo -> ()
$crnf :: TickInfo -> ()
NFData, Get TickInfo
[TickInfo] -> Put
TickInfo -> Put
(TickInfo -> Put)
-> Get TickInfo -> ([TickInfo] -> Put) -> Binary TickInfo
forall t. (t -> Put) -> Get t -> ([t] -> Put) -> Binary t
putList :: [TickInfo] -> Put
$cputList :: [TickInfo] -> Put
get :: Get TickInfo
$cget :: Get TickInfo
put :: TickInfo -> Put
$cput :: TickInfo -> Put
Binary)

instance Ord TickInfo where
  compare :: TickInfo -> TickInfo -> Ordering
compare (SrcSpan SrcSpan
s1) (SrcSpan SrcSpan
s2) = SrcSpan -> SrcSpan -> Ordering
leftmost_smallest SrcSpan
s1 SrcSpan
s2
  compare (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
forall a. Ord a => a -> a -> Ordering
compare Type
t1 Type
t2
  compare 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)
    where
      getRank :: TickInfo -> Word
      getRank :: TickInfo -> Word
getRank SrcSpan{} = Word
0
      getRank NameMod{} = Word
1
      getRank TickInfo
DeDup     = Word
2
      getRank TickInfo
NoDeDup   = Word
3

-- | Tag to indicate which instance/register name modifier was used
data NameMod
  = PrefixName
  -- ^ @Clash.Magic.prefixName@
  | SuffixName
  -- ^ @Clash.Magic.suffixName@
  | SuffixNameP
  -- ^ @Clash.Magic.suffixNameP@
  | SetName
  -- ^ @Clash.Magic.setName@
  deriving (NameMod -> NameMod -> Bool
(NameMod -> NameMod -> Bool)
-> (NameMod -> NameMod -> Bool) -> Eq NameMod
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: NameMod -> NameMod -> Bool
$c/= :: NameMod -> NameMod -> Bool
== :: NameMod -> NameMod -> Bool
$c== :: NameMod -> NameMod -> Bool
Eq,Eq NameMod
Eq NameMod
-> (NameMod -> NameMod -> Ordering)
-> (NameMod -> NameMod -> Bool)
-> (NameMod -> NameMod -> Bool)
-> (NameMod -> NameMod -> Bool)
-> (NameMod -> NameMod -> Bool)
-> (NameMod -> NameMod -> NameMod)
-> (NameMod -> NameMod -> NameMod)
-> Ord NameMod
NameMod -> NameMod -> Bool
NameMod -> NameMod -> Ordering
NameMod -> NameMod -> NameMod
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: NameMod -> NameMod -> NameMod
$cmin :: NameMod -> NameMod -> NameMod
max :: NameMod -> NameMod -> NameMod
$cmax :: NameMod -> NameMod -> NameMod
>= :: NameMod -> NameMod -> Bool
$c>= :: NameMod -> NameMod -> Bool
> :: NameMod -> NameMod -> Bool
$c> :: NameMod -> NameMod -> Bool
<= :: NameMod -> NameMod -> Bool
$c<= :: NameMod -> NameMod -> Bool
< :: NameMod -> NameMod -> Bool
$c< :: NameMod -> NameMod -> Bool
compare :: NameMod -> NameMod -> Ordering
$ccompare :: NameMod -> NameMod -> Ordering
$cp1Ord :: Eq NameMod
Ord,Int -> NameMod -> ShowS
[NameMod] -> ShowS
NameMod -> String
(Int -> NameMod -> ShowS)
-> (NameMod -> String) -> ([NameMod] -> ShowS) -> Show NameMod
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [NameMod] -> ShowS
$cshowList :: [NameMod] -> ShowS
show :: NameMod -> String
$cshow :: NameMod -> String
showsPrec :: Int -> NameMod -> ShowS
$cshowsPrec :: Int -> NameMod -> ShowS
Show,(forall x. NameMod -> Rep NameMod x)
-> (forall x. Rep NameMod x -> NameMod) -> Generic NameMod
forall x. Rep NameMod x -> NameMod
forall x. NameMod -> Rep NameMod x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep NameMod x -> NameMod
$cfrom :: forall x. NameMod -> Rep NameMod x
Generic,NameMod -> ()
(NameMod -> ()) -> NFData NameMod
forall a. (a -> ()) -> NFData a
rnf :: NameMod -> ()
$crnf :: NameMod -> ()
NFData,Eq NameMod
Eq NameMod
-> (Int -> NameMod -> Int) -> (NameMod -> Int) -> Hashable NameMod
Int -> NameMod -> Int
NameMod -> Int
forall a. Eq a -> (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: NameMod -> Int
$chash :: NameMod -> Int
hashWithSalt :: Int -> NameMod -> Int
$chashWithSalt :: Int -> NameMod -> Int
$cp1Hashable :: Eq NameMod
Hashable,Get NameMod
[NameMod] -> Put
NameMod -> Put
(NameMod -> Put)
-> Get NameMod -> ([NameMod] -> Put) -> Binary NameMod
forall t. (t -> Put) -> Get t -> ([t] -> Put) -> Binary t
putList :: [NameMod] -> Put
$cputList :: [NameMod] -> Put
get :: Get NameMod
$cget :: Get NameMod
put :: NameMod -> Put
$cput :: NameMod -> Put
Binary)

data IsMultiPrim
  = SingleResult
  | MultiResult
  deriving (Int -> IsMultiPrim -> ShowS
[IsMultiPrim] -> ShowS
IsMultiPrim -> String
(Int -> IsMultiPrim -> ShowS)
-> (IsMultiPrim -> String)
-> ([IsMultiPrim] -> ShowS)
-> Show IsMultiPrim
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [IsMultiPrim] -> ShowS
$cshowList :: [IsMultiPrim] -> ShowS
show :: IsMultiPrim -> String
$cshow :: IsMultiPrim -> String
showsPrec :: Int -> IsMultiPrim -> ShowS
$cshowsPrec :: Int -> IsMultiPrim -> ShowS
Show, (forall x. IsMultiPrim -> Rep IsMultiPrim x)
-> (forall x. Rep IsMultiPrim x -> IsMultiPrim)
-> Generic IsMultiPrim
forall x. Rep IsMultiPrim x -> IsMultiPrim
forall x. IsMultiPrim -> Rep IsMultiPrim x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep IsMultiPrim x -> IsMultiPrim
$cfrom :: forall x. IsMultiPrim -> Rep IsMultiPrim x
Generic, IsMultiPrim -> ()
(IsMultiPrim -> ()) -> NFData IsMultiPrim
forall a. (a -> ()) -> NFData a
rnf :: IsMultiPrim -> ()
$crnf :: IsMultiPrim -> ()
NFData, IsMultiPrim -> IsMultiPrim -> Bool
(IsMultiPrim -> IsMultiPrim -> Bool)
-> (IsMultiPrim -> IsMultiPrim -> Bool) -> Eq IsMultiPrim
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: IsMultiPrim -> IsMultiPrim -> Bool
$c/= :: IsMultiPrim -> IsMultiPrim -> Bool
== :: IsMultiPrim -> IsMultiPrim -> Bool
$c== :: IsMultiPrim -> IsMultiPrim -> Bool
Eq, Eq IsMultiPrim
Eq IsMultiPrim
-> (Int -> IsMultiPrim -> Int)
-> (IsMultiPrim -> Int)
-> Hashable IsMultiPrim
Int -> IsMultiPrim -> Int
IsMultiPrim -> Int
forall a. Eq a -> (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: IsMultiPrim -> Int
$chash :: IsMultiPrim -> Int
hashWithSalt :: Int -> IsMultiPrim -> Int
$chashWithSalt :: Int -> IsMultiPrim -> Int
$cp1Hashable :: Eq IsMultiPrim
Hashable, Get IsMultiPrim
[IsMultiPrim] -> Put
IsMultiPrim -> Put
(IsMultiPrim -> Put)
-> Get IsMultiPrim -> ([IsMultiPrim] -> Put) -> Binary IsMultiPrim
forall t. (t -> Put) -> Get t -> ([t] -> Put) -> Binary t
putList :: [IsMultiPrim] -> Put
$cputList :: [IsMultiPrim] -> Put
get :: Get IsMultiPrim
$cget :: Get IsMultiPrim
put :: IsMultiPrim -> Put
$cput :: IsMultiPrim -> Put
Binary)

data PrimInfo = PrimInfo
  { PrimInfo -> Text
primName :: !Text
  , PrimInfo -> Type
primType :: !Type
  , PrimInfo -> WorkInfo
primWorkInfo :: !WorkInfo
  , PrimInfo -> IsMultiPrim
primMultiResult :: !IsMultiPrim
  -- ^ Primitive with multiple return values. Useful for primitives that cannot
  -- return their results as a single product type, due to limitation of
  -- synthesis tooling. It will be applied to its normal arguments, followed by
  -- the variables it should assign its results to.
  --
  -- See: 'Clash.Normalize.Transformations.setupMultiResultPrim'
  , PrimInfo -> PrimUnfolding
primUnfolding :: !PrimUnfolding
  } deriving (Int -> PrimInfo -> ShowS
[PrimInfo] -> ShowS
PrimInfo -> String
(Int -> PrimInfo -> ShowS)
-> (PrimInfo -> String) -> ([PrimInfo] -> ShowS) -> Show PrimInfo
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PrimInfo] -> ShowS
$cshowList :: [PrimInfo] -> ShowS
show :: PrimInfo -> String
$cshow :: PrimInfo -> String
showsPrec :: Int -> PrimInfo -> ShowS
$cshowsPrec :: Int -> PrimInfo -> ShowS
Show, (forall x. PrimInfo -> Rep PrimInfo x)
-> (forall x. Rep PrimInfo x -> PrimInfo) -> Generic PrimInfo
forall x. Rep PrimInfo x -> PrimInfo
forall x. PrimInfo -> Rep PrimInfo x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep PrimInfo x -> PrimInfo
$cfrom :: forall x. PrimInfo -> Rep PrimInfo x
Generic, PrimInfo -> ()
(PrimInfo -> ()) -> NFData PrimInfo
forall a. (a -> ()) -> NFData a
rnf :: PrimInfo -> ()
$crnf :: PrimInfo -> ()
NFData, Get PrimInfo
[PrimInfo] -> Put
PrimInfo -> Put
(PrimInfo -> Put)
-> Get PrimInfo -> ([PrimInfo] -> Put) -> Binary PrimInfo
forall t. (t -> Put) -> Get t -> ([t] -> Put) -> Binary t
putList :: [PrimInfo] -> Put
$cputList :: [PrimInfo] -> Put
get :: Get PrimInfo
$cget :: Get PrimInfo
put :: PrimInfo -> Put
$cput :: PrimInfo -> Put
Binary)

data PrimUnfolding
  = NoUnfolding
  | Unfolding !Id
  deriving (Int -> PrimUnfolding -> ShowS
[PrimUnfolding] -> ShowS
PrimUnfolding -> String
(Int -> PrimUnfolding -> ShowS)
-> (PrimUnfolding -> String)
-> ([PrimUnfolding] -> ShowS)
-> Show PrimUnfolding
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PrimUnfolding] -> ShowS
$cshowList :: [PrimUnfolding] -> ShowS
show :: PrimUnfolding -> String
$cshow :: PrimUnfolding -> String
showsPrec :: Int -> PrimUnfolding -> ShowS
$cshowsPrec :: Int -> PrimUnfolding -> ShowS
Show, (forall x. PrimUnfolding -> Rep PrimUnfolding x)
-> (forall x. Rep PrimUnfolding x -> PrimUnfolding)
-> Generic PrimUnfolding
forall x. Rep PrimUnfolding x -> PrimUnfolding
forall x. PrimUnfolding -> Rep PrimUnfolding x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep PrimUnfolding x -> PrimUnfolding
$cfrom :: forall x. PrimUnfolding -> Rep PrimUnfolding x
Generic, PrimUnfolding -> ()
(PrimUnfolding -> ()) -> NFData PrimUnfolding
forall a. (a -> ()) -> NFData a
rnf :: PrimUnfolding -> ()
$crnf :: PrimUnfolding -> ()
NFData, PrimUnfolding -> PrimUnfolding -> Bool
(PrimUnfolding -> PrimUnfolding -> Bool)
-> (PrimUnfolding -> PrimUnfolding -> Bool) -> Eq PrimUnfolding
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PrimUnfolding -> PrimUnfolding -> Bool
$c/= :: PrimUnfolding -> PrimUnfolding -> Bool
== :: PrimUnfolding -> PrimUnfolding -> Bool
$c== :: PrimUnfolding -> PrimUnfolding -> Bool
Eq, Eq PrimUnfolding
Eq PrimUnfolding
-> (Int -> PrimUnfolding -> Int)
-> (PrimUnfolding -> Int)
-> Hashable PrimUnfolding
Int -> PrimUnfolding -> Int
PrimUnfolding -> Int
forall a. Eq a -> (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: PrimUnfolding -> Int
$chash :: PrimUnfolding -> Int
hashWithSalt :: Int -> PrimUnfolding -> Int
$chashWithSalt :: Int -> PrimUnfolding -> Int
$cp1Hashable :: Eq PrimUnfolding
Hashable, Get PrimUnfolding
[PrimUnfolding] -> Put
PrimUnfolding -> Put
(PrimUnfolding -> Put)
-> Get PrimUnfolding
-> ([PrimUnfolding] -> Put)
-> Binary PrimUnfolding
forall t. (t -> Put) -> Get t -> ([t] -> Put) -> Binary t
putList :: [PrimUnfolding] -> Put
$cputList :: [PrimUnfolding] -> Put
get :: Get PrimUnfolding
$cget :: Get PrimUnfolding
put :: PrimUnfolding -> Put
$cput :: PrimUnfolding -> Put
Binary)

data MultiPrimInfo = MultiPrimInfo
  { MultiPrimInfo -> PrimInfo
mpi_primInfo :: PrimInfo
  , MultiPrimInfo -> DataCon
mpi_resultDc :: DataCon
  , MultiPrimInfo -> [Type]
mpi_resultTypes :: [Type]
  }

data WorkInfo
  = WorkConstant
  -- ^ Ignores its arguments, and outputs a constant
  | WorkNever
  -- ^ Never adds any work
  | WorkVariable
  -- ^ Does work when the arguments are variable
  | WorkAlways
  -- ^ Performs work regardless of whether the variables are constant or
  -- variable; these are things like clock or reset generators
  | WorkIdentity Int [Int]
  -- ^ A more restrictive version of 'WorkNever', where the value is the
  -- argument at the given position if all arguments for the given list of
  -- positions are also 'WorkIdentity'
  deriving (WorkInfo -> WorkInfo -> Bool
(WorkInfo -> WorkInfo -> Bool)
-> (WorkInfo -> WorkInfo -> Bool) -> Eq WorkInfo
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: WorkInfo -> WorkInfo -> Bool
$c/= :: WorkInfo -> WorkInfo -> Bool
== :: WorkInfo -> WorkInfo -> Bool
$c== :: WorkInfo -> WorkInfo -> Bool
Eq,Int -> WorkInfo -> ShowS
[WorkInfo] -> ShowS
WorkInfo -> String
(Int -> WorkInfo -> ShowS)
-> (WorkInfo -> String) -> ([WorkInfo] -> ShowS) -> Show WorkInfo
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [WorkInfo] -> ShowS
$cshowList :: [WorkInfo] -> ShowS
show :: WorkInfo -> String
$cshow :: WorkInfo -> String
showsPrec :: Int -> WorkInfo -> ShowS
$cshowsPrec :: Int -> WorkInfo -> ShowS
Show,(forall x. WorkInfo -> Rep WorkInfo x)
-> (forall x. Rep WorkInfo x -> WorkInfo) -> Generic WorkInfo
forall x. Rep WorkInfo x -> WorkInfo
forall x. WorkInfo -> Rep WorkInfo x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep WorkInfo x -> WorkInfo
$cfrom :: forall x. WorkInfo -> Rep WorkInfo x
Generic,WorkInfo -> ()
(WorkInfo -> ()) -> NFData WorkInfo
forall a. (a -> ()) -> NFData a
rnf :: WorkInfo -> ()
$crnf :: WorkInfo -> ()
NFData,Eq WorkInfo
Eq WorkInfo
-> (Int -> WorkInfo -> Int)
-> (WorkInfo -> Int)
-> Hashable WorkInfo
Int -> WorkInfo -> Int
WorkInfo -> Int
forall a. Eq a -> (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: WorkInfo -> Int
$chash :: WorkInfo -> Int
hashWithSalt :: Int -> WorkInfo -> Int
$chashWithSalt :: Int -> WorkInfo -> Int
$cp1Hashable :: Eq WorkInfo
Hashable,Get WorkInfo
[WorkInfo] -> Put
WorkInfo -> Put
(WorkInfo -> Put)
-> Get WorkInfo -> ([WorkInfo] -> Put) -> Binary WorkInfo
forall t. (t -> Put) -> Get t -> ([t] -> Put) -> Binary t
putList :: [WorkInfo] -> Put
$cputList :: [WorkInfo] -> Put
get :: Get WorkInfo
$cget :: Get WorkInfo
put :: WorkInfo -> Put
$cput :: WorkInfo -> Put
Binary)

-- | Term reference
type TmName     = Name Term
-- | Binding in a LetRec construct
type LetBinding = (Id, Term)

data Bind a
  = NonRec Id a
  | Rec [(Id, a)]
  deriving (Bind a -> Bind a -> Bool
(Bind a -> Bind a -> Bool)
-> (Bind a -> Bind a -> Bool) -> Eq (Bind a)
forall a. Eq a => Bind a -> Bind a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Bind a -> Bind a -> Bool
$c/= :: forall a. Eq a => Bind a -> Bind a -> Bool
== :: Bind a -> Bind a -> Bool
$c== :: forall a. Eq a => Bind a -> Bind a -> Bool
Eq, Int -> Bind a -> ShowS
[Bind a] -> ShowS
Bind a -> String
(Int -> Bind a -> ShowS)
-> (Bind a -> String) -> ([Bind a] -> ShowS) -> Show (Bind a)
forall a. Show a => Int -> Bind a -> ShowS
forall a. Show a => [Bind a] -> ShowS
forall a. Show a => Bind a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Bind a] -> ShowS
$cshowList :: forall a. Show a => [Bind a] -> ShowS
show :: Bind a -> String
$cshow :: forall a. Show a => Bind a -> String
showsPrec :: Int -> Bind a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Bind a -> ShowS
Show, (forall x. Bind a -> Rep (Bind a) x)
-> (forall x. Rep (Bind a) x -> Bind a) -> Generic (Bind a)
forall x. Rep (Bind a) x -> Bind a
forall x. Bind a -> Rep (Bind a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (Bind a) x -> Bind a
forall a x. Bind a -> Rep (Bind a) x
$cto :: forall a x. Rep (Bind a) x -> Bind a
$cfrom :: forall a x. Bind a -> Rep (Bind a) x
Generic, Bind a -> ()
(Bind a -> ()) -> NFData (Bind a)
forall a. NFData a => Bind a -> ()
forall a. (a -> ()) -> NFData a
rnf :: Bind a -> ()
$crnf :: forall a. NFData a => Bind a -> ()
NFData, Eq (Bind a)
Eq (Bind a)
-> (Int -> Bind a -> Int) -> (Bind a -> Int) -> Hashable (Bind a)
Int -> Bind a -> Int
Bind a -> Int
forall a. Eq a -> (Int -> a -> Int) -> (a -> Int) -> Hashable a
forall a. Hashable a => Eq (Bind a)
forall a. Hashable a => Int -> Bind a -> Int
forall a. Hashable a => Bind a -> Int
hash :: Bind a -> Int
$chash :: forall a. Hashable a => Bind a -> Int
hashWithSalt :: Int -> Bind a -> Int
$chashWithSalt :: forall a. Hashable a => Int -> Bind a -> Int
$cp1Hashable :: forall a. Hashable a => Eq (Bind a)
Hashable, Get (Bind a)
[Bind a] -> Put
Bind a -> Put
(Bind a -> Put)
-> Get (Bind a) -> ([Bind a] -> Put) -> Binary (Bind a)
forall a. Binary a => Get (Bind a)
forall a. Binary a => [Bind a] -> Put
forall a. Binary a => Bind a -> Put
forall t. (t -> Put) -> Get t -> ([t] -> Put) -> Binary t
putList :: [Bind a] -> Put
$cputList :: forall a. Binary a => [Bind a] -> Put
get :: Get (Bind a)
$cget :: forall a. Binary a => Get (Bind a)
put :: Bind a -> Put
$cput :: forall a. Binary a => Bind a -> Put
Binary, a -> Bind b -> Bind a
(a -> b) -> Bind a -> Bind b
(forall a b. (a -> b) -> Bind a -> Bind b)
-> (forall a b. a -> Bind b -> Bind a) -> Functor Bind
forall a b. a -> Bind b -> Bind a
forall a b. (a -> b) -> Bind a -> Bind b
forall (f :: Type -> Type).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> Bind b -> Bind a
$c<$ :: forall a b. a -> Bind b -> Bind a
fmap :: (a -> b) -> Bind a -> Bind b
$cfmap :: forall a b. (a -> b) -> Bind a -> Bind b
Functor)
  -- Structural equivalence instead of alpha equivalance

-- | Patterns in the LHS of a case-decomposition
data Pat
  = DataPat !DataCon [TyVar] [Id]
  -- ^ Datatype pattern, '[TyVar]' bind existentially-quantified
  -- type-variables of a DataCon
  | LitPat  !Literal
  -- ^ Literal pattern
  | DefaultPat
  -- ^ Default pattern
  deriving (Pat -> Pat -> Bool
(Pat -> Pat -> Bool) -> (Pat -> Pat -> Bool) -> Eq Pat
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Pat -> Pat -> Bool
$c/= :: Pat -> Pat -> Bool
== :: Pat -> Pat -> Bool
$c== :: Pat -> Pat -> Bool
Eq, Eq Pat
Eq Pat
-> (Pat -> Pat -> Ordering)
-> (Pat -> Pat -> Bool)
-> (Pat -> Pat -> Bool)
-> (Pat -> Pat -> Bool)
-> (Pat -> Pat -> Bool)
-> (Pat -> Pat -> Pat)
-> (Pat -> Pat -> Pat)
-> Ord Pat
Pat -> Pat -> Bool
Pat -> Pat -> Ordering
Pat -> Pat -> Pat
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Pat -> Pat -> Pat
$cmin :: Pat -> Pat -> Pat
max :: Pat -> Pat -> Pat
$cmax :: Pat -> Pat -> Pat
>= :: Pat -> Pat -> Bool
$c>= :: Pat -> Pat -> Bool
> :: Pat -> Pat -> Bool
$c> :: Pat -> Pat -> Bool
<= :: Pat -> Pat -> Bool
$c<= :: Pat -> Pat -> Bool
< :: Pat -> Pat -> Bool
$c< :: Pat -> Pat -> Bool
compare :: Pat -> Pat -> Ordering
$ccompare :: Pat -> Pat -> Ordering
$cp1Ord :: Eq Pat
Ord, Int -> Pat -> ShowS
[Pat] -> ShowS
Pat -> String
(Int -> Pat -> ShowS)
-> (Pat -> String) -> ([Pat] -> ShowS) -> Show Pat
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Pat] -> ShowS
$cshowList :: [Pat] -> ShowS
show :: Pat -> String
$cshow :: Pat -> String
showsPrec :: Int -> Pat -> ShowS
$cshowsPrec :: Int -> Pat -> ShowS
Show, (forall x. Pat -> Rep Pat x)
-> (forall x. Rep Pat x -> Pat) -> Generic Pat
forall x. Rep Pat x -> Pat
forall x. Pat -> Rep Pat x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Pat x -> Pat
$cfrom :: forall x. Pat -> Rep Pat x
Generic, Pat -> ()
(Pat -> ()) -> NFData Pat
forall a. (a -> ()) -> NFData a
rnf :: Pat -> ()
$crnf :: Pat -> ()
NFData, Get Pat
[Pat] -> Put
Pat -> Put
(Pat -> Put) -> Get Pat -> ([Pat] -> Put) -> Binary Pat
forall t. (t -> Put) -> Get t -> ([t] -> Put) -> Binary t
putList :: [Pat] -> Put
$cputList :: [Pat] -> Put
get :: Get Pat
$cget :: Get Pat
put :: Pat -> Put
$cput :: Pat -> Put
Binary)

type Alt = (Pat,Term)

-- | Get the list of term-binders out of a DataType pattern
patIds :: Pat -> ([TyVar],[Id])
patIds :: Pat -> ([TyVar], [Id])
patIds (DataPat DataCon
_  [TyVar]
tvs [Id]
ids) = ([TyVar]
tvs,[Id]
ids)
patIds Pat
_                    = ([],[])

patVars :: Pat -> [Var a]
patVars :: Pat -> [Var a]
patVars (DataPat DataCon
_ [TyVar]
tvs [Id]
ids) = [TyVar] -> [Var a]
coerce [TyVar]
tvs [Var a] -> [Var a] -> [Var a]
forall a. [a] -> [a] -> [a]
++ [Id] -> [Var a]
coerce [Id]
ids
patVars Pat
_ = []

-- | Abstract a term over a list of term and type variables
mkAbstraction :: Term -> [Either Id TyVar] -> Term
mkAbstraction :: Term -> [Either Id TyVar] -> Term
mkAbstraction = (Either Id TyVar -> Term -> Term)
-> Term -> [Either Id TyVar] -> Term
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Id -> Term -> Term)
-> (TyVar -> Term -> Term) -> Either Id TyVar -> Term -> Term
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either Id -> Term -> Term
Lam TyVar -> Term -> Term
TyLam)

-- | Abstract a term over a list of type variables
mkTyLams :: Term -> [TyVar] -> Term
mkTyLams :: Term -> [TyVar] -> Term
mkTyLams Term
tm = Term -> [Either Id TyVar] -> Term
mkAbstraction Term
tm ([Either Id TyVar] -> Term)
-> ([TyVar] -> [Either Id TyVar]) -> [TyVar] -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TyVar -> Either Id TyVar) -> [TyVar] -> [Either Id TyVar]
forall a b. (a -> b) -> [a] -> [b]
map TyVar -> Either Id TyVar
forall a b. b -> Either a b
Right

-- | Abstract a term over a list of variables
mkLams :: Term -> [Id] -> Term
mkLams :: Term -> [Id] -> Term
mkLams Term
tm = Term -> [Either Id TyVar] -> Term
mkAbstraction Term
tm ([Either Id TyVar] -> Term)
-> ([Id] -> [Either Id TyVar]) -> [Id] -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Id -> Either Id TyVar) -> [Id] -> [Either Id TyVar]
forall a b. (a -> b) -> [a] -> [b]
map Id -> Either Id TyVar
forall a b. a -> Either a b
Left

-- | Apply a list of types and terms to a term
mkApps :: Term -> [Either Term Type] -> Term
mkApps :: Term -> [Either Term Type] -> Term
mkApps = (Term -> Either Term Type -> Term)
-> Term -> [Either Term Type] -> Term
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Term
e Either Term Type
a -> (Term -> Term) -> (Type -> Term) -> Either Term Type -> Term
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Term -> Term -> Term
App Term
e) (Term -> Type -> Term
TyApp Term
e) Either Term Type
a)

-- | Apply a list of terms to a term
mkTmApps :: Term -> [Term] -> Term
mkTmApps :: Term -> [Term] -> Term
mkTmApps = (Term -> Term -> Term) -> Term -> [Term] -> Term
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Term -> Term -> Term
App

-- | Apply a list of types to a term
mkTyApps :: Term -> [Type] -> Term
mkTyApps :: Term -> [Type] -> Term
mkTyApps = (Term -> Type -> Term) -> Term -> [Type] -> Term
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Term -> Type -> Term
TyApp

mkTicks :: Term -> [TickInfo] -> Term
mkTicks :: Term -> [TickInfo] -> Term
mkTicks Term
tm [TickInfo]
ticks = (Term -> TickInfo -> Term) -> Term -> [TickInfo] -> Term
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Term
e TickInfo
s -> TickInfo -> Term -> Term
Tick TickInfo
s Term
e) Term
tm ([TickInfo] -> [TickInfo]
forall a. Eq a => [a] -> [a]
nub [TickInfo]
ticks)

-- | Context in which a term appears
data CoreContext
  = AppFun
  -- ^ Function position of an application
  | AppArg (Maybe (Text, Int, Int))
  -- ^ Argument position of an application. If this is an argument applied to
  -- a primitive, a tuple is defined containing (name of the primitive, #type
  -- args, #term args)
  | TyAppC
  -- ^ Function position of a type application
  | LetBinding Id [Id]
  -- ^ RHS of a Let-binder with the sibling LHS'
  | LetBody [LetBinding]
  -- ^ Body of a Let-binding with the bound LHS'
  | LamBody Id
  -- ^ Body of a lambda-term with the abstracted variable
  | TyLamBody TyVar
  -- ^ Body of a TyLambda-term with the abstracted type-variable
  | CaseAlt Pat
  -- ^ RHS of a case-alternative with the bound pattern on the LHS
  | CaseScrut
  -- ^ Subject of a case-decomposition
  | CastBody
  -- ^ Body of a Cast
  | TickC TickInfo
  -- ^ Body of a Tick
  deriving (Int -> CoreContext -> ShowS
[CoreContext] -> ShowS
CoreContext -> String
(Int -> CoreContext -> ShowS)
-> (CoreContext -> String)
-> ([CoreContext] -> ShowS)
-> Show CoreContext
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CoreContext] -> ShowS
$cshowList :: [CoreContext] -> ShowS
show :: CoreContext -> String
$cshow :: CoreContext -> String
showsPrec :: Int -> CoreContext -> ShowS
$cshowsPrec :: Int -> CoreContext -> ShowS
Show, (forall x. CoreContext -> Rep CoreContext x)
-> (forall x. Rep CoreContext x -> CoreContext)
-> Generic CoreContext
forall x. Rep CoreContext x -> CoreContext
forall x. CoreContext -> Rep CoreContext x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep CoreContext x -> CoreContext
$cfrom :: forall x. CoreContext -> Rep CoreContext x
Generic, CoreContext -> ()
(CoreContext -> ()) -> NFData CoreContext
forall a. (a -> ()) -> NFData a
rnf :: CoreContext -> ()
$crnf :: CoreContext -> ()
NFData, Get CoreContext
[CoreContext] -> Put
CoreContext -> Put
(CoreContext -> Put)
-> Get CoreContext -> ([CoreContext] -> Put) -> Binary CoreContext
forall t. (t -> Put) -> Get t -> ([t] -> Put) -> Binary t
putList :: [CoreContext] -> Put
$cputList :: [CoreContext] -> Put
get :: Get CoreContext
$cget :: Get CoreContext
put :: CoreContext -> Put
$cput :: CoreContext -> Put
Binary)

-- | A list of @CoreContext@ describes the complete navigation path from the
-- top-level to a specific sub-expression.
type Context = [CoreContext]

-- [Note] Custom @Eq@ instance for @CoreContext@
--
-- We need a manual equality instance here, due to the argument of `AppArg`.
-- Specifically, it is the only piece of information kept in `CoreContext`,
-- which references information about its children, breaking the invariant
-- that contexts represent a navigation to a specific sub-expression.
--
-- One would expect equal contexts to navigate to the same place, but if
-- these navigate to an argument position that contains different children,
-- we will get inequality from the derived `Eq`.
instance Eq CoreContext where
  CoreContext
c == :: CoreContext -> CoreContext -> Bool
== CoreContext
c' = case (CoreContext
c, CoreContext
c') of
    (CoreContext
AppFun,          CoreContext
AppFun)            -> Bool
True
    (AppArg Maybe (Text, Int, Int)
_,        AppArg Maybe (Text, Int, Int)
_)          -> Bool
True
    -- NB: we do not see inside the argument here
    (CoreContext
TyAppC,          CoreContext
TyAppC)            -> Bool
True
    (LetBinding Id
i [Id]
is, LetBinding Id
i' [Id]
is') -> Id
i Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
i' Bool -> Bool -> Bool
&& [Id]
is [Id] -> [Id] -> Bool
forall a. Eq a => a -> a -> Bool
== [Id]
is'
    (LetBody [LetBinding]
is,      LetBody [LetBinding]
is')       -> (LetBinding -> Id) -> [LetBinding] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map LetBinding -> Id
forall a b. (a, b) -> a
fst [LetBinding]
is [Id] -> [Id] -> Bool
forall a. Eq a => a -> a -> Bool
== (LetBinding -> Id) -> [LetBinding] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map LetBinding -> Id
forall a b. (a, b) -> a
fst [LetBinding]
is'
    (LamBody Id
i,       LamBody Id
i')        -> Id
i Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
i'
    (TyLamBody TyVar
tv,    TyLamBody TyVar
tv')     -> TyVar
tv TyVar -> TyVar -> Bool
forall a. Eq a => a -> a -> Bool
== TyVar
tv'
    (CaseAlt Pat
p,       CaseAlt Pat
p')        -> Pat
p Pat -> Pat -> Bool
forall a. Eq a => a -> a -> Bool
== Pat
p'
    (CoreContext
CaseScrut,       CoreContext
CaseScrut)         -> Bool
True
    (CoreContext
CastBody,        CoreContext
CastBody)          -> Bool
True
    (TickC TickInfo
sp,        TickC TickInfo
sp')         -> TickInfo
sp TickInfo -> TickInfo -> Bool
forall a. Eq a => a -> a -> Bool
== TickInfo
sp'
    (CoreContext
_,               CoreContext
_)                 -> Bool
False

-- | Is the Context a Lambda/Term-abstraction context?
isLambdaBodyCtx :: CoreContext -> Bool
isLambdaBodyCtx :: CoreContext -> Bool
isLambdaBodyCtx (LamBody Id
_) = Bool
True
isLambdaBodyCtx CoreContext
_           = Bool
False

-- | Is the Context a Tick context?
isTickCtx :: CoreContext -> Bool
isTickCtx :: CoreContext -> Bool
isTickCtx (TickC TickInfo
_) = Bool
True
isTickCtx CoreContext
_         = Bool
False

stripTicks :: Term -> Term
stripTicks :: Term -> Term
stripTicks (Tick TickInfo
_ Term
e) = Term -> Term
stripTicks Term
e
stripTicks Term
e = Term
e

-- | Like 'stripTicks' but removes all ticks from subexpressions.
stripAllTicks :: Term -> Term
stripAllTicks :: Term -> Term
stripAllTicks = Term -> Term
go
 where
  go :: Term -> Term
go (Lam Id
i Term
x) = Id -> Term -> Term
Lam Id
i (Term -> Term
go Term
x)
  go (TyLam TyVar
i Term
x) = TyVar -> Term -> Term
TyLam TyVar
i (Term -> Term
go Term
x)
  go (App Term
f Term
x) = Term -> Term -> Term
App (Term -> Term
go Term
f) (Term -> Term
go Term
x)
  go (TyApp Term
f Type
a) = Term -> Type -> Term
TyApp (Term -> Term
go Term
f) Type
a
  go (Let Bind Term
bs Term
x) = Bind Term -> Term -> Term
Let (Bind Term -> Bind Term
goBinds Bind Term
bs) (Term -> Term
go Term
x)
  go (Case Term
x Type
ty [Alt]
alts) = Term -> Type -> [Alt] -> Term
Case (Term -> Term
go Term
x) Type
ty ((Term -> Term) -> Alt -> Alt
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Term -> Term
go (Alt -> Alt) -> [Alt] -> [Alt]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> [Alt]
alts)
  go (Cast Term
x Type
a Type
b) = Term -> Type -> Type -> Term
Cast (Term -> Term
go Term
x) Type
a Type
b
  go (Tick TickInfo
_ Term
x) = Term -> Term
go Term
x
  go Term
x = Term
x

  goBinds :: Bind Term -> Bind Term
goBinds (NonRec Id
i Term
x) = Id -> Term -> Bind Term
forall a. Id -> a -> Bind a
NonRec Id
i (Term -> Term
go Term
x)
  goBinds (Rec [LetBinding]
ixs) = [LetBinding] -> Bind Term
forall a. [(Id, a)] -> Bind a
Rec ((Term -> Term) -> LetBinding -> LetBinding
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Term -> Term
go (LetBinding -> LetBinding) -> [LetBinding] -> [LetBinding]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> [LetBinding]
ixs)

-- | Split a (Type)Application in the applied term and it arguments
collectArgs :: Term -> (Term, [Either Term Type])
collectArgs :: Term -> (Term, [Either Term Type])
collectArgs = [Either Term Type] -> Term -> (Term, [Either Term Type])
go []
  where
    go :: [Either Term Type] -> Term -> (Term, [Either Term Type])
go [Either Term Type]
args (App Term
e1 Term
e2) = [Either Term Type] -> Term -> (Term, [Either Term Type])
go (Term -> Either Term Type
forall a b. a -> Either a b
Left Term
e2Either Term Type -> [Either Term Type] -> [Either Term Type]
forall a. a -> [a] -> [a]
:[Either Term Type]
args) Term
e1
    go [Either Term Type]
args (TyApp Term
e Type
t) = [Either Term Type] -> Term -> (Term, [Either Term Type])
go (Type -> Either Term Type
forall a b. b -> Either a b
Right Type
tEither Term Type -> [Either Term Type] -> [Either Term Type]
forall a. a -> [a] -> [a]
:[Either Term Type]
args) Term
e
    go [Either Term Type]
args (Tick TickInfo
_ Term
e)  = [Either Term Type] -> Term -> (Term, [Either Term Type])
go [Either Term Type]
args Term
e
    go [Either Term Type]
args Term
e           = (Term
e, [Either Term Type]
args)

collectTicks :: Term -> (Term, [TickInfo])
collectTicks :: Term -> (Term, [TickInfo])
collectTicks = [TickInfo] -> Term -> (Term, [TickInfo])
go []
 where
  go :: [TickInfo] -> Term -> (Term, [TickInfo])
go [TickInfo]
ticks (Tick TickInfo
s Term
e) = [TickInfo] -> Term -> (Term, [TickInfo])
go (TickInfo
sTickInfo -> [TickInfo] -> [TickInfo]
forall a. a -> [a] -> [a]
:[TickInfo]
ticks) Term
e
  go [TickInfo]
ticks Term
e          = (Term
e,[TickInfo]
ticks)

collectArgsTicks :: Term -> (Term, [Either Term Type], [TickInfo])
collectArgsTicks :: Term -> (Term, [Either Term Type], [TickInfo])
collectArgsTicks = [Either Term Type]
-> [TickInfo] -> Term -> (Term, [Either Term Type], [TickInfo])
go [] []
 where
  go :: [Either Term Type]
-> [TickInfo] -> Term -> (Term, [Either Term Type], [TickInfo])
go [Either Term Type]
args [TickInfo]
ticks (App Term
e1 Term
e2) = [Either Term Type]
-> [TickInfo] -> Term -> (Term, [Either Term Type], [TickInfo])
go (Term -> Either Term Type
forall a b. a -> Either a b
Left Term
e2Either Term Type -> [Either Term Type] -> [Either Term Type]
forall a. a -> [a] -> [a]
:[Either Term Type]
args) [TickInfo]
ticks     Term
e1
  go [Either Term Type]
args [TickInfo]
ticks (TyApp Term
e Type
t) = [Either Term Type]
-> [TickInfo] -> Term -> (Term, [Either Term Type], [TickInfo])
go (Type -> Either Term Type
forall a b. b -> Either a b
Right Type
tEither Term Type -> [Either Term Type] -> [Either Term Type]
forall a. a -> [a] -> [a]
:[Either Term Type]
args) [TickInfo]
ticks     Term
e
  go [Either Term Type]
args [TickInfo]
ticks (Tick TickInfo
s Term
e)  = [Either Term Type]
-> [TickInfo] -> Term -> (Term, [Either Term Type], [TickInfo])
go [Either Term Type]
args           (TickInfo
sTickInfo -> [TickInfo] -> [TickInfo]
forall a. a -> [a] -> [a]
:[TickInfo]
ticks) Term
e
  go [Either Term Type]
args [TickInfo]
ticks Term
e           = (Term
e, [Either Term Type]
args, [TickInfo]
ticks)

-- | Split a (Type)Abstraction in the bound variables and the abstracted term
collectBndrs :: Term -> ([Either Id TyVar], Term)
collectBndrs :: Term -> ([Either Id TyVar], Term)
collectBndrs = [Either Id TyVar] -> Term -> ([Either Id TyVar], Term)
go []
 where
  go :: [Either Id TyVar] -> Term -> ([Either Id TyVar], Term)
go [Either Id TyVar]
bs (Lam Id
v Term
e')    = [Either Id TyVar] -> Term -> ([Either Id TyVar], Term)
go (Id -> Either Id TyVar
forall a b. a -> Either a b
Left Id
vEither Id TyVar -> [Either Id TyVar] -> [Either Id TyVar]
forall a. a -> [a] -> [a]
:[Either Id TyVar]
bs) Term
e'
  go [Either Id TyVar]
bs (TyLam TyVar
tv Term
e') = [Either Id TyVar] -> Term -> ([Either Id TyVar], Term)
go (TyVar -> Either Id TyVar
forall a b. b -> Either a b
Right TyVar
tvEither Id TyVar -> [Either Id TyVar] -> [Either Id TyVar]
forall a. a -> [a] -> [a]
:[Either Id TyVar]
bs) Term
e'
  go [Either Id TyVar]
bs Term
e'            = ([Either Id TyVar] -> [Either Id TyVar]
forall a. [a] -> [a]
reverse [Either Id TyVar]
bs,Term
e')

-- | Given a function application, find the primitive it's applied. Yields
-- Nothing if given term is not an application or if it is not a primitive.
primArg
  :: Term
  -- ^ Function application
  -> Maybe (Text, Int, Int)
  -- ^ If @Term@ was a primitive: (name of primitive, #type args, #term args)
primArg :: Term -> Maybe (Text, Int, Int)
primArg (Term -> (Term, [Either Term Type])
collectArgs -> (Term, [Either Term Type])
t) =
  case (Term, [Either Term Type])
t of
    (Prim PrimInfo
p, [Either Term Type]
args) ->
      (Text, Int, Int) -> Maybe (Text, Int, Int)
forall a. a -> Maybe a
Just (PrimInfo -> Text
primName PrimInfo
p, [Type] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length ([Either Term Type] -> [Type]
forall a b. [Either a b] -> [b]
rights [Either Term Type]
args), [Term] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length ([Either Term Type] -> [Term]
forall a b. [Either a b] -> [a]
lefts [Either Term Type]
args))
    (Term, [Either Term Type])
_ ->
      Maybe (Text, Int, Int)
forall a. Maybe a
Nothing

-- | Partition ticks in source ticks and nameMod ticks
partitionTicks
  :: [TickInfo]
  -> ([TickInfo], [TickInfo])
  -- ^ (source ticks, nameMod ticks)
partitionTicks :: [TickInfo] -> ([TickInfo], [TickInfo])
partitionTicks = (TickInfo -> Bool) -> [TickInfo] -> ([TickInfo], [TickInfo])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (\case {SrcSpan {} -> Bool
True; TickInfo
_ -> Bool
False})

-- | Visit all terms in a term, testing it with a predicate, and returning
-- a list of predicate yields.
walkTerm :: forall a . (Term -> Maybe a) -> Term -> [a]
walkTerm :: (Term -> Maybe a) -> Term -> [a]
walkTerm Term -> Maybe a
f = [Maybe a] -> [a]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe a] -> [a]) -> (Term -> [Maybe a]) -> Term -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DList (Maybe a) -> [Maybe a]
forall a. DList a -> [a]
DList.toList (DList (Maybe a) -> [Maybe a])
-> (Term -> DList (Maybe a)) -> Term -> [Maybe a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> DList (Maybe a)
go
 where
  go :: Term -> DList.DList (Maybe a)
  go :: Term -> DList (Maybe a)
go Term
t = Maybe a -> DList (Maybe a) -> DList (Maybe a)
forall a. a -> DList a -> DList a
DList.cons (Term -> Maybe a
f Term
t) (DList (Maybe a) -> DList (Maybe a))
-> DList (Maybe a) -> DList (Maybe a)
forall a b. (a -> b) -> a -> b
$ case Term
t of
    Var Id
_ -> DList (Maybe a)
forall a. Monoid a => a
mempty
    Data DataCon
_ -> DList (Maybe a)
forall a. Monoid a => a
mempty
    Literal Literal
_ -> DList (Maybe a)
forall a. Monoid a => a
mempty
    Prim PrimInfo
_ -> DList (Maybe a)
forall a. Monoid a => a
mempty
    Lam Id
_ Term
t1 -> Term -> DList (Maybe a)
go Term
t1
    TyLam TyVar
_ Term
t1 -> Term -> DList (Maybe a)
go Term
t1
    App Term
t1 Term
t2 -> Term -> DList (Maybe a)
go Term
t1 DList (Maybe a) -> DList (Maybe a) -> DList (Maybe a)
forall a. Semigroup a => a -> a -> a
<> Term -> DList (Maybe a)
go Term
t2
    TyApp Term
t1 Type
_ -> Term -> DList (Maybe a)
go Term
t1
    Let (NonRec Id
_ Term
x) Term
t1 -> Term -> DList (Maybe a)
go Term
t1 DList (Maybe a) -> DList (Maybe a) -> DList (Maybe a)
forall a. Semigroup a => a -> a -> a
<> Term -> DList (Maybe a)
go Term
x
    Let (Rec [LetBinding]
bndrs) Term
t1 -> Term -> DList (Maybe a)
go Term
t1 DList (Maybe a) -> DList (Maybe a) -> DList (Maybe a)
forall a. Semigroup a => a -> a -> a
<> [DList (Maybe a)] -> DList (Maybe a)
forall a. Monoid a => [a] -> a
mconcat ((LetBinding -> DList (Maybe a))
-> [LetBinding] -> [DList (Maybe a)]
forall a b. (a -> b) -> [a] -> [b]
map (Term -> DList (Maybe a)
go (Term -> DList (Maybe a))
-> (LetBinding -> Term) -> LetBinding -> DList (Maybe a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LetBinding -> Term
forall a b. (a, b) -> b
snd) [LetBinding]
bndrs)
    Case Term
t1 Type
_ [Alt]
alts -> Term -> DList (Maybe a)
go Term
t1 DList (Maybe a) -> DList (Maybe a) -> DList (Maybe a)
forall a. Semigroup a => a -> a -> a
<> [DList (Maybe a)] -> DList (Maybe a)
forall a. Monoid a => [a] -> a
mconcat ((Alt -> DList (Maybe a)) -> [Alt] -> [DList (Maybe a)]
forall a b. (a -> b) -> [a] -> [b]
map (Term -> DList (Maybe a)
go (Term -> DList (Maybe a))
-> (Alt -> Term) -> Alt -> DList (Maybe a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alt -> Term
forall a b. (a, b) -> b
snd) [Alt]
alts)
    Cast Term
t1 Type
_ Type
_ -> Term -> DList (Maybe a)
go Term
t1
    Tick TickInfo
_ Term
t1 -> Term -> DList (Maybe a)
go Term
t1

-- Collect all term ids mentioned in a term
collectTermIds :: Term -> [Id]
collectTermIds :: Term -> [Id]
collectTermIds = [[Id]] -> [Id]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat ([[Id]] -> [Id]) -> (Term -> [[Id]]) -> Term -> [Id]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Term -> Maybe [Id]) -> Term -> [[Id]]
forall a. (Term -> Maybe a) -> Term -> [a]
walkTerm ([Id] -> Maybe [Id]
forall a. a -> Maybe a
Just ([Id] -> Maybe [Id]) -> (Term -> [Id]) -> Term -> Maybe [Id]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> [Id]
go)
 where
  go :: Term -> [Id]
  go :: Term -> [Id]
go (Var Id
i) = [Id
i]
  go (Lam Id
i Term
_) = [Id
i]
  go (Let (NonRec Id
i Term
_) Term
_) = [Id
i]
  go (Let (Rec [LetBinding]
bndrs) Term
_) = (LetBinding -> Id) -> [LetBinding] -> [Id]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap LetBinding -> Id
forall a b. (a, b) -> a
fst [LetBinding]
bndrs
  go (Case Term
_ Type
_ [Alt]
alts) = (Alt -> [Id]) -> [Alt] -> [Id]
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> [b]) -> t a -> [b]
concatMap (Pat -> [Id]
pat (Pat -> [Id]) -> (Alt -> Pat) -> Alt -> [Id]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alt -> Pat
forall a b. (a, b) -> a
fst) [Alt]
alts
  go (Data DataCon
_) = []
  go (Literal Literal
_) = []
  go (Prim PrimInfo
_) = []
  go (TyLam TyVar
_ Term
_) = []
  go (App Term
_ Term
_) = []
  go (TyApp Term
_ Type
_) = []
  go (Cast Term
_ Type
_ Type
_) = []
  go (Tick TickInfo
_ Term
_) = []

  pat :: Pat -> [Id]
  pat :: Pat -> [Id]
pat (DataPat DataCon
_ [TyVar]
_ [Id]
ids) = [Id]
ids
  pat (LitPat Literal
_) = []
  pat Pat
DefaultPat = []

-- | Make a term variable out of a variable reference or ticked variable
-- reference
varToId :: Term -> Id
varToId :: Term -> Id
varToId = \case
  Var Id
i    -> Id
i
  Tick TickInfo
_ Term
e -> Term -> Id
varToId Term
e
  Term
e        -> String -> Id
forall a. HasCallStack => String -> a
error (String -> Id) -> String -> Id
forall a b. (a -> b) -> a -> b
$ $(String
curLoc) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"varToId: not a var: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Term -> String
forall a. Show a => a -> String
show Term
e