{-# LANGUAGE BlockArguments, LambdaCase, NoImplicitPrelude #-}

-- | Functions involving lists of 'Either'.

module Data.List.EitherFunctions
  (

    {- * Map       -}  partlyMap,
    {- * Group     -}  groupEither,
    {- * Partition -}  partition,
    {- * Span      -}  spanLeft, spanLeft', spanRight, spanRight',
    {- * Lead      -}  leadLeft, leadLeft', leadRight, leadRight',
    {- * Branch    -}  branchLeft, branchRight, BranchComparison

  ) where

import Data.Bool                   ( Bool (..) )
import Data.Either                 ( Either (..) )
import Data.Function               ( fix )
import Data.Functor.Contravariant  ( Comparison (..), contramap )
import Data.List                   ( foldr, map, span )
import Data.Tree                   ( Tree (..), Forest )
import Data.Maybe                  ( Maybe (..), maybe )
import Data.Ord                    ( Ordering (..) )

-- |
-- >>> import Prelude (even, show)
--
-- >>> partlyMap (\x -> if even x then Just (show x) else Nothing) [1..5]
-- [Left 1,Right "2",Left 3,Right "4",Left 5]

partlyMap :: (a -> Maybe b) -> [a] -> [Either a b]

partlyMap :: forall a b. (a -> Maybe b) -> [a] -> [Either a b]
partlyMap a -> Maybe b
f = forall a b. (a -> b) -> [a] -> [b]
map \a
x -> forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall a b. a -> Either a b
Left a
x) forall a b. b -> Either a b
Right (a -> Maybe b
f a
x)

-- |
-- >>> groupEither [Left 1, Left 2, Right 'a', Left 3, Right 'b', Right 'c']
-- [Left [1,2],Right "a",Left [3],Right "bc"]

groupEither :: [Either a b] -> [Either [a] [b]]

groupEither :: forall a b. [Either a b] -> [Either [a] [b]]
groupEither = forall a. (a -> a) -> a
fix \[Either a b] -> [Either [a] [b]]
r -> \case
    []            ->  []
    Left  a
x : [Either a b]
xs  ->  forall a b. a -> Either a b
Left  (a
x forall a. a -> [a] -> [a]
: [a]
ys) forall a. a -> [a] -> [a]
: [Either a b] -> [Either [a] [b]]
r [Either a b]
zs  where ([a]
ys, [Either a b]
zs) = forall a b. [Either a b] -> ([a], [Either a b])
spanLeft  [Either a b]
xs
    Right b
x : [Either a b]
xs  ->  forall a b. b -> Either a b
Right (b
x forall a. a -> [a] -> [a]
: [b]
ys) forall a. a -> [a] -> [a]
: [Either a b] -> [Either [a] [b]]
r [Either a b]
zs  where ([b]
ys, [Either a b]
zs) = forall a b. [Either a b] -> ([b], [Either a b])
spanRight [Either a b]
xs

-- |
-- >>> leadLeft [Right 'a', Right 'b', Left 1, Right 'c', Right 'd', Left 2, Right 'e', Right 'f']
-- ("ab",[(1,"cd"),(2,"ef")])
--
-- >>> leadLeft [Left 1, Left 2, Right 'a', Left 3, Right 'b', Right 'c']
-- ("",[(1,""),(2,"a"),(3,"bc")])

leadLeft :: [Either a b] -> ([b], [(a, [b])])

leadLeft :: forall a b. [Either a b] -> ([b], [(a, [b])])
leadLeft = forall a b. [Either a b] -> ([b], [(a, [b])])
f
  where
    f :: [Either a b] -> ([b], [(a, [b])])
f [Either a b]
xs = ([b]
unledItems, [(a, [b])]
ledGroups)
      where
        ([b]
unledItems, Maybe (a, [Either a b])
ysMaybe) = forall a b. [Either a b] -> ([b], Maybe (a, [Either a b]))
spanRight' [Either a b]
xs
        ledGroups :: [(a, [b])]
ledGroups = case Maybe (a, [Either a b])
ysMaybe of
            Maybe (a, [Either a b])
