CESK.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}

-- |
-- SPDX-License-Identifier: BSD-3-Clause
-- Description: State machine of Swarm's interpreter
--
-- The Swarm interpreter uses a technique known as a
-- <https://matt.might.net/articles/cesk-machines/ CESK machine> (if
-- you want to read up on them, you may want to start by reading about
-- <https://matt.might.net/articles/cek-machines/ CEK machines>
-- first).  Execution happens simply by iterating a step function,
-- sending one state of the CESK machine to the next. In addition to
-- being relatively efficient, this means we can easily run a bunch of
-- robots synchronously, in parallel, without resorting to any threads
-- (by stepping their machines in a round-robin fashion); pause and
-- single-step the game; save and resume, and so on.
--
-- Essentially, a CESK machine state has four components:
--
-- - The __C__ontrol is the thing we are currently focused on:
--   either a 'Term' to evaluate, or a 'Value' that we have
--   just finished evaluating.
-- - The __E__nvironment ('Env') is a mapping from variables that might
--   occur free in the Control to their values.
-- - The __S__tore ('Store') is a mapping from abstract integer
--   /locations/ to values.  We use it to store delayed (lazy) values,
--   so they will be computed at most once.
-- - The __K__ontinuation ('Cont') is a stack of 'Frame's,
--   representing the evaluation context, /i.e./ what we are supposed
--   to do after we finish with the currently focused thing.  When we
--   reduce the currently focused term to a value, the top frame on
--   the stack tells us how to proceed.
--
-- You can think of a CESK machine as a defunctionalization of a
-- recursive big-step interpreter, where we explicitly keep track of
-- the call stack and the environments that would be in effect at
-- various places in the recursion.  One could probably even derive
-- this mechanically, by writing a recursive big-step interpreter,
-- then converting it to CPS, then defunctionalizing the
-- continuations.
--
-- The slightly confusing thing about CESK machines is how we
-- have to pass around environments everywhere.  Basically,
-- anywhere there can be unevaluated terms containing free
-- variables (in values, in continuation stack frames, ...), we
-- have to store the proper environment alongside so that when
-- we eventually get around to evaluating it, we will be able to
-- pull out the environment to use.
module Swarm.Game.CESK (
  -- * Frames and continuations
  Frame (..),
  Cont,

  -- ** Wrappers for creating delayed change of state
  WorldUpdate (..),
  RobotUpdate (..),

  -- * Store
  Store,
  Addr,
  emptyStore,
  allocate,
  resolveValue,
  lookupStore,
  setStore,

  -- * CESK machine states
  CESK (..),

  -- ** Construction
  initMachine,
  continue,
  cancel,
  prepareTerm,

  -- ** Extracting information
  finalValue,
  suspendedEnv,
  store,
  cont,
) where

import Control.Lens (Lens', Traversal', lens, traversal, (^.))
import Data.Aeson (FromJSON (..), ToJSON (..), genericParseJSON, genericToJSON)
import Data.IntMap.Strict (IntMap)
import Data.IntMap.Strict qualified as IM
import GHC.Generics (Generic)
import Prettyprinter (Doc, Pretty (..), encloseSep, hsep, (<+>))
import Swarm.Game.Entity (Entity)
import Swarm.Game.Exception
import Swarm.Game.Ingredients (Count)
import Swarm.Game.Tick
import Swarm.Game.World (WorldUpdate (..))
import Swarm.Language.Elaborate (insertSuspend)
import Swarm.Language.Requirements.Type (Requirements)
import Swarm.Language.Syntax
import Swarm.Language.Types
import Swarm.Language.Value as V
import Swarm.Pretty (PrettyPrec (..), pparens, ppr)
import Swarm.Util.JSON (optionsMinimize)

------------------------------------------------------------
-- Frames and continuations
------------------------------------------------------------

