-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

{-# LANGUAGE DerivingStrategies, InstanceSigs #-}

{- |
This module contains the core of Indigo language:
the 'IndigoState' monad, a datatype that represents its state.
It also includes some convenient functions to work with the state in IndigoM,
to provide rebindable syntax.

The 'IndigoState' monad implements the functionality of a symbolic interpreter.
During its execution Lorentz code is being generated.
-}

module Indigo.Internal.State
  ( -- * Indigo State
    IndigoState (..)
  , usingIndigoState
  , (>>=)
  , (=<<)
  , (>>)
  , (<$>)
  , return
  , iget
  , iput

  , RefId
  , StkEl (..)
  , StackVars
  , GenCode (..)
  , MetaData (..)
  , emptyMetadata
  , cleanGenCode
  , DefaultStack
  ) where

import qualified Data.Kind as Kind
import Data.Type.Equality (TestEquality (..))
import Data.Typeable (eqT)

import Indigo.Backend.Prelude
import Indigo.Lorentz
import qualified Lorentz.Instr as L

{-# ANN module ("HLint: ignore Reduce duplication" :: Text) #-}

----------------------------------------------------------------------------
-- Indigo State
----------------------------------------------------------------------------

-- | IndigoState monad. It's basically
-- [Control.Monad.Indexed.State](https://hackage.haskell.org/package/category-extras-0.53.5/docs/Control-Monad-Indexed-State.html)
-- , however this package is not in the used lts and it doesn't compile.
--
-- It takes as input a 'MetaData' (for the initial state) and returns a
-- 'GenCode' (for the resulting state and the generated Lorentz code).
--
-- IndigoState has to be used to write backend typed Lorentz code
-- from the corresponding frontend constructions.
newtype IndigoState inp out a =
  IndigoState {runIndigoState :: MetaData inp -> GenCode inp out a}
  deriving stock Functor

usingIndigoState :: MetaData inp -> IndigoState inp out a -> GenCode inp out a
usingIndigoState = flip runIndigoState

-- | Return for rebindable syntax.
return :: a -> IndigoState inp inp a
return a = IndigoState $ \md -> GenCode a md L.nop L.nop

-- | Bind for rebindable syntax.
--
-- It's basically like the bind for the 'State' monad, but it also composes the
-- generated code from @m a@ and @a -> m b@.
(>>=) :: forall inp out out1 a b . IndigoState inp out a -> (a -> IndigoState out out1 b) -> IndigoState inp out1 b
(>>=) m f = IndigoState $ \md ->
  let GenCode a md1 cd1 cl1 = runIndigoState m md in
  let GenCode b md2 cd2 cl2 = runIndigoState (f a) md1 in
  GenCode b md2 (cd1 ## cd2) (cl2 ## cl1)

(=<<) :: (a -> IndigoState out out1 b) -> IndigoState inp out a -> IndigoState inp out1 b
(=<<) = flip (>>=)

-- | Then for rebindable syntax.
(>>) :: IndigoState inp out a -> IndigoState out out1 b -> IndigoState inp out1 b
(>>) a b = a >>= const b

-- | Get current 'MetaData'.
iget :: IndigoState inp inp (MetaData inp)
iget = IndigoState $ \md -> GenCode md md L.nop L.nop

-- | Put new 'GenCode'.
iput :: GenCode inp out a -> IndigoState inp out a
iput gc = IndigoState $ \_ -> gc

----------------------------------------------------------------------------
-- Indigo stack and code gen primitives
----------------------------------------------------------------------------

-- | Reference id to a stack cell
newtype RefId = RefId Word
  deriving stock (Show, Generic)
  deriving newtype (Eq, Ord, Real, Num)

-- | Stack element of the symbolic interpreter.
--
-- It holds either a reference index that refers to this element
-- or just 'NoRef', indicating that there are no references
-- to this element.
data StkEl a where
  NoRef :: KnownValue a => StkEl a
  Ref :: KnownValue a => RefId -> StkEl a

instance TestEquality StkEl where
  testEquality NoRef NoRef = eqT
  testEquality (Ref _) (Ref _) = eqT
  testEquality (Ref _) NoRef = eqT
  testEquality NoRef (Ref _) = eqT

-- | Stack of the symbolic interpreter.
type StackVars (stk :: [Kind.Type]) = Rec StkEl stk

-- | Initial state of 'IndigoState'.
data MetaData stk = MetaData
  { mdStack :: StackVars stk
  -- ^ Stack of the symbolic interpreter.
  , mdRefCount :: RefId
  -- ^ Number of allocated variables.
  }

emptyMetadata :: MetaData '[]
emptyMetadata = MetaData RNil 0

type DefaultStack stk = Default (MetaData stk)

instance Default (MetaData '[]) where
  def = emptyMetadata

instance (KnownValue x, Default (MetaData xs)) => Default (MetaData (x ': xs)) where
  def = MetaData (NoRef :& mdStack def) 0

-- | Resulting state of IndigoM.
data GenCode inp out a = GenCode
  { gcOut :: ~a
  -- ^ Interpreter output value
  , gcMeta :: ~(MetaData out)
  -- ^ Interpreter meta data.
  , gcCode  :: inp :-> out
  -- ^ Generated Lorentz code.
  , gcClear :: out :-> inp
  -- ^ Clearing Lorentz code.
  } deriving stock Functor

-- | Produces the generated Lorentz code that cleans after itself, leaving the
-- same stack as the input one
cleanGenCode :: GenCode inp out a -> inp :-> inp
cleanGenCode GenCode {..} = gcCode ## gcClear