{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}

-- | This module contains a collection of generalized graph search algorithms,
-- for when you don't want to explicitly represent your data as a graph. The
-- general idea is to provide these algorithms with a way of generating "next"
-- states, a way of generating associated information, a way of determining
-- when you have found a solution, and an initial state.
module Algorithm.Search (
  -- * Searches
  bfs,
  dfs,
  dijkstra,
  dijkstraAssoc,
  aStar,
  aStarAssoc,
  -- * Monadic Searches
  -- $monadic
  bfsM,
  dfsM,
  dijkstraM,
  aStarM,
  -- * Utility
  incrementalCosts,
  incrementalCostsM,
  pruning,
  pruningM
  ) where

import qualified Data.Map as Map
import qualified Data.Sequence as Seq
import qualified Data.Set as Set
import qualified Data.List as List
import qualified Data.Foldable as Foldable
import Data.Functor.Identity (Identity(..))
import Control.Monad (filterM, zipWithM)

-- | @bfs next found initial@ performs a breadth-first search over a set of
-- states, starting with @initial@, and generating neighboring states with
-- @next@. It returns a path to a state for which @found@ returns 'True'.
-- Returns 'Nothing' if no path is possible.
--
-- === Example: Making change problem
--
-- >>> :{
-- countChange target = bfs (add_one_coin `pruning` (> target)) (== target) 0
--   where
--     add_one_coin amt = map (+ amt) coins
--     coins = [25, 10, 5, 1]
-- :}
--
-- >>> countChange 67
-- Just [25,50,60,65,66,67]
bfs :: (Foldable f, Ord state)
  => (state -> f state)
  -- ^ Function to generate "next" states given a current state
  -> (state -> Bool)
  -- ^ Predicate to determine if solution found. 'bfs' returns a path to the
  -- first state for which this predicate returns 'True'.
  -> state
  -- ^ Initial state
  -> Maybe [state]
  -- ^ First path found to a state matching the predicate, or 'Nothing' if no
  -- such path exists.
bfs :: (state -> f state) -> (state -> Bool) -> state -> Maybe [state]
bfs =
  -- BFS is a generalized search using a queue, which directly compares states,
  -- and which always uses the first path found to a state
  Seq state
-> (state -> state)
-> ([state] -> [state] -> Bool)
-> (state -> f state)
-> (state -> Bool)
-> state
-> Maybe [state]
forall (f :: * -> *) container stateKey state.
(Foldable f, SearchContainer container, Ord stateKey,
 Elem container ~ state) =>
container
-> (state -> stateKey)
-> ([state] -> [state] -> Bool)
-> (state -> f state)
-> (state -> Bool)
-> state
-> Maybe [state]
generalizedSearch Seq state
forall a. Seq a
Seq.empty state -> state
forall a. a -> a
id (\[state]
_ [state]
_ -> Bool
False)


-- | @dfs next found initial@ performs a depth-first search over a set
-- of states, starting with @initial@ and generating neighboring states with
-- @next@. It returns a depth-first path to a state for which @found@ returns
-- 'True'. Returns 'Nothing' if no path is possible.
--
-- === Example: Simple directed graph search
--
-- >>> import qualified Data.Map as Map
--
-- >>> graph = Map.fromList [(1, [2, 3]), (2, [4]), (3, [4]), (4, [])]
--
-- >>> dfs (graph Map.!) (== 4) 1
-- Just [3,4]
dfs :: (Foldable f, Ord state)
  => (state -> f state)
  -- ^ Function to generate "next" states given a current state. These should be
  -- given in the order in which states should be pushed onto the stack, i.e.
  -- the "last" state in the Foldable will be the first one visited.
  -> (state -> Bool)
  -- ^ Predicate to determine if solution found. 'dfs' returns a path to the
  -- first state for which this predicate returns 'True'.
  -> state
  -- ^ Initial state
  -> Maybe [state]
  -- ^ First path found to a state matching the predicate, or 'Nothing' if no
  -- such path exists.
dfs :: (state -> f state) -> (state -> Bool) -> state -> Maybe [state]
dfs =
  -- DFS is a generalized search using a stack, which directly compares states,
  -- and which always uses the most recent path found to a state
  [state]
-> (state -> state)
-> ([state] -> [state] -> Bool)
-> (state -> f state)
-> (state -> Bool)
-> state
-> Maybe [state]
forall (f :: * -> *) container stateKey state.
(Foldable f, SearchContainer container, Ord stateKey,
 Elem container ~ state) =>
container
-> (state -> stateKey)
-> ([state] -> [state] -> Bool)
-> (state -> f state)
-> (state -> Bool)
-> state
-> Maybe [state]
generalizedSearch [] state -> state
forall a. a -> a
id (\[state]
_ [state]
_ -> Bool
True)


-- | @dijkstra next cost found initial@ performs a shortest-path search over
-- a set of states using Dijkstra's algorithm, starting with @initial@,
-- generating neighboring states with @next@, and their incremental costs with
-- @costs@. This will find the least-costly path from an initial state to a
-- state for which @found@ returns 'True'. Returns 'Nothing' if no path to a
-- solved state is possible.
--
-- === Example: Making change problem, with a twist
--
-- >>> :{
-- -- Twist: dimes have a face value of 10 cents, but are actually rare
-- -- misprints which are worth 10 dollars
-- countChange target =
--   dijkstra (add_coin `pruning` (> target)) true_cost  (== target) 0
--   where
--     coin_values = [(25, 25), (10, 1000), (5, 5), (1, 1)]
--     add_coin amt = map ((+ amt) . snd) coin_values
--     true_cost low high =
--       case lookup (high - low) coin_values of
--         Just val -> val
--         Nothing -> error $ "invalid costs: " ++ show high ++ ", " ++ show low
-- :}
--
-- >>> countChange 67
-- Just (67,[1,2,7,12,17,42,67])
dijkstra :: (Foldable f, Num cost, Ord cost, Ord state)
  => (state -> f state)
  -- ^ Function to generate list of neighboring states given the current state
  -> (state -> state -> cost)
  -- ^ Function to generate transition costs between neighboring states. This is
  -- only called for adjacent states, so it is safe to have this function be
  -- partial for non-neighboring states.
  -> (state -> Bool)
  -- ^ Predicate to determine if solution found. 'dijkstra' returns the shortest
  -- path to the first state for which this predicate returns 'True'.
  -> state
  -- ^ Initial state
  -> Maybe (cost, [state])
  -- ^ (Total cost, list of steps) for the first path found which
  -- satisfies the given predicate
dijkstra :: (state -> f state)
-> (state -> state -> cost)
-> (state -> Bool)
-> state
-> Maybe (cost, [state])
dijkstra state -> f state
next state -> state -> cost
cost state -> Bool
found state
initial =
  -- This API to Dijkstra's algorithm is useful when the state transition
  -- function and the cost function are logically separate.
  -- It is implemented by using @dijkstraAssoc@ with appropriate mapping of
  -- arguments.
  (state -> [(state, cost)])
-> (state -> Bool) -> state -> Maybe (cost, [state])
forall cost state.
(Num cost, Ord cost, Ord state) =>
(state -> [(state, cost)])
-> (state -> Bool) -> state -> Maybe (cost, [state])
dijkstraAssoc state -> [(state, cost)]
next' state -> Bool
found state
initial
  where
    next' :: state -> [(state, cost)]
next' state
st = (state -> (state, cost)) -> [state] -> [(state, cost)]
forall a b. (a -> b) -> [a] -> [b]
map (\state
new_st -> (state
new_st, state -> state -> cost
cost state
st state
new_st)) ([state] -> [(state, cost)]) -> [state] -> [(state, cost)]
forall a b. (a -> b) -> a -> b
$
               f state -> [state]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList (state -> f state
next state
st)

