-- SPDX-FileCopyrightText: Copyright Preetham Gujjula
-- SPDX-License-Identifier: BSD-3-Clause
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE NoFieldSelectors #-}

module ApplyMerge.IntSet (applyMerge, applyMergeNonEmpty) where

import Control.Arrow ((>>>))
import Control.Monad (guard)
import Data.IntSet (IntSet)
import Data.IntSet qualified as IntSet
import Data.List (unfoldr)
import Data.List.NonEmpty (NonEmpty ((:|)), nonEmpty)
import Data.List.NonEmpty qualified as NonEmpty
import Data.Maybe (fromMaybe)
import Data.PQueue.Prio.Min (MinPQueue)
import Data.PQueue.Prio.Min qualified as MinPQueue

data Node a b c = Node
  { forall a b c. Node a b c -> (Int, Int)
position :: (Int, Int),
    forall a b c. Node a b c -> c
value :: c,
    forall a b c. Node a b c -> NonEmpty a
as :: NonEmpty a,
    forall a b c. Node a b c -> NonEmpty b
bs :: NonEmpty b
  }

data Frontier a b c = Frontier
  { forall a b c. Frontier a b c -> MinPQueue c (Node a b c)
queue :: MinPQueue c (Node a b c),
    forall a b c. Frontier a b c -> IntSet
indexSetA :: IntSet,
    forall a b c. Frontier a b c -> IntSet
indexSetB :: IntSet
  }

applyMerge :: (Ord c) => (a -> b -> c) -> [a] -> [b] -> [c]
applyMerge :: forall c a b. Ord c => (a -> b -> c) -> [a] -> [b] -> [c]
applyMerge a -> b -> c
f [a]
as [b]
bs = [c] -> Maybe [c] -> [c]
forall a. a -> Maybe a -> a
fromMaybe [] (Maybe [c] -> [c]) -> Maybe [c] -> [c]
forall a b. (a -> b) -> a -> b
$ do
  NonEmpty a
as' <- [a] -> Maybe (NonEmpty a)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty [a]
as
  NonEmpty b
