module AOC.Common.Search (
    aStar
  , bfs
  , binarySearch
  , exponentialSearch
  , binaryMinSearch
  , exponentialMinSearch
  , binaryFindMin
  , exponentialFindMin
  ) where

import           Data.Bifunctor
import           Data.Map       (Map)
import           Data.OrdPSQ    (OrdPSQ)
import           Data.Sequence  (Seq(..))
import           Data.Set       (Set)
import qualified Data.Map       as M
import qualified Data.OrdPSQ    as Q
import qualified Data.Sequence  as Seq
import qualified Data.Set       as S

data AStarState n p = AS { forall n p. AStarState n p -> Map n (Maybe n)
_asClosed  :: !(Map n (Maybe n))         -- ^ map of item to "parent"
                         , forall n p. AStarState n p -> OrdPSQ n p (p, Maybe n)
_asOpen    :: !(OrdPSQ n p (p, Maybe n))    -- ^ map of item to "parent", and cost-so-far
                         }

-- | A* Search
aStar
    :: forall n p. (Ord n, Ord p, Num p)
    => (n -> p)         -- ^ heuristic
    -> (n -> Map n p)   -- ^ neighborhood
    -> n                -- ^ start
    -> (n -> Bool)      -- ^ target
    -> Maybe (p, [n])   -- ^ the shortest path, if it exists, and its cost
aStar :: forall n p.
(Ord n, Ord p, Num p) =>
(n -> p) -> (n -> Map n p) -> n -> (n -> Bool) -> Maybe (p, [n])
aStar n -> p
h n -> Map n p
ex n
x0 n -> Bool
dest = ((n, Map n (Maybe n)) -> [n])
-> (p, (n, Map n (Maybe n))) -> (p, [n])
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (n, Map n (Maybe n)) -> [n]
reconstruct ((p, (n, Map n (Maybe n))) -> (p, [n]))
-> Maybe (p, (n, Map n (Maybe n))) -> Maybe (p, [n])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> AStarState n p -> Maybe (p, (n, Map n (Maybe n)))
go (n -> p -> Maybe n -> AStarState n p -> AStarState n p
addBack n
x0 p
0 Maybe n
forall a. Maybe a
Nothing (Map n (Maybe n) -> OrdPSQ n p (p, Maybe n) -> AStarState n p
forall n p.
Map n (Maybe n) -> OrdPSQ n p (p, Maybe n) -> AStarState n p
AS Map n (Maybe n)
forall k a. Map k a
M.empty OrdPSQ n p (p, Maybe n)
forall k p v. OrdPSQ k p v
Q.empty))
  where
    reconstruct :: (n, Map n (Maybe n)) -> [n]
    reconstruct :: (n, Map n (Maybe n)) -> [n]
reconstruct (n
goal, Map n (Maybe n)
mp) = [n] -> [n]
forall a. [a] -> [a]
reverse ([n] -> [n]) -> [n] -> [n]
forall a b. (a -> b) -> a -> b
$ n -> [n]
goreco n
goal
      where
        goreco :: n -> [n]
goreco n
n = n
n n -> [n] -> [n]
forall a. a -> [a] -> [a]
: [n] -> (n -> [n]) -> Maybe n -> [n]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] n -> [n]
goreco (Map n (Maybe n)
mp Map n (Maybe n) -> n -> Maybe n
forall k a. Ord k => Map k a -> k -> a
M.! n
n)
    go :: AStarState n p -> Maybe (p, (n, Map n (Maybe n)))
    go :: AStarState n p -> Maybe (p, (n, Map n (Maybe n)))