Nothing            ->  []
            Just (a
leader, [Either a b]
ys)  ->  forall {a} {b}. a -> [Either a b] -> [(a, [b])]
r a
leader [Either a b]
ys

    r :: a -> [Either a b] -> [(a, [b])]
r a
leader [Either a b]
xs = (a, [b])
firstGroup forall a. a -> [a] -> [a]
: [(a, [b])]
moreGroups
      where
        firstGroup :: (a, [b])
firstGroup = (a
leader, [b]
followers)
        ([b]
followers, Maybe (a, [Either a b])
ysMaybe) = forall a b. [Either a b] -> ([b], Maybe (a, [Either a b]))
spanRight' [Either a b]
xs
        moreGroups :: [(a, [b])]
moreGroups = case Maybe (a, [Either a b])
ysMaybe of
            Maybe (a, [Either a b])
Nothing             ->  []
            Just (a
leader', [Either a b]
ys)  ->  a -> [Either a b] -> [(a, [b])]
r a
leader' [Either a b]
ys

-- |
-- >>> leadLeft' 0 [Right 'a', Right 'b', Left 1, Right 'c', Right 'd', Left 2, Right 'e', Right 'f']
-- [(0,"ab"),(1,"cd"),(2,"ef")]
--
-- >>> leadLeft' 0 [Left 1, Left 2, Right 'a', Left 3, Right 'b', Right 'c']
-- [(1,""),(2,"a"),(3,"bc")]

leadLeft' ::
    a -- ^ Leader to use for the first group in case the list does not begin with a 'Left'.
    -> [Either a b] -> [(a, [b])]

leadLeft' :: forall {a} {b}. a -> [Either a b] -> [(a, [b])]
leadLeft' a
leader [Either a b]
xs = forall a b. a -> ([b], [(a, [b])]) -> [(a, [b])]
addMissingLeader a
leader (forall a b. [Either a b] -> ([b], [(a, [b])])
leadLeft [Either a b]
xs)

-- |
-- >>> leadRight [Left 1, Left 2, Right 'a', Left 3, Left 4, Right 'b', Left 5, Left 6]
-- ([1,2],[('a',[3,4]),('b',[5,6])])
--
-- >>> leadRight [Right 'a', Left 3, Left 4, Right 'b', Right 'c', Left 5, Left 6]
-- ([],[('a',[3,4]),('b',[]),('c',[5,6])])

leadRight :: [Either a b] -> ([a], [(b, [a])])

leadRight :: forall a b. [Either a b] -> ([a], [(b, [a])])
leadRight = forall a b. [Either a b] -> ([a], [(b, [a])])
f
  where
    f :: [Either a b] -> ([a], [(b, [a])])
f [Either a b]
xs = ([a]
unledItems, [(b, [a])]
ledGroups)
      where
        ([a]
unledItems, Maybe (b, [Either a b])
ysMaybe) = forall a b. [Either a b] -> ([a], Maybe (b, [Either a b]))
spanLeft' [Either a b]
xs
        ledGroups :: [(b, [a])]
ledGroups = case Maybe (b, [Either a b])
ysMaybe of
            Maybe (b, [Either a b])
Nothing            ->  []
            Just (b
leader, [Either a b]
ys)  ->  forall {b} {a}. b -> [Either a b] -> [(b, [a])]
r b
leader [Either a b]
ys

    r :: b -> [Either a b] -> [(b, [a])]
r b
leader [Either a b]
xs = (b, [a])
firstGroup forall a. a -> [a] -> [a]
: [(b, [a])]
moreGroups
      where
        firstGroup :: (b, [a])
firstGroup = (b
leader, [a]
followers)
        ([a]
followers, Maybe (b, [Either a b])
ysMaybe) = forall a b. [Either a b] -> ([a], Maybe (b, [Either a b]))
spanLeft' [Either a b]
xs
        moreGroups :: [(b, [a])]
moreGroups = case Maybe (b, [Either a b])
ysMaybe of
            Maybe (b, [Either a b])
