State.hs

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

{-# LANGUAGE InstanceSigs #-}

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

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

Functionally, it's the same as having Lorentz instruction that can access and
modify a 'StackVars', referring to values on the stack with a 'RefId'.
-}

module Indigo.Internal.State
  ( -- * Indigo State
    IndigoState (..)
  , usingIndigoState
  , (>>)
  , (<$>)
  , iput
  , nopState
  , assignTopVar
  , withObject
  , withObjectState
  , withStackVars

  , DecomposedObjects
  , MetaData (..)
  , replStkMd
  , alterStkMd
  , pushRefMd
  , pushNoRefMd
  , popNoRefMd

  , GenCode (..)
  , cleanGenCode
  ) where

import qualified Data.Map as M
import Data.Typeable ((:~:)(..), eqT)

import Indigo.Internal.Object
import Indigo.Internal.Var
import Indigo.Backend.Prelude
import Indigo.Lorentz
import qualified Lorentz.Instr as L
import Util.Peano

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

-- | IndigoState data type.
--
-- It takes as input a 'StackVars' (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.
--
-- It has no return type, IndigoState instruction may take one or more
-- "return variables", that they assign to values produced during their execution.
newtype IndigoState inp out = IndigoState {
    runIndigoState :: MetaData inp -> GenCode inp out
  }

-- | Inverse of 'runIndigoState' for utility.
usingIndigoState :: MetaData inp -> IndigoState inp out -> GenCode inp out
usingIndigoState md act = runIndigoState act md

-- | Then for rebindable syntax.
(>>) :: IndigoState inp out -> IndigoState out out1 -> IndigoState inp out1
(>>) a b = IndigoState $ \md ->
  let GenCode st1 cd1 cl1 = runIndigoState a md in
  let GenCode st2 cd2 cl2 = runIndigoState b (replStkMd md st1) in
  GenCode st2 (cd1 ## cd2) (cl2 ## cl1)

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

-- | The simplest 'IndigoState', it does not modify the stack, nor the produced
-- code.
nopState :: IndigoState inp inp
nopState = IndigoState $ \md -> GenCode (mdStack md) L.nop L.nop

-- | Assigns a variable to reference the element on top of the stack.
assignTopVar :: KnownValue x => Var x -> IndigoState (x & inp) (x & inp)
assignTopVar var = IndigoState $ \md ->
  GenCode (assignVarAt var (mdStack md) SZ) L.nop L.nop

withObject
  :: forall a r .  KnownValue a
  => DecomposedObjects
  -> Var a
  -> (Object a -> r)
  -> r
withObject objs (Var refId) f = case M.lookup refId objs of
  Nothing -> f (Cell refId)
  Just so -> case so of
    SomeObject (obj :: Object a1) -> case eqT @a @a1 of
      Just Refl -> f obj
      Nothing ->
        error $ "unexpectedly SomeObject with by reference #" <> show refId <> " has different type"

withObjectState
  :: forall a inp out . KnownValue a
  => Var a
  -> (Object a -> IndigoState inp out)
  -> IndigoState inp out
withObjectState v f = IndigoState $ \md -> usingIndigoState md (withObject (mdObjects md) v f)

-- | Utility function to create 'IndigoState' that need access to the current 'StackVars'.
withStackVars :: (StackVars inp -> IndigoState inp out) -> IndigoState inp out
withStackVars fIs = IndigoState $ \md -> usingIndigoState md (fIs $ mdStack md)

----------------------------------------------------------------------------
-- MetaData primitives
----------------------------------------------------------------------------

type DecomposedObjects = Map RefId SomeObject

data MetaData inp = MetaData
  { mdStack   :: StackVars inp
  , mdObjects :: DecomposedObjects
  }

replStkMd :: MetaData inp -> StackVars inp1 -> MetaData inp1
replStkMd md = alterStkMd md . const

alterStkMd :: MetaData inp -> (StackVars inp -> StackVars inp1) -> MetaData inp1
alterStkMd (MetaData stk objs) f = MetaData (f stk) objs

-- | 'pushRef' version for 'MetaData'
pushRefMd :: KnownValue a => Var a -> MetaData inp -> MetaData (a & inp)
pushRefMd var md = alterStkMd md (pushRef var)

-- | 'pushNoRef' version for 'MetaData'
pushNoRefMd :: KnownValue a => MetaData inp -> MetaData (a & inp)
pushNoRefMd md = alterStkMd md pushNoRef

-- | 'popNoRef' version for 'MetaData'
popNoRefMd :: MetaData (a & inp) -> MetaData inp
popNoRefMd md = alterStkMd md popNoRef

----------------------------------------------------------------------------
-- Code generation primitives
----------------------------------------------------------------------------

-- | Resulting state of IndigoM.
data GenCode inp out = GenCode
  { gcStack :: ~(StackVars out)
  -- ^ Stack of the symbolic interpreter.
  , gcCode  :: inp :-> out
  -- ^ Generated Lorentz code.
  , gcClear :: out :-> inp
  -- ^ Clearing Lorentz code.
  }

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