Use type family for robot CESK machine field (#1729)

Towards #1715 and #1043.

This refactoring step is a prerequisite for #1719 to extricate references to the `CESK` module from the base `RobotR` definition.

In this PR:
* `Swarm.Game.CESK` is imported qualified to more easily track usages
* A new `RobotMachine` type family is added to parameterize the `_machine` field.
* `CESK` is a new parameter to `addTRobot`

app/doc/Swarm/Doc/Gen.hs

@@ -136,7 +136,7 @@ generateSpecialKeyNames =
 generateRecipe :: IO String
 generateRecipe = simpleErrorHandle $ do
   (classic, (worlds, entities, recipes)) <- loadStandaloneScenario "data/scenarios/classic.yaml"
-  baseRobot <- getBaseRobot classic
+  baseRobot <- instantiateBaseRobot classic
   return . Dot.showDot $ recipesToDot baseRobot (worlds ! "classic") entities recipes
 
 recipesToDot :: Robot -> Some (TTerm '[]) -> EntityMap -> [Recipe Entity] -> Dot ()

src/Swarm/Doc/Util.hs

@@ -48,7 +48,7 @@ commands = filter Syntax.isCmd Syntax.allConst
 constSyntax :: Const -> Text
 constSyntax = Syntax.syntax . Syntax.constInfo
 
-getBaseRobot :: Has (Throw SystemFailure) sig m => Scenario -> m Robot
-getBaseRobot s = case listToMaybe $ view scenarioRobots s of
-  Just r -> pure $ instantiateRobot 0 r
+instantiateBaseRobot :: Has (Throw SystemFailure) sig m => Scenario -> m Robot
+instantiateBaseRobot s = case listToMaybe $ view scenarioRobots s of
+  Just r -> pure $ instantiateRobot Nothing 0 r
   Nothing -> throwError $ CustomFailure "Scenario contains no robots"

src/swarm-engine/Swarm/Game/Robot.hs

@@ -21,7 +21,7 @@ module Swarm.Game.Robot (
   TRobot,
 
   -- ** Runtime robot update
-  RobotUpdate (..),
+  C.RobotUpdate (..),
 
   -- * Robot context
   RobotContext,
@@ -97,13 +97,14 @@ import GHC.Generics (Generic)
 import Linear
 import Servant.Docs (ToSample)
 import Servant.Docs qualified as SD
-import Swarm.Game.CESK
+import Swarm.Game.CESK qualified as C
 import Swarm.Game.Display (Display, curOrientation, defaultRobotDisplay, invisible)
 import Swarm.Game.Entity hiding (empty)
 import Swarm.Game.Location (Heading, Location, toDirection, toHeading)
 import Swarm.Game.Universe
 import Swarm.Language.Capability (Capability)
 import Swarm.Language.Context qualified as Ctx
+import Swarm.Language.Pipeline (ProcessedTerm)
 import Swarm.Language.Pipeline.QQ (tmQ)
 import Swarm.Language.Requirement (ReqCtx)
 import Swarm.Language.Syntax (Const, Syntax)
@@ -129,7 +130,7 @@ data RobotContext = RobotContext
   -- ^ Map definition names to their values. Note that since
   --   definitions are delayed, the values will just consist of
   --   'VRef's pointing into the store.
-  , _defStore :: Store
+  , _defStore :: C.Store
   -- ^ A store containing memory cells allocated to hold
   --   definitions.
   }
@@ -138,7 +139,7 @@ data RobotContext = RobotContext
 makeLenses ''RobotContext
 
 emptyRobotContext :: RobotContext
-emptyRobotContext = RobotContext Ctx.empty Ctx.empty Ctx.empty emptyStore
+emptyRobotContext = RobotContext Ctx.empty Ctx.empty Ctx.empty C.emptyStore
 
 type instance Index RobotContext = Ctx.Var
 type instance IxValue RobotContext = Typed Value
@@ -181,7 +182,7 @@ data ActivityCounts = ActivityCounts
   , _tangibleCommandCount :: Int
   , _commandsHistogram :: Map Const Int
   , _lifetimeStepCount :: Int
-  , _activityWindow :: WindowedCounter TickNumber
+  , _activityWindow :: WindowedCounter C.TickNumber
   }
   deriving (Eq, Show, Generic, Ae.FromJSON, Ae.ToJSON)
 
@@ -253,7 +254,7 @@ commandsHistogram :: Lens' ActivityCounts (Map Const Int)
 lifetimeStepCount :: Lens' ActivityCounts Int
 
 -- | Sliding window over a span of ticks indicating ratio of activity
-activityWindow :: Lens' ActivityCounts (WindowedCounter TickNumber)
+activityWindow :: Lens' ActivityCounts (WindowedCounter C.TickNumber)
 
 -- | With a robot template, we may or may not have a location.  With a
 --   concrete robot we must have a location.
@@ -266,6 +267,10 @@ type family RobotID (phase :: RobotPhase) :: Data.Kind.Type where
   RobotID 'TemplateRobot = ()
   RobotID 'ConcreteRobot = RID
 
+type family RobotMachine (phase :: RobotPhase) :: Data.Kind.Type
+type instance RobotMachine 'TemplateRobot = Maybe ProcessedTerm
+type instance RobotMachine 'ConcreteRobot = C.CESK
+
 -- | A value of type 'RobotR' is a record representing the state of a
 --   single robot.  The @f@ parameter is for tracking whether or not
 --   the robot has been assigned a unique ID.
@@ -282,7 +287,7 @@ data RobotR (phase :: RobotPhase) = RobotR
   , _robotID :: RobotID phase
   , _robotParentID :: Maybe RID
   , _robotHeavy :: Bool
-  , _machine :: CESK
+  , _machine :: RobotMachine phase
   , _systemRobot :: Bool
   , _selfDestruct :: Bool
   , _activityCounts :: ActivityCounts
@@ -292,8 +297,8 @@ data RobotR (phase :: RobotPhase) = RobotR
   }
   deriving (Generic)
 
