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

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

{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell #-}

module Clash.Core.Term
  ( Term (..)
  , mkAbstraction
  , mkTyLams
  , mkLams
  , mkApps
  , mkTyApps
  , mkTmApps
  , mkTicks
  , TmName
  , idToVar
  , varToId
  , LetBinding
  , Pat (..)
  , patIds
  , patVars
  , Alt
  , TickInfo (..)
  , stripTicks
  , partitionTicks
  , NameMod (..)
  , PrimInfo (..)
  , 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)
import Data.Foldable                           (foldl')
import Data.Maybe                              (catMaybes)
import Data.Hashable                           (Hashable)
import Data.List                               (nub, partition)
import Data.Text                               (Text)
import GHC.Generics
import SrcLoc                                  (SrcSpan)

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

-- | 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
  | Letrec  [LetBinding] 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,Int -> Term -> Int
Term -> Int
(Int -> Term -> Int) -> (Term -> Int) -> Hashable Term
forall a. (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: Term -> Int
$chash :: Term -> Int
hashWithSalt :: Int -> Term -> Int
$chashWithSalt :: Int -> Term -> Int
Hashable,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)

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,Int -> TickInfo -> Int
TickInfo -> Int
(Int -> TickInfo -> Int) -> (TickInfo -> Int) -> Hashable TickInfo
forall a. (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: TickInfo -> Int
$chash :: TickInfo -> Int
hashWithSalt :: Int -> TickInfo -> Int
$chashWithSalt :: Int -> TickInfo -> Int
Hashable,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)

-- | 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,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,Int -> NameMod -> Int
NameMod -> Int
(Int -> NameMod -> Int) -> (NameMod -> Int) -> Hashable NameMod
forall a. (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: NameMod -> Int
$chash :: NameMod -> Int
hashWithSalt :: Int -> NameMod -> Int
$chashWithSalt :: Int -> NameMod -> Int
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 PrimInfo = PrimInfo
  { PrimInfo -> Text
primName     :: !Text
  , PrimInfo -> Type
primType     :: !Type
  , PrimInfo -> WorkInfo
primWorkInfo :: !WorkInfo
  } 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,Int -> PrimInfo -> Int
PrimInfo -> Int
(Int -> PrimInfo -> Int) -> (PrimInfo -> Int) -> Hashable PrimInfo
forall a. (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: PrimInfo -> Int
$chash :: PrimInfo -> Int
hashWithSalt :: Int -> PrimInfo -> Int
$chashWithSalt :: Int -> PrimInfo -> Int
Hashable,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 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
  deriving (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,Int -> WorkInfo -> Int
WorkInfo -> Int
(Int -> WorkInfo -> Int) -> (WorkInfo -> Int) -> Hashable WorkInfo
forall a. (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: WorkInfo -> Int
$chash :: WorkInfo -> Int
hashWithSalt :: Int -> WorkInfo -> Int
$chashWithSalt :: Int -> WorkInfo -> Int
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)

-- | 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,Int -> Pat -> Int
Pat -> Int
(Int -> Pat -> Int) -> (Pat -> Int) -> Hashable Pat
forall a. (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: Pat -> Int
$chash :: Pat -> Int
hashWithSalt :: Int -> Pat -> Int
$chashWithSalt :: Int -> Pat -> Int
Hashable,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 term 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 type 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 [Id]
  -- ^ 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, Int -> CoreContext -> Int
CoreContext -> Int
(Int -> CoreContext -> Int)
-> (CoreContext -> Int) -> Hashable CoreContext
forall a. (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: CoreContext -> Int
$chash :: CoreContext -> Int
hashWithSalt :: Int -> CoreContext -> Int
$chashWithSalt :: Int -> CoreContext -> Int
Hashable, 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 [Id]
is,      LetBody [Id]
is')       -> [Id]
is [Id] -> [Id] -> Bool
forall a. Eq a => a -> a -> Bool
== [Id]
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

-- | 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
    Letrec [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 (Letrec [LetBinding]
bndrs Term
_) = (LetBinding -> Id) -> [LetBinding] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map 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 variable reference out of term variable
idToVar :: Id -> Term
idToVar :: Id -> Term
idToVar i :: Id
i@(Id {}) = Id -> Term
Var Id
i
idToVar Id
tv        = String -> Term
forall a. HasCallStack => String -> a
error (String -> Term) -> String -> Term
forall a b. (a -> b) -> a -> b
$ $(String
curLoc) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"idToVar: tyVar: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Id -> String
forall a. Show a => a -> String
show Id
tv

-- | Make a term variable out of a variable reference
varToId :: Term -> Id
varToId :: Term -> Id
varToId (Var Id
i) = Id
i
varToId 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