{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Test.Util.Split (
    spanLeft
  , spanLeft'
  , splitAtJust
  ) where

import           Data.Bifunctor (first)
import           Data.Word (Word64)

{-------------------------------------------------------------------------------
  spanLeft
-------------------------------------------------------------------------------}

-- | The returned @b@ is the first in the list.
--
-- INVARIANT The output data is a segmentation of the given list.
spanLeft ::
     forall x a b.
     (x -> Either a b) -> [x] -> ([a], Maybe (b, [x]))
spanLeft :: forall x a b. (x -> Either a b) -> [x] -> ([a], Maybe (b, [x]))
spanLeft x -> Either a b
prj [x]
xs = ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
acc, Maybe (b, [x])
mbBxs)
  where
    ([a]
acc, Maybe (b, [x])
mbBxs) = (x -> Either a b) -> [x] -> ([a], Maybe (b, [x]))
forall x a b. (x -> Either a b) -> [x] -> ([a], Maybe (b, [x]))
spanLeft' x -> Either a b
prj [x]
xs

-- | As 'spanLeft', but the @[a]@ is reversed.
spanLeft'
  :: forall x a b.
     (x -> Either a b) -> [x] -> ([a], Maybe (b, [x]))
spanLeft' :: forall x a b. (x -> Either a b) -> [x] -> ([a], Maybe (b, [x]))
spanLeft' x -> Either a b
prj = [a] -> [x] -> ([a], Maybe (b, [x]))
go []
  where
    go :: [a] -> [x] -> ([a], Maybe (b, [x]))
go [a]
acc = \case
      []     -> ([a]
acc, Maybe (b, [x])
forall a. Maybe a
Nothing)
      x
x : [x]
xs -> case x -> Either a b
prj x
x of
        Left a
a  -> [a] -> [x] -> ([a], Maybe (b, [x]))
go (a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc) [x]
xs
        Right b
b -> ([a]
acc, (b, [x]) -> Maybe (b, [x])
forall a. a -> Maybe a
Just (b
b, [x]
xs))

{-------------------------------------------------------------------------------
  splitAtJust
-------------------------------------------------------------------------------}

-- | INVARIANT: The second is a function of the first.
data Prj a b = Prj !a !b

-- | The returned @b@ is either the @n@th @b@ or else the last in the given
-- list.
--
-- INVARIANT The output data is a segmentation of the given list.
splitAtJust ::
     forall x b.
     (x -> Maybe b) -> Word64 -> [x] -> (Maybe ([x], b), [x])
splitAtJust :: forall x b.
(x -> Maybe b) -> Word64 -> [x] -> (Maybe ([x], b), [x])
splitAtJust x -> Maybe b
prj = \Word64
n [x]
xs ->
  if Word64
0 Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
n then (Maybe ([x], b)
forall a. Maybe a
Nothing, [x]
xs)
  else case [x] -> ([x], Maybe (Prj x b, [x]))
peel [x]
xs of
    ([x]
pre, Just (Prj x b
xb, [x]
xs')) -> ([x], b) -> Maybe ([x], b)
forall a. a -> Maybe a
Just (([x], b) -> Maybe ([x], b))
-> (([x], b), [x]) -> (Maybe ([x], b), [x])
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
`first` [x] -> Prj x b -> Word64 -> [x] -> (([x], b), [x])
forall {t}.
(Eq t, Num t) =>
[x] -> Prj x b -> t -> [x] -> (([x], b), [x])
go [x]
pre Prj x b
xb (Word64
n Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1) [x]
xs'
    ([x]
_, Maybe (Prj x b, [x])
Nothing)          -> (Maybe ([x], b)
forall a. Maybe a
Nothing, [x]
xs)
  where
    peel :: [x] -> ([x], Maybe (Prj x b, [x]))
    peel :: [x] -> ([x], Maybe (Prj x b, [x]))
peel = (x -> Either x (Prj x b)) -> [x] -> ([x], Maybe (Prj x b, [x]))
forall x a b. (x -> Either a b) -> [x] -> ([a], Maybe (b, [x]))
spanLeft' x -> Either x (Prj x b)
prj'
      where
        prj' :: x -> Either x (Prj x b)
prj' x
x = case x -> Maybe b
prj x
x of
          Maybe b
Nothing -> x -> Either x (Prj x b)
forall a b. a -> Either a b
Left x
x
          Just b
b  -> Prj x b -> Either x (Prj x b)
forall a b. b -> Either a b
Right (x -> b -> Prj x b
forall a b. a -> b -> Prj a b
Prj x
x b
b)

    go :: [x] -> Prj x b -> t -> [x] -> (([x], b), [x])
go [x]
pre (Prj x
x b
b) t
n [x]
xs
      | t
0 t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
n    = (([x] -> [x]
forall a. [a] -> [a]
reverse [x]
pre, b
b), [x]
xs)
      | Bool
otherwise = case [x] -> ([x], Maybe (Prj x b, [x]))
peel [x]
xs of
      ([x]
pre', Maybe (Prj x b, [x])
Nothing       ) -> (([x] -> [x]
forall a. [a] -> [a]
reverse [x]
pre, b
b), [x] -> [x]
forall a. [a] -> [a]
reverse [x]
pre')
      ([x]
pre', Just (Prj x b
xb, [x]
xs')) -> [x] -> Prj x b -> t -> [x] -> (([x], b), [x])
go ([x]
pre' [x] -> [x] -> [x]
forall a. [a] -> [a] -> [a]
++ x
x x -> [x] -> [x]
forall a. a -> [a] -> [a]
: [x]
pre) Prj x b
xb (t
n t -> t -> t
forall a. Num a => a -> a -> a
- t
1) [x]
xs'