{-|
  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 (..)
  , TmName
  , LetBinding
  , Pat (..)
  , Alt
  , TickInfo (..), NameMod (..)
  , PrimInfo (..)
  , WorkInfo (..)
  , CoreContext (..), Context, isLambdaBodyCtx, isTickCtx, walkTerm
  , collectArgs, collectArgsTicks, collectTicks, collectTermIds, primArg
  , partitionTicks
  )
where

-- External Modules
import Control.DeepSeq
import Data.Binary                             (Binary)
import qualified Data.DList                    as DList
import Data.Either                             (lefts, rights)
import Data.Maybe                              (catMaybes)
import Data.Hashable                           (Hashable)
import Data.List                               (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          (Type)
import Clash.Core.Var                          (Id, TyVar)

-- | 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)

-- | 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
  c :: CoreContext
c == :: CoreContext -> CoreContext -> Bool
== c' :: CoreContext
c' = case (CoreContext
c, CoreContext
c') of
    (AppFun,          AppFun)            -> Bool
True
    (AppArg _,        AppArg _)          -> Bool
True
    -- NB: we do not see inside the argument here
    (TyAppC,          TyAppC)            -> Bool
True
    (LetBinding i :: Id
i is :: [Id]
is, LetBinding i' :: Id
i' is' :: [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 is :: [Id]
is,      LetBody is' :: [Id]
is')       -> [Id]
is [Id] -> [Id] -> Bool
forall a. Eq a => a -> a -> Bool
== [Id]
is'
    (LamBody i :: Id
i,       LamBody i' :: Id
i')        -> Id
i Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
i'
    (TyLamBody tv :: TyVar
tv,    TyLamBody tv' :: TyVar
tv')     -> TyVar
tv TyVar -> TyVar -> Bool
forall a. Eq a => a -> a -> Bool
== TyVar
tv'
    (CaseAlt p :: Pat
p,       CaseAlt p' :: Pat
p')        -> Pat
p Pat -> Pat -> Bool
forall a. Eq a => a -> a -> Bool
== Pat
p'
    (CaseScrut,       CaseScrut)         -> Bool
True
    (CastBody,        CastBody)          -> Bool
True
    (TickC sp :: TickInfo
sp,        TickC sp' :: TickInfo
sp')         -> TickInfo
sp TickInfo -> TickInfo -> Bool
forall a. Eq a => a -> a -> Bool
== TickInfo
sp'
    (_,               _)                 -> Bool
False

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

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

-- | 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 args :: [Either Term Type]
args (App e1 :: Term
e1 e2 :: 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 args :: [Either Term Type]
args (TyApp e :: Term
e t :: 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 args :: [Either Term Type]
args (Tick _ e :: Term
e)  = [Either Term Type] -> Term -> (Term, [Either Term Type])
go [Either Term Type]
args Term
e
    go args :: [Either Term Type]
args e :: 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 ticks :: [TickInfo]
ticks (Tick s :: TickInfo
s e :: Term
e) = [TickInfo] -> Term -> (Term, [TickInfo])
go (TickInfo
sTickInfo -> [TickInfo] -> [TickInfo]
forall a. a -> [a] -> [a]
:[TickInfo]
ticks) Term
e
  go ticks :: [TickInfo]
ticks e :: 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 args :: [Either Term Type]
args ticks :: [TickInfo]
ticks (App e1 :: Term
e1 e2 :: 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 args :: [Either Term Type]
args ticks :: [TickInfo]
ticks (TyApp e :: Term
e t :: 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 args :: [Either Term Type]
args ticks :: [TickInfo]
ticks (Tick s :: TickInfo
s e :: 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 args :: [Either Term Type]
args ticks :: [TickInfo]
ticks e :: Term
e           = (Term
e, [Either Term Type]
args, [TickInfo]
ticks)

-- | 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 p :: PrimInfo
p, args :: [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))
    _ ->
      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; _ -> 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 f :: 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 t :: 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 _ -> DList (Maybe a)
forall a. Monoid a => a
mempty
    Data _ -> DList (Maybe a)
forall a. Monoid a => a
mempty
    Literal _ -> DList (Maybe a)
forall a. Monoid a => a
mempty
    Prim _ -> DList (Maybe a)
forall a. Monoid a => a
mempty
    Lam _ t1 :: Term
t1 -> Term -> DList (Maybe a)
go Term
t1
    TyLam _ t1 :: Term
t1 -> Term -> DList (Maybe a)
go Term
t1
    App t1 :: Term
t1 t2 :: 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 t1 :: Term
t1 _ -> Maybe a -> DList (Maybe a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> Maybe a
f Term
t1)
    Letrec bndrs :: [LetBinding]
bndrs t1 :: Term
t1 -> Maybe a -> DList (Maybe a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> Maybe a
f 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 t1 :: Term
t1 _ alts :: [Alt]
alts -> Maybe a -> DList (Maybe a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> Maybe a
f 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 t1 :: Term
t1 _ _ -> Term -> DList (Maybe a)
go Term
t1
    Tick _ t1 :: Term
t1 -> Maybe a -> DList (Maybe a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term -> Maybe a
f 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 i :: Id
i) = [Id
i]
  go (Lam i :: Id
i _) = [Id
i]
  go (Letrec bndrs :: [LetBinding]
bndrs _) = (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 _ _ alts :: [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 _) = []
  go (Literal _) = []
  go (Prim _) = []
  go (TyLam _ _) = []
  go (App _ _) = []
  go (TyApp _ _) = []
  go (Cast _ _ _) = []
  go (Tick _ _) = []

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