{-# LANGUAGE GADTs #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- | This module contains useful functions for working with 'Either's.
module Data.Either.Linear
  ( Either (..),
    either,
    lefts,
    rights,
    fromLeft,
    fromRight,
    partitionEithers,
  )
where

import Data.Unrestricted.Linear
import Prelude (Either (..))

-- XXX Design Notes
-- Functions like isLeft do not make sense in a linear program.
--------------------------------------------------------------------------------

-- | Linearly consume an @Either@ by applying the first linear function on a
-- value constructed with @Left@ and the second linear function on a value
-- constructed with @Right@.
either :: (a %1 -> c) -> (b %1 -> c) -> Either a b %1 -> c
either f _ (Left x) = f x
either _ g (Right y) = g y

-- | Get all the left elements in order, and consume the right ones.
lefts :: Consumable b => [Either a b] %1 -> [a]
lefts [] = []
lefts (Left a : xs) = a : lefts xs
lefts (Right b : xs) = lseq b (lefts xs)

-- | Get all the right elements in order, and consume the left ones.
rights :: Consumable a => [Either a b] %1 -> [b]
rights [] = []
rights (Left a : xs) = lseq a (rights xs)
rights (Right b : xs) = b : rights xs

-- | Get the left element of a consumable @Either@ with a default
fromLeft :: (Consumable a, Consumable b) => a %1 -> Either a b %1 -> a
fromLeft x (Left a) = lseq x a
fromLeft x (Right b) = lseq b x

-- | Get the right element of a consumable @Either@ with a default
fromRight :: (Consumable a, Consumable b) => b %1 -> Either a b %1 -> b
fromRight x (Left a) = lseq a x
fromRight x (Right b) = lseq x b

-- | Partition and consume a list of @Either@s into two lists with all the
-- lefts in one and the rights in the second, in the order they appeared in the
-- initial list.
partitionEithers :: [Either a b] %1 -> ([a], [b])
partitionEithers [] = ([], [])
partitionEithers (x : xs) = fromRecur x (partitionEithers xs)
  where
    fromRecur :: Either a b %1 -> ([a], [b]) %1 -> ([a], [b])
    fromRecur (Left a) (as, bs) = (a : as, bs)
    fromRecur (Right b) (as, bs) = (as, b : bs)