-- | @dijkstraAssoc next found initial@ performs a shortest-path search over
-- a set of states using Dijkstra's algorithm, starting with @initial@,
-- generating neighboring states with associated incremenal costs with
-- @next@. This will find the least-costly path from an initial state to a
-- state for which @found@ returns 'True'. Returns 'Nothing' if no path to a
-- solved state is possible.
dijkstraAssoc :: (Num cost, Ord cost, Ord state)
  => (state -> [(state, cost)])
  -- ^ function to generate list of neighboring states with associated
  -- transition costs given the current state
  -> (state -> Bool)
  -- ^ Predicate to determine if solution found. 'dijkstraAssoc' returns the
  -- shortest path to the first state for which this predicate returns 'True'.
  -> state
  -- ^ Initial state
  -> Maybe (cost, [state])
  -- ^ (Total cost, list of steps) for the first path found which
  -- satisfies the given predicate
dijkstraAssoc :: (state -> [(state, cost)])
-> (state -> Bool) -> state -> Maybe (cost, [state])
dijkstraAssoc state -> [(state, cost)]
next state -> Bool
found state
initial =
  -- This API to Dijkstra's algoritm is useful in the common case when next
  -- states and their associated transition costs are generated together.
  --
  -- Dijkstra's algorithm can be viewed as a generalized search, with the search
  -- container being a heap, with the states being compared without regard to
  -- cost, with the shorter paths taking precedence over longer ones, and with
  -- the stored state being (cost so far, state).
  -- This implementation makes that transformation, then transforms that result
  -- back into the desired result from @dijkstraAssoc@
  [(cost, state)] -> (cost, [state])
forall a a. Num a => [(a, a)] -> (a, [a])
unpack ([(cost, state)] -> (cost, [state]))
-> Maybe [(cost, state)] -> Maybe (cost, [state])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
    LIFOHeap cost state
-> ((cost, state) -> state)
-> ([(cost, state)] -> [(cost, state)] -> Bool)
-> ((cost, state) -> [(cost, state)])
-> ((cost, state) -> Bool)
-> (cost, state)
-> Maybe [(cost, state)]
forall (f :: * -> *) container stateKey state.
(Foldable f, SearchContainer container, Ord stateKey,
 Elem container ~ state) =>
container
-> (state -> stateKey)
-> ([state] -> [state] -> Bool)
-> (state -> f state)
-> (state -> Bool)
-> state
-> Maybe [state]
generalizedSearch LIFOHeap cost state
forall k a. LIFOHeap k a
emptyLIFOHeap (cost, state) -> state
forall a b. (a, b) -> b
snd [(cost, state)] -> [(cost, state)] -> Bool
forall a b. Ord a => [(a, b)] -> [(a, b)] -> Bool
leastCostly (cost, state) -> [(cost, state)]
next' (state -> Bool
found (state -> Bool)
-> ((cost, state) -> state) -> (cost, state) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (cost, state) -> state
forall a b. (a, b) -> b
snd)
      (cost
0, state
initial)
  where
    next' :: (cost, state) -> [(cost, state)]
next' (cost
old_cost, state
st) =
      (\(state
new_st, cost
new_cost) -> (cost
new_cost cost -> cost -> cost
forall a. Num a => a -> a -> a
+ cost
old_cost, state
new_st))
        ((state, cost) -> (cost, state))
-> [(state, cost)] -> [(cost, state)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (state -> [(state, cost)]
next state
st)
    unpack :: [(a, a)] -> (a, [a])
unpack [] = (a
0, [])
    unpack [(a, a)]
packed_states = ((a, a) -> a
forall a b. (a, b) -> a
fst ((a, a) -> a) -> ([(a, a)] -> (a, a)) -> [(a, a)] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(a, a)] -> (a, a)
forall a. [a] -> a
last ([(a, a)] -> a) -> [(a, a)] -> a
forall a b. (a -> b) -> a -> b
$ [(a, a)]
packed_states, ((a, a) -> a) -> [(a, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, a) -> a
forall a b. (a, b) -> b
snd [(a, a)]
packed_states)


-- | @aStar next cost remaining found initial@ performs a best-first search
-- using the A* search algorithm, starting with the state @initial@, generating
-- neighboring states with @next@, their cost with @cost@, and an estimate of
-- the remaining cost with @remaining@. This returns a path to a state for which
-- @found@ returns 'True'. If @remaining@ is strictly a lower bound on the
-- remaining cost to reach a solved state, then the returned path is the
-- shortest path. Returns 'Nothing' if no path to a solved state is possible.
--
-- === Example: Path finding in taxicab geometry
--
-- >>> :{
-- neighbors (x, y) = [(x, y + 1), (x - 1, y), (x + 1, y), (x, y - 1)]
-- dist (x1, y1) (x2, y2) = abs (y2 - y1) + abs (x2 - x1)
-- start = (0, 0)
-- end = (0, 2)
-- isWall = (== (0, 1))
-- :}
--
-- >>> aStar (neighbors `pruning` isWall) dist (dist end) (== end) start
-- Just (4,[(1,0),(1,1),(1,2),(0,2)])
aStar :: (Foldable f, Num cost, Ord cost, Ord state)
  => (state -> f state)
  -- ^ Function to generate list of neighboring states given the current state
  -> (state -> state -> cost)
  -- ^ Function to generate transition costs between neighboring states. This is
  -- only called for adjacent states, so it is safe to have this function be
  -- partial for non-neighboring states.
  -> (state -> cost)
  -- ^ Estimate on remaining cost given a state
  -> (state -> Bool)
  -- ^ Predicate to determine if solution found. 'aStar' returns the shortest
  -- path to the first state for which this predicate returns 'True'.
  -> state
  -- ^ Initial state
  -> Maybe (cost, [state])
  -- ^ (Total cost, list of steps) for the first path found which satisfies the
  -- given predicate
aStar :: (state -> f state)
-> (state -> state -> cost)
-> (state -> cost)
-> (state -> Bool)
-> state
-> Maybe (cost, [state])
aStar state -> f state
next state -> state -> cost
cost state -> cost
remaining state -> Bool
found state
initial =
  -- This API to A* search is useful when the state transition
  -- function and the cost function are logically separate.
  -- It is implemented by using @aStarAssoc@ with appropriate mapping of
  -- arguments.
  (state -> [(state, cost)])
-> (state -> cost)
-> (state -> Bool)
-> state
-> Maybe (cost, [state])
forall cost state.
(Num cost, Ord cost, Ord state) =>
(state -> [(state, cost)])
-> (state -> cost)
-> (state -> Bool)
-> state
-> Maybe (cost, [state])
aStarAssoc state -> [(state, cost)]
next' state -> cost
remaining state -> Bool
found state
initial
  where
    next' :: state -> [(state, cost)]
next' state
st = (state -> (state, cost)) -> [state] -> [(state, cost)]
forall a b. (a -> b) -> [a] -> [b]
map (\state
new_st -> (state
new_st, state -> state -> cost
cost state
st state
new_st)) ([state] -> [(state, cost)]) -> [state] -> [(state, cost)]
forall a b. (a -> b) -> a -> b
$
               f state -> [state]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList (state -> f state
next state
st)

-- | @aStarAssoc next remaining found initial@ performs a best-first search
-- using the A* search algorithm, starting with the state @initial@, generating
-- neighboring states and their associated costs with @next@, and an estimate of
-- the remaining cost with @remaining@. This returns a path to a state for which
-- @found@ returns 'True'. If @remaining@ is strictly a lower bound on the
-- remaining cost to reach a solved state, then the returned path is the
-- shortest path. Returns 'Nothing' if no path to a solved state is possible.
aStarAssoc :: (Num cost, Ord cost, Ord state)
  => (state -> [(state, cost)])
  -- ^ function to generate list of neighboring states with associated
  -- transition costs given the current state
  -> (state -> cost)
  -- ^ Estimate on remaining cost given a state
  -> (state -> Bool)
  -- ^ Predicate to determine if solution found. 'aStar' returns the shortest
  -- path to the first state for which this predicate returns 'True'.
  -> state
  -- ^ Initial state
  -> Maybe (cost, [state])
  -- ^ (Total cost, list of steps) for the first path found which satisfies the
  -- given predicate
aStarAssoc :: (state -> [(state, cost)])
-> (state -> cost)
-> (state -> Bool)
-> state
-> Maybe (cost, [state])
aStarAssoc state -> [(state, cost)]
next state -> cost
remaining state -> Bool
found state
initial =
  -- This API to A* search is useful in the common case when next
  -- states and their associated transition costs are generated together.
  --
  -- A* can be viewed as a generalized search, with the search container being a
  -- heap, with the states being compared without regard to cost, with the
  -- shorter paths taking precedence over longer ones, and with
  -- the stored state being (total cost estimate, (cost so far, state)).
  -- This implementation makes that transformation, then transforms that result
  -- back into the desired result from @aStarAssoc@
  [(cost, (cost, state))] -> (cost, [state])
forall a a a. Num a => [(a, (a, a))] -> (a, [a])
unpack ([(cost, (cost, state))] -> (cost, [state]))
-> Maybe [(cost, (cost, state))] -> Maybe (cost, [state])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> LIFOHeap cost (cost, state)
-> ((cost, (cost, state)) -> state)
-> ([(cost, (cost, state))] -> [(cost, (cost, state))] -> Bool)
-> ((cost, (cost, state)) -> [(cost, (cost, state))])
-> ((cost, (cost, state)) -> Bool)
-> (cost, (cost, state))
-> Maybe [(cost, (cost, state))]
forall (f :: * -> *) container stateKey state.
(Foldable f, SearchContainer container, Ord stateKey,
 Elem container ~ state) =>
container
-> (state -> stateKey)
-> ([state] -> [state] -> Bool)
-> (state -> f state)
-> (state -> Bool)
-> state
-> Maybe [state]
generalizedSearch LIFOHeap cost (cost, state)
forall k a. LIFOHeap k a
emptyLIFOHeap (cost, (cost, state)) -> state
forall a a c. (a, (a, c)) -> c
snd2 [(cost, (cost, state))] -> [(cost, (cost, state))] -> Bool
forall a b. Ord a => [(a, b)] -> [(a, b)] -> Bool
leastCostly (cost, (cost, state)) -> [(cost, (cost, state))]
next'
    (state -> Bool
found (state -> Bool)
-> ((cost, (cost, state)) -> state)
-> (cost, (cost, state))
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (cost, (cost, state)) -> state
forall a a c. (a, (a, c)) -> c
snd2) (state -> cost
remaining state
initial, (cost
0, state
initial))
  where
    next' :: (cost, (cost, state)) -> [(cost, (cost, state))]
next' (cost
_, (cost
old_cost, state
old_st)) =
      (state, cost) -> (cost, (cost, state))