Nothing             ->  []
            Just (b
leader', [Either a b]
ys)  ->  b -> [Either a b] -> [(b, [a])]
r b
leader' [Either a b]
ys

-- |
-- >>> leadRight' 'z' [Left 1, Left 2, Right 'a', Left 3, Left 4, Right 'b', Left 5, Left 6]
-- [('z',[1,2]),('a',[3,4]),('b',[5,6])]
--
-- >>> leadRight' 'z' [Right 'a', Left 3, Left 4, Right 'b', Right 'c', Left 5, Left 6]
-- [('a',[3,4]),('b',[]),('c',[5,6])]

leadRight' ::
    b -- ^ Leader to use for the first group in case the list does not begin with a 'Right'.
    -> [Either a b] -> [(b, [a])]

leadRight' :: forall {b} {a}. b -> [Either a b] -> [(b, [a])]
leadRight' b
leader [Either a b]
xs = forall a b. a -> ([b], [(a, [b])]) -> [(a, [b])]
addMissingLeader b
leader (forall a b. [Either a b] -> ([a], [(b, [a])])
leadRight [Either a b]
xs)

addMissingLeader :: a -> ([b], [(a, [b])]) -> [(a, [b])]

addMissingLeader :: forall a b. a -> ([b], [(a, [b])]) -> [(a, [b])]
addMissingLeader a
_      ( []         , [(a, [b])]
groups ) =                        [(a, [b])]
groups
addMissingLeader a
leader ( [b]
unledIntro , [(a, [b])]
groups ) = (a
leader, [b]
unledIntro) forall a. a -> [a] -> [a]
: [(a, [b])]
groups

-- |
-- >>> spanLeft [Left 1, Left 2, Right 'a', Left 3, Right 'b', Right 'c']
-- ([1,2],[Right 'a',Left 3,Right 'b',Right 'c'])
--
-- >>> spanLeft [Right 'a', Left 3, Right 'b', Right 'c']
-- ([],[Right 'a',Left 3,Right 'b',Right 'c'])

spanLeft :: [Either a b] -> ([a], [Either a b])

spanLeft :: forall a b. [Either a b] -> ([a], [Either a b])
spanLeft = forall a. (a -> a) -> a
fix \[Either a b] -> ([a], [Either a b])
r -> \case
    []           ->  ( []     , [] )
    Left a
x : [Either a b]
xs  ->  ( a
x forall a. a -> [a] -> [a]
: [a]
ys , [Either a b]
zs )  where ([a]
ys, [Either a b]
zs) = [Either a b] -> ([a], [Either a b])
r [Either a b]
xs
    [Either a b]
xs           ->  ( []     , [Either a b]
xs )

-- | Similar to 'spanLeft', but preserves a little more information in the return type: if the remainder of the list is non-empty, then it necessarily begins with a 'Right', and so we can go ahead and unwrap that and return it as a value of type `b`.
--
-- >>> spanLeft' [Left 1, Left 2, Right 'a', Left 3, Right 'b', Right 'c']
-- ([1,2],Just ('a',[Left 3,Right 'b',Right 'c']))
--
-- >>> spanLeft' [Right 'a', Left 3, Right 'b', Right 'c']
-- ([],Just ('a',[Left 3,Right 'b',Right 'c']))
--
-- >>> spanLeft' [Left 1, Left 2, Left 3]
-- ([1,2,3],Nothing)

spanLeft' :: [Either a b] -> ([a], Maybe (b, [Either a b]))

spanLeft' :: forall a b. [Either a b] -> ([a], Maybe (b, [Either a b]))
spanLeft' = forall a. (a -> a) -> a
fix \[Either a b] -> ([a], Maybe (b, [Either a b]))
r -> \case
    []            ->  ( []     , forall a. Maybe a
Nothing      )
    Left  a
x : [Either a b]
xs  ->  ( a
x forall a. a -> [a] -> [a]
: [a]
ys , Maybe (b, [Either a b])
zs           )  where ([a]
ys, Maybe (b, [Either a b])
zs) = [Either a b] -> ([a], Maybe (b, [Either a b]))
r [Either a b]
xs
    Right b
