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

module Data.List.ApplyMerge.IntSet (applyMerge) where

import Control.Monad (guard)
import Data.Function ((&))
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 ((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) ((a -> b -> c) -> NonEmpty a -> NonEmpty b -> Frontier a b c
forall a b c.
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> Frontier a b c
initialFrontier a -> b -> c
f NonEmpty a
as' NonEmpty b
bs'))

initialFrontier :: (a -> b -> c) -> NonEmpty a -> NonEmpty b -> Frontier a b c
initialFrontier :: forall a b c.
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> Frontier a b c
initialFrontier a -> b -> c
f NonEmpty a
as NonEmpty b
bs =
  let node :: Node a b c
node = (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
   in Frontier
        { $sel:queue:Frontier :: MinPQueue c (Node a b c)
queue = c -> Node a b c -> MinPQueue c (Node a b c)
forall k a. k -> a -> MinPQueue k a
MinPQueue.singleton Node a b c
node.value Node a b c
node,
          $sel:indexSetA:Frontier :: IntSet
indexSetA = Int -> IntSet
IntSet.singleton Int
0,
          $sel:indexSetB:Frontier :: IntSet
indexSetB = Int -> IntSet
IntSet.singleton Int
0
        }

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 Frontier a b c
frontier = do
  (Node a b c
node, Frontier a b c
frontier') <- 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 Frontier a b c
frontier
  let frontier'' :: Frontier a b c
frontier'' =
        Frontier a b c
frontier'
          Frontier a b c
-> (Frontier a b c -> Frontier a b c) -> Frontier a b c
forall a b. a -> (a -> b) -> b
& (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) -> Frontier a b c
forall a b. a -> (a -> b) -> b
& (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
  (c, Frontier a b c) -> Maybe (c, Frontier a b c)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Node a b c
node.value, Frontier a b c
frontier'')

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')

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
    }