State.hs

-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

{-# 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.Common.State
  ( -- * Indigo State
    IndigoState (..)
  , usingIndigoState
  , (>>)
  , (<$>)
  , iput
  , nopState
  , assignTopVar
  , withObject
  , withObjectState
  , withStackVars

  , DecomposedObjects
  , GenCodeHooks (..)
  , emptyGenCodeHooks
  , MetaData (..)
  , stmtHook
  , stmtHookState
  , auxiliaryHook
  , auxiliaryHookState
  , exprHook
  , exprHookState
  , replStkMd
  , alterStkMd
  , pushRefMd
  , pushNoRefMd
  , popNoRefMd

  , GenCode (..)
  , cleanGenCode

  , (##)
  ) where

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

import Indigo.Backend.Prelude
import Indigo.Common.Object
import Indigo.Common.Var
import Indigo.Lorentz
import Lorentz.Instr qualified as L
import Morley.Michelson.Typed qualified as M
import Morley.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 #" <> pretty 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
  , mdHooks   :: GenCodeHooks
  }

data GenCodeHooks = GenCodeHooks
  { gchStmtHook      :: forall inp out . Text -> (inp :-> out) -> (inp :-> out)
  , gchAuxiliaryHook :: forall inp out . Text -> (inp :-> out) -> (inp :-> out)
  , gchExprHook      :: forall inp out . Text -> (inp :-> out) -> (inp :-> out)
  -- pva701: dunno whether this level of verbosity is needed
  --, csSubExpr    :: forall a inp out . Expr a -> (inp :-> out) -> (inp :-> out)
  }

instance Semigroup GenCodeHooks where
  GenCodeHooks a b c <> GenCodeHooks a1 b1 c1 = GenCodeHooks
    { gchStmtHook = \t cd -> a1 t (a t cd)
    , gchAuxiliaryHook = \t cd -> b1 t (b t cd)
    , gchExprHook = \t cd -> c1 t (c t cd)
    }

instance Monoid GenCodeHooks where
  mempty = emptyGenCodeHooks

emptyGenCodeHooks :: GenCodeHooks
emptyGenCodeHooks = GenCodeHooks (const id) (const id) (const id)

stmtHook :: forall inp out any . MetaData any -> Text -> (inp :-> out) -> (inp :-> out)
stmtHook MetaData{..} tx cd = (gchStmtHook mdHooks) tx cd

stmtHookState :: Text -> IndigoState inp out -> IndigoState inp out
stmtHookState tx cd = IndigoState $ \md ->
  let GenCode st c cl = usingIndigoState md cd in
  GenCode st (stmtHook md tx c) cl

auxiliaryHook :: forall inp out any . MetaData any -> Text -> (inp :-> out) -> (inp :-> out)
auxiliaryHook MetaData{..} tx cd = (gchAuxiliaryHook mdHooks) tx cd

auxiliaryHookState :: Text -> IndigoState inp out -> IndigoState inp out
auxiliaryHookState tx cd = IndigoState $ \md ->
  let GenCode st c cl = usingIndigoState md cd in
  GenCode st (auxiliaryHook md tx c) cl

exprHook :: forall inp out any . MetaData any -> Text -> (inp :-> out) -> (inp :-> out)
exprHook MetaData{..} exTx cd = (gchExprHook mdHooks) exTx cd

exprHookState :: Text -> IndigoState inp out -> IndigoState inp out
exprHookState tx cd = IndigoState $ \md ->
  let GenCode st c cl = usingIndigoState md cd in
  GenCode st (exprHook md tx c) cl


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 cm) f = MetaData (f stk) objs cm

-- | '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

----------------------------------------------------------------------------
-- Helpers
----------------------------------------------------------------------------

-- | Version of '#' which performs some optimizations immediately.
--
-- In particular, this avoids glueing @Nop@s.
(##) :: (a :-> b) -> (b :-> c) -> (a :-> c)
l ## r =
  -- We are very verbose about cases to avoid
  -- significant compilation time increase
  case l of
    I M.Nop -> case r of
      I x -> I x
      _   -> l # r
    I x -> case r of
      I M.Nop -> I x
      _       -> l # r
    _ -> l # r