-- | A frame is a single component of a continuation stack, explaining
--   what to do next after we finish evaluating the currently focused
--   term.
data Frame
  = -- | We were evaluating the first component of a pair; next, we
    --   should evaluate the second component which was saved in this
    --   frame (and push a 'FFst' frame on the stack to save the first component).
    FSnd Term Env
  | -- | We were evaluating the second component of a pair; when done,
    --   we should combine it with the value of the first component saved
    --   in this frame to construct a fully evaluated pair.
    FFst Value
  | -- | @FArg t e@ says that we were evaluating the left-hand side of
    -- an application, so the next thing we should do is evaluate the
    -- term @t@ (the right-hand side, /i.e./ argument of the
    -- application) in environment @e@.  We will also push an 'FApp'
    -- frame on the stack.
    FArg Term Env
  | -- | @FVArg v@ says that we were evaluating the left-hand side of
    --   an application, and the next thing we should do is apply it
    --   to the given value.  This does not normally occur as part of
    --   the usual evaluation process for applications, which instead
    --   uses FArg.  However, it is sometimes useful when reducing
    --   other constructs---for example, the pair eliminator 'match'.
    FVArg Value
  | -- | @FApp v@ says that we were evaluating the right-hand side of
    -- an application; once we are done, we should pass the resulting
    -- value as an argument to @v@.
    FApp Value
  | -- | @FLet x ty t2 e@ says that we were evaluating a term @t1@ of
    -- type @ty@ in an expression of the form @let x = t1 in t2@, that
    -- is, we were evaluating the definition of @x@; the next thing we
    -- should do is evaluate @t2@ in the environment @e@ extended with
    -- a binding for @x@.
    FLet Var (Maybe (Polytype, Requirements)) Term Env
  | -- | We are executing inside a 'Try' block.  If an exception is
    --   raised, we will execute the stored term (the "catch" block).
    FTry Value
  | -- | An @FExec@ frame means the focused value is a command, which
    -- we should now execute.
    FExec
  | -- | We are in the process of executing the first component of a
    --   bind; once done, we should also execute the second component
    --   in the given environment (extended by binding the variable,
    --   if there is one, to the output of the first command).
    FBind (Maybe Var) (Maybe (Polytype, Requirements)) Term Env
  | -- | Apply specific updates to the world and current robot.
    --
    -- The 'Const' is used to track the original command for error messages.
    FImmediate Const [WorldUpdate Entity] [RobotUpdate]
  | -- | Update the cell at a certain location in the store with the computed value.
    FUpdate Addr
  | -- | Signal that we are done with an atomic computation.
    FFinishAtomic
  | -- | We are in the middle of evaluating a record: some fields have
    --   already been evaluated; we are focusing on evaluating one
    --   field; and some fields have yet to be evaluated.
    FRcd Env [(Var, Value)] Var [(Var, Maybe Term)]
  | -- | We are in the middle of evaluating a record field projection.
    FProj Var
  | -- | We should suspend with the given environment once we finish
    --   the current evaluation.
    FSuspend Env
  | -- | If an exception bubbles all the way up to this frame, then
    --   switch to Suspended mode with this saved top-level context.
    FRestoreEnv Env
  deriving (Generic)

instance ToJSON Frame where
  toJSON = genericToJSON optionsMinimize

-- | A continuation is just a stack of frames.
type Cont = [Frame]

------------------------------------------------------------
-- Store
------------------------------------------------------------

type Addr = Int

-- | 'Store' represents a store, /i.e./ memory, indexing integer
--   locations to 'Value's.
data Store = Store {next :: Addr, mu :: IntMap Value}
  deriving (Generic, ToJSON)

emptyStore :: Store
emptyStore = Store 0 IM.empty

-- | Allocate a new memory cell containing a given value.  Return the
--   index of the allocated cell.
allocate :: Value -> Store -> (Addr, Store)
allocate v (Store n m) = (n, Store (n + 1) (IM.insert n v m))

-- | Resolve a value, recursively looking up any indirections in the
--   store.
resolveValue :: Store -> Value -> Either Addr Value
resolveValue s = \case
  VIndir loc -> lookupStore s loc
  v -> Right v