update_state ((state, cost) -> (cost, (cost, state)))
-> [(state, cost)] -> [(cost, (cost, state))]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (state -> [(state, cost)]
next state
old_st)
      where
        update_state :: (state, cost) -> (cost, (cost, state))
update_state (state
new_st, cost
cost) =
          let new_cost :: cost
new_cost = cost
old_cost cost -> cost -> cost
forall a. Num a => a -> a -> a
+ cost
cost
              new_est :: cost
new_est = cost
new_cost cost -> cost -> cost
forall a. Num a => a -> a -> a
+ state -> cost
remaining state
new_st
          in (cost
new_est, (cost
new_cost, state
new_st))
    unpack :: [(a, (a, a))] -> (a, [a])
unpack [] = (a
0, [])
    unpack [(a, (a, a))]
packed_states =
      ((a, a) -> a
forall a b. (a, b) -> a
fst ((a, a) -> a) -> ([(a, (a, a))] -> (a, a)) -> [(a, (a, a))] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, (a, a)) -> (a, a)
forall a b. (a, b) -> b
snd ((a, (a, a)) -> (a, a))
-> ([(a, (a, a))] -> (a, (a, a))) -> [(a, (a, a))] -> (a, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(a, (a, a))] -> (a, (a, a))
forall a. [a] -> a
last ([(a, (a, a))] -> a) -> [(a, (a, a))] -> a
forall a b. (a -> b) -> a -> b
$ [(a, (a, a))]
packed_states, ((a, (a, a)) -> a) -> [(a, (a, a))] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, (a, a)) -> a
forall a a c. (a, (a, c)) -> c
snd2 [(a, (a, a))]
packed_states)
    snd2 :: (a, (a, c)) -> c
snd2 = (a, c) -> c
forall a b. (a, b) -> b
snd ((a, c) -> c) -> ((a, (a, c)) -> (a, c)) -> (a, (a, c)) -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, (a, c)) -> (a, c)
forall a b. (a, b) -> b
snd

-- $monadic
-- Note that for all monadic searches, it is up to the user to ensure that
-- side-effecting monads do not logically change the structure of the graph.
-- For example, if the list of neighbors is being read from a file, the user
-- must ensure that those values do not change between reads.

-- | @bfsM@ is a monadic version of 'bfs': it has support for monadic @next@ and
-- @found@ parameters.
bfsM :: (Monad m, Foldable f, Ord state)
  => (state -> m (f state))
  -- ^ Function to generate "next" states given a current state
  -> (state -> m Bool)
  -- ^ Predicate to determine if solution found. 'bfsM' returns a path to the
  -- first state for which this predicate returns 'True'.
  -> state
  -- ^ Initial state
  -> m (Maybe [state])
  -- ^ First path found to a state matching the predicate, or 'Nothing' if no
  -- such path exists.
bfsM :: (state -> m (f state))
-> (state -> m Bool) -> state -> m (Maybe [state])
bfsM = Seq state
-> (state -> state)
-> ([state] -> [state] -> Bool)
-> (state -> m (f state))
-> (state -> m Bool)
-> state
-> m (Maybe [state])
forall (m :: * -> *) (f :: * -> *) container stateKey state.
(Monad m, Foldable f, SearchContainer container, Ord stateKey,
 Elem container ~ state) =>
container
-> (state -> stateKey)
-> ([state] -> [state] -> Bool)
-> (state -> m (f state))
-> (state -> m Bool)
-> state
-> m (Maybe [state])
generalizedSearchM Seq state
forall a. Seq a
Seq.empty state -> state
forall a. a -> a
id (\[state]
_ [state]
_ -> Bool
False)


-- | @dfsM@ is a monadic version of 'dfs': it has support for monadic @next@ and
-- @found@ parameters.
dfsM :: (Monad m, Foldable f, Ord state)
  => (state -> m (f state))
  -- ^ Function to generate "next" states given a current state
  -> (state -> m Bool)
  -- ^ Predicate to determine if solution found. 'dfsM' returns a path to the
  -- first state for which this predicate returns 'True'.
  -> state
  -- ^ Initial state
  -> m (Maybe [state])
  -- ^ First path found to a state matching the predicate, or 'Nothing' if no
  -- such path exists.
dfsM :: (state -> m (f state))
-> (state -> m Bool) -> state -> m (Maybe [state])
dfsM =
  [state]
-> (state -> state)
-> ([state] -> [state] -> Bool)
-> (state -> m (f state))
-> (state -> m Bool)
-> state
-> m (Maybe [state])
forall (m :: * -> *) (f :: * -> *) container stateKey state.
(Monad m, Foldable f, SearchContainer container, Ord stateKey,
 Elem container ~ state) =>
container
-> (state -> stateKey)
-> ([state] -> [state] -> Bool)
-> (state -> m (f state))
-> (state -> m Bool)
-> state
-> m (Maybe [state])
generalizedSearchM [] state -> state
forall a. a -> a
id (\[state]
_ [state]
_ -> Bool
True)

