{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}

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

  Types for the Partial Evaluator
-}
module Clash.Core.Evaluator.Types where

import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap (insert, lookup)
#if !MIN_VERSION_base(4,20,0)
import Data.List (foldl')
#endif
import Data.Maybe (fromMaybe, isJust)

#if MIN_VERSION_prettyprinter(1,7,0)
import Prettyprinter (hsep)
#else
import Data.Text.Prettyprint.Doc (hsep)
#endif

import Clash.Core.DataCon (DataCon, dcType)
import Clash.Core.HasType
import Clash.Core.Literal (Literal(CharLiteral))
import Clash.Core.Pretty (fromPpr, ppr, showPpr)
import Clash.Core.Term (Term(..), PrimInfo(..), TickInfo, Alt, mkApps)
import Clash.Core.TyCon (TyConMap)
import Clash.Core.Type (Type (..), mkFunTy)
import Clash.Core.Var (Id, IdScope(..), TyVar)
import Clash.Core.VarEnv
import Clash.Driver.Types (BindingMap, bindingTerm)
import Clash.Pretty (ClashPretty(..), fromPretty, showDoc)
import Clash.Util.Supply (Supply)

whnf'
  :: Evaluator
  -> BindingMap
  -> VarEnv Term
  -> TyConMap
  -> PrimHeap
  -> Supply
  -> InScopeSet
  -> Bool
  -> Term
  -> (PrimHeap, PureHeap, Term)
whnf' eval bm lh tcm ph ids is isSubj e =
  toResult $ whnf eval tcm isSubj m
 where
  toResult x = (mHeapPrim x, mHeapLocal x, mTerm x)

  m  = Machine ph gh lh [] ids is e
  gh = mapVarEnv bindingTerm bm

-- | Evaluate to WHNF given an existing Heap and Stack
whnf
  :: Evaluator
  -> TyConMap
  -> Bool
  -> Machine
  -> Machine
whnf eval tcm isSubj m
  | isSubj =
      -- See [Note: empty case expressions]
      let ty = inferCoreTypeOf tcm (mTerm m)
       in go (stackPush (Scrutinise ty []) m)
  | otherwise = go m
  where
    go :: Machine -> Machine
    go s = case step eval s tcm of
      Just s' -> go s'
      Nothing -> fromMaybe (error . showDoc . ppr $ mTerm m) (unwindStack s)


-- | An evaluator is a collection of basic building blocks which are used to
-- define partial evaluation. In this implementation, it consists of two types
-- of function:
--
--   * steps, which applies the reduction realtion to the current term
--   * unwindings, which pop the stack and evaluate the stack frame
--
-- Variants of these functions also exist for evalauting primitive operations.
-- This is because there may be multiple frontends to the compiler which can
-- reuse a common step and unwind, but have different primitives.
--
data Evaluator = Evaluator
  { step        :: Step
  , unwind      :: Unwind
  , primStep    :: PrimStep
  , primUnwind  :: PrimUnwind
  }

-- | Completely unwind the stack to get back the complete term
unwindStack :: Machine -> Maybe Machine
unwindStack m
  | stackNull m = Just m
  | otherwise = do
      (m', kf) <- stackPop m

      case kf of
        PrimApply p tys vs tms ->
          let term = foldl' App
                       (foldl' App
                         (foldl' TyApp (Prim p) tys)
                         (fmap valToTerm vs))
                       (mTerm m' : tms)
           in unwindStack (setTerm term m')

        Instantiate ty ->
          let term = TyApp (getTerm m') ty
           in unwindStack (setTerm term m')

        Apply n ->
          case heapLookup LocalId n m' of
            Just e ->
              let term = App (getTerm m') e
               in unwindStack (setTerm term m')

            Nothing -> error $ unlines $
              [ "Clash.Core.Evaluator.unwindStack:"
              , "Stack:"
              ] <>
              [ "  " <> showDoc (clashPretty frame) | frame <- mStack m] <>
              [ ""
              , "Expression:"
              , showPpr (mTerm m)
              , ""
              , "Heap:"
              , showDoc (clashPretty $ mHeapLocal m)
              ]

        Scrutinise _ [] ->
          unwindStack m'

        Scrutinise ty alts ->
          let term = Case (getTerm m') ty alts
           in unwindStack (setTerm term m')

        Update LocalId x ->
          unwindStack (heapInsert LocalId x (mTerm m') m')

        Update GlobalId _ ->
          unwindStack m'

        Tickish sp ->
          let term = Tick sp (getTerm m')
           in unwindStack (setTerm term m')

-- | A single step in the partial evaluator. The result is the new heap and
-- stack, and the next expression to be reduced.
--
type Step = Machine -> TyConMap -> Maybe Machine

type Unwind = Value -> Step

type PrimStep
  =  TyConMap
  -> Bool
  -> PrimInfo
  -> [Type]
  -> [Value]
  -> Machine
  -> Maybe Machine

type PrimUnwind
  =  TyConMap
  -> PrimInfo
  -> [Type]
  -> [Value]
  -> Value
  -> [Term]
  -> Machine
  -> Maybe Machine

-- | A machine represents the current state of the abstract machine used to
-- evaluate terms. A machine has a term under evaluation, a stack, and three
-- heaps:
--
--  * a primitive heap to store IO values from primitives (like ByteArrays)
--  * a global heap to store top-level bindings in scope
--  * a local heap to store local bindings in scope
--
-- Machines also include a unique supply and InScopeSet. These are needed when
-- new heap bindings are created, and are just an implementation detail.
--
data Machine = Machine
  { mHeapPrim   :: PrimHeap
  , mHeapGlobal :: PureHeap
  , mHeapLocal  :: PureHeap
  , mStack      :: Stack
  , mSupply     :: Supply
  , mScopeNames :: InScopeSet
  , mTerm       :: Term
  }

instance Show Machine where
  show (Machine ph gh lh s _ _ x) =
    unlines
      [ "Machine:"
      , ""
      , "Heap (Prim):"
      , show ph
      , ""
      , "Heap (Globals):"
      , show gh
      , ""
      , "Heap (Locals):"
      , show lh
      , ""
      , "Stack:"
      , show (fmap clashPretty s)
      , ""
      , "Term:"
      , show x
      ]

type PrimHeap = (IntMap Term, Int)
type PureHeap = VarEnv Term

type Stack = [StackFrame]

data StackFrame
  = Update IdScope Id
  | Apply  Id
  | Instantiate Type
  | PrimApply  PrimInfo [Type] [Value] [Term]
  | Scrutinise Type [Alt]
  | Tickish TickInfo
  deriving Show

instance ClashPretty StackFrame where
  clashPretty (Update GlobalId i) = hsep ["Update(Global)", fromPpr i]
  clashPretty (Update LocalId i)  = hsep ["Update(Local)", fromPpr i]
  clashPretty (Apply i) = hsep ["Apply", fromPpr i]
  clashPretty (Instantiate t) = hsep ["Instantiate", fromPpr t]
  clashPretty (PrimApply p tys vs ts) =
    hsep ["PrimApply", fromPretty (primName p), "::", fromPpr (coreTypeOf p),
          "; type args=", fromPpr tys,
          "; val args=", fromPpr (map valToTerm vs),
          "term args=", fromPpr ts]
  clashPretty (Scrutinise a b) =
    hsep ["Scrutinise ", fromPpr a,
          fromPpr (Case (Literal (CharLiteral '_')) a b)]
  clashPretty (Tickish sp) =
    hsep ["Tick", fromPpr sp]

-- Values
data Value
  = Lambda Id Term
  -- ^ Functions
  | TyLambda TyVar Term
  -- ^ Type abstractions
  | DC DataCon [Either Term Type]
  -- ^ Data constructors
  | Lit Literal
  -- ^ Literals
  | PrimVal  PrimInfo [Type] [Value]
  -- ^ Clash's number types are represented by their "fromInteger#" primitive
  -- function. So some primitives are values.
  | Suspend Term
  -- ^ Used by lazy primitives
  | TickValue TickInfo Value
  -- ^ Preserve ticks from Terms in Values
  | CastValue Value Type Type
  -- ^ Preserve casts from Terms in Values
  deriving Show

instance InferType Value where
  inferCoreTypeOf tcm = go
   where
    go = \case
      Lambda i t -> mkFunTy (coreTypeOf i) (inferCoreTypeOf tcm t)
      TyLambda v t -> ForAllTy v (inferCoreTypeOf tcm t)
      DC dc args -> applyTypeToArgs (mkApps (Data dc) args) tcm (dcType dc) args
      Lit l -> coreTypeOf l
      PrimVal p tys vals ->
        let args = map Right tys ++ map (Left . valToTerm) vals
         in applyTypeToArgs
              (mkApps (Prim p) args)
              tcm
              (primType p)
              args
      Suspend t -> inferCoreTypeOf tcm t
      TickValue _ v -> go v
      CastValue _ _ t -> t

valToTerm :: Value -> Term
valToTerm v = case v of
  Lambda x e           -> Lam x e
  TyLambda x e         -> TyLam x e
  DC dc pxs            -> foldl' (\e a -> either (App e) (TyApp e) a)
                                 (Data dc) pxs
  Lit l                -> Literal l
  PrimVal ty tys vs    -> foldl' App (foldl' TyApp (Prim ty) tys)
                                 (map valToTerm vs)
  Suspend e            -> e
  TickValue t x        -> Tick t (valToTerm x)
  CastValue x t1 t2    -> Cast (valToTerm x) t1 t2

-- Collect all the ticks from a value, exposing the ticked value.
--
collectValueTicks
  :: Value
  -> (Value, [TickInfo])
collectValueTicks = go []
 where
  go ticks (TickValue t v) = go (t:ticks) v
  go ticks v = (v, ticks)

-- | Are we in a context where special primitives must be forced.
--
-- See [Note: forcing special primitives]
forcePrims :: Machine -> Bool
forcePrims = go . mStack
 where
  -- When do we need to force the compile-time evaluation of a primitive?
  --
  -- 1. When they are the subject of a case-expression
  go (Scrutinise{}:_) = True
  -- 2. When they are in the argument position of another primitive:
  --    primitives are assumed to be strict in their arguments
  go (PrimApply{}:_)  = True
  -- We look through ticks
  go (Tickish{}:xs)   = go xs
  -- We are in a context where we dereferenced a heap-binding, hence the
  -- update fram on the stack. So now we need to check whether that variable
  -- reference was in a position where the result must be evaluated to WHNF
  go (Update{}:xs)    = go xs
  go _                = False

primCount :: Machine -> Int
primCount = snd . mHeapPrim

primLookup :: Int -> Machine -> Maybe Term
primLookup i = IntMap.lookup i . fst . mHeapPrim

primInsert :: Int -> Term -> Machine -> Machine
primInsert i x m =
  let (gh, c) = mHeapPrim m
   in m { mHeapPrim = (IntMap.insert i x gh, c + 1) }

primUpdate :: Int -> Term -> Machine -> Machine
primUpdate i x m =
  let (gh, c) = mHeapPrim m
   in m { mHeapPrim = (IntMap.insert i x gh, c) }

heapLookup :: IdScope -> Id -> Machine -> Maybe Term
heapLookup GlobalId i m =
  lookupVarEnv i $ mHeapGlobal m
heapLookup LocalId i m =
  lookupVarEnv i $ mHeapLocal m

heapContains :: IdScope -> Id -> Machine -> Bool
heapContains scope i = isJust . heapLookup scope i

heapInsert :: IdScope -> Id -> Term -> Machine -> Machine
heapInsert GlobalId i x m =
  m { mHeapGlobal = extendVarEnv i x (mHeapGlobal m) }
heapInsert LocalId i x m =
  m { mHeapLocal = extendVarEnv i x (mHeapLocal m) }

heapDelete :: IdScope -> Id -> Machine -> Machine
heapDelete GlobalId i m =
  m { mHeapGlobal = delVarEnv (mHeapGlobal m) i }
heapDelete LocalId i m =
  m { mHeapLocal = delVarEnv (mHeapLocal m) i }

stackPush :: StackFrame -> Machine -> Machine
stackPush f m = m { mStack = f : mStack m }

stackPop :: Machine -> Maybe (Machine, StackFrame)
stackPop m = case mStack m of
  [] -> Nothing
  (x:xs) -> Just (m { mStack = xs }, x)

stackClear :: Machine -> Machine
stackClear m = m { mStack = [] }

stackNull :: Machine -> Bool
stackNull = null . mStack

getTerm :: Machine -> Term
getTerm = mTerm

setTerm :: Term -> Machine -> Machine
setTerm x m = m { mTerm = x }
