{-# Language LambdaCase, ImportQualifiedPost, BangPatterns #-}
{-|
Module      : Main
Description : Generalized graph search
Copyright   : (c) Eric Mertens, 2019-2021
License     : ISC
Maintainer  : emertens@gmail.com

These implementations provide a lazily-generated list of visited
states with the order defined by the search strategy.

-}
module Advent.Search (
  -- * Depth-first search
  dfs, dfsOn,

  -- * Breadth-first search
  bfs, bfsOn, bfsOnN,

  -- * A* search
  AStep(..),
  astar, astarOn,

  ) where

import Advent.PQueue qualified as PQueue
import Advent.Queue qualified as Queue
import Data.Foldable (foldl')
import Data.Set qualified as Set
import Data.IntSet qualified as IntSet

-- | Shortcut for @'dfsOn' 'id'@
dfs :: Ord a => (a -> [a]) -> a -> [a]
dfs :: forall a. Ord a => (a -> [a]) -> a -> [a]
dfs = (a -> a) -> (a -> [a]) -> a -> [a]
forall r a. Ord r => (a -> r) -> (a -> [a]) -> a -> [a]
dfsOn a -> a
forall a. a -> a
id
{-# INLINE dfs #-}

-- | Depth-first search.
--
-- Generates the list of unique visited states from a
-- given starting state. States are unique up to the
-- characterizing function.
dfsOn ::
  Ord r =>
  (a -> r)   {- ^ state characterization              -} ->
  (a -> [a]) {- ^ successors function                 -} ->
  a          {- ^ initial state                       -} ->
  [a]        {- ^ visited states in depth-first order -}
dfsOn :: forall r a. Ord r => (a -> r) -> (a -> [a]) -> a -> [a]
dfsOn a -> r
rep a -> [a]
next a
start = Set r -> [a] -> [a]
loop Set r
forall a. Set a
Set.empty [a
start]
  where
    loop :: Set r -> [a] -> [a]
loop !Set r
seen = \case
      [] -> []
      a
x:[a]
xs
        | r -> Set r -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member r
r Set r
seen ->     Set r -> [a] -> [a]
loop Set r
seen [a]
xs
        | Bool
otherwise         -> a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Set r -> [a] -> [a]
loop Set r
seen1 (a -> [a]
next a
x [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
xs)
        where
          r :: r
r     = a -> r
rep a
x
          seen1 :: Set r
seen1 = r -> Set r -> Set r
forall a. Ord a => a -> Set a -> Set a
Set.insert r
r Set r
seen

-- | Shortcut for @'bfsOn' 'id'@
bfs :: Ord a => (a -> [a]) -> a -> [a]
bfs :: forall a. Ord a => (a -> [a]) -> a -> [a]
bfs = (a -> a) -> (a -> [a]) -> a -> [a]
forall r a. Ord r => (a -> r) -> (a -> [a]) -> a -> [a]
bfsOn a -> a
forall a. a -> a
id
{-# INLINE bfs #-}

-- | Enumerate the reachable states in breadth-first order
-- given a successor state function and initial state.
--
-- States are compared for equality using the representative
-- function. If the representatives are equal the state is
-- considered already visited.
{-# INLINE [0] bfsOn #-}
bfsOn ::
  Ord r =>
  (a -> r)   {- ^ representative function   -} ->
  (a -> [a]) {- ^ successor state generator -} ->
  a          {- ^ initial state             -} ->
  [a]        {- ^ reachable states          -}
bfsOn :: forall r a. Ord r => (a -> r) -> (a -> [a]) -> a -> [a]
bfsOn a -> r
rep a -> [a]
next a
start = (a -> r) -> (a -> [a]) -> [a] -> [a]
forall r a. Ord r => (a -> r) -> (a -> [a]) -> [a] -> [a]
bfsOnN a -> r
rep a -> [a]
next [a
start]

-- | Generalization of 'bfsOn' allowing multiple
-- initial states to be considered.
bfsOnN ::
  Ord r =>
  (a -> r)   {- ^ representative function   -} ->
  (a -> [a]) {- ^ successor state generator -} ->
  [a]        {- ^ initial states            -} ->
  [a]        {- ^ reachable states          -}
bfsOnN :: forall r a. Ord r => (a -> r) -> (a -> [a]) -> [a] -> [a]
bfsOnN a -> r
rep a -> [a]
next [a]
start = Set r -> Queue a -> [a]
loop Set r
forall a. Set a
Set.empty ([a] -> Queue a
forall a. [a] -> Queue a
Queue.fromList [a]
start)
  where
    loop :: Set r -> Queue a -> [a]
loop !Set r
seen = \case
      Queue a
Queue.Empty -> []
      a
x Queue.:<| Queue a
q
        | r -> Set r -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member r
r Set r
seen ->     Set r -> Queue a -> [a]
loop Set r
seen  Queue a
q
        | Bool
otherwise         -> a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Set r -> Queue a -> [a]
loop Set r
seen' Queue a
q'
        where
          r :: r
r     = a -> r
rep a
x
          seen' :: Set r
seen' = r -> Set r -> Set r
forall a. Ord a => a -> Set a -> Set a
Set.insert r
r Set r
seen
          q' :: Queue a
q'    = [a] -> Queue a -> Queue a
forall a. [a] -> Queue a -> Queue a
Queue.appendList (a -> [a]
next a
x) Queue a
q
{-# INLINE [0] bfsOnN #-}

{-# RULES "bfsOn/Int" bfsOn = bfsOnInt #-}
{-# INLINE bfsOnInt #-}
bfsOnInt :: (a -> Int) -> (a -> [a]) -> a -> [a]
bfsOnInt :: forall a. (a -> Int) -> (a -> [a]) -> a -> [a]
bfsOnInt a -> Int
rep a -> [a]
next a
start = IntSet -> Queue a -> [a]
loop IntSet
IntSet.empty (a -> Queue a
forall a. a -> Queue a
Queue.singleton a
start)
  where
    loop :: IntSet -> Queue a -> [a]
loop !IntSet
seen = \case
      Queue a
Queue.Empty -> []
      a
x Queue.:<| Queue a
q
        | Int -> IntSet -> Bool
IntSet.member Int
r IntSet
seen ->     IntSet -> Queue a -> [a]
loop IntSet
seen  Queue a
q
        | Bool
otherwise            -> a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: IntSet -> Queue a -> [a]
loop IntSet
seen' Queue a
q'
        where
          r :: Int
r     = a -> Int
rep a
x
          seen' :: IntSet
seen' = Int -> IntSet -> IntSet
IntSet.insert Int
r IntSet
seen
          q' :: Queue a
q'    = [a] -> Queue a -> Queue a
forall a. [a] -> Queue a -> Queue a
Queue.appendList (a -> [a]
next a
x) Queue a
q

-- | Shortcut for @'astarOn' 'id'@
astar :: Ord a => (a -> [AStep a]) -> a -> [(a,Int)]
astar :: forall a. Ord a => (a -> [AStep a]) -> a -> [(a, Int)]
astar = (a -> a) -> (a -> [AStep a]) -> a -> [(a, Int)]
forall b a.
Ord b =>
(a -> b) -> (a -> [AStep a]) -> a -> [(a, Int)]
astarOn a -> a
forall a. a -> a
id
{-# INLINE astar #-}

-- | A* graph search producing a list of reached states and the
-- minimum cost of reaching that state.
--
-- Returned states will be unique up to the characterization function.
-- This allows extra information of a node to be ignored for the
-- purposes of the search. For example, a node might remember the
-- path used to reach it while for the search the particular path
-- taken might not matter.
astarOn ::
  Ord b =>
  (a -> b)         {- ^ state characterization                                   -} ->
  (a -> [AStep a]) {- ^ step function (new state, step cost, distance heuristic) -} ->
  a                {- ^ starting state                                           -} ->
  [(a,Int)]        {- ^ list of states visited                                   -}
astarOn :: forall b a.
Ord b =>
(a -> b) -> (a -> [AStep a]) -> a -> [(a, Int)]
astarOn a -> b
rep a -> [AStep a]
nexts a
start = Set b -> PQueue (WithCost a) -> [(a, Int)]
go Set b
forall a. Set a
Set.empty (Int -> WithCost a -> PQueue (WithCost a)
forall a. Int -> a -> PQueue a
PQueue.singleton Int
0 (Int -> a -> WithCost a
forall a. Int -> a -> WithCost a
WC Int
0 a
start))
  where
    go :: Set b -> PQueue (WithCost a) -> [(a, Int)]
go !Set b
seen = \case
      PQueue (WithCost a)
PQueue.Empty -> []
      WC Int
cost a
x PQueue.:<| PQueue (WithCost a)
work
        | b -> Set b -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member b
r Set b
seen -> Set b -> PQueue (WithCost a) -> [(a, Int)]
go Set b
seen PQueue (WithCost a)
work
        | Bool
otherwise         -> (a
x,Int
cost) (a, Int) -> [(a, Int)] -> [(a, Int)]
forall a. a -> [a] -> [a]
: Set b -> PQueue (WithCost a) -> [(a, Int)]
go Set b
seen' PQueue (WithCost a)
work'
        where
          r :: b
r     = a -> b
rep a
x
          seen' :: Set b
seen' = b -> Set b -> Set b
forall a. Ord a => a -> Set a -> Set a
Set.insert b
r Set b
seen
          work' :: PQueue (WithCost a)
work' = (PQueue (WithCost a) -> AStep a -> PQueue (WithCost a))
-> PQueue (WithCost a) -> [AStep a] -> PQueue (WithCost a)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' PQueue (WithCost a) -> AStep a -> PQueue (WithCost a)
forall {a}. PQueue (WithCost a) -> AStep a -> PQueue (WithCost a)
addWork PQueue (WithCost a)
work (a -> [AStep a]
nexts a
x)
          addWork :: PQueue (WithCost a) -> AStep a -> PQueue (WithCost a)
addWork PQueue (WithCost a)
w (AStep a
x' Int
stepcost Int
heuristic) =
            Int -> WithCost a -> PQueue (WithCost a) -> PQueue (WithCost a)
forall a. Int -> a -> PQueue a -> PQueue a
PQueue.insert (Int
cost' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
heuristic) (Int -> a -> WithCost a
forall a. Int -> a -> WithCost a
WC Int
cost' a
x') PQueue (WithCost a)
w
            where
              cost' :: Int
cost' = Int
cost Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stepcost
{-# INLINE astarOn #-}

-- Helper type to unpack the cost value in the A* priority queue
data WithCost a = WC !Int a

-- | A step in the A* graph search annotated with its cost and an
-- estimate of the distance remaining to the goal. The estimate
-- must be an underapproximation to ensure the search finds the
-- optimal solution
data AStep a = AStep {
  forall a. AStep a -> a
astepNext      :: a,    -- ^ successor node
  forall a. AStep a -> Int
astepCost      :: !Int, -- ^ cost of edge
  forall a. AStep a -> Int
astepHeuristic :: !Int  -- ^ heuristic cost to goal from this new node
  } deriving Int -> AStep a -> ShowS
[AStep a] -> ShowS
AStep a -> String
(Int -> AStep a -> ShowS)
-> (AStep a -> String) -> ([AStep a] -> ShowS) -> Show (AStep a)
forall a. Show a => Int -> AStep a -> ShowS
forall a. Show a => [AStep a] -> ShowS
forall a. Show a => AStep a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AStep a] -> ShowS
$cshowList :: forall a. Show a => [AStep a] -> ShowS
show :: AStep a -> String
$cshow :: forall a. Show a => AStep a -> String
showsPrec :: Int -> AStep a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> AStep a -> ShowS
Show