-- | @dijkstraM@ is a monadic version of 'dijkstra': it has support for monadic
-- @next@, @cost@, and @found@ parameters.
dijkstraM :: (Monad m, Foldable f, Num cost, Ord cost, Ord state)
  => (state -> m (f state))
  -- ^ Function to generate list of neighboring states given the current state
  -> (state -> state -> m cost)
  -- ^ Function to generate list of costs between neighboring states. This is
  -- only called for adjacent states, so it is safe to have this function be
  -- partial for non-neighboring states.
  -> (state -> m Bool)
  -- ^ Predicate to determine if solution found. 'dijkstraM' returns the
  -- shortest path to the first state for which this predicate returns 'True'.
  -> state
  -- ^ Initial state
  -> m (Maybe (cost, [state]))
  -- ^ (Total cost, list of steps) for the first path found which
  -- satisfies the given predicate
dijkstraM :: (state -> m (f state))
-> (state -> state -> m cost)
-> (state -> m Bool)
-> state
-> m (Maybe (cost, [state]))
dijkstraM state -> m (f state)
nextM state -> state -> m cost
costM state -> m Bool
foundM state
initial =
  ([(cost, state)] -> (cost, [state]))
-> m (Maybe [(cost, state)]) -> m (Maybe (cost, [state]))
forall (f1 :: * -> *) (f2 :: * -> *) a b.
(Functor f1, Functor f2) =>
(a -> b) -> f1 (f2 a) -> f1 (f2 b)
fmap2 [(cost, state)] -> (cost, [state])
forall a a. Num a => [(a, a)] -> (a, [a])
unpack (m (Maybe [(cost, state)]) -> m (Maybe (cost, [state])))
-> m (Maybe [(cost, state)]) -> m (Maybe (cost, [state]))
forall a b. (a -> b) -> a -> b
$ LIFOHeap cost state
-> ((cost, state) -> state)
-> ([(cost, state)] -> [(cost, state)] -> Bool)
-> ((cost, state) -> m [(cost, state)])
-> ((cost, state) -> m Bool)
-> (cost, state)
-> m (Maybe [(cost, state)])
forall (m :: * -> *) (f :: * -> *) container stateKey state.
(Monad m, Foldable f, SearchContainer container, Ord stateKey,
 Elem container ~ state) =>
container
-> (state -> stateKey)
-> ([state] -> [state] -> Bool)
-> (state -> m (f state))
-> (state -> m Bool)
-> state
-> m (Maybe [state])
generalizedSearchM LIFOHeap cost state
forall k a. LIFOHeap k a
emptyLIFOHeap (cost, state) -> state
forall a b. (a, b) -> b
snd [(cost, state)] -> [(cost, state)] -> Bool
forall a b. Ord a => [(a, b)] -> [(a, b)] -> Bool
leastCostly (cost, state) -> m [(cost, state)]
nextM'
    (state -> m Bool
foundM (state -> m Bool)
-> ((cost, state) -> state) -> (cost, state) -> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (cost, state) -> state
forall a b. (a, b) -> b
snd) (cost
0, state
initial)
  where
    nextM' :: (cost, state) -> m [(cost, state)]
nextM' (cost
old_cost, state
old_st) = do
      [state]
new_states <- f state -> [state]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList (f state -> [state]) -> m (f state) -> m [state]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> state -> m (f state)
nextM state
old_st
      [cost]
incr_costs <- [m cost] -> m [cost]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence ([m cost] -> m [cost]) -> [m cost] -> m [cost]
forall a b. (a -> b) -> a -> b
$ state -> state -> m cost
costM state
old_st (state -> m cost) -> [state] -> [m cost]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [state]
new_states
      let new_costs :: [cost]
new_costs = (cost -> cost -> cost
forall a. Num a => a -> a -> a
+ cost
old_cost) (cost -> cost) -> [cost] -> [cost]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [cost]
incr_costs
      [(cost, state)] -> m [(cost, state)]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(cost, state)] -> m [(cost, state)])
-> [(cost, state)] -> m [(cost, state)]
forall a b. (a -> b) -> a -> b
$ [cost] -> [state] -> [(cost, state)]
forall a b. [a] -> [b] -> [(a, b)]
zip [cost]
new_costs [state]
new_states
    unpack :: [(a, a)] -> (a, [a])
unpack [] = (a
0, [])
    unpack [(a, a)]
packed_states = ((a, a) -> a
forall a b. (a, b) -> a
fst ((a, a) -> a) -> ([(a, a)] -> (a, a)) -> [(a, a)] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(a, a)] -> (a, a)
forall a. [a] -> a
last ([(a, a)] -> a) -> [(a, a)] -> a
forall a b. (a -> b) -> a -> b
$ [(a, a)]
packed_states, ((a, a) -> a) -> [(a, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, a) -> a
forall a b. (a, b) -> b
snd [(a, a)]
packed_states)


-- | @aStarM@ is a monadic version of 'aStar': it has support for monadic
-- @next@, @cost@, @remaining@, and @found@ parameters.
aStarM :: (Monad m, Foldable f, Num cost, Ord cost, Ord state)
  => (state -> m (f state))
  -- ^ Function which, when given the current state, produces a list whose
  -- elements are (incremental cost to reach neighboring state,
  -- estimate on remaining cost from said state, said state).
  -> (state -> state -> m cost)
  -- ^ Function to generate list of costs between neighboring states. This is
  -- only called for adjacent states, so it is safe to have this function be
  -- partial for non-neighboring states.
  -> (state -> m cost)
  -- ^ Estimate on remaining cost given a state
  -> (state -> m Bool)
  -- ^ Predicate to determine if solution found. 'aStarM' returns the shortest
  -- path to the first state for which this predicate returns 'True'.
  -> state
  -- ^ Initial state
  -> m (Maybe (cost, [state]))
  -- ^ (Total cost, list of steps) for the first path found which satisfies the
  -- given predicate
aStarM :: (state -> m (f state))
-> (state -> state -> m cost)
-> (state -> m cost)
-> (state -> m Bool)
-> state
-> m (Maybe (cost, [state]))
aStarM state -> m (f state)
nextM state -> state -> m cost
costM state -> m cost
remainingM state -> m Bool
foundM state
initial = do
  cost
remaining_init <- state -> m cost
remainingM state
initial
  ([(cost, (cost, state))] -> (cost, [state]))
-> m (Maybe [(cost, (cost, state))]) -> m (Maybe (cost, [state]))
forall (f1 :: * -> *) (f2 :: * -> *) a b.
(Functor f1, Functor f2) =>
(a -> b) -> f1 (f2 a) -> f1 (f2 b)
fmap2 [(cost, (cost, state))] -> (cost, [state])
forall a a a. Num a => [(a, (a, a))] -> (a, [a])
unpack (m (Maybe [(cost, (cost, state))]) -> m (Maybe (cost, [state])))
-> m (Maybe [(cost, (cost, state))]) -> m (Maybe (cost, [state]))
forall a b. (a -> b) -> a -> b
$ LIFOHeap cost (cost, state)
-> ((cost, (cost, state)) -> state)
-> ([(cost, (cost, state))] -> [(cost, (cost, state))] -> Bool)
-> ((cost, (cost, state)) -> m [(cost, (cost, state))])
-> ((cost, (cost, state)) -> m Bool)
-> (cost, (cost, state))
-> m (Maybe [(cost, (cost, state))])
forall (m :: * -> *) (f :: * -> *) container stateKey state.
(Monad m, Foldable f, SearchContainer container, Ord stateKey,
 Elem container ~ state) =>
