    {-|
    Copyright   : (C) 2020-2021, QBayLogic B.V.,
                      2022     , Google Inc.
    License     : BSD2 (see the file LICENSE)
    Maintainer  : QBayLogic B.V. <devops@qbaylogic.com>

    Normal forms for the partial evaluator. These provide a restricted model of
    how terms can be constructed (compared to the more liberal Term type) which
    give a stronger guarantee that evaluation does not produce invalid results.
    This module is only needed to define new evaluators, for calling an existing
    evaluator see Clash.Core.PartialEval.
    -}

    {-# LANGUAGE LambdaCase #-}
    {-# LANGUAGE RankNTypes #-}

    module Clash.Core.PartialEval.NormalForm
      ( Arg
      , Args
      , Neutral(..)
      , Value(..)
      , mkValueTicks
      , stripValue
      , collectValueTicks
      , isUndefined
      , isUndefinedX
      , Normal(..)
      , LocalEnv(..)
      , GlobalEnv(..)
      , workFreeCache
      ) where

    import Control.Lens (Lens', lens)
    import Data.IntMap.Strict (IntMap)
    import Data.Map.Strict (Map)

    import Clash.Core.DataCon (DataCon)
    import Clash.Core.Literal
    import Clash.Core.Term (Bind, Term(..), PrimInfo(primName), TickInfo, Pat)
    import Clash.Core.TyCon (TyConMap)
    import Clash.Core.Type (Type, TyVar)
    import Clash.Core.Util (undefinedPrims, undefinedXPrims)
    import Clash.Core.Var (Id)
    import Clash.Core.VarEnv (VarEnv, InScopeSet)
    import Clash.Driver.Types (Binding(..))
    import Clash.Util.Supply (Supply)

    type Args a
      = [Arg a]

    -- | An argument applied to a function / data constructor / primitive.
    --
    type Arg a
      = Either a Type

    -- | Neutral terms cannot be reduced, as they represent things like variables
    -- which are unknown, partially applied functions, or case expressions where
    -- the subject cannot be inspected. Consider:
    --
    -- v              Stuck if "v" is a free variable
    -- p x1 ... xn    Stuck if "p" is a primitive that cannot be reduced
    -- x $ y          Stuck if "x" is not known to be a lambda
    -- x @ A          Stuck if "x" is not known to be a type lambda
    -- case x of ...  Stuck if "x" is neutral (cannot choose an alternative)
    --
    -- Neutral terms can also be let expressions which preserve required bindings
    -- in the normal form representation. Examples of bindings that may be kept are
    -- bindings which perform work (and should not be copied) or bindings that
    -- are recursive and are still referred to by the body of the let expression.
    --
    -- let ... in ... Preserved bindings are needed by the body
    --
    data Neutral a
      = NeVar    !Id
      | NePrim   !PrimInfo !(Args a)
      | NeApp    !(Neutral a) !a
      | NeTyApp  !(Neutral a) !Type
      | NeLet    !(Bind a) !a
      | NeCase   !a !Type ![(Pat, a)]
      deriving (Int -> Neutral a -> ShowS
    [Neutral a] -> ShowS
    Neutral a -> String
    (Int -> Neutral a -> ShowS)
    -> (Neutral a -> String)
    -> ([Neutral a] -> ShowS)
    -> Show (Neutral a)
    forall a. Show a => Int -> Neutral a -> ShowS
    forall a. Show a => [Neutral a] -> ShowS
    forall a. Show a => Neutral a -> String
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: forall a. Show a => Int -> Neutral a -> ShowS
    showsPrec :: Int -> Neutral a -> ShowS
    $cshow :: forall a. Show a => Neutral a -> String
    show :: Neutral a -> String
    $cshowList :: forall a. Show a => [Neutral a] -> ShowS
    showList :: [Neutral a] -> ShowS
    Show)

    -- TODO Write an instance (InferType a) => InferType (Neutral a)
    -- TODO Write an instance (HasFreeVars a) => HasFreeVars (Neutral a)

    -- | A term which has been potentially evaluated to WHNF. If evaluation has
    -- occurred, then there will be no redexes at the head of the Value, but
    -- sub-terms may still have redexes. Data constructors are only considered to
    -- be values when fully applied, if partially applied they should be
    -- eta-expanded during evaluation.
    --
    -- Thunks are included so that lazy evaluation can be modelled without needing
    -- to store Either Term Value in the environment. This makes the presentation
    -- simpler, with the caveat that values must be forced when they are required
    -- to not be thunks.
    --
    data Value
      = VNeutral  !(Neutral Value)
      | VLiteral  !Literal
      | VData     !DataCon !(Args Value) !LocalEnv
      | VLam      !Id !Term !LocalEnv
      | VTyLam    !TyVar !Term !LocalEnv
      | VCast     !Value !Type !Type
      | VTick     !Value !TickInfo
      | VThunk    !Term !LocalEnv
      deriving (Int -> Value -> ShowS
    [Value] -> ShowS
    Value -> String
    (Int -> Value -> ShowS)
    -> (Value -> String) -> ([Value] -> ShowS) -> Show Value
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> Value -> ShowS
    showsPrec :: Int -> Value -> ShowS
    $cshow :: Value -> String
    show :: Value -> String
    $cshowList :: [Value] -> ShowS
    showList :: [Value] -> ShowS
    Show)

    -- TODO Write an instance InferType Value
    -- TODO Write an instance HasFreeVars Value

    mkValueTicks :: Value -> [TickInfo] -> Value
    mkValueTicks :: Value -> [TickInfo] -> Value
    mkValueTicks = (Value -> TickInfo -> Value) -> Value -> [TickInfo] -> Value
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: Type -> Type) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl Value -> TickInfo -> Value
    VTick

    stripValue :: Value -> Value
    stripValue :: Value -> Value
    stripValue = (Value, [TickInfo]) -> Value
    forall a b. (a, b) -> a
    fst ((Value, [TickInfo]) -> Value)
    -> (Value -> (Value, [TickInfo])) -> Value -> Value
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Value -> (Value, [TickInfo])
    collectValueTicks

    collectValueTicks :: Value -> (Value, [TickInfo])
    collectValueTicks :: Value -> (Value, [TickInfo])
    collectValueTicks = [TickInfo] -> Value -> (Value, [TickInfo])
    go []
     where
      go :: [TickInfo] -> Value -> (Value, [TickInfo])
    go ![TickInfo]
    acc = \case
        VTick Value
    v TickInfo
    tick -> [TickInfo] -> Value -> (Value, [TickInfo])
    go (TickInfo
    tick TickInfo -> [TickInfo] -> [TickInfo]
    forall a. a -> [a] -> [a]
    : [TickInfo]
    acc) Value
    v
        Value
    value -> (Value
    value, [TickInfo]
    acc)

    isUndefined :: Value -> Bool
    isUndefined :: Value -> Bool
    isUndefined = \case
      VNeutral (NePrim PrimInfo
    pr Args Value
    _) ->
        PrimInfo -> Text
    primName PrimInfo
    pr Text -> [Text] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: Type -> Type) a.
    (Foldable t, Eq a) =>
    a -> t a -> Bool
    `elem` [Text]
    undefinedPrims

      Value
    _ -> Bool
    False

    isUndefinedX :: Value -> Bool
    isUndefinedX :: Value -> Bool
    isUndefinedX = \case
      VNeutral (NePrim PrimInfo
    pr Args Value
    _) ->
        PrimInfo -> Text
    primName PrimInfo
    pr Text -> [Text] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: Type -> Type) a.
    (Foldable t, Eq a) =>
    a -> t a -> Bool
    `elem` [Text]
    undefinedXPrims

      Value
    _ -> Bool
    False

    -- | A term which is in beta-normal eta-long form (NF). This has no redexes,
    -- and all partially applied functions in sub-terms are eta-expanded.
    --
    -- While not strictly necessary, NLam includes the environment at the point the
    -- original term was evaluated. This makes it easier for the AsTerm instance
    -- for Normal to reintroduce let expressions before lambdas without
    -- accidentally floating a let using a lambda bound variable outwards.
    --
    data Normal
      = NNeutral  !(Neutral Normal)
      | NLiteral  !Literal
      | NData     !DataCon !(Args Normal)
      | NLam      !Id !Normal !LocalEnv
      | NTyLam    !TyVar !Normal !LocalEnv
      | NCast     !Normal !Type !Type
      | NTick     !Normal !TickInfo
      deriving (Int -> Normal -> ShowS
    [Normal] -> ShowS
    Normal -> String
    (Int -> Normal -> ShowS)
    -> (Normal -> String) -> ([Normal] -> ShowS) -> Show Normal
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> Normal -> ShowS
    showsPrec :: Int -> Normal -> ShowS
    $cshow :: Normal -> String
    show :: Normal -> String
    $cshowList :: [Normal] -> ShowS
    showList :: [Normal] -> ShowS
    Show)

    data LocalEnv = LocalEnv
      { LocalEnv -> Id
    lenvContext :: Id
        -- ^ The id of the term currently under evaluation.
      , LocalEnv -> Map TyVar Type
    lenvTypes :: Map TyVar Type
        -- ^ Local type environment. These are types that are introduced while
        -- evaluating the current term (i.e. by type applications)
      , LocalEnv -> Map Id Value
    lenvValues :: Map Id Value
        -- ^ Local term environment. These are WHNF terms or unevaluated thunks
        -- introduced while evaluating the current term (i.e. by applications)
      , LocalEnv -> Word
    lenvFuel :: Word
        -- ^ The amount of fuel left in the local environment when the previous
        -- head was reached. This is needed so resuming evaluation does not lead
        -- to additional fuel being available.
      , LocalEnv -> Bool
    lenvKeepLifted :: Bool
        -- ^ When evaluating, keep data constructors for boxed data types (e.g. I#)
        -- instead of converting these back to their corresponding primitive. This
        -- is used when evaluating terms where the result is subject of a case
        -- expression (see note: lifted data types).
      } deriving (Int -> LocalEnv -> ShowS
    [LocalEnv] -> ShowS
    LocalEnv -> String
    (Int -> LocalEnv -> ShowS)
    -> (LocalEnv -> String) -> ([LocalEnv] -> ShowS) -> Show LocalEnv
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> LocalEnv -> ShowS
    showsPrec :: Int -> LocalEnv -> ShowS
    $cshow :: LocalEnv -> String
    show :: LocalEnv -> String
    $cshowList :: [LocalEnv] -> ShowS
    showList :: [LocalEnv] -> ShowS
    Show)

    -- TODO Add recursion info to the global environment. Until then we are forced
    -- to spend fuel on non-recursive (terminating) terms.

    data GlobalEnv = GlobalEnv
      { GlobalEnv -> VarEnv (Binding Value)
    genvBindings :: VarEnv (Binding Value)
        -- ^ Global term environment. These are the potentially evaluated bodies
        -- of the top level definitions which are forced on lookup.
      , GlobalEnv -> TyConMap
    genvTyConMap :: TyConMap
        -- ^ The type constructors known about by Clash.
      , GlobalEnv -> InScopeSet
    genvInScope :: InScopeSet
        -- ^ The set of in scope variables during partial evaluation. This includes
        -- new variables introduced by the evaluator (such as the ids of binders
        -- introduced during eta expansion.)
      , GlobalEnv -> Supply
    genvSupply :: Supply
        -- ^ The supply of fresh names for generating identifiers.
      , GlobalEnv -> Word
    genvFuel :: Word
        -- ^ The remaining fuel which can be spent inlining global variables. This
        -- is saved in the local environment, so when evaluation resumes from WHNF
        -- the amount of fuel used is preserved.
      , GlobalEnv -> IntMap Value
    genvHeap :: IntMap Value
        -- ^ The heap containing the results of any evaluated IO primitives.
      , GlobalEnv -> Int
    genvAddr :: Int
        -- ^ The address of the next element to be inserted into the heap.
      , GlobalEnv -> VarEnv Bool
    genvWorkCache :: VarEnv Bool
        -- ^ Cache for the results of isWorkFree. This is required to use
        -- Clash.Rewrite.WorkFree.isWorkFree.
      }

    workFreeCache :: Lens' GlobalEnv (VarEnv Bool)
    workFreeCache :: Lens' GlobalEnv (VarEnv Bool)
    workFreeCache = (GlobalEnv -> VarEnv Bool)
    -> (GlobalEnv -> VarEnv Bool -> GlobalEnv)
    -> Lens' GlobalEnv (VarEnv Bool)
    forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
    lens GlobalEnv -> VarEnv Bool
    genvWorkCache (\GlobalEnv
    env VarEnv Bool
    x -> GlobalEnv
    env { genvWorkCache = x })