bs' <- [b] -> Maybe (NonEmpty b)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty [b]
bs
  [c] -> Maybe [c]
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NonEmpty c -> [c]
forall a. NonEmpty a -> [a]
NonEmpty.toList ((a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
forall c a b.
Ord c =>
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
applyMergeNonEmpty a -> b -> c
f NonEmpty a
as' NonEmpty b
bs'))

applyMergeNonEmpty ::
  (Ord c) => (a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
applyMergeNonEmpty :: forall c a b.
Ord c =>
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
applyMergeNonEmpty a -> b -> c
f NonEmpty a
as NonEmpty b
bs =
  let (c
c, Frontier a b c
frontier) = (a -> b -> c) -> NonEmpty a -> NonEmpty b -> (c, Frontier a b c)
forall a b c.
Ord c =>
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> (c, Frontier a b c)
initialState a -> b -> c
f NonEmpty a
as NonEmpty b
bs
   in c
c c -> [c] -> NonEmpty c
forall a. a -> [a] -> NonEmpty a
:| (Frontier a b c -> Maybe (c, Frontier a b c))
-> Frontier a b c -> [c]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr ((a -> b -> c) -> Frontier a b c -> Maybe (c, Frontier a b c)
forall c a b.
Ord c =>
(a -> b -> c) -> Frontier a b c -> Maybe (c, Frontier a b c)
step a -> b -> c
f) Frontier a b c
frontier

initialState ::
  forall a b c.
  (Ord c) =>
  (a -> b -> c) ->
  NonEmpty a ->
  NonEmpty b ->
  (c, Frontier a b c)
initialState :: forall a b c.
Ord c =>
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> (c, Frontier a b c)
initialState a -> b -> c
f NonEmpty a
as NonEmpty b
bs =
  let initialNode :: Node a b c
      initialNode :: Node a b c
initialNode = (a -> b -> c)
-> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
forall a b c.
(a -> b -> c)
-> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
mkNode a -> b -> c
f (Int
0, Int
0) NonEmpty a
as NonEmpty b
bs

      emptyFrontier :: Frontier a b c
      emptyFrontier :: Frontier a b c
emptyFrontier =
        Frontier
          { $sel:queue:Frontier :: MinPQueue c (Node a b c)
queue = MinPQueue c (Node a b c)
forall k a. MinPQueue k a
MinPQueue.empty,
            $sel:indexSetA:Frontier :: IntSet
indexSetA = IntSet
IntSet.empty,
            $sel:indexSetB:Frontier :: IntSet
indexSetB = IntSet
IntSet.empty
          }
   in (a -> b -> c)
-> Node a b c -> Frontier a b c -> (c, Frontier a b c)
forall c a b.
Ord c =>
(a -> b -> c)
-> Node a b c -> Frontier a b c -> (c, Frontier a b c)
peekInsertChildren a -> b -> c
f Node a b c
initialNode Frontier a b c
emptyFrontier

step :: (Ord c) => (a -> b -> c) -> Frontier a b c -> Maybe (c, Frontier a b c)
step :: forall c a b.
Ord c =>
(a -> b -> c) -> Frontier a b c -> Maybe (c, Frontier a b c)
step a -> b -> c
f = ((Node a b c, Frontier a b c) -> (c, Frontier a b c))
-> Maybe (Node a b c, Frontier a b c) -> Maybe (c, Frontier a b c)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Node a b c -> Frontier a b c -> (c, Frontier a b c))
-> (Node a b c, Frontier a b c) -> (c, Frontier a b c)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((a -> b -> c)
-> Node a b c -> Frontier a b c -> (c, Frontier a b c)
forall c a b.
Ord c =>
(a -> b -> c)
-> Node a b c -> Frontier a b c -> (c, Frontier a b c)
peekInsertChildren a -> b -> c
f)) (Maybe (Node a b c, Frontier a b c) -> Maybe (c, Frontier a b c))
-> (Frontier a b c -> Maybe (Node a b c, Frontier a b c))
-> Frontier a b c
-> Maybe (c, Frontier a b c)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Frontier a b c -> Maybe (Node a b c, Frontier a b c)
forall c a b.
Ord c =>
Frontier a b c -> Maybe (Node a b c, Frontier a b c)
deleteMinNode

deleteMinNode :: (Ord c) => Frontier a b c -> Maybe (Node a b c, Frontier a b c)
deleteMinNode :: forall c a b.
Ord c =>
Frontier a b c -> Maybe (Node a b c, Frontier a b c)
deleteMinNode Frontier a b c
frontier = do
  (Node a b c
node, MinPQueue c (Node a b c)
queue') <- MinPQueue c (Node a b c)
-> Maybe (Node a b c, MinPQueue c (Node a b c))
forall k a. Ord k => MinPQueue k a -> Maybe (a, MinPQueue k a)
MinPQueue.minView Frontier a b c
frontier.queue
  let (Int
ia, Int
ib) = Node a b c
node.position
      frontier' :: Frontier a b c
frontier' =
        Frontier
          { $sel:queue:Frontier :: MinPQueue c (Node a b c)
queue = MinPQueue c (Node a b c)
queue',
            $sel:indexSetA:Frontier :: IntSet
indexSetA = Int -> IntSet -> IntSet
IntSet.delete Int
ia Frontier a b c
frontier.indexSetA,
            $sel:indexSetB:Frontier :: IntSet
indexSetB = Int -> IntSet -> IntSet
IntSet.delete Int
ib Frontier a b c
frontier.indexSetB
          }
  (Node a b c, Frontier a b c) -> Maybe (Node a b c, Frontier a b c)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Node a b c
node, Frontier a b c
frontier')

peekInsertChildren ::
  (Ord c) =>
  (a -> b -> c) ->
  Node a b c ->
  Frontier a b c ->
  (c, Frontier a b c)
peekInsertChildren :: forall c a b.
Ord c =>
(a -> b -> c)
-> Node a b c -> Frontier a b c -> (c, Frontier a b c)
peekInsertChildren a -> b -> c
f Node a b c
node =
  (a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
forall c a b.
Ord c =>
(a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
insertChildA a -> b -> c
f Node a b c
node
    (Frontier a b c -> Frontier a b c)
-> (Frontier a b c -> (c, Frontier a b c))
-> Frontier a b c
-> (c, Frontier a b c)
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> (a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
forall c a b.
Ord c =>
(a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
insertChildB a -> b -> c
f Node a b c
node
    (Frontier a b c -> Frontier a b c)
-> (Frontier a b c -> (c, Frontier a b c))
-> Frontier a b c
-> (c, Frontier a b c)
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> (Node a b c
node.value,)

insertChildA ::
  (Ord c) => (a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
insertChildA :: forall c a b.
Ord c =>
(a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
insertChildA a -> b -> c
f (Node (Int
ia, Int
ib) c
_ NonEmpty a
as NonEmpty b
bs) Frontier a b c
frontier = Frontier a b c -> Maybe (Frontier a b c) -> Frontier a b c
forall a. a -> Maybe a -> a
fromMaybe Frontier a b c
frontier (Maybe (Frontier a b c) -> Frontier a b c)
-> Maybe (Frontier a b c) -> Frontier a b c
forall a b. (a -> b) -> a -> b
$ do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Bool
not (Int -> IntSet -> Bool
IntSet.member (Int
ia Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Frontier a b c
frontier.indexSetA))
  NonEmpty a
as' <- [a] -> Maybe (NonEmpty a)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty (NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NonEmpty.tail NonEmpty a
as)
  let childA :: Node a b c
childA = (a -> b -> c)
-> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
forall a b c.
(a -> b -> c)
-> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
mkNode a -> b -> c
f (Int
ia Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
ib) NonEmpty a
as' NonEmpty b
bs
  Frontier a b c -> Maybe (Frontier a b c)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Frontier a b c -> Maybe (Frontier a b c))
-> Frontier a b c -> Maybe (Frontier a b c)
forall a b. (a -> b) -> a -> b
$ Node a b c -> Frontier a b c -> Frontier a b c
forall c a b.
Ord c =>
Node a b c -> Frontier a b c -> Frontier a b c
insertNode Node a b c
childA Frontier a b c
frontier

insertChildB ::
  (Ord c) => (a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
insertChildB :: forall c a b.
Ord c =>
(a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
insertChildB a -> b -> c
f (Node (Int
ia, Int
ib) c
_ NonEmpty a
as NonEmpty b
bs) Frontier a b c
frontier = Frontier a b c -> Maybe (Frontier a b c) -> Frontier a b c
forall a. a -> Maybe a -> a
fromMaybe Frontier a b c
frontier (Maybe (Frontier a b c) -> Frontier a b c)
-> Maybe (Frontier a b c) -> Frontier a b c
forall a b. (a -> b) -> a -> b
$ do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Bool
not (Int -> IntSet -> Bool
IntSet.member (Int
ib Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Frontier a b c
frontier.indexSetB))
  NonEmpty b
bs' <- [b] -> Maybe (NonEmpty b)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty (NonEmpty b -> [b]
forall a. NonEmpty a -> [a]
NonEmpty.tail NonEmpty b
bs)
  let childB :: Node a b c
childB = (a -> b -> c)
-> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
forall a b c.
(a -> b -> c)
-> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
mkNode a -> b -> c
f (Int
ia, Int
ib Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) NonEmpty a
as NonEmpty b
bs'
  Frontier a b c -> Maybe (Frontier a b c)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Frontier a b c -> Maybe (Frontier a b c))
-> Frontier a b c -> Maybe (Frontier a b c)
forall a b. (a -> b) -> a -> b
$ Node a b c -> Frontier a b c -> Frontier a b c
forall c a b.
Ord c =>
Node a b c -> Frontier a b c -> Frontier a b c
insertNode Node a b c
childB Frontier a b c
frontier

insertNode :: (Ord c) => Node a b c -> Frontier a b c -> Frontier a b c
insertNode :: forall c a b.
Ord c =>
Node a b c -> Frontier a b c -> Frontier a b c
insertNode Node a b c
node Frontier a b c
frontier =
  let (Int
ia, Int
ib) = Node a b c
node.position
   in Frontier
        { $sel:queue:Frontier :: MinPQueue c (Node a b c)
queue = c
-> Node a b c
-> MinPQueue c (Node a b c)
-> MinPQueue c (Node a b c)
forall k a. Ord k => k -> a -> MinPQueue k a -> MinPQueue k a
MinPQueue.insert Node a b c
node.value Node a b c
node Frontier a b c
frontier.queue,
          $sel:indexSetA:Frontier :: IntSet
indexSetA = Int -> IntSet -> IntSet
IntSet.insert Int
ia Frontier a b c
frontier.indexSetA,
          $sel:indexSetB:Frontier :: IntSet
indexSetB = Int -> IntSet -> IntSet
IntSet.insert Int
ib Frontier a b c
frontier.indexSetB
        }

mkNode :: (a -> b -> c) -> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
mkNode :: forall a b c.
(a -> b -> c)
-> (Int, Int) -> NonEmpty a -> NonEmpty b -> Node a b c
mkNode a -> b -> c
f (Int
ia, Int
ib) NonEmpty a
as NonEmpty b
bs =
  Node
    { $sel:position:Node :: (Int, Int)
position = (Int
ia, Int
ib),
      $sel:value:Node :: c
value = a -> b -> c
f (NonEmpty a -> a
forall a. NonEmpty a -> a
NonEmpty.head NonEmpty a
as) (NonEmpty b -> b
forall a. NonEmpty a -> a
NonEmpty.head NonEmpty b
bs),
      $sel:as:Node :: NonEmpty a
as = NonEmpty a
as,
      $sel:bs:Node :: NonEmpty b
bs = NonEmpty b
bs
    }