container
-> (state -> stateKey)
-> ([state] -> [state] -> Bool)
-> (state -> m (f state))
-> (state -> m Bool)
-> state
-> m (Maybe [state])
generalizedSearchM LIFOHeap cost (cost, state)
forall k a. LIFOHeap k a
emptyLIFOHeap (cost, (cost, state)) -> state
forall a a c. (a, (a, c)) -> c
snd2 [(cost, (cost, state))] -> [(cost, (cost, state))] -> Bool
forall a b. Ord a => [(a, b)] -> [(a, b)] -> Bool
leastCostly (cost, (cost, state)) -> m [(cost, (cost, state))]
nextM'
    (state -> m Bool
foundM (state -> m Bool)
-> ((cost, (cost, state)) -> state)
-> (cost, (cost, state))
-> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (cost, (cost, state)) -> state
forall a a c. (a, (a, c)) -> c
snd2) (cost
remaining_init, (cost
0, state
initial))
  where
    nextM' :: (cost, (cost, state)) -> m [(cost, (cost, state))]
nextM' (cost
_, (cost
old_cost, state
old_st)) = do
      [state]
new_states <- f state -> [state]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList (f state -> [state]) -> m (f state) -> m [state]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> state -> m (f state)
nextM state
old_st
      [m (cost, (cost, state))] -> m [(cost, (cost, state))]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence ([m (cost, (cost, state))] -> m [(cost, (cost, state))])
-> [m (cost, (cost, state))] -> m [(cost, (cost, state))]
forall a b. (a -> b) -> a -> b
$ state -> m (cost, (cost, state))
update_stateM (state -> m (cost, (cost, state)))
-> [state] -> [m (cost, (cost, state))]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [state]
new_states
      where
        update_stateM :: state -> m (cost, (cost, state))
update_stateM state
new_st = do
          cost
remaining <- state -> m cost
remainingM state
new_st
          cost
cost <- state -> state -> m cost
costM state
old_st state
new_st
          let new_cost :: cost
new_cost = cost
old_cost cost -> cost -> cost
forall a. Num a => a -> a -> a
+ cost
cost
              new_est :: cost
new_est = cost
new_cost cost -> cost -> cost
forall a. Num a => a -> a -> a
+ cost
remaining
          (cost, (cost, state)) -> m (cost, (cost, state))
forall (m :: * -> *) a. Monad m => a -> m a
return (cost
new_est, (cost
new_cost, state
new_st))
    unpack :: [(a, (a, a))] -> (a, [a])
unpack [] = (a
0, [])
    unpack [(a, (a, a))]
packed_states =
      ((a, a) -> a
forall a b. (a, b) -> a
fst ((a, a) -> a) -> ([(a, (a, a))] -> (a, a)) -> [(a, (a, a))] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, (a, a)) -> (a, a)
forall a b. (a, b) -> b
snd ((a, (a, a)) -> (a, a))
-> ([(a, (a, a))] -> (a, (a, a))) -> [(a, (a, a))] -> (a, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(a, (a, a))] -> (a, (a, a))
forall a. [a] -> a
last ([(a, (a, a))] -> a) -> [(a, (a, a))] -> a
forall a b. (a -> b) -> a -> b
$ [(a, (a, a))]
packed_states, ((a, (a, a)) -> a) -> [(a, (a, a))] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, (a, a)) -> a
forall a a c. (a, (a, c)) -> c
snd2 [(a, (a, a))]
packed_states)
    snd2 :: (a, (a, c)) -> c
snd2 = (a, c) -> c
forall a b. (a, b) -> b
snd ((a, c) -> c) -> ((a, (a, c)) -> (a, c)) -> (a, (a, c)) -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, (a, c)) -> (a, c)
forall a b. (a, b) -> b
snd


-- | @incrementalCosts cost_fn states@ gives a list of the incremental costs
-- going from state to state along the path given in @states@, using the cost
-- function given by @cost_fn@. Note that the paths returned by the searches
-- in this module do not include the initial state, so if you want the
-- incremental costs along a @path@ returned by one of these searches, you
-- want to use @incrementalCosts cost_fn (initial : path)@.
--
-- === Example: Getting incremental costs from dijkstra
--
-- >>> import Data.Maybe (fromJust)
--
-- >>> :{
-- cyclicWeightedGraph :: Map.Map Char [(Char, Int)]
-- cyclicWeightedGraph = Map.fromList [
--   ('a', [('b', 1), ('c', 2)]),
--   ('b', [('a', 1), ('c', 2), ('d', 5)]),
--   ('c', [('a', 1), ('d', 2)]),
--   ('d', [])
--   ]
-- start = (0, 0)
-- end = (0, 2)
-- cost a b = fromJust . lookup b $ cyclicWeightedGraph Map.! a
-- :}
--
-- >>> incrementalCosts cost ['a', 'b', 'd']
-- [1,5]
incrementalCosts ::
  (state -> state -> cost)
  -- ^ Function to generate list of costs between neighboring states. This is
  -- only called for adjacent states in the `states` list, so it is safe to have
  -- this function be partial for non-neighboring states.
  -> [state]
  -- ^ A path, given as a list of adjacent states, along which to find the
  -- incremental costs
  -> [cost]
  -- ^ List of incremental costs along given path
incrementalCosts :: (state -> state -> cost) -> [state] -> [cost]
incrementalCosts state -> state -> cost
cost_fn [state]
states = (state -> state -> cost) -> [state] -> [state] -> [cost]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith state -> state -> cost
cost_fn [state]
states ([state] -> [state]
forall a. [a] -> [a]
tail [state]
states)

-- | @incrementalCostsM@ is a monadic version of 'incrementalCosts': it has
-- support for a monadic @const_fn@ parameter.
incrementalCostsM ::
  (Monad m) =>
  (state -> state -> m cost)
  -- ^ Function to generate list of costs between neighboring states. This is
  -- only called for adjacent states in the `states` list, so it is safe to have
  -- this function be partial for non-neighboring states.
  -> [state]
  -- ^ A path, given as a list of adjacent states, along which to find the
  -- incremental costs
  -> m [cost]
  -- ^ List of incremental costs along given path
incrementalCostsM :: (state -> state -> m cost) -> [state] -> m [cost]
incrementalCostsM state -> state -> m cost
costM [state]
states = (state -> state -> m cost) -> [state] -> [state] -> m [cost]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM state -> state -> m cost
costM [state]
states ([state] -> [state]
forall a. [a] -> [a]
tail [state]
states)