-- | Look up the value at a given index, but keep following
--   indirections until encountering a value that is not a 'VIndir'.
lookupStore :: Store -> Addr -> Either Addr Value
lookupStore s = go
 where
  go loc = case IM.lookup loc (mu s) of
    Nothing -> Left loc
    Just v -> case v of
      VIndir loc' -> go loc'
      _ -> Right v

-- | Set the value at a given index.
setStore :: Addr -> Value -> Store -> Store
setStore n c (Store nxt m) = Store nxt (IM.insert n c m)

------------------------------------------------------------
-- CESK machine
------------------------------------------------------------

-- | The overall state of a CESK machine, which can actually be one of
--   four kinds of states. The CESK machine is named after the first
--   kind of state, and it would probably be possible to inline a
--   bunch of things and get rid of the second state, but I find it
--   much more natural and elegant this way.  Most tutorial
--   presentations of CEK/CESK machines only have one kind of state, but
--   then again, most tutorial presentations only deal with the bare
--   lambda calculus, so one can tell whether a term is a value just
--   by seeing whether it is syntactically a lambda.  I learned this
--   approach from Harper's Practical Foundations of Programming
--   Languages.
data CESK
  = -- | When we are on our way "in/down" into a term, we have a
    --   currently focused term to evaluate in the environment, a store,
    --   and a continuation.  In this mode we generally pattern-match on the
    --   'Term' to decide what to do next.
    In Term Env Store Cont
  | -- | Once we finish evaluating a term, we end up with a 'Value'
    --   and we switch into "out" mode, bringing the value back up
    --   out of the depths to the context that was expecting it.  In
    --   this mode we generally pattern-match on the 'Cont' to decide
    --   what to do next.
    --
    --   Note that there is no 'Env', because we don't have anything
    --   with variables to evaluate at the moment, and we maintain the
    --   invariant that any unevaluated terms buried inside a 'Value'
    --   or 'Cont' must carry along their environment with them.
    Out Value Store Cont
  | -- | An exception has been raised.  Keep unwinding the
    --   continuation stack (until finding an enclosing 'Try' in the
    --   case of a command failure or a user-generated exception, or
    --   until the stack is empty in the case of a fatal exception).
    Up Exn Store Cont
  | -- | The machine is waiting for the game to reach a certain time
    --   to resume its execution.
    Waiting TickNumber CESK
  | -- | The machine is suspended, i.e. waiting for another term to
    --   evaluate.  This happens after we have evaluated whatever the
    --   user entered at the REPL and we are waiting for them to type
    --   something else.  Conceptually, this is like a combination of
    --   'Out' and 'In': we store a 'Value' that was just yielded by
    --   evaluation, and otherwise it is just like 'In' with a hole
    --   for the 'Term' we are going to evaluate.
    Suspended Value Env Store Cont
  deriving (Generic)

instance ToJSON CESK where
  toJSON = genericToJSON optionsMinimize