x : [Either a b]
xs  ->  ( []     , forall a. a -> Maybe a
Just (b
x, [Either a b]
xs) )

-- | Similar to 'spanRight', but preserves a little more information in the return type: if the remainder of the list is non-empty, then it necessarily begins with a 'Left', and so we can go ahead and unwrap that and return it as a value of type `a`.
--
-- >>> spanRight [Left 1, Left 2, Right 'a', Left 3, Right 'b', Right 'c']
-- ("",[Left 1,Left 2,Right 'a',Left 3,Right 'b',Right 'c'])
--
-- >>> spanRight [Right 'a', Left 3, Right 'b', Right 'c']
-- ("a",[Left 3,Right 'b',Right 'c'])

spanRight :: [Either a b] -> ([b], [Either a b])

spanRight :: forall a b. [Either a b] -> ([b], [Either a b])
spanRight = forall a. (a -> a) -> a
fix \[Either a b] -> ([b], [Either a b])
r -> \case
    []            ->  ( []     , [] )
    Right b
x : [Either a b]
xs  ->  ( b
x forall a. a -> [a] -> [a]
: [b]
ys , [Either a b]
zs )  where ([b]
ys, [Either a b]
zs) = [Either a b] -> ([b], [Either a b])
r [Either a b]
xs
    [Either a b]
xs            ->  ( []     , [Either a b]
xs )

-- |
-- >>> spanRight' [Left 1, Left 2, Right 'a', Left 3, Right 'b', Right 'c']
-- ("",Just (1,[Left 2,Right 'a',Left 3,Right 'b',Right 'c']))
--
-- >>> spanRight' [Right 'a', Left 3, Right 'b', Right 'c']
-- ("a",Just (3,[Right 'b',Right 'c']))
--
-- >>> spanRight' [Right 'a', Right 'b', Right 'c']
-- ("abc",Nothing)

spanRight' :: [Either a b] -> ([b], Maybe (a, [Either a b]))

spanRight' :: forall a b. [Either a b] -> ([b], Maybe (a, [Either a b]))
spanRight' = forall a. (a -> a) -> a
fix \[Either a b] -> ([b], Maybe (a, [Either a b]))
r -> \case
    []            ->  ( []     , forall a. Maybe a
Nothing      )
    Right b
x : [Either a b]
xs  ->  ( b
x forall a. a -> [a] -> [a]
: [b]
ys , Maybe (a, [Either a b])
zs           )  where ([b]
ys, Maybe (a, [Either a b])
zs) = [Either a b] -> ([b], Maybe (a, [Either a b]))
r [Either a b]
xs
    Left  a
x : [Either a b]
xs  ->  ( []     , forall a. a -> Maybe a
Just (a
x, [Either a b]
xs) )

-- |
-- >>> partition [Left 1, Left 2, Right 'a', Left 3, Right 'b', Right 'c']
-- ([1,2,3],"abc")

partition :: [Either a b] -> ([a], [b])

partition :: forall a b. [Either a b] -> ([a], [b])
partition = forall a. (a -> a) -> a
fix \[Either a b] -> ([a], [b])
r -> \case
    []            ->  ( []     , []     )
    Left  a
a : [Either a b]
xs  ->  ( a
a forall a. a -> [a] -> [a]
: [a]
as , [b]
bs     )  where ([a]
as, [b]
bs) = [Either a b] -> ([a], [b])
r [Either a b]
xs
    Right b
b : [Either a b]
xs  ->  ( [a]
as     , b
b forall a. a -> [a] -> [a]
: [b]
bs )  where ([a]
as, [b]
bs) = [Either a b] -> ([a], [b])
r [Either a b]
xs

-- | The relative significance of branches (greater values are closer to the root).
type BranchComparison a = Comparison a

-- |
-- >>> import Prelude
--
-- >>> heading level title = Left (level, title)
-- >>> chapter = heading 1
-- >>> section = heading 2
-- >>> p text = Right text
--
-- >>> :{
-- >>> list =
-- >>>     [ p "Copyright"
-- >>>     , p "Preface"
-- >>>     , chapter "Animals"
-- >>>     , p "The kingdom animalia"
-- >>>     , section "Vertebrates"
-- >>>     , p "Cats"
-- >>>     , p "Snakes"
-- >>>     , section "Invertebrates"
-- >>>     , p "Worms"
-- >>>     , p "Jellyfishes"
-- >>>     , chapter "Fungus"
-- >>>     , p "Yeast"
-- >>>     , p "Truffles"
-- >>>     , p "Morels"
-- >>>     ]
-- >>> :}
--
-- >>> import Data.Functor.Contravariant
-- >>> flipComparison (Comparison f) = Comparison (flip f)
-- >>> headingComparison = contramap fst (flipComparison defaultComparison)
--
-- >>> (frontMatter, mainMatter) = branchLeft headingComparison list
--
-- >>> frontMatter
-- ["Copyright","Preface"]
--
-- >>> import Data.List
-- >>> showContent ((_, x), ys) = x ++ ": " ++ intercalate ", " ys
--
-- >>> import Data.Tree
-- >>> putStrLn $ drawForest $ map (fmap showContent) mainMatter
-- Animals: The kingdom animalia
-- |
-- +- Vertebrates: Cats, Snakes
-- |
-- `- Invertebrates: Worms, Jellyfishes
-- <BLANKLINE>
-- Fungus: Yeast, Truffles, Morels
-- <BLANKLINE>
-- <BLANKLINE>

branchLeft :: BranchComparison a -> [Either a b] -> ([b], Forest (a, [b]))

branchLeft :: forall a b.
BranchComparison a -> [Either a b] -> ([b], Forest (a, [b]))
branchLeft BranchComparison a
c [Either a b]
xs = ([b]
rejects, Forest (a, [b])
forest)
  where
    ([b]
rejects, [(a, [b])]
nodes) = forall a b. [Either a b] -> ([b], [(a, [b])])
leadLeft [Either a b]
xs
    forest :: Forest (a, [b])
forest = forall a. BranchComparison a -> [a] -> Forest a
makeForest forall {b}. Comparison (a, b)
c' [(a, [b])]
nodes
    c' :: Comparison (a, b)
c' = forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap (\(a
x, b
_) -> a
x) BranchComparison a
c

-- | Same as 'branchLeft', but with the types flipped; here, 'Right' is the case that indicates a branch.

branchRight :: BranchComparison b -> [Either a b] -> ([a], Forest (b, [a]))

branchRight :: forall b a.
BranchComparison b -> [Either a b] -> ([a], Forest (b, [a]))
branchRight BranchComparison b
c [Either a b]
xs = ([a]
rejects, Forest (b, [a])
forest)
  where
    ([a]
rejects, [(b, [a])]
nodes) = forall a b. [Either a b] -> ([a], [(b, [a])])
leadRight [Either a b]
xs
    forest :: Forest (b, [a])
forest = forall a. BranchComparison a -> [a] -> Forest a
makeForest forall {b}. Comparison (b, b)
c' [(b, [a])]
nodes
    c' :: Comparison (b, b)
c' = forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap (\(b
x, b
_) -> b
x) BranchComparison b
c

makeForest :: BranchComparison a -> [a] -> Forest a

makeForest :: forall a. BranchComparison a -> [a] -> Forest a
makeForest BranchComparison a
c = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> [Tree a] -> [Tree a]
f []
  where
    f :: a -> [Tree a] -> [Tree a]
f a
x [Tree a]
xs = forall a. a -> [Tree a] -> Tree a
Node a
x [Tree a]
chomped forall a. a -> [a] -> [a]
: [Tree a]
remainder
      where
        ([Tree a]
chomped, [Tree a]
remainder) = forall a. (a -> Bool) -> [a] -> ([a], [a])
span (\(Node a
y [Tree a]
_) -> a
x a -> a -> Bool
> a
y) [Tree a]
xs

    a
x > :: a -> a -> Bool
> a
y = case forall a. Comparison a -> a -> a -> Ordering
getComparison BranchComparison a
c a
x a
y of Ordering
GT -> Bool
True; Ordering
_ -> Bool
False