-deriving instance (Show (RobotLocation phase), Show (RobotID phase)) => Show (RobotR phase)
-deriving instance (Eq (RobotLocation phase), Eq (RobotID phase)) => Eq (RobotR phase)
+deriving instance (Show (RobotLocation phase), Show (RobotID phase), Show (RobotMachine phase)) => Show (RobotR phase)
+deriving instance (Eq (RobotLocation phase), Eq (RobotID phase), Eq (RobotMachine phase)) => Eq (RobotR phase)
 
 -- See https://byorgey.wordpress.com/2021/09/17/automatically-updated-cached-views-with-lens/
 -- for the approach used here with lenses.
@@ -320,7 +325,7 @@ instance ToSample Robot where
         defaultCosmicLocation
         zero
         defaultRobotDisplay
-        (initMachine [tmQ| move |] mempty emptyStore)
+        (C.initMachine [tmQ| move |] mempty C.emptyStore)
         []
         []
         False
@@ -421,11 +426,15 @@ robotID :: Getter Robot RID
 --    if the robot template didn't have a location already, just set
 --    the location to (0,0) by default.  If you want a different location,
 --    set it via 'trobotLocation' before calling 'instantiateRobot'.
-instantiateRobot :: RID -> TRobot -> Robot
-instantiateRobot i r =
+--
+-- If a machine is not supplied (i.e. 'Nothing'), will fallback to any
+-- program specified in the template robot.
+instantiateRobot :: Maybe C.CESK -> RID -> TRobot -> Robot
+instantiateRobot maybeMachine i r =
   r
     { _robotID = i
     , _robotLocation = fromMaybe defaultCosmicLocation $ _robotLocation r
+    , _machine = fromMaybe (mkMachine $ _machine r) maybeMachine
     }
 
 -- | The ID number of the robot's parent, that is, the robot that
@@ -495,7 +504,7 @@ robotCapabilities :: Getter Robot (Set Capability)
 robotCapabilities = to _robotCapabilities
 
 -- | The robot's current CEK machine state.
-machine :: Lens' Robot CESK
+machine :: Lens' Robot C.CESK
 
 -- | Is this robot a "system robot"?  System robots are generated by
 --   the system (as opposed to created by the user) and are not
@@ -523,6 +532,10 @@ walkabilityContext :: Getter Robot WalkabilityContext
 walkabilityContext = to $
   \x -> WalkabilityContext (_robotCapabilities x) (_unwalkableEntities x)
 
+mkMachine :: Maybe ProcessedTerm -> C.CESK
+mkMachine Nothing = C.Out VUnit C.emptyStore []
+mkMachine (Just pt) = C.initMachine pt mempty C.emptyStore
+
 -- | A general function for creating robots.
 mkRobot ::
   -- | ID number of the robot.
@@ -540,7 +553,7 @@ mkRobot ::
   -- | Robot display.
   Display ->
   -- | Initial CESK machine.
-  CESK ->
+  RobotMachine phase ->
   -- | Equipped devices.
   [Entity] ->
   -- | Initial inventory.
@@ -603,16 +616,13 @@ instance FromJSONE EntityMap TRobot where
       <*> liftE (v .:? "loc")
       <*> liftE (fmap getHeading $ v .:? "dir" .!= HeadingSpec zero)
       <*> localE (const defDisplay) (v ..:? "display" ..!= defDisplay)
-      <*> liftE (mkMachine <$> (v .:? "program"))
+      <*> liftE (v .:? "program")
       <*> v ..:? "devices" ..!= []
       <*> v ..:? "inventory" ..!= []
       <*> pure sys
       <*> liftE (v .:? "heavy" .!= False)
       <*> liftE (v .:? "unwalkable" ..!= mempty)
       <*> pure 0
-   where
-    mkMachine Nothing = Out VUnit emptyStore []
-    mkMachine (Just pt) = initMachine pt mempty emptyStore
 
 (.=?) :: (Ae.KeyValue a, Ae.ToJSON v, Eq v) => Ae.Key -> v -> v -> Maybe a
 (.=?) n v defaultVal = if defaultVal /= v then Just $ n Ae..= v else Nothing
@@ -657,22 +667,22 @@ isActive = isNothing . getResult
 -- | "Active" robots include robots that are waiting; 'wantsToStep' is
 --   true if the robot actually wants to take another step right now
 --   (this is a /subset/ of active robots).
-wantsToStep :: TickNumber -> Robot -> Bool
+wantsToStep :: C.TickNumber -> Robot -> Bool
 wantsToStep now robot
   | not (isActive robot) = False
   | otherwise = maybe True (now >=) (waitingUntil robot)
 
 -- | The time until which the robot is waiting, if any.
-waitingUntil :: Robot -> Maybe TickNumber
+waitingUntil :: Robot -> Maybe C.TickNumber
 waitingUntil robot =
   case _machine robot of
-    Waiting time _ -> Just time
+    C.Waiting time _ -> Just time
     _ -> Nothing
 
 -- | Get the result of the robot's computation if it is finished.
-getResult :: Robot -> Maybe (Value, Store)
+getResult :: Robot -> Maybe (Value, C.Store)
 {-# INLINE getResult #-}
-getResult = finalValue . view machine
+getResult = C.finalValue . view machine
 
 hearingDistance :: (Num i) => i
 hearingDistance = 32

src/swarm-engine/Swarm/Game/State.hs

@@ -677,8 +677,11 @@ pureScenarioToGameState scenario theSeed now toRun gsc =
 
   initialCodeToRun = getCodeToRun <$> toRun
 
+  robotListRaw =
+    zipWith (instantiateRobot Nothing) [baseID ..] robotsByBasePrecedence
+
   robotList =
-    zipWith instantiateRobot [baseID ..] robotsByBasePrecedence
+    robotListRaw
       -- If the  --run flag was used, use it to replace the CESK machine of the
       -- robot whose id is 0, i.e. the first robot listed in the scenario.
       -- Note that this *replaces* any program the base robot otherwise

src/swarm-engine/Swarm/Game/State/Robot.hs

@@ -231,10 +231,10 @@ viewCenterRule = lens getter setter
 --   First, generate a unique ID number for it.  Then, add it to the
 --   main robot map, the active robot set, and to to the index of
 --   robots by location. Return the updated robot.
-addTRobot :: (Has (State Robots) sig m) => TRobot -> m Robot
-addTRobot r = do
+addTRobot :: (Has (State Robots) sig m) => CESK -> TRobot -> m Robot
+addTRobot initialMachine r = do
   rid <- robotNaming . gensym <+= 1
-  let r' = instantiateRobot rid r
+  let r' = instantiateRobot (Just initialMachine) rid r
   addRobot r'
   return r'
 

src/swarm-engine/Swarm/Game/Step/Combustion.hs

@@ -97,7 +97,7 @@ addCombustionBot inputEntity combustibility ts loc = do
   let combustionProg = combustionProgram combustionDurationRand combustibility
   void
     . zoomRobots
-    . addTRobot
+    . addTRobot (initMachine combustionProg empty emptyStore)
     $ mkRobot
       ()
       Nothing
@@ -109,7 +109,7 @@ addCombustionBot inputEntity combustibility ts loc = do
           & displayAttr .~ AWorld "fire"
           & displayPriority .~ 0
       )
-      (initMachine combustionProg empty emptyStore)
+      Nothing
       []
       botInventory
       True
@@ -212,7 +212,7 @@ addIgnitionBot ::
   m ()
 addIgnitionBot ignitionDelay inputEntity ts loc =
   void $
-    addTRobot $
+    addTRobot (initMachine (ignitionProgram ignitionDelay) empty emptyStore) $
       mkRobot
         ()
         Nothing
@@ -223,7 +223,7 @@ addIgnitionBot ignitionDelay inputEntity ts loc =
         ( defaultEntityDisplay '*'
             & invisible .~ True
         )
-        (initMachine (ignitionProgram ignitionDelay) empty emptyStore)
+        Nothing
         []
         []
         True

src/swarm-engine/Swarm/Game/Step/Const.hs

@@ -1061,7 +1061,7 @@ execConst runChildProg c vs s k = do
         -- Construct the new robot and add it to the world.
         parentCtx <- use robotContext
         newRobot <-
-          zoomRobots . addTRobot . (trobotContext .~ parentCtx) $
+          zoomRobots . addTRobot (In cmd e s [FExec]) . (trobotContext .~ parentCtx) $
             mkRobot
               ()
               (Just pid)
@@ -1074,7 +1074,7 @@ execConst runChildProg c vs s k = do
               ( defaultRobotDisplay
                   & inherit displayAttr (r ^. robotDisplay)
               )
-              (In cmd e s [FExec])
+              Nothing
               []
               []
               isSystemRobot

src/swarm-engine/Swarm/Game/Step/Util/Command.hs

@@ -375,7 +375,7 @@ addSeedBot ::
 addSeedBot e (minT, maxT) loc ts =
   void
     . zoomRobots
-    . addTRobot
+    . addTRobot (initMachine (seedProgram minT (maxT - minT) (e ^. entityName)) empty emptyStore)
     $ mkRobot
       ()
       Nothing
@@ -387,7 +387,7 @@ addSeedBot e (minT, maxT) loc ts =
           & displayAttr .~ (e ^. entityDisplay . displayAttr)
           & displayPriority .~ 0
       )
-      (initMachine (seedProgram minT (maxT - minT) (e ^. entityName)) empty emptyStore)
+      Nothing
       []
       [(1, e)]
       True

test/bench/Benchmark.hs

@@ -113,16 +113,16 @@ waveProgram manualInline =
 
 -- | Initializes a robot with program prog at location loc facing north.
 initRobot :: ProcessedTerm -> Location -> TRobot
-initRobot prog loc = mkRobot () Nothing "" mempty (Just $ Cosmic DefaultRootSubworld loc) north defaultRobotDisplay (initMachine prog Context.empty emptyStore) [] [] False False mempty 0
+initRobot prog loc = mkRobot () Nothing "" mempty (Just $ Cosmic DefaultRootSubworld loc) north defaultRobotDisplay (Just prog) [] [] False False mempty 0
 
 -- | Creates a GameState with numRobot copies of robot on a blank map, aligned
 --   in a row starting at (0,0) and spreading east.
-mkGameState :: (Location -> TRobot) -> Int -> IO GameState
-mkGameState robotMaker numRobots = do
+mkGameState :: ProcessedTerm -> (Location -> TRobot) -> Int -> IO GameState
+mkGameState prog robotMaker numRobots = do
   let robots = [robotMaker (Location (fromIntegral x) 0) | x <- [0 .. numRobots - 1]]
   Right initAppState <- runExceptT classicGame0
   execStateT
-    (zoomRobots $ mapM addTRobot robots)
+    (zoomRobots $ mapM (addTRobot $ initMachine prog Context.empty emptyStore) robots)
     ( (initAppState ^. gameState)
         & creativeMode .~ True
         & landscape . multiWorld .~ M.singleton DefaultRootSubworld (newWorld (WF $ const (fromEnum DirtT, ENothing)))
@@ -162,7 +162,7 @@ main = do
   robotNumbers = [10, 20 .. 40]
 
   mkGameStates :: ProcessedTerm -> IO [(Int, GameState)]
-  mkGameStates prog = zip robotNumbers <$> mapM (mkGameState (initRobot prog)) robotNumbers
+  mkGameStates prog = zip robotNumbers <$> mapM (mkGameState prog $ initRobot prog) robotNumbers
 
   toBenchmarks :: [(Int, GameState)] -> [Benchmark]
   toBenchmarks gameStates =