content: {-|
    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 (Show, Generic, NFData, Binary)
  
  -- TODO When it is possible, remove this pattern.
  pattern Letrec :: [LetBinding] -> Term -> Term
  pattern Letrec bs x <- Let (bindToList -> bs) x
   where
    Letrec bs x = Let (Rec bs) x
  
  bindToList :: Bind a -> [(Id, a)]
  bindToList (NonRec i x) = [(i, x)]
  bindToList (Rec xs) = 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 (Eq, Show, Generic, NFData, Binary)
  
  instance Ord TickInfo where
    compare (SrcSpan s1) (SrcSpan s2) = leftmost_smallest s1 s2
    compare (NameMod m1 t1) (NameMod m2 t2) =
      compare m1 m2 `thenCompare` compare t1 t2
    compare t1 t2 = compare (getRank t1) (getRank t2)
      where
        getRank :: TickInfo -> Word
        getRank SrcSpan{} = 0
        getRank NameMod{} = 1
        getRank DeDup     = 2
        getRank NoDeDup   = 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 (Eq,Ord,Show,Generic,NFData,Hashable,Binary)
  
  data IsMultiPrim
    = SingleResult
    | MultiResult
    deriving (Show, Generic, NFData, Eq, Hashable, Binary)
  
  data PrimInfo = PrimInfo
    { primName :: !Text
    , primType :: !Type
    , primWorkInfo :: !WorkInfo
    , 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'
    , primUnfolding :: !PrimUnfolding
    } deriving (Show, Generic, NFData, Binary)
  
  data PrimUnfolding
    = NoUnfolding
    | Unfolding !Id
    deriving (Show, Generic, NFData, Eq, Hashable, Binary)
  
  data MultiPrimInfo = MultiPrimInfo
    { mpi_primInfo :: PrimInfo
    , mpi_resultDc :: DataCon
    , 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 (Eq,Show,Generic,NFData,Hashable,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 (Eq, Show, Generic, NFData, Hashable, Binary, 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 (Eq, Ord, Show, Generic, NFData, Binary)
  
  type Alt = (Pat,Term)
  
  -- | Get the list of term-binders out of a DataType pattern
  patIds :: Pat -> ([TyVar],[Id])
  patIds (DataPat _  tvs ids) = (tvs,ids)
  patIds _                    = ([],[])
  
  patVars :: Pat -> [Var a]
  patVars (DataPat _ tvs ids) = coerce tvs ++ coerce ids
  patVars _ = []
  
  -- | Abstract a term over a list of term and type variables
  mkAbstraction :: Term -> [Either Id TyVar] -> Term
  mkAbstraction = foldr (either Lam TyLam)
  
  -- | Abstract a term over a list of type variables
  mkTyLams :: Term -> [TyVar] -> Term
  mkTyLams tm = mkAbstraction tm . map Right
  
  -- | Abstract a term over a list of variables
  mkLams :: Term -> [Id] -> Term
  mkLams tm = mkAbstraction tm . map Left
  
  -- | Apply a list of types and terms to a term
  mkApps :: Term -> [Either Term Type] -> Term
  mkApps = foldl' (\e a -> either (App e) (TyApp e) a)
  
  -- | Apply a list of terms to a term
  mkTmApps :: Term -> [Term] -> Term
  mkTmApps = foldl' App
  
  -- | Apply a list of types to a term
  mkTyApps :: Term -> [Type] -> Term
  mkTyApps = foldl' TyApp
  
  mkTicks :: Term -> [TickInfo] -> Term
  mkTicks tm ticks = foldl' (\e s -> Tick s e) tm (nub 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 (Show, Generic, NFData, 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 == c' = case (c, c') of
      (AppFun,          AppFun)            -> True
      (AppArg _,        AppArg _)          -> True
      -- NB: we do not see inside the argument here
      (TyAppC,          TyAppC)            -> True
      (LetBinding i is, LetBinding i' is') -> i == i' && is == is'
      (LetBody is,      LetBody is')       -> map fst is == map fst is'
      (LamBody i,       LamBody i')        -> i == i'
      (TyLamBody tv,    TyLamBody tv')     -> tv == tv'
      (CaseAlt p,       CaseAlt p')        -> p == p'
      (CaseScrut,       CaseScrut)         -> True
      (CastBody,        CastBody)          -> True
      (TickC sp,        TickC sp')         -> sp == sp'
      (_,               _)                 -> False
  
  -- | Is the Context a Lambda/Term-abstraction context?
  isLambdaBodyCtx :: CoreContext -> Bool
  isLambdaBodyCtx (LamBody _) = True
  isLambdaBodyCtx _           = False
  
  -- | Is the Context a Tick context?
  isTickCtx :: CoreContext -> Bool
  isTickCtx (TickC _) = True
  isTickCtx _         = False
  
  stripTicks :: Term -> Term
  stripTicks (Tick _ e) = stripTicks e
  stripTicks e = e
  
  -- | Like 'stripTicks' but removes all ticks from subexpressions.
  stripAllTicks :: Term -> Term
  stripAllTicks = go
   where
    go (Lam i x) = Lam i (go x)
    go (TyLam i x) = TyLam i (go x)
    go (App f x) = App (go f) (go x)
    go (TyApp f a) = TyApp (go f) a
    go (Let bs x) = Let (goBinds bs) (go x)
    go (Case x ty alts) = Case (go x) ty (fmap go <$> alts)
    go (Cast x a b) = Cast (go x) a b
    go (Tick _ x) = go x
    go x = x
  
    goBinds (NonRec i x) = NonRec i (go x)
    goBinds (Rec ixs) = Rec (fmap go <$> ixs)
  
  -- | Split a (Type)Application in the applied term and it arguments
  collectArgs :: Term -> (Term, [Either Term Type])
  collectArgs = go []
    where
      go args (App e1 e2) = go (Left e2:args) e1
      go args (TyApp e t) = go (Right t:args) e
      go args (Tick _ e)  = go args e
      go args e           = (e, args)
  
  collectTicks :: Term -> (Term, [TickInfo])
  collectTicks = go []
   where
    go ticks (Tick s e) = go (s:ticks) e
    go ticks e          = (e,ticks)
  
  collectArgsTicks :: Term -> (Term, [Either Term Type], [TickInfo])
  collectArgsTicks = go [] []
   where
    go args ticks (App e1 e2) = go (Left e2:args) ticks     e1
    go args ticks (TyApp e t) = go (Right t:args) ticks     e
    go args ticks (Tick s e)  = go args           (s:ticks) e
    go args ticks e           = (e, args, ticks)
  
  -- | Split a (Type)Abstraction in the bound variables and the abstracted term
  collectBndrs :: Term -> ([Either Id TyVar], Term)
  collectBndrs = go []
   where
    go bs (Lam v e')    = go (Left v:bs) e'
    go bs (TyLam tv e') = go (Right tv:bs) e'
    go bs e'            = (reverse bs,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 (collectArgs -> t) =
    case t of
      (Prim p, args) ->
        Just (primName p, length (rights args), length (lefts args))
      _ ->
        Nothing
  
  -- | Partition ticks in source ticks and nameMod ticks
  partitionTicks
    :: [TickInfo]
    -> ([TickInfo], [TickInfo])
    -- ^ (source ticks, nameMod ticks)
  partitionTicks = partition (\case {SrcSpan {} -> True; _ -> 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 f = catMaybes . DList.toList . go
   where
    go :: Term -> DList.DList (Maybe a)
    go t = DList.cons (f t) $ case t of
      Var _ -> mempty
      Data _ -> mempty
      Literal _ -> mempty
      Prim _ -> mempty
      Lam _ t1 -> go t1
      TyLam _ t1 -> go t1
      App t1 t2 -> go t1 <> go t2
      TyApp t1 _ -> go t1
      Let (NonRec _ x) t1 -> go t1 <> go x
      Let (Rec bndrs) t1 -> go t1 <> mconcat (map (go . snd) bndrs)
      Case t1 _ alts -> go t1 <> mconcat (map (go . snd) alts)
      Cast t1 _ _ -> go t1
      Tick _ t1 -> go t1
  
  -- Collect all term ids mentioned in a term
  collectTermIds :: Term -> [Id]
  collectTermIds = concat . walkTerm (Just . go)
   where
    go :: Term -> [Id]
    go (Var i) = [i]
    go (Lam i _) = [i]
    go (Let (NonRec i _) _) = [i]
    go (Let (Rec bndrs) _) = fmap fst bndrs
    go (Case _ _ alts) = concatMap (pat . fst) alts
    go (Data _) = []
    go (Literal _) = []
    go (Prim _) = []
    go (TyLam _ _) = []
    go (App _ _) = []
    go (TyApp _ _) = []
    go (Cast _ _ _) = []
    go (Tick _ _) = []
  
    pat :: Pat -> [Id]
    pat (DataPat _ _ ids) = ids
    pat (LitPat _) = []
    pat DefaultPat = []
  
  -- | Make a term variable out of a variable reference or ticked variable
  -- reference
  varToId :: Term -> Id
  varToId = \case
    Var i    -> i
    Tick _ e -> varToId e
    e        -> error $ $(curLoc) ++ "varToId: not a var: " ++ show e
  