-- | @next \`pruning\` predicate@ streams the elements generate by @next@ into a
-- list, removing elements which satisfy @predicate@. This is useful for the
-- common case when you want to logically separate your search's `next` function
-- from some way of determining when you've reached a dead end.
--
-- === Example: Pruning a Set
--
-- >>> import qualified Data.Set as Set
--
-- >>> ((\x -> Set.fromList [0..x]) `pruning` even) 10
-- [1,3,5,7,9]
--
-- === Example: depth-first search, avoiding certain nodes
--
-- >>> import qualified Data.Map as Map
--
-- >>> :{
-- graph = Map.fromList [
--   ('a', ['b', 'c', 'd']),
--   ('b', [undefined]),
--   ('c', ['e']),
--   ('d', [undefined]),
--   ('e', [])
--   ]
-- :}
--
-- >>> dfs ((graph Map.!) `pruning` (`elem` "bd")) (== 'e') 'a'
-- Just "ce"
pruning ::
  (Foldable f)
  => (a -> f a)
  -- ^ Function to generate next states
  -> (a -> Bool)
  -- ^ Predicate to prune on
  -> (a -> [a])
  -- ^ Version of @next@ which excludes elements satisfying @predicate@
a -> f a
next pruning :: (a -> f a) -> (a -> Bool) -> a -> [a]
`pruning` a -> Bool
predicate =
  ((a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (a -> Bool) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
predicate) ([a] -> [a]) -> (f a -> [a]) -> f a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList) (f a -> [a]) -> (a -> f a) -> a -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f a
next


-- | @pruningM@ is a monadic version of 'pruning': it has support for monadic
-- @next@ and @predicate@ parameters
pruningM ::
  (Monad m, Foldable f)
  => (a -> m (f a))
  -- ^ Function to generate next states
  -> (a -> m Bool)
  -- ^ Predicate to prune on
  -> (a -> m [a])
  -- ^ Version of @next@ which excludes elements satisfying @predicate@
pruningM :: (a -> m (f a)) -> (a -> m Bool) -> a -> m [a]
pruningM a -> m (f a)
nextM a -> m Bool
predicateM a
a = do
  f a
next_states <- a -> m (f a)
nextM a
a
  (a -> m Bool) -> [a] -> m [a]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM ((Bool -> Bool) -> m Bool -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Bool -> Bool
not(m Bool -> m Bool) -> (a -> m Bool) -> a -> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> m Bool
predicateM) ([a] -> m [a]) -> [a] -> m [a]
forall a b. (a -> b) -> a -> b
$ f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList f a
next_states


-- | A @SearchState@ represents the state of a generalized search at a given
-- point in an algorithms execution. The advantage of this abstraction is that
-- it can be used for things like bidirectional searches, where you want to
-- stop and start a search part-way through.
data SearchState container stateKey state = SearchState {
  SearchState container stateKey state -> state
current :: state,
  SearchState container stateKey state -> container
queue :: container,
  SearchState container stateKey state -> Set stateKey
visited :: Set.Set stateKey,
  SearchState container stateKey state -> Map stateKey [state]
paths :: Map.Map stateKey [state]
  }

-- | Workhorse simple search algorithm, generalized over search container
-- and path-choosing function. The idea here is that many search algorithms are
-- at their core the same, with these details substituted. By writing these
-- searches in terms of this function, we reduce the chances of errors sneaking
-- into each separate implementation.
generalizedSearch ::
  (Foldable f, SearchContainer container, Ord stateKey, Elem container ~ state)
  => container
  -- ^ Empty @SearchContainer@
  -> (state -> stateKey)
  -- ^ Function to turn a @state@ into a key by which states will be compared
  -- when determining whether a state has be enqueued and / or visited
  -> ([state] -> [state] -> Bool)
  -- ^ Function @better old new@, which when given a choice between an @old@ and
  -- a @new@ path to a state, returns True when @new@ is a "better" path than
  -- old and should thus be inserted
  -> (state -> f state)
  -- ^ Function to generate "next" states given a current state
  -> (state -> Bool)
  -- ^ Predicate to determine if solution found. @generalizedSearch@ returns a
  -- path to the first state for which this predicate returns 'True'.
  -> state
  -- ^ Initial state
  -> Maybe [state]
  -- ^ First path found to a state matching the predicate, or 'Nothing' if no
  -- such path exists.
generalizedSearch :: container
-> (state -> stateKey)
-> ([state] -> [state] -> Bool)
-> (state -> f state)
-> (state -> Bool)
-> state
-> Maybe [state]
generalizedSearch container
empty state -> stateKey
mk_key [state] -> [state] -> Bool
better state -> f state
next state -> Bool
found state
initial = Identity (Maybe [state]) -> Maybe [state]
forall a. Identity a -> a
runIdentity (Identity (Maybe [state]) -> Maybe [state])
-> Identity (Maybe [state]) -> Maybe [state]
forall a b. (a -> b) -> a -> b
$
  container
-> (state -> stateKey)
-> ([state] -> [state] -> Bool)
-> (state -> Identity (f state))
-> (state -> Identity Bool)
-> state
-> Identity (Maybe [state])
forall (m :: * -> *) (f :: * -> *) container stateKey state.
(Monad m, Foldable f, SearchContainer container, Ord stateKey,
 Elem container ~ state) =>
container
-> (state -> stateKey)
-> ([state] -> [state] -> Bool)
-> (state -> m (f state))
-> (state -> m Bool)
-> state
-> m (Maybe [state])
generalizedSearchM container
empty state -> stateKey
mk_key [state] -> [state] -> Bool
better (f state -> Identity (f state)
forall a. a -> Identity a
Identity (f state -> Identity (f state))
-> (state -> f state) -> state -> Identity (f state)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. state -> f state
next) (Bool -> Identity Bool
forall a. a -> Identity a
Identity (Bool -> Identity Bool)
-> (state -> Bool) -> state -> Identity Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. state -> Bool
found) state
initial

-- | @nextSearchState@ moves from one @searchState@ to the next in the
-- generalized search algorithm
nextSearchStateM ::
  (Monad m, Foldable f, SearchContainer container, Ord stateKey,
   Elem container ~ state)
  => ([state] -> [state] -> Bool)
  -> (state -> stateKey)
  -> (state -> m (f state))
  -> SearchState container stateKey state
  -> m (Maybe (SearchState container stateKey state))
nextSearchStateM :: ([state] -> [state] -> Bool)
-> (state -> stateKey)
-> (state -> m (f state))
-> SearchState container stateKey state
-> m (Maybe (SearchState container stateKey state))
nextSearchStateM [state] -> [state] -> Bool
better state -> stateKey
mk_key state -> m (f state)
nextM SearchState container stateKey state
old = do
  (container
new_queue, Map stateKey [state]
new_paths) <- m (container, Map stateKey [state])
new_queue_paths_M
  let new_state_May :: Maybe (SearchState container stateKey state)
new_state_May = Map stateKey [state]
-> (state, container) -> SearchState container stateKey state
mk_search_state Map stateKey [state]
new_paths ((state, container) -> SearchState container stateKey state)
-> Maybe (state, container)
-> Maybe (SearchState container stateKey state)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> container -> Maybe (Elem container, container)
forall container.
SearchContainer container =>
container -> Maybe (Elem container, container)
pop container
new_queue
  case Maybe (SearchState container stateKey state)
new_state_May of
    Just SearchState container stateKey state
new_state ->
      if state -> stateKey
mk_key (SearchState container stateKey state -> state
forall container stateKey state.
SearchState container stateKey state -> state
current SearchState container stateKey state
new_state) stateKey -> Set stateKey -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` SearchState container stateKey state -> Set stateKey
forall container stateKey state.
SearchState container stateKey state -> Set stateKey
visited SearchState container stateKey state
old
      then ([state] -> [state] -> Bool)
-> (state -> stateKey)
-> (state -> m (f state))
-> SearchState container stateKey state
-> m (Maybe (SearchState container stateKey state))
forall (m :: * -> *) (f :: * -> *) container stateKey state.
(Monad m, Foldable f, SearchContainer container, Ord stateKey,
 Elem container ~ state) =>
([state] -> [state] -> Bool)
-> (state -> stateKey)
-> (state -> m (f state))
-> SearchState container stateKey state
-> m (Maybe (SearchState container stateKey state))
nextSearchStateM [state] -> [state] -> Bool
better state -> stateKey
mk_key state -> m (f state)
nextM SearchState container stateKey state
new_state
      else Maybe (SearchState container stateKey state)
-> m (Maybe (SearchState container stateKey state))
forall (m :: * -> *) a. Monad m => a -> m a
return (SearchState container stateKey state
-> Maybe (SearchState container stateKey state)
forall a. a -> Maybe a
Just SearchState container stateKey state
new_state)
    Maybe (SearchState container stateKey state)
Nothing -> Maybe (SearchState container stateKey state)
-> m (Maybe (SearchState container stateKey state))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (SearchState container stateKey state)
forall a. Maybe a
Nothing
  where
    mk_search_state :: Map stateKey [state]
-> (state, container) -> SearchState container stateKey state
mk_search_state Map stateKey [state]
new_paths (state
new_current, container
remaining_queue) = SearchState :: forall container stateKey state.
state
-> container
-> Set stateKey
-> Map stateKey [state]
-> SearchState container stateKey state
SearchState {
      current :: state
current = state
new_current,
      queue :: container
queue = container
remaining_queue,
      visited :: Set stateKey
visited = stateKey -> Set stateKey -> Set stateKey
forall a. Ord a => a -> Set a -> Set a
Set.insert (state -> stateKey
mk_key state
new_current) (SearchState container stateKey state -> Set stateKey
forall container stateKey state.
SearchState container stateKey state -> Set stateKey
visited SearchState container stateKey state
old),
      paths :: Map stateKey [state]
paths = Map stateKey [state]
new_paths
      }
    new_queue_paths_M :: m (container, Map stateKey [state])
