{-|
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 OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TemplateHaskell #-}

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

import Control.Concurrent.Supply (Supply)
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(..))

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
showList :: [Neutral a] -> ShowS
$cshowList :: forall a. Show a => [Neutral a] -> ShowS
show :: Neutral a -> String
$cshow :: forall a. Show a => Neutral a -> String
showsPrec :: Int -> Neutral a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> 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
showList :: [Value] -> ShowS
$cshowList :: [Value] -> ShowS
show :: Value -> String
$cshow :: Value -> String
showsPrec :: Int -> Value -> ShowS
$cshowsPrec :: Int -> 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 (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 (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 (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
showList :: [Normal] -> ShowS
$cshowList :: [Normal] -> ShowS
show :: Normal -> String
$cshow :: Normal -> String
showsPrec :: Int -> Normal -> ShowS
$cshowsPrec :: Int -> 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
showList :: [LocalEnv] -> ShowS
$cshowList :: [LocalEnv] -> ShowS
show :: LocalEnv -> String
$cshow :: LocalEnv -> String
showsPrec :: Int -> LocalEnv -> ShowS
$cshowsPrec :: Int -> 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 :: (VarEnv Bool -> f (VarEnv Bool)) -> GlobalEnv -> f GlobalEnv
workFreeCache = (GlobalEnv -> VarEnv Bool)
-> (GlobalEnv -> VarEnv Bool -> GlobalEnv)
-> Lens GlobalEnv GlobalEnv (VarEnv Bool) (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 :: VarEnv Bool
genvWorkCache = VarEnv Bool
x })