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

The main API of the partial evaluator. This exposes the main functions needed
to call the evaluator, and the type of evaluators. A concrete implementation
of an evaluator is required to use this module: this can be imported from the
library for the compiler front-end, e.g. Clash.GHC.PartialEval in clash-ghc.
-}

module Clash.Core.PartialEval where

import Control.Concurrent.Supply (Supply)
import Data.IntMap.Strict (IntMap)

import Clash.Core.PartialEval.AsTerm
import Clash.Core.PartialEval.Monad
import Clash.Core.PartialEval.NormalForm
import Clash.Core.Term (Term)
import Clash.Core.TyCon (TyConMap)
import Clash.Core.Var (Id)
import Clash.Core.VarEnv (InScopeSet)
import Clash.Driver.Types (Binding(..), BindingMap)

-- | An evaluator for Clash core. This consists of two functions: one to
-- evaluate a term to weak-head normal form (WHNF) and another to recursively
-- evaluate sub-terms to obtain beta-normal eta-long form (NF).
--
data Evaluator = Evaluator
  { Evaluator -> Term -> Eval Value
evalWhnf :: Term  -> Eval Value
  , Evaluator -> Value -> Eval Normal
quoteNf  :: Value -> Eval Normal
  }

-- | Evaluate a term to WHNF, converting the result back to a Term.
-- The global environment at the end of evaluation is also returned, callers
-- should preserve any parts of the global environment needed for later calls.
--
whnf
  :: Evaluator
  -- ^ The evaluator implementation to use.
  -> GlobalEnv
  -- ^ The initial global environment.
  -> Bool
  -- ^ Whether evaluation should keep lifted data constructors.
  -- See NOTE [Lifted Constructors] in Clash.Core.PartialEval.NormalForm.
  -> Id
  -- ^ The id of the term under evaluation.
  -> Term
  -- ^ The term under evaluation.
  -> IO (Term, GlobalEnv)
  -- ^ The term evalated to WHNF, and the final global environment.
whnf :: Evaluator
-> GlobalEnv -> Bool -> Id -> Term -> IO (Term, GlobalEnv)
whnf Evaluator
e GlobalEnv
g Bool
isSubj Id
i Term
x =
  let l :: LocalEnv
l = Id -> Map TyVar Type -> Map Id Value -> Word -> Bool -> LocalEnv
LocalEnv Id
i Map TyVar Type
forall a. Monoid a => a
mempty Map Id Value
forall a. Monoid a => a
mempty (GlobalEnv -> Word
genvFuel GlobalEnv
g) Bool
isSubj
   in GlobalEnv -> LocalEnv -> Eval Term -> IO (Term, GlobalEnv)
forall a. GlobalEnv -> LocalEnv -> Eval a -> IO (a, GlobalEnv)
runEval GlobalEnv
g LocalEnv
l (Value -> Term
forall a. AsTerm a => a -> Term
asTerm (Value -> Term) -> Eval Value -> Eval Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Evaluator -> Term -> Eval Value
evalWhnf Evaluator
e Term
x)

-- | Evaluate a term to NF, converting the result back to a Term.
-- See `whnf` for more details.
--
nf
  :: Evaluator
  -- ^ The evaluator implementation to use.
  -> GlobalEnv
  -- ^ The initial global environment.
  -> Bool
  -- ^ Whether evaluation should keep lifted data constructors.
  -- See NOTE [Lifted Constructors] in Clash.Core.PartialEval.NormalForm.
  -> Id
  -- ^ The id of the term under evaluation.
  -> Term
  -- ^ The term under evaluation.
  -> IO (Term, GlobalEnv)
  -- ^ The term evalated to NF, and the final global environment.
nf :: Evaluator
-> GlobalEnv -> Bool -> Id -> Term -> IO (Term, GlobalEnv)
nf Evaluator
e GlobalEnv
g Bool
isSubj Id
i Term
x =
  let l :: LocalEnv
l = Id -> Map TyVar Type -> Map Id Value -> Word -> Bool -> LocalEnv
LocalEnv Id
i Map TyVar Type
forall a. Monoid a => a
mempty Map Id Value
forall a. Monoid a => a
mempty (GlobalEnv -> Word
genvFuel GlobalEnv
g) Bool
isSubj
   in GlobalEnv -> LocalEnv -> Eval Term -> IO (Term, GlobalEnv)
forall a. GlobalEnv -> LocalEnv -> Eval a -> IO (a, GlobalEnv)
runEval GlobalEnv
g LocalEnv
l (Normal -> Term
forall a. AsTerm a => a -> Term
asTerm (Normal -> Term) -> Eval Normal -> Eval Term
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (Evaluator -> Term -> Eval Value
evalWhnf Evaluator
e Term
x Eval Value -> (Value -> Eval Normal) -> Eval Normal
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= Evaluator -> Value -> Eval Normal
quoteNf Evaluator
e))

mkGlobalEnv
  :: BindingMap
  -- ^ Global bindings available to the evaluator.
  -> TyConMap
  -- ^ The type constructors known by Clash.
  -> InScopeSet
  -- ^ The set of variables in scope during evaluation.
  -> Supply
  -- ^ The supply of fresh names for variables.
  -> Word
  -- ^ The initial supply of fuel.
  -> IntMap Value
  -- ^ The initial IO heap.
  -> Int
  -- ^ The address of the next heap element.
  -> GlobalEnv
mkGlobalEnv :: BindingMap
-> TyConMap
-> InScopeSet
-> Supply
-> Word
-> IntMap Value
-> Int
-> GlobalEnv
mkGlobalEnv BindingMap
bm TyConMap
tcm InScopeSet
iss Supply
ids Word
fuel IntMap Value
heap Int
addr =
  VarEnv (Binding Value)
-> TyConMap
-> InScopeSet
-> Supply
-> Word
-> IntMap Value
-> Int
-> VarEnv Bool
-> GlobalEnv
GlobalEnv ((Binding Term -> Binding Value)
-> BindingMap -> VarEnv (Binding Value)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Binding Term -> Binding Value
asThunk BindingMap
bm) TyConMap
tcm InScopeSet
iss Supply
ids Word
fuel IntMap Value
heap Int
addr VarEnv Bool
forall a. Monoid a => a
mempty
 where
  asThunk :: Binding Term -> Binding Value
asThunk b :: Binding Term
b@Binding{bindingId :: forall a. Binding a -> Id
bindingId=Id
i,bindingTerm :: forall a. Binding a -> a
bindingTerm=Term
t} =
    Binding Term
b { bindingTerm :: Value
bindingTerm = Term -> LocalEnv -> Value
VThunk Term
t (Id -> Map TyVar Type -> Map Id Value -> Word -> Bool -> LocalEnv
LocalEnv Id
i Map TyVar Type
forall a. Monoid a => a
mempty Map Id Value
forall a. Monoid a => a
mempty Word
fuel Bool
False) }