new_queue_paths_M =
      ((container, Map stateKey [state])
 -> state -> (container, Map stateKey [state]))
-> (container, Map stateKey [state])
-> f state
-> (container, Map stateKey [state])
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (container, Map stateKey [state])
-> state -> (container, Map stateKey [state])
update_queue_paths (SearchState container stateKey state -> container
forall container stateKey state.
SearchState container stateKey state -> container
queue SearchState container stateKey state
old, SearchState container stateKey state -> Map stateKey [state]
forall container stateKey state.
SearchState container stateKey state -> Map stateKey [state]
paths SearchState container stateKey state
old)
        (f state -> (container, Map stateKey [state]))
-> m (f state) -> m (container, Map stateKey [state])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> state -> m (f state)
nextM (SearchState container stateKey state -> state
forall container stateKey state.
SearchState container stateKey state -> state
current SearchState container stateKey state
old)
    update_queue_paths :: (container, Map stateKey [state])
-> state -> (container, Map stateKey [state])
update_queue_paths (container
old_queue, Map stateKey [state]
old_paths) state
st =
      if state -> stateKey
mk_key state
st stateKey -> Set stateKey -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` SearchState container stateKey state -> Set stateKey
forall container stateKey state.
SearchState container stateKey state -> Set stateKey
visited SearchState container stateKey state
old
      then (container
old_queue, Map stateKey [state]
old_paths)
      else
        case stateKey -> Map stateKey [state] -> Maybe [state]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (state -> stateKey
mk_key state
st) Map stateKey [state]
old_paths of
          Just [state]
old_path ->
            if [state] -> [state] -> Bool
better [state]
old_path (state
st state -> [state] -> [state]
forall a. a -> [a] -> [a]
: [state]
steps_so_far)
            then (container
q', Map stateKey [state]
ps')
            else (container
old_queue, Map stateKey [state]
old_paths)
          Maybe [state]
Nothing -> (container
q', Map stateKey [state]
ps')
        where
          steps_so_far :: [state]
steps_so_far = SearchState container stateKey state -> Map stateKey [state]
forall container stateKey state.
SearchState container stateKey state -> Map stateKey [state]
paths SearchState container stateKey state
old Map stateKey [state] -> stateKey -> [state]
forall k a. Ord k => Map k a -> k -> a
Map.! state -> stateKey
mk_key (SearchState container stateKey state -> state
forall container stateKey state.
SearchState container stateKey state -> state
current SearchState container stateKey state
old)
          q' :: container
q' = container -> Elem container -> container
forall container.
SearchContainer container =>
container -> Elem container -> container
push container
old_queue state
Elem container
st
          ps' :: Map stateKey [state]
ps' = stateKey -> [state] -> Map stateKey [state] -> Map stateKey [state]
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (state -> stateKey
mk_key state
st) (state
st state -> [state] -> [state]
forall a. a -> [a] -> [a]
: [state]
steps_so_far) Map stateKey [state]
old_paths


-- | @generalizedSearchM@ is a monadic version of generalizedSearch
generalizedSearchM ::
  (Monad m, Foldable f, SearchContainer container, Ord stateKey,
   Elem container ~ state)
  => container
  -- ^ Empty @SearchContainer@
  -> (state -> stateKey)
  -- ^ Function to turn a @state@ into a key by which states will be compared
  -- when determining whether a state has be enqueued and / or visited
  -> ([state] -> [state] -> Bool)
  -- ^ Function @better old new@, which when given a choice between an @old@ and
  -- a @new@ path to a state, returns True when @new@ is a "better" path than
  -- old and should thus be inserted
  -> (state -> m (f state))
  -- ^ Function to generate "next" states given a current state
  -> (state -> m Bool)
  -- ^ Predicate to determine if solution found. @generalizedSearch@ returns a
  -- path to the first state for which this predicate returns 'True'.
  -> state
  -- ^ Initial state
  -> m (Maybe [state])
  -- ^ First path found to a state matching the predicate, or 'Nothing' if no
  -- such path exists.
generalizedSearchM :: container
-> (state -> stateKey)
-> ([state] -> [state] -> Bool)
-> (state -> m (f state))
-> (state -> m Bool)
-> state
-> m (Maybe [state])
generalizedSearchM container
empty state -> stateKey
mk_key [state] -> [state] -> Bool
better state -> m (f state)
nextM state -> m Bool
foundM state
initial = do
  let initial_state :: SearchState container stateKey state
initial_state =
        state
-> container
-> Set stateKey
-> Map stateKey [state]
-> SearchState container stateKey state
forall container stateKey state.
state
-> container
-> Set stateKey
-> Map stateKey [state]
-> SearchState container stateKey state
SearchState state
initial container
empty (stateKey -> Set stateKey
forall a. a -> Set a
Set.singleton (stateKey -> Set stateKey) -> stateKey -> Set stateKey
forall a b. (a -> b) -> a -> b
$ state -> stateKey
mk_key state
initial)
        (stateKey -> [state] -> Map stateKey [state]
forall k a. k -> a -> Map k a
Map.singleton (state -> stateKey
mk_key state
initial) [])
  Maybe (SearchState container stateKey state)
end_May <- (SearchState container stateKey state
 -> m (Maybe (SearchState container stateKey state)))
-> (SearchState container stateKey state -> m Bool)
-> SearchState container stateKey state
-> m (Maybe (SearchState container stateKey state))
forall (m :: * -> *) a.
Monad m =>
(a -> m (Maybe a)) -> (a -> m Bool) -> a -> m (Maybe a)
findIterateM (([state] -> [state] -> Bool)
-> (state -> stateKey)
-> (state -> m (f state))
-> SearchState container stateKey state
-> m (Maybe (SearchState container stateKey state))
forall (m :: * -> *) (f :: * -> *) container stateKey state.
(Monad m, Foldable f, SearchContainer container, Ord stateKey,
 Elem container ~ state) =>
([state] -> [state] -> Bool)
-> (state -> stateKey)
-> (state -> m (f state))
-> SearchState container stateKey state
-> m (Maybe (SearchState container stateKey state))
nextSearchStateM [state] -> [state] -> Bool
better state -> stateKey
mk_key state -> m (f state)
nextM)
    (state -> m Bool
foundM (state -> m Bool)
-> (SearchState container stateKey state -> state)
-> SearchState container stateKey state
-> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SearchState container stateKey state -> state
forall container stateKey state.
SearchState container stateKey state -> state
current) SearchState container stateKey state
initial_state
  Maybe [state] -> m (Maybe [state])
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe [state] -> m (Maybe [state]))
-> Maybe [state] -> m (Maybe [state])
forall a b. (a -> b) -> a -> b
$ (SearchState container stateKey state -> [state])
-> Maybe (SearchState container stateKey state) -> Maybe [state]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([state] -> [state]
forall a. [a] -> [a]
reverse ([state] -> [state])
-> (SearchState container stateKey state -> [state])
-> SearchState container stateKey state
-> [state]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SearchState container stateKey state -> [state]
get_steps) Maybe (SearchState container stateKey state)
end_May
  where
    get_steps :: SearchState container stateKey state -> [state]
get_steps SearchState container stateKey state
search_st = SearchState container stateKey state -> Map stateKey [state]
forall container stateKey state.
SearchState container stateKey state -> Map stateKey [state]
paths SearchState container stateKey state
search_st Map stateKey [state] -> stateKey -> [state]
forall k a. Ord k => Map k a -> k -> a
Map.! state -> stateKey
mk_key (SearchState container stateKey state -> state
forall container stateKey state.
SearchState container stateKey state -> state
current SearchState container stateKey state
search_st)


newtype LIFOHeap k a = LIFOHeap (Map.Map k [a])


emptyLIFOHeap :: LIFOHeap k a
emptyLIFOHeap :: LIFOHeap k a
emptyLIFOHeap = Map k [a] -> LIFOHeap k a
forall k a. Map k [a] -> LIFOHeap k a
LIFOHeap Map k [a]
forall k a. Map k a
Map.empty


-- | The @SearchContainer@ class abstracts the idea of a container to be used in
-- @generalizedSearch@
class SearchContainer container where
  type Elem container
  pop :: container -> Maybe (Elem container, container)
  push :: container -> Elem container -> container

instance SearchContainer (Seq.Seq a) where
  type Elem (Seq.Seq a) = a
  pop :: Seq a -> Maybe (Elem (Seq a), Seq a)
pop Seq a
s =
    case Seq a -> ViewL a
forall a. Seq a -> ViewL a
Seq.viewl Seq a
s of
      ViewL a
Seq.EmptyL -> Maybe (Elem (Seq a), Seq a)
forall a. Maybe a
Nothing
      (a
x Seq.:< Seq a
xs) -> (a, Seq a) -> Maybe (a, Seq a)
forall a. a -> Maybe a
Just (a
x, Seq a
xs)
  push :: Seq a -> Elem (Seq a) -> Seq a
push Seq a
s Elem (Seq a)
a = Seq a
s Seq a -> a -> Seq a
forall a. Seq a -> a -> Seq a
Seq.|> a
Elem (Seq a)
a

instance SearchContainer [a] where
  type Elem [a] = a
  pop :: [a] -> Maybe (Elem [a], [a])
pop [a]
list =
    case [a]
list of
      [] -> Maybe (Elem [a], [a])
forall a. Maybe a
Nothing
      (a
x : [a]
xs) -> (a, [a]) -> Maybe (a, [a])
forall a. a -> Maybe a
Just (a
x, [a]
xs)
  push :: [a] -> Elem [a] -> [a]
push [a]
list Elem [a]
a = a
Elem [a]
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
list

instance Ord k => SearchContainer (LIFOHeap k a) where
  type Elem (LIFOHeap k a) = (k, a)
  pop :: LIFOHeap k a -> Maybe (Elem (LIFOHeap k a), LIFOHeap k a)
pop (LIFOHeap Map k [a]
inner)
    | Map k [a] -> Bool
forall k a. Map k a -> Bool
Map.null Map k [a]
inner = Maybe (Elem (LIFOHeap k a), LIFOHeap k a)
forall a. Maybe a
Nothing
    | Bool
otherwise = case Map k [a] -> (k, [a])
forall k a. Map k a -> (k, a)
Map.findMin Map k [a]
inner of
      (k
k, [a
a]) -> ((k, a), LIFOHeap k a) -> Maybe ((k, a), LIFOHeap k a)
forall a. a -> Maybe a
Just ((k
k, a
a), Map k [a] -> LIFOHeap k a
forall k a. Map k [a] -> LIFOHeap k a
LIFOHeap (Map k [a] -> LIFOHeap k a) -> Map k [a] -> LIFOHeap k a
forall a b. (a -> b) -> a -> b
$ Map k [a] -> Map k [a]
forall k a. Map k a -> Map k a
Map.deleteMin Map k [a]
inner)
      (k
k, a
a : [a]
_) -> ((k, a), LIFOHeap k a) -> Maybe ((k, a), LIFOHeap k a)
forall a. a -> Maybe a
Just ((k
k, a
a), Map k [a] -> LIFOHeap k a
forall k a. Map k [a] -> LIFOHeap k a
LIFOHeap (Map k [a] -> LIFOHeap k a) -> Map k [a] -> LIFOHeap k a
forall a b. (a -> b) -> a -> b
$ ([a] -> Maybe [a]) -> Map k [a] -> Map k [a]
forall a k. (a -> Maybe a) -> Map k a -> Map k a
Map.updateMin ([a] -> Maybe [a]
forall a. a -> Maybe a
Just ([a] -> Maybe [a]) -> ([a] -> [a]) -> [a] -> Maybe [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. [a] -> [a]
tail) Map k [a]
inner)
      (k
_, []) -> LIFOHeap k a -> Maybe (Elem (LIFOHeap k a), LIFOHeap k a)
forall container.
SearchContainer container =>
container -> Maybe (Elem container, container)
pop (Map k [a] -> LIFOHeap k a
forall k a. Map k [a] -> LIFOHeap k a
LIFOHeap (Map k [a] -> LIFOHeap k a) -> Map k [a] -> LIFOHeap k a
forall a b. (a -> b) -> a -> b
$ Map k [a] -> Map k [a]
forall k a. Map k a -> Map k a
Map.deleteMin Map k [a]
inner)
                 -- Logically, this should never happen
  push :: LIFOHeap k a -> Elem (LIFOHeap k a) -> LIFOHeap k a
push (LIFOHeap Map k [a]
inner) (k, a) = Map k [a] -> LIFOHeap k a
forall k a. Map k [a] -> LIFOHeap k a
LIFOHeap (Map k [a] -> LIFOHeap k a) -> Map k [a] -> LIFOHeap k a
forall a b. (a -> b) -> a -> b
$ ([a] -> [a] -> [a]) -> k -> [a] -> Map k [a] -> Map k [a]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
(++) k
k [a
a] Map k [a]
inner


-- | @findIterateM@ is a monadic version of @findIterate@
findIterateM :: Monad m => (a -> m (Maybe a)) -> (a -> m Bool) -> a -> m (Maybe a)
findIterateM :: (a -> m (Maybe a)) -> (a -> m Bool) -> a -> m (Maybe a)
findIterateM a -> m (Maybe a)
nextM a -> m Bool
foundM a
initial = do
  Bool
found <- a -> m Bool
foundM a
initial
  if Bool
found
  then Maybe a -> m (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> m (Maybe a)) -> Maybe a -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just a
initial
  else a -> m (Maybe a)
nextM a
initial m (Maybe a) -> (Maybe a -> m (Maybe a)) -> m (Maybe a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= m (Maybe a) -> (a -> m (Maybe a)) -> Maybe a -> m (Maybe a)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Maybe a -> m (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing) ((a -> m (Maybe a)) -> (a -> m Bool) -> a -> m (Maybe a)
forall (m :: * -> *) a.
Monad m =>
(a -> m (Maybe a)) -> (a -> m Bool) -> a -> m (Maybe a)
findIterateM a -> m (Maybe a)
nextM a -> m Bool
foundM)


-- | @leastCostly paths_a paths_b@ is a utility function to be used with
-- 'dijkstra'-like functions. It returns True when the cost of @paths_a@
-- is less than the cost of @paths_b@, where the total costs are the first
-- elements in each tuple in each path
leastCostly :: Ord a => [(a, b)] -> [(a, b)] -> Bool
leastCostly :: [(a, b)] -> [(a, b)] -> Bool
leastCostly ((a
cost_a, b
_):[(a, b)]
_) ((a
cost_b, b
_):[(a, b)]
_) = a
cost_b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
cost_a
-- logically this never happens, because if you have a
-- zero-length path a point, you already visited it
-- and thus do not consider other paths to it
leastCostly [] [(a, b)]
_ = Bool
False
-- logically this never happens, because you cannot find
-- a new zero-length path to a point
leastCostly [(a, b)]
_ [] = Bool
True


-- | This is just a convenience function which @fmap@s two deep
fmap2 :: (Functor f1, Functor f2) => (a -> b) -> f1 (f2 a) -> f1 (f2 b)
fmap2 :: (a -> b) -> f1 (f2 a) -> f1 (f2 b)
fmap2 = (f2 a -> f2 b) -> f1 (f2 a) -> f1 (f2 b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((f2 a -> f2 b) -> f1 (f2 a) -> f1 (f2 b))
-> ((a -> b) -> f2 a -> f2 b) -> (a -> b) -> f1 (f2 a) -> f1 (f2 b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> f2 a -> f2 b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap