{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeFamilies #-}

module Algorithm.TwoPointers where

import Control.Monad.Primitive
import Data.Function
import qualified Data.Vector.Fusion.Bundle.Monadic as MBundle
import qualified Data.Vector.Fusion.Bundle.Size as Bundle
import qualified Data.Vector.Fusion.Stream.Monadic as MS
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Unboxed as U
import GHC.Generics

-- | @[l, r)@
data Window a = Window
  { forall a. Window a -> Int
leftW :: !Int
  , forall a. Window a -> Int
rightW :: !Int
  , forall a. Window a -> a
contextW :: !a
  }
  deriving (Window a -> Window a -> Bool
(Window a -> Window a -> Bool)
-> (Window a -> Window a -> Bool) -> Eq (Window a)
forall a. Eq a => Window a -> Window a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Window a -> Window a -> Bool
== :: Window a -> Window a -> Bool
$c/= :: forall a. Eq a => Window a -> Window a -> Bool
/= :: Window a -> Window a -> Bool
Eq, Eq (Window a)
Eq (Window a) =>
(Window a -> Window a -> Ordering)
-> (Window a -> Window a -> Bool)
-> (Window a -> Window a -> Bool)
-> (Window a -> Window a -> Bool)
-> (Window a -> Window a -> Bool)
-> (Window a -> Window a -> Window a)
-> (Window a -> Window a -> Window a)
-> Ord (Window a)
Window a -> Window a -> Bool
Window a -> Window a -> Ordering
Window a -> Window a -> Window a
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. Ord a => Eq (Window a)
forall a. Ord a => Window a -> Window a -> Bool
forall a. Ord a => Window a -> Window a -> Ordering
forall a. Ord a => Window a -> Window a -> Window a
$ccompare :: forall a. Ord a => Window a -> Window a -> Ordering
compare :: Window a -> Window a -> Ordering
$c< :: forall a. Ord a => Window a -> Window a -> Bool
< :: Window a -> Window a -> Bool
$c<= :: forall a. Ord a => Window a -> Window a -> Bool
<= :: Window a -> Window a -> Bool
$c> :: forall a. Ord a => Window a -> Window a -> Bool
> :: Window a -> Window a -> Bool
$c>= :: forall a. Ord a => Window a -> Window a -> Bool
>= :: Window a -> Window a -> Bool
$cmax :: forall a. Ord a => Window a -> Window a -> Window a
max :: Window a -> Window a -> Window a
$cmin :: forall a. Ord a => Window a -> Window a -> Window a
min :: Window a -> Window a -> Window a
Ord, Int -> Window a -> ShowS
[Window a] -> ShowS
Window a -> String
(Int -> Window a -> ShowS)
-> (Window a -> String) -> ([Window a] -> ShowS) -> Show (Window a)
forall a. Show a => Int -> Window a -> ShowS
forall a. Show a => [Window a] -> ShowS
forall a. Show a => Window a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Window a -> ShowS
showsPrec :: Int -> Window a -> ShowS
$cshow :: forall a. Show a => Window a -> String
show :: Window a -> String
$cshowList :: forall a. Show a => [Window a] -> ShowS
showList :: [Window a] -> ShowS
Show, (forall x. Window a -> Rep (Window a) x)
-> (forall x. Rep (Window a) x -> Window a) -> Generic (Window a)
forall x. Rep (Window a) x -> Window a
forall x. Window a -> Rep (Window a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (Window a) x -> Window a
forall a x. Window a -> Rep (Window a) x
$cfrom :: forall a x. Window a -> Rep (Window a) x
from :: forall x. Window a -> Rep (Window a) x
$cto :: forall a x. Rep (Window a) x -> Window a
to :: forall x. Rep (Window a) x -> Window a
Generic)

instance U.IsoUnbox (Window a) (Int, Int, a)

newtype instance U.MVector s (Window a) = MV_Window (U.MVector s (Int, Int, a))
newtype instance U.Vector (Window a) = V_Window (U.Vector (Int, Int, a))
deriving via (Window a `U.As` (Int, Int, a)) instance (U.Unbox a) => GM.MVector U.MVector (Window a)
deriving via (Window a `U.As` (Int, Int, a)) instance (U.Unbox a) => G.Vector U.Vector (Window a)
instance (U.Unbox a) => U.Unbox (Window a)

enumerateTwoPointers ::
  (Monad m) =>
  -- | size
  Int ->
  -- | shrinkL (@l < r@)
  (Window a -> m a) ->
  -- | tryExtendR (@l <= r@)
  (Window a -> m (Maybe a)) ->
  -- | context for the empty
  a ->
  m Int
enumerateTwoPointers :: forall (m :: * -> *) a.
Monad m =>
Int -> (Window a -> m a) -> (Window a -> m (Maybe a)) -> a -> m Int
enumerateTwoPointers Int
n Window a -> m a
shrinkL Window a -> m (Maybe a)
tryExtendR a
x0 = do
  ((Int -> Window a -> m Int) -> Int -> Window a -> m Int)
-> Int -> Window a -> m Int
forall a. (a -> a) -> a
fix
    ( \Int -> Window a -> m Int
loop !Int
acc w :: Window a
w@(Window Int
l Int
r a
_) -> do
        if Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
          then do
            Window a -> m (Maybe a)
tryExtendR Window a
w m (Maybe a) -> (Maybe a -> m Int) -> m Int
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
              Maybe a
Nothing
                | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
r -> Window a -> m a
shrinkL Window a
w m a -> (a -> m Int) -> m Int
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Window a -> m Int
loop (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l) (Window a -> m Int) -> (a -> Window a) -> a -> m Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> a -> Window a
forall a. Int -> Int -> a -> Window a
Window (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
r
                | Bool
otherwise -> Int -> Window a -> m Int
loop Int
acc (Int -> Int -> a -> Window a
forall a. Int -> Int -> a -> Window a
Window (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
x0)
              Just a
w' -> Int -> Window a -> m Int
loop Int
acc (Int -> Int -> a -> Window a
forall a. Int -> Int -> a -> Window a
Window Int
l (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
w')
          else Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l) Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2
    )
    Int
0
    (Int -> Int -> a -> Window a
forall a. Int -> Int -> a -> Window a
Window Int
0 Int
0 a
x0)
{-# INLINE enumerateTwoPointers #-}

maxLengthTwoPointers ::
  (Monad m) =>
  -- | size
  Int ->
  -- | shrinkL (@l < r@)
  (Window a -> m a) ->
  -- | tryExtendR (@l <= r@)
  (Window a -> m (Maybe a)) ->
  -- | context for the empty
  a ->
  m Int
maxLengthTwoPointers :: forall (m :: * -> *) a.
Monad m =>
Int -> (Window a -> m a) -> (Window a -> m (Maybe a)) -> a -> m Int
maxLengthTwoPointers Int
n Window a -> m a
shrinkL Window a -> m (Maybe a)
tryExtendR a
x0 = do
  ((Int -> Window a -> m Int) -> Int -> Window a -> m Int)
-> Int -> Window a -> m Int
forall a. (a -> a) -> a
fix
    ( \Int -> Window a -> m Int
loop !Int
acc w :: Window a
w@(Window Int
l Int
r a
_) -> do
        if Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
          then do
            Window a -> m (Maybe a)
tryExtendR Window a
w m (Maybe a) -> (Maybe a -> m Int) -> m Int
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
              Maybe a
Nothing
                | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
r -> Window a -> m a
shrinkL Window a
w m a -> (a -> m Int) -> m Int
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Window a -> m Int
loop (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
acc (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l)) (Window a -> m Int) -> (a -> Window a) -> a -> m Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> a -> Window a
forall a. Int -> Int -> a -> Window a
Window (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
r
                | Bool
otherwise -> Int -> Window a -> m Int
loop Int
acc (Int -> Int -> a -> Window a
forall a. Int -> Int -> a -> Window a
Window (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
x0)
              Just a
w' -> Int -> Window a -> m Int
loop Int
acc (Int -> Int -> a -> Window a
forall a. Int -> Int -> a -> Window a
Window Int
l (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
w')
          else Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
acc (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l)
    )
    Int
0
    (Int -> Int -> a -> Window a
forall a. Int -> Int -> a -> Window a
Window Int
0 Int
0 a
x0)
{-# INLINE maxLengthTwoPointers #-}

runTwoPointersStream ::
  (Monad m) =>
  -- | start
  Int ->
  -- | end
  Int ->
  -- | shrinkL (@l < r@)
  (Window a -> m a) ->
  -- | tryExtendR (@l <= r@)
  (Window a -> m (Maybe a)) ->
  a ->
  -- | context for the empty
  MS.Stream m (Window a)
runTwoPointersStream :: forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Window a -> m a)
-> (Window a -> m (Maybe a))
-> a
-> Stream m (Window a)
runTwoPointersStream Int
l0 Int
r0 Window a -> m a
shrinkL Window a -> m (Maybe a)
tryExtendR a
x0 = (Window a -> m (Step (Window a) (Window a)))
-> Window a -> Stream m (Window a)
forall (m :: * -> *) a s. (s -> m (Step s a)) -> s -> Stream m a
MS.Stream Window a -> m (Step (Window a) (Window a))
step (Int -> Int -> a -> Window a
forall a. Int -> Int -> a -> Window a
Window Int
l0 Int
l0 a
x0)
  where
    step :: Window a -> m (Step (Window a) (Window a))
step w :: Window a
w@(Window Int
l Int
r a
_)
      | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
r0 = do
          Window a -> m (Maybe a)
tryExtendR Window a
w m (Maybe a)
-> (Maybe a -> m (Step (Window a) (Window a)))
-> m (Step (Window a) (Window a))
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Maybe a
Nothing
              | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
r -> do
                  Window a -> Window a -> Step (Window a) (Window a)
forall a s. a -> s -> Step s a
MS.Yield Window a
w (Window a -> Step (Window a) (Window a))
-> (a -> Window a) -> a -> Step (Window a) (Window a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> a -> Window a
forall a. Int -> Int -> a -> Window a
Window (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
r (a -> Step (Window a) (Window a))
-> m a -> m (Step (Window a) (Window a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Window a -> m a
shrinkL Window a
w
              | Bool
otherwise -> do
                  Step (Window a) (Window a) -> m (Step (Window a) (Window a))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Window a) (Window a) -> m (Step (Window a) (Window a)))
-> Step (Window a) (Window a) -> m (Step (Window a) (Window a))
forall a b. (a -> b) -> a -> b
$ Window a -> Window a -> Step (Window a) (Window a)
forall a s. a -> s -> Step s a
MS.Yield Window a
w (Int -> Int -> a -> Window a
forall a. Int -> Int -> a -> Window a
Window (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
x0)
            Just a
x' -> Step (Window a) (Window a) -> m (Step (Window a) (Window a))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Window a) (Window a) -> m (Step (Window a) (Window a)))
-> Step (Window a) (Window a) -> m (Step (Window a) (Window a))
forall a b. (a -> b) -> a -> b
$ Window a -> Step (Window a) (Window a)
forall s a. s -> Step s a
MS.Skip (Int -> Int -> a -> Window a
forall a. Int -> Int -> a -> Window a
Window Int
l (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
x')
      | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
r0 = Window a -> Window a -> Step (Window a) (Window a)
forall a s. a -> s -> Step s a
MS.Yield Window a
w (Window a -> Step (Window a) (Window a))
-> (a -> Window a) -> a -> Step (Window a) (Window a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> a -> Window a
forall a. Int -> Int -> a -> Window a
Window (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
r (a -> Step (Window a) (Window a))
-> m a -> m (Step (Window a) (Window a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Window a -> m a
shrinkL Window a
w
      | Bool
otherwise = Step (Window a) (Window a) -> m (Step (Window a) (Window a))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (Window a) (Window a)
forall s a. Step s a
MS.Done
    {-# INLINE [0] step #-}
{-# INLINE [1] runTwoPointersStream #-}

runTwoPointersM ::
  (PrimMonad m, G.Vector v (Window a)) =>
  -- | start
  Int ->
  -- | end
  Int ->
  -- | shrinkL (@l < r@)
  (Window a -> m a) ->
  -- | tryExtendR (@l <= r@)
  (Window a -> m (Maybe a)) ->
  a ->
  -- | context for the empty
  m (v (Window a))
runTwoPointersM :: forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v (Window a)) =>
Int
-> Int
-> (Window a -> m a)
-> (Window a -> m (Maybe a))
-> a
-> m (v (Window a))
runTwoPointersM Int
l0 Int
r0 Window a -> m a
shrinkL Window a -> m (Maybe a)
tryExtendR a
x0 =
  MBundle m (ZonkAny 0) (Window a)
-> m (Mutable v (PrimState m) (Window a))
forall (m :: * -> *) (v :: * -> * -> *) a (u :: * -> *).
(PrimMonad m, MVector v a) =>
MBundle m u a -> m (v (PrimState m) a)
GM.munstream
    ( Stream m (Window a) -> Size -> MBundle m (ZonkAny 0) (Window a)
forall (m :: * -> *) a (v :: * -> *).
Monad m =>
Stream m a -> Size -> Bundle m v a
MBundle.fromStream
        (Int
-> Int
-> (Window a -> m a)
-> (Window a -> m (Maybe a))
-> a
-> Stream m (Window a)
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Window a -> m a)
-> (Window a -> m (Maybe a))
-> a
-> Stream m (Window a)
runTwoPointersStream Int
l0 Int
r0 Window a -> m a
shrinkL Window a -> m (Maybe a)
tryExtendR a
x0)
        (Int -> Size
Bundle.Exact (Int
r0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l0))
    )
    m (Mutable v (PrimState m) (Window a))
-> (Mutable v (PrimState m) (Window a) -> m (v (Window a)))
-> m (v (Window a))
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Mutable v (PrimState m) (Window a) -> m (v (Window a))
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze
{-# INLINE runTwoPointersM #-}

runTwoPointers ::
  (G.Vector v (Window a)) =>
  -- | start
  Int ->
  -- | end
  Int ->
  -- | shrinkL (@l < r@)
  (Window a -> a) ->
  -- | tryExtendR (@l <= r@)
  (Window a -> Maybe a) ->
  -- | context for the empty
  a ->
  v (Window a)
runTwoPointers :: forall (v :: * -> *) a.
Vector v (Window a) =>
Int
-> Int
-> (Window a -> a)
-> (Window a -> Maybe a)
-> a
-> v (Window a)
runTwoPointers Int
l0 Int
r0 Window a -> a
shrinkL Window a -> Maybe a
tryExtendR a
x0 =
  Bundle v (Window a) -> v (Window a)
forall (v :: * -> *) a. Vector v a => Bundle v a -> v a
G.unstream
    ( Stream Id (Window a) -> Size -> Bundle v (Window a)
forall (m :: * -> *) a (v :: * -> *).
Monad m =>
Stream m a -> Size -> Bundle m v a
MBundle.fromStream
        (Int
-> Int
-> (Window a -> Id a)
-> (Window a -> Id (Maybe a))
-> a
-> Stream Id (Window a)
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Window a -> m a)
-> (Window a -> m (Maybe a))
-> a
-> Stream m (Window a)
runTwoPointersStream Int
l0 Int
r0 (a -> Id a
forall a. a -> Id a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Id a) -> (Window a -> a) -> Window a -> Id a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Window a -> a
shrinkL) (Maybe a -> Id (Maybe a)
forall a. a -> Id a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> Id (Maybe a))
-> (Window a -> Maybe a) -> Window a -> Id (Maybe a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Window a -> Maybe a
tryExtendR) a
x0)
        (Int -> Size
Bundle.Exact (Int
r0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l0))
    )
{-# INLINE runTwoPointers #-}