-- | Is the CESK machine in a final (finished) state?  If so, extract
--   the final value and store.
finalValue :: CESK -> Maybe Value
{-# INLINE finalValue #-}
finalValue (Out v _ []) = Just v
finalValue (Suspended v _ _ []) = Just v
finalValue _ = Nothing

-- | Extract the environment from a suspended CESK machine (/e.g./ to
--   use for typechecking).
suspendedEnv :: Traversal' CESK Env
suspendedEnv = traversal go
 where
  go :: Applicative f => (Env -> f Env) -> CESK -> f CESK
  go f (Suspended v e s k) = Suspended v <$> f e <*> pure s <*> pure k
  go _ cesk = pure cesk

-- | Lens focusing on the store of a CESK machine.
store :: Lens' CESK Store
store = lens get set
 where
  get = \case
    In _ _ s _ -> s
    Out _ s _ -> s
    Up _ s _ -> s
    Waiting _ c -> get c
    Suspended _ _ s _ -> s
  set cesk s = case cesk of
    In t e _ k -> In t e s k
    Out v _ k -> Out v s k
    Up x _ k -> Up x s k
    Waiting t c -> Waiting t (set c s)
    Suspended v e _ k -> Suspended v e s k

-- | Lens focusing on the continuation of a CESK machine.
cont :: Lens' CESK Cont
cont = lens get set
 where
  get = \case
    In _ _ _ k -> k
    Out _ _ k -> k
    Up _ _ k -> k
    Waiting _ c -> get c
    Suspended _ _ _ k -> k
  set cesk k = case cesk of
    In t e s _ -> In t e s k
    Out v s _ -> Out v s k
    Up x s _ -> Up x s k
    Waiting t c -> Waiting t (set c k)
    Suspended v e s _ -> Suspended v e s k

-- | Create a brand new CESK machine, with empty environment and
--   store, to evaluate a given term.  We always initialize the
--   machine with a single FExec frame as the continuation; if the
--   given term does not have a command type, we wrap it in @pure@.
initMachine :: TSyntax -> CESK
initMachine t = In (prepareTerm V.emptyEnv t) V.emptyEnv emptyStore [FExec]

-- | Load a program into an existing robot CESK machine: either
--   continue from a suspended state, or, as a fallback, start from
--   scratch with an empty environment but the same store.
--
--   Also insert a @suspend@ primitive at the end, so the resulting
--   term is suitable for execution by the base (REPL) robot.
continue :: TSyntax -> CESK -> CESK
continue t = \case
  -- The normal case is when we are continuing from a suspended state. We:
  --
  --   (1) insert a suspend call at the end of the term, so that in
  --   the normal case after executing the entire term we will suspend
  --   in the innermost scope, to continue executing another term
  --   within that scope later.
  --
  --   (2) insert a failsafe FRestoreEnv frame into the continuation
  --   stack, in case execution of the term throws an exception.  In
  --   that case we will fall back to suspending with the original
  --   environment e (any names brought into scope by executing the
  --   term will be discarded).  If the term succeeds, the extra
  --   FRestoreEnv frame will be discarded.
  Suspended _ e s k -> In (insertSuspend $ prepareTerm e t) e s (FExec : FRestoreEnv e : k)
  -- In any other state, just start with an empty environment.  This
  -- happens e.g. when running a program on the base robot for the
  -- very first time.
  cesk -> In (insertSuspend $ prepareTerm V.emptyEnv t) V.emptyEnv (cesk ^. store) (FExec : (cesk ^. cont))

-- | Prepare a term for evaluation by a CESK machine in the given
--   environment: erase all type annotations, and optionally wrap it
--   in @pure@ if it does not have a command type.  Note that since
--   the environment might contain type aliases, we have to be careful
--   to expand them before concluding whether the term has a command
--   type or not.
prepareTerm :: Env -> TSyntax -> Term
prepareTerm e t = case whnfType (e ^. envTydefs) (ptBody (t ^. sType)) of
  TyCmd _ -> t'
  _ -> TApp (TConst Pure) t'
 where
  t' = eraseS t

-- | Cancel the currently running computation.
cancel :: CESK -> CESK
cancel cesk = Up Cancel (cesk ^. store) (cesk ^. cont)

------------------------------------------------------------
-- Pretty printing CESK machine states
------------------------------------------------------------

instance PrettyPrec CESK where
  prettyPrec _ = \case
    In c _ _ k -> prettyCont k (11, "▶" <> ppr c <> "◀")
    Out v _ k -> prettyCont k (11, "◀" <> ppr (valueToTerm v) <> "▶")
    Up e _ k -> prettyCont k (11, "!" <> (pretty (formatExn mempty e) <> "!"))
    Waiting t cesk -> "🕑" <> pretty t <> "(" <> ppr cesk <> ")"
    Suspended v _ _ k -> prettyCont k (11, "◀" <> ppr (valueToTerm v) <> "...▶")

-- | Take a continuation, and the pretty-printed expression which is
--   the focus of the continuation (i.e. the expression whose value
--   will be given to the continuation) along with its top-level
--   precedence, and pretty-print the whole thing.
--
--   As much as possible, we try to print to look like an *expression*
--   with a currently focused part, that is, we print the continuation
--   from the inside out instead of as a list of frames.  This makes
--   it much more intuitive to read.
prettyCont :: Cont -> (Int, Doc ann) -> Doc ann
prettyCont [] (_, inner) = inner
prettyCont (f : k) inner = prettyCont k (prettyFrame f inner)

-- | Pretty-print a single continuation frame, given its already
--   pretty-printed focus.  In particular, given a frame and its
--   "inside" (i.e. the expression or other frames being focused on,
--   whose value will eventually be passed to this frame), with the
--   precedence of the inside's top-level construct, return a
--   pretty-printed version of the entire frame along with its
--   top-level precedence.
prettyFrame :: Frame -> (Int, Doc ann) -> (Int, Doc ann)
prettyFrame f (p, inner) = case f of
  FSnd t _ -> (11, "(" <> inner <> "," <+> ppr t <> ")")
  FFst v -> (11, "(" <> ppr (valueToTerm v) <> "," <+> inner <> ")")
  FArg t _ -> (10, pparens (p < 10) inner <+> prettyPrec 11 t)
  FVArg v -> (10, pparens (p < 10) inner <+> prettyPrec 11 (valueToTerm v))
  FApp v -> (10, prettyPrec 10 (valueToTerm v) <+> pparens (p < 11) inner)
  FLet x _ t _ -> (11, hsep ["let", ppr x, "=", inner, "in", ppr t])
  FTry v -> (10, "try" <+> pparens (p < 11) inner <+> prettyPrec 11 (valueToTerm v))
  FExec -> prettyPrefix "E·" (p, inner)
  FBind Nothing _ t _ -> (0, pparens (p < 1) inner <+> ";" <+> ppr t)
  FBind (Just x) _ t _ -> (0, hsep [ppr x, "<-", pparens (p < 1) inner, ";", ppr t])
  FImmediate c _worldUpds _robotUpds -> prettyPrefix ("I[" <> ppr c <> "]·") (p, inner)
  FUpdate {} -> (p, inner)
  FFinishAtomic -> prettyPrefix "A·" (p, inner)
  FRcd _ done foc rest -> (11, encloseSep "[" "]" ", " (pDone ++ [pFoc] ++ pRest))
   where
    pDone = map (\(x, v) -> ppr x <+> "=" <+> ppr (valueToTerm v)) (reverse done)
    pFoc = ppr foc <+> "=" <+> inner
    pRest = map pprEq rest
    pprEq (x, Nothing) = ppr x
    pprEq (x, Just t) = ppr x <+> "=" <+> ppr t
  FProj x -> (11, pparens (p < 11) inner <> "." <> ppr x)
  FSuspend _ -> (10, "suspend" <+> pparens (p < 11) inner)
  FRestoreEnv _ -> (10, "restore" <+> pparens (p < 11) inner)

-- | Pretty-print a special "prefix application" frame, i.e. a frame
--   formatted like @X· inner@.  Unlike typical applications, these
--   associate to the *right*, so that we can print something like @X·
--   Y· Z· inner@ with no parens.
prettyPrefix :: Doc ann -> (Int, Doc ann) -> (Int, Doc ann)
prettyPrefix pre (p, inner) = (11, pre <+> pparens (p < 11) inner)

--------------------------------------------------------------
-- Runtime robot update
--------------------------------------------------------------

-- | Enumeration of robot updates.  This type is used for changes by
--   /e.g./ the @drill@ command which must be carried out at a later
--   tick. Using a first-order representation (as opposed to /e.g./
--   just a @Robot -> Robot@ function) allows us to serialize and
--   inspect the updates.
--
--   Note that this can not be in 'Swarm.Game.Robot' as it would create
--   a cyclic dependency.
data RobotUpdate
  = -- | Add copies of an entity to the robot's inventory.
    AddEntity Count Entity
  | -- | Make the robot learn about an entity.
    LearnEntity Entity
  deriving (Eq, Ord, Show, Generic)

instance ToJSON RobotUpdate where
  toJSON = genericToJSON optionsMinimize

instance FromJSON RobotUpdate where
  parseJSON = genericParseJSON optionsMinimize