go as0 :: AStarState n p
as0@AS{Map n (Maybe n)
OrdPSQ n p (p, Maybe n)
_asOpen :: OrdPSQ n p (p, Maybe n)
_asClosed :: Map n (Maybe n)
_asOpen :: forall n p. AStarState n p -> OrdPSQ n p (p, Maybe n)
_asClosed :: forall n p. AStarState n p -> Map n (Maybe n)
..} = OrdPSQ n p (p, Maybe n)
-> Maybe (n, p, (p, Maybe n), OrdPSQ n p (p, Maybe n))
forall k p v.
(Ord k, Ord p) =>
OrdPSQ k p v -> Maybe (k, p, v, OrdPSQ k p v)
Q.minView OrdPSQ n p (p, Maybe n)
_asOpen Maybe (n, p, (p, Maybe n), OrdPSQ n p (p, Maybe n))
-> ((n, p, (p, Maybe n), OrdPSQ n p (p, Maybe n))
    -> Maybe (p, (n, Map n (Maybe n))))
-> Maybe (p, (n, Map n (Maybe n)))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(n
n,p
p,(p
g,Maybe n
up),OrdPSQ n p (p, Maybe n)
queue') ->
      let closed' :: Map n (Maybe n)
closed' = n -> Maybe n -> Map n (Maybe n) -> Map n (Maybe n)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert n
n Maybe n
up Map n (Maybe n)
_asClosed
      in  if n -> Bool
dest n
n
            then (p, (n, Map n (Maybe n))) -> Maybe (p, (n, Map n (Maybe n)))
forall a. a -> Maybe a
Just (p
p, (n
n, Map n (Maybe n)
closed'))
            else AStarState n p -> Maybe (p, (n, Map n (Maybe n)))
go (AStarState n p -> Maybe (p, (n, Map n (Maybe n))))
-> (Map n p -> AStarState n p)
-> Map n p
-> Maybe (p, (n, Map n (Maybe n)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AStarState n p -> n -> p -> AStarState n p)
-> AStarState n p -> Map n p -> AStarState n p
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey' (n -> p -> AStarState n p -> n -> p -> AStarState n p
processNeighbor n
n p
g) (AStarState n p
as0 { _asOpen :: OrdPSQ n p (p, Maybe n)
_asOpen = OrdPSQ n p (p, Maybe n)
queue', _asClosed :: Map n (Maybe n)
_asClosed = Map n (Maybe n)
closed'  })
                    (Map n p -> Maybe (p, (n, Map n (Maybe n))))
-> Map n p -> Maybe (p, (n, Map n (Maybe n)))
forall a b. (a -> b) -> a -> b
$ n -> Map n p
ex n
n
    addBack :: n -> p -> Maybe n -> AStarState n p -> AStarState n p
    addBack :: n -> p -> Maybe n -> AStarState n p -> AStarState n p
addBack n
x p
g Maybe n
up AStarState n p
as0 = AStarState n p
as0 { _asOpen :: OrdPSQ n p (p, Maybe n)
_asOpen = n
-> p
-> (p, Maybe n)
-> OrdPSQ n p (p, Maybe n)
-> OrdPSQ n p (p, Maybe n)
forall k p v.
(Ord k, Ord p) =>
k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
insertIfBetter n
x (p
g p -> p -> p
forall a. Num a => a -> a -> a
+ n -> p
h n
x) (p
g, Maybe n
up) (OrdPSQ n p (p, Maybe n) -> OrdPSQ n p (p, Maybe n))
-> (AStarState n p -> OrdPSQ n p (p, Maybe n))
-> AStarState n p
-> OrdPSQ n p (p, Maybe n)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AStarState n p -> OrdPSQ n p (p, Maybe n)
forall n p. AStarState n p -> OrdPSQ n p (p, Maybe n)
_asOpen (AStarState n p -> OrdPSQ n p (p, Maybe n))
-> AStarState n p -> OrdPSQ n p (p, Maybe n)
forall a b. (a -> b) -> a -> b
$ AStarState n p
as0 }
    processNeighbor :: n -> p -> AStarState n p -> n -> p -> AStarState n p
    processNeighbor :: n -> p -> AStarState n p -> n -> p -> AStarState n p
processNeighbor n
curr p
currCost as0 :: AStarState n p
as0@AS{Map n (Maybe n)
OrdPSQ n p (p, Maybe n)
_asOpen :: OrdPSQ n p (p, Maybe n)
_asClosed :: Map n (Maybe n)
_asOpen :: forall n p. AStarState n p -> OrdPSQ n p (p, Maybe n)
_asClosed :: forall n p. AStarState n p -> Map n (Maybe n)
..} n
neighb p
moveCost
      --     | neighb `Q.member` _asOpen || neighb `M.member` _asClosed = as0
      | n
neighb n -> Map n (Maybe n) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map n (Maybe n)
_asClosed = AStarState n p
as0
      | Bool
otherwise = n -> p -> Maybe n -> AStarState n p -> AStarState n p
addBack n
neighb (p
currCost p -> p -> p
forall a. Num a => a -> a -> a
+ p
moveCost) (n -> Maybe n
forall a. a -> Maybe a
Just n
curr) AStarState n p
as0
        -- addBack neighb (currCost + moveCost) (Just curr) as0

insertIfBetter :: (Ord k, Ord p) => k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
insertIfBetter :: forall k p v.
(Ord k, Ord p) =>
k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
insertIfBetter k
k p
p v
x OrdPSQ k p v
q = case k -> OrdPSQ k p v -> Maybe (p, v)
forall k p v. Ord k => k -> OrdPSQ k p v -> Maybe (p, v)
Q.lookup k
k OrdPSQ k p v
q of
    Maybe (p, v)
Nothing       -> k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
forall k p v.
(Ord k, Ord p) =>
k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
Q.insert k
k p
p v
x OrdPSQ k p v
q
    Just (p
p', v
_)
      | p
p p -> p -> Bool
forall a. Ord a => a -> a -> Bool
< p
p'    -> k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
forall k p v.
(Ord k, Ord p) =>
k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
Q.insert k
k p
p v
x OrdPSQ k p v
q
      | Bool
otherwise -> OrdPSQ k p v
q

data BFSState n = BS { forall n. BFSState n -> Map n (Maybe n)
_bsClosed  :: !(Map n (Maybe n))  -- ^ map of item to "parent"
                     , forall n. BFSState n -> Seq n
_bsOpen    :: !(Seq n          ) -- ^ queue
                     }

-- | Breadth-first search, with loop detection
bfs :: forall n. Ord n
    => (n -> Set n)   -- ^ neighborhood
    -> n              -- ^ start
    -> (n -> Bool)    -- ^ target
    -> Maybe [n]      -- ^ the shortest path, if it exists
bfs :: forall n. Ord n => (n -> Set n) -> n -> (n -> Bool) -> Maybe [n]
bfs n -> Set n
ex n
x0 n -> Bool
dest = (n, Map n (Maybe n)) -> [n]
reconstruct ((n, Map n (Maybe n)) -> [n])
-> Maybe (n, Map n (Maybe n)) -> Maybe [n]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BFSState n -> Maybe (n, Map n (Maybe n))
go (n -> Maybe n -> BFSState n -> BFSState n
addBack n
x0 Maybe n
forall a. Maybe a
Nothing (Map n (Maybe n) -> Seq n -> BFSState n
forall n. Map n (Maybe n) -> Seq n -> BFSState n
BS Map n (Maybe n)
forall k a. Map k a
M.empty Seq n
forall a. Seq a
Seq.empty))
  where
    reconstruct :: (n, Map n (Maybe n)) -> [n]
    reconstruct :: (n, Map n (Maybe n)) -> [n]
reconstruct (n
goal, Map n (Maybe n)
mp) = Int -> [n] -> [n]
forall a. Int -> [a] -> [a]
drop Int
1 ([n] -> [n]) -> ([n] -> [n]) -> [n] -> [n]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [n] -> [n]
forall a. [a] -> [a]
reverse ([n] -> [n]) -> [n] -> [n]
forall a b. (a -> b) -> a -> b
$ n -> [n]
goreco n
goal
      where
        goreco :: n -> [n]
goreco n
n = n
n n -> [n] -> [n]
forall a. a -> [a] -> [a]
: [n] -> (n -> [n]) -> Maybe n -> [n]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] n -> [n]
goreco (Map n (Maybe n)
mp Map n (Maybe n) -> n -> Maybe n
forall k a. Ord k => Map k a -> k -> a
M.! n
n)
    go :: BFSState n -> Maybe (n, Map n (Maybe n))
    go :: BFSState n -> Maybe (n, Map n (Maybe n))
go BS{Map n (Maybe n)
Seq n
_bsOpen :: Seq n
_bsClosed :: Map n (Maybe n)
_bsOpen :: forall n. BFSState n -> Seq n
_bsClosed :: forall n. BFSState n -> Map n (Maybe n)
..} = case Seq n
_bsOpen of
      Seq n
Empty    -> Maybe (n, Map n (Maybe n))
forall a. Maybe a
Nothing
      n
n :<| Seq n
ns
        | n -> Bool
dest n
n    -> (n, Map n (Maybe n)) -> Maybe (n, Map n (Maybe n))
forall a. a -> Maybe a
Just (n
n, Map n (Maybe n)
_bsClosed)
        | Bool
otherwise -> BFSState n -> Maybe (n, Map n (Maybe n))
go (BFSState n -> Maybe (n, Map n (Maybe n)))
-> (Set n -> BFSState n) -> Set n -> Maybe (n, Map n (Maybe n))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (BFSState n -> n -> BFSState n)
-> BFSState n -> Set n -> BFSState n
forall a b. (a -> b -> a) -> a -> Set b -> a
S.foldl' (n -> BFSState n -> n -> BFSState n
processNeighbor n
n) (Map n (Maybe n) -> Seq n -> BFSState n
forall n. Map n (Maybe n) -> Seq n -> BFSState n
BS Map n (Maybe n)
_bsClosed Seq n
ns) (Set n -> Maybe (n, Map n (Maybe n)))
-> Set n -> Maybe (n, Map n (Maybe n))
forall a b. (a -> b) -> a -> b
$ n -> Set n
ex n
n
    addBack :: n -> Maybe n -> BFSState n -> BFSState n
    addBack :: n -> Maybe n -> BFSState n -> BFSState n
addBack n
x Maybe n
up BS{Map n (Maybe n)
Seq n
_bsOpen :: Seq n
_bsClosed :: Map n (Maybe n)
_bsOpen :: forall n. BFSState n -> Seq n
_bsClosed :: forall n. BFSState n -> Map n (Maybe n)
..} = BS :: forall n. Map n (Maybe n) -> Seq n -> BFSState n
BS
      { _bsClosed :: Map n (Maybe n)
_bsClosed = n -> Maybe n -> Map n (Maybe n) -> Map n (Maybe n)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert n
x Maybe n
up Map n (Maybe n)
_bsClosed
      , _bsOpen :: Seq n
_bsOpen   = Seq n
_bsOpen Seq n -> n -> Seq n
forall a. Seq a -> a -> Seq a
:|> n
x
      }
    processNeighbor :: n -> BFSState n -> n -> BFSState n
    processNeighbor :: n -> BFSState n -> n -> BFSState n
processNeighbor n
curr bs0 :: BFSState n
bs0@BS{Map n (Maybe n)
Seq n
_bsOpen :: Seq n
_bsClosed :: Map n (Maybe n)
_bsOpen :: forall n. BFSState n -> Seq n
_bsClosed :: forall n. BFSState n -> Map n (Maybe n)
..} n
neighb
      | n
neighb n -> Map n (Maybe n) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map n (Maybe n)
_bsClosed = BFSState n
bs0
      | Bool
otherwise                   = n -> Maybe n -> BFSState n -> BFSState n
addBack n
neighb (n -> Maybe n
forall a. a -> Maybe a
Just n
curr) BFSState n
bs0

-- -- | Breadth-first search, with loop detection, stopping at first result
-- bfs :: forall n. Ord n
--     => (n -> Set n)   -- ^ neighborhood
--     -> n              -- ^ start
--     -> (n -> Bool)    -- ^ target
--     -> Maybe [n]      -- ^ the shortest path, if it exists
-- bfs ex x0 dest = fmap snd . M.lookupMin $ bfsAll ex x0 (mfilter dest . Just) ((> 0) . S.size)

-- data BFSState n a = BS { _bsClosed  :: !(Map n (Maybe n))  -- ^ map of item to "parent"
--                        , _bsOpen    :: !(Seq n          )  -- ^ queue
--                        , _bsFound   :: !(Map a n        )  -- ^ found items
--                        }

-- -- | Breadth-first search, with loop detection, to find all matches.
-- bfsAll :: forall n a. (Ord n, Ord a)
--     => (n -> Set n)             -- ^ neighborhood
--     -> n                        -- ^ start
--     -> (n -> Maybe a)              -- ^ keep me when True
--     -> (Set a -> Bool)          -- ^ stop when True
--     -> Map a [n]
-- bfsAll ex x0 isGood stopper = reconstruct <$> founds
-- -- M.fromSet reconstruct founds
--   where
--     (founds, parentMap) = go . addBack x0 Nothing $ BS M.empty Seq.empty M.empty
--     reconstruct :: n -> [n]
--     reconstruct goal = drop 1 . reverse $ goreco goal
--       where
--         goreco n = n : maybe [] goreco (parentMap M.! n)
--     go :: BFSState n a -> (Map a n, Map n (Maybe n))
--     go BS{..} = case _bsOpen of
--       Empty    -> (_bsFound, _bsClosed)
--       (!n) :<| ns ->
--         let (found', updated) = case isGood n of
--               Just x
--                 | x `M.notMember` _bsFound -> (M.insert x n _bsFound, True)
--               _   -> (_bsFound, False)
--             stopHere = updated && stopper (M.keysSet found')
--         in  if stopHere
--               then (found', _bsClosed)
--               else go . S.foldl' (processNeighbor n) (BS _bsClosed ns found') $ ex n
--     addBack :: n -> Maybe n -> BFSState n a -> BFSState n a
--     addBack !x !up BS{..} = BS
--       { _bsClosed = M.insert x up _bsClosed
--       , _bsOpen   = _bsOpen :|> x
--       , _bsFound  = _bsFound
--       }
--     processNeighbor :: n -> BFSState n a -> n -> BFSState n a
--     processNeighbor !curr bs0@BS{..} neighb
--       | neighb `M.member` _bsClosed = bs0
--       | otherwise                   = addBack neighb (Just curr) bs0

binarySearch
    :: (Int -> Ordering)        -- LT: Too small, GT: Too big
    -> Int
    -> Int
    -> Maybe Int
binarySearch :: (Int -> Ordering) -> Int -> Int -> Maybe Int
binarySearch Int -> Ordering
p = Int -> Int -> Maybe Int
go
  where
    go :: Int -> Int -> Maybe Int
go !Int
x !Int
y
        | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y    = if Int -> Ordering
p Int
x Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ then Int -> Maybe Int
forall a. a -> Maybe a
Just Int
x else Maybe Int
forall a. Maybe a
Nothing
        | Bool
otherwise = case Int -> Ordering
p Int
mid of
            Ordering
LT -> Int -> Int -> Maybe Int
go Int
mid Int
y
            Ordering
EQ -> Int -> Maybe Int
forall a. a -> Maybe a
Just Int
mid
            Ordering
GT -> Int -> Int -> Maybe Int
go Int
x Int
mid
      where
        mid :: Int
mid = ((Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x

exponentialSearch
    :: (Int -> Ordering)        -- LT: Too small, GT: Too big
    -> Int
    -> Maybe Int
exponentialSearch :: (Int -> Ordering) -> Int -> Maybe Int
exponentialSearch Int -> Ordering
p = Int -> Maybe Int
go
  where
    go :: Int -> Maybe Int
go !Int
x = case Int -> Ordering
p Int
x of
      Ordering
LT -> Int -> Maybe Int
go (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2)
      Ordering
EQ -> Int -> Maybe Int
forall a. a -> Maybe a
Just Int
x
      Ordering
GT -> (Int -> Ordering) -> Int -> Int -> Maybe Int
binarySearch Int -> Ordering
p (Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) Int
x

-- | Find the lowest value where the predicate is satisfied within the
-- given bounds.
binaryMinSearch
    :: (Int -> Bool)
    -> Int
    -> Int
    -> Maybe Int
binaryMinSearch :: (Int -> Bool) -> Int -> Int -> Maybe Int
binaryMinSearch Int -> Bool
p = Int -> Int -> Maybe Int
go
  where
    go :: Int -> Int -> Maybe Int
go !Int
x !Int
y
        | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
mid Bool -> Bool -> Bool
|| Int
y Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
mid = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Int -> Bool
p Int
mid                = Int -> Int -> Maybe Int
go Int
x Int
mid
        | Bool
otherwise            = Int -> Int -> Maybe Int
go Int
mid Int
y
      where
        mid :: Int
mid = ((Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x

-- | Find the lowest value where the predicate is satisfied above a given
-- bound.
exponentialMinSearch
    :: (Int -> Bool)
    -> Int
    -> Maybe Int
exponentialMinSearch :: (Int -> Bool) -> Int -> Maybe Int
exponentialMinSearch Int -> Bool
p = Int -> Maybe Int
go
  where
    go :: Int -> Maybe Int
go !Int
x
      | Int -> Bool
p Int
x       = (Int -> Bool) -> Int -> Int -> Maybe Int
binaryMinSearch Int -> Bool
p (Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) Int
x
      | Bool
otherwise = Int -> Maybe Int
go (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2)

-- | Find the lowest value where the predicate is 'Just' within the
-- given bounds.
binaryFindMin
    :: (Int -> Maybe a)
    -> Int
    -> Int
    -> Maybe a
binaryFindMin :: forall a. (Int -> Maybe a) -> Int -> Int -> Maybe a
binaryFindMin Int -> Maybe a
p Int
x0 Int
y0 = (Int -> Maybe a) -> Maybe a -> Int -> Int -> Maybe a
forall a. (Int -> Maybe a) -> Maybe a -> Int -> Int -> Maybe a
binaryFindMin_ Int -> Maybe a
p (Int -> Maybe a
p Int
y0) Int
x0 Int
y0

binaryFindMin_
    :: (Int -> Maybe a)
    -> Maybe a          -- p y0
    -> Int
    -> Int
    -> Maybe a
binaryFindMin_ :: forall a. (Int -> Maybe a) -> Maybe a -> Int -> Int -> Maybe a
binaryFindMin_ Int -> Maybe a
p = Maybe a -> Int -> Int -> Maybe a
go
  where
    go :: Maybe a -> Int -> Int -> Maybe a
go Maybe a
found !Int
x !Int
y
      | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
mid Bool -> Bool -> Bool
|| Int
y Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
mid = Maybe a
found
      | Bool
otherwise            = case Int -> Maybe a
p Int
mid of
          Maybe a
Nothing    -> Maybe a -> Int -> Int -> Maybe a
go Maybe a
found Int
mid Int
y
          f :: Maybe a
f@(Just a
_) -> Maybe a -> Int -> Int -> Maybe a
go Maybe a
f     Int
x   Int
mid
      where
        mid :: Int
mid = ((Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x

-- | Find the lowest value where the predicate is 'Just' above a given
-- bound.
exponentialFindMin
    :: (Int -> Maybe a)
    -> Int
    -> Maybe a
exponentialFindMin :: forall a. (Int -> Maybe a) -> Int -> Maybe a
exponentialFindMin Int -> Maybe a
p = Int -> Maybe a
go
  where
    go :: Int -> Maybe a
go !Int
x = case Int -> Maybe a
p Int
x of
      Maybe a
Nothing -> Int -> Maybe a
go (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2)
      f :: Maybe a
f@(Just a
_) -> (Int -> Maybe a) -> Maybe a -> Int -> Int -> Maybe a
forall a. (Int -> Maybe a) -> Maybe a -> Int -> Int -> Maybe a
binaryFindMin_ Int -> Maybe a
p Maybe a
f (Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) Int
x