{-# LANGUAGE TemplateHaskell #-}

-- |
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Implementation of the 'Swarm.Language.Syntax.Volume' command for robots.
--
-- Note: If the robot is currently on an unwalkable cell (which may happen in
-- the case of teleportation or if an entity is placed or pushed into its cell),
-- the volume shall be zero.
module Swarm.Game.Step.Flood (
  floodFill,
) where

import Control.Effect.Lens
import Control.Lens (makeLenses, (%~), (&))
import Data.HashSet (HashSet)
import Data.HashSet qualified as HashSet
import Swarm.Game.Location
import Swarm.Game.Step.RobotStepState
import Swarm.Game.Step.Util (checkMoveFailureUnprivileged)
import Swarm.Game.Step.Util.Inspect (getNeighborLocs)
import Swarm.Game.Universe

data FloodParms = FloodParms
  { FloodParms -> SubworldName
theSubworld :: SubworldName
  , FloodParms -> Int
maxVisits :: Int
  }

data Tracking = Tracking
  { Tracking -> HashSet Location
visited :: HashSet Location
  , Tracking -> FloodPartition
floodPartition :: FloodPartition
  }

-- | We annotate each visited cell as
-- being part of the boundary or the interior.
-- This lets us:
--
-- 1. Use the interior cell count as a termination condition
-- 2. Handle (eventual) cache invalidation differently for boundary
--    members than interior members.
data FloodPartition = FloodPartition
  { FloodPartition -> HashSet Location
_boundary :: HashSet Location
  , FloodPartition -> HashSet Location
_interior :: HashSet Location
  }

makeLenses ''FloodPartition

-- |
-- == Algorithm
--
-- Explore via DFS using a list as a stack.
-- Each iteration examines a single cell.
--
-- 1. Mark the popped cell as visited, regardless of walkability.
-- 2. Check popped cell for walkability
-- 3. Add all neighbors that aren't already visited, regardless of walkability, to the stack.
--    But unwalkable cells shall not produce neighbors and shall be marked with a boundary/interior distinction.
floodRecursive ::
  HasRobotStepState sig m =>
  Tracking ->
  [Location] ->
  FloodParms ->
  m (Maybe Int)
floodRecursive :: forall (sig :: (* -> *) -> * -> *) (m :: * -> *).
HasRobotStepState sig m =>
Tracking -> [Location] -> FloodParms -> m (Maybe Int)
floodRecursive Tracking
tracking [Location]
pending FloodParms
params =
  case [Location]
pending of
    Location
nextLoc : [Location]
otherLocs ->
      if Int
interiorCount Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> FloodParms -> Int
maxVisits FloodParms
params
        then Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Int
forall a. Maybe a
Nothing
        else Location -> [Location] -> m (Maybe Int)
forall {m :: * -> *} {sig :: (* -> *) -> * -> *}.
(Algebra sig m, Member (Throw Exn) sig, Member (State Robot) sig,
 Member (State GameState) sig) =>
Location -> [Location] -> m (Maybe Int)
checkNeighbors Location
nextLoc [Location]
otherLocs
    [] -> Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Int -> m (Maybe Int)) -> Maybe Int -> m (Maybe Int)
forall a b. (a -> b) -> a -> b
$ Int -> Maybe Int
forall a. a -> Maybe a
Just Int
interiorCount
 where
  interiorCount :: Int
interiorCount = HashSet Location -> Int
forall a. HashSet a -> Int
HashSet.size (HashSet Location -> Int) -> HashSet Location -> Int
forall a b. (a -> b) -> a -> b
$ FloodPartition -> HashSet Location
_interior (FloodPartition -> HashSet Location)
-> FloodPartition -> HashSet Location
forall a b. (a -> b) -> a -> b
$ Tracking -> FloodPartition
floodPartition Tracking
tracking
  checkNeighbors :: Location -> [Location] -> m (Maybe Int)
checkNeighbors Location
nextLoc [Location]
otherLocs = do
    Bool
isWalkable <- Maybe MoveFailureMode -> Bool
forall a. Maybe a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Maybe MoveFailureMode -> Bool)
-> m (Maybe MoveFailureMode) -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Cosmic Location -> m (Maybe MoveFailureMode)
forall (sig :: (* -> *) -> * -> *) (m :: * -> *).
HasRobotStepState sig m =>
Cosmic Location -> m (Maybe MoveFailureMode)
checkMoveFailureUnprivileged Cosmic Location
cosmicLoc
    let candidateNeighbors :: [Location]
candidateNeighbors =
          if Bool
isWalkable
            then (Cosmic Location -> Location) -> [Cosmic Location] -> [Location]
forall a b. (a -> b) -> [a] -> [b]
map (Getting Location (Cosmic Location) Location
-> Cosmic Location -> Location
forall r a (sig :: (* -> *) -> * -> *) (m :: * -> *).
Has (Reader r) sig m =>
Getting a r a -> m a
view Getting Location (Cosmic Location) Location
forall a1 a2 (f :: * -> *).
Functor f =>
(a1 -> f a2) -> Cosmic a1 -> f (Cosmic a2)
planar) ([Cosmic Location] -> [Location])
-> [Cosmic Location] -> [Location]
forall a b. (a -> b) -> a -> b
$ Cosmic Location -> [Cosmic Location]
getNeighborLocs Cosmic Location
cosmicLoc
            else []
        visitableNeighbors :: [Location]
visitableNeighbors = (Location -> Bool) -> [Location] -> [Location]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Location -> Bool) -> Location -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Location -> HashSet Location -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`HashSet.member` Tracking -> HashSet Location
visited Tracking
tracking)) [Location]
candidateNeighbors

        -- It's cheaper to prepend the "visitableNeighbors" list because
        -- it should in general be a shorter list than the "pending" list.
        newPending :: [Location]
newPending = [Location]
visitableNeighbors [Location] -> [Location] -> [Location]
forall a. Semigroup a => a -> a -> a
<> [Location]
otherLocs

        partitionMutator :: (HashSet Location -> Identity (HashSet Location))
-> FloodPartition -> Identity FloodPartition
partitionMutator =
          if Bool
isWalkable
            then (HashSet Location -> Identity (HashSet Location))
-> FloodPartition -> Identity FloodPartition
Lens' FloodPartition (HashSet Location)
interior
            else (HashSet Location -> Identity (HashSet Location))
-> FloodPartition -> Identity FloodPartition
Lens' FloodPartition (HashSet Location)
boundary
        newPartition :: FloodPartition
newPartition = Tracking -> FloodPartition
floodPartition Tracking
tracking FloodPartition
-> (FloodPartition -> FloodPartition) -> FloodPartition
forall a b. a -> (a -> b) -> b
& (HashSet Location -> Identity (HashSet Location))
-> FloodPartition -> Identity FloodPartition
partitionMutator ((HashSet Location -> Identity (HashSet Location))
 -> FloodPartition -> Identity FloodPartition)
-> (HashSet Location -> HashSet Location)
-> FloodPartition
-> FloodPartition
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ Location -> HashSet Location -> HashSet Location
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
HashSet.insert Location
nextLoc

        newTracking :: Tracking
newTracking =
          Tracking
tracking
            { visited = newVisited
            , floodPartition = newPartition
            }
    Tracking -> [Location] -> FloodParms -> m (Maybe Int)
forall (sig :: (* -> *) -> * -> *) (m :: * -> *).
HasRobotStepState sig m =>
Tracking -> [Location] -> FloodParms -> m (Maybe Int)
floodRecursive Tracking
newTracking [Location]
newPending FloodParms
params
   where
    newVisited :: HashSet Location
newVisited = Location -> HashSet Location -> HashSet Location
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
HashSet.insert Location
nextLoc (HashSet Location -> HashSet Location)
-> HashSet Location -> HashSet Location
forall a b. (a -> b) -> a -> b
$ Tracking -> HashSet Location
visited Tracking
tracking
    cosmicLoc :: Cosmic Location
cosmicLoc = SubworldName -> Location -> Cosmic Location
forall a. SubworldName -> a -> Cosmic a
Cosmic (FloodParms -> SubworldName
theSubworld FloodParms
params) Location
nextLoc

floodFill ::
  HasRobotStepState sig m =>
  Cosmic Location ->
  Int ->
  m (Maybe Int)
floodFill :: forall (sig :: (* -> *) -> * -> *) (m :: * -> *).
HasRobotStepState sig m =>
Cosmic Location -> Int -> m (Maybe Int)
floodFill (Cosmic SubworldName
swn Location
curLoc) =
  Tracking -> [Location] -> FloodParms -> m (Maybe Int)
forall (sig :: (* -> *) -> * -> *) (m :: * -> *).
HasRobotStepState sig m =>
Tracking -> [Location] -> FloodParms -> m (Maybe Int)
floodRecursive Tracking
emptyTracking [Location
curLoc] (FloodParms -> m (Maybe Int))
-> (Int -> FloodParms) -> Int -> m (Maybe Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubworldName -> Int -> FloodParms
FloodParms SubworldName
swn
 where
  emptyTracking :: Tracking
emptyTracking = HashSet Location -> FloodPartition -> Tracking
Tracking HashSet Location
forall a. Monoid a => a
mempty (FloodPartition -> Tracking) -> FloodPartition -> Tracking
forall a b. (a -> b) -> a -> b
$ HashSet Location -> HashSet Location -> FloodPartition
FloodPartition HashSet Location
forall a. Monoid a => a
mempty HashSet Location
forall a. Monoid a => a
mempty