Plus.hs

{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TypeFamilies    #-}

-- |
-- Module      : Data.List.Kleene.Star
-- Description : A list type isomorphic to 'Data.List.NonEmpty.NonEmpty' based on the Kleene plus.
-- Copyright   : (c) Donnacha Oisín Kidney, 2020
-- License     : Apache
-- Maintainer  : mail@doisinkidney.com
-- Stability   : experimental
-- Portability : ghc
--
-- This module provides a simple list type isomorphic to 'Data.List.NonEmpty.NonEmpty',
-- but which is defined in terms of a possibly empty list
-- type. This can make moving between one type and another
-- easier.

module Data.List.Kleene.Plus
  ( -- * The list type
   Plus(..)
  ,pattern (:+)
  ,pattern One
  ,Star(..)
  -- * Utility functions
  ,last
  ,filter
  ,reverse
  ,take
  ,(!!)
  -- * Building lists
  ,unfoldr
  ,iterate
  ,cycle
  -- * scans
  ,scanr
  ,scanl
  ,prescanr
  ,prescanl
  -- * Sorting
  ,sortBy
  ,sortOn
  ,sort
  )
  where

import           Data.List.Kleene.Internal

import           Data.Ord

import           Prelude                   hiding (cycle, filter, head, iterate,
                                            last, reverse, scanl, scanr, take,
                                            (!!))

-- | Sort given a comparison function. Stable.
-- \(\mathcal{O}(n \log n)\).
--
-- >>> sortBy (\x y -> compare (fst x) (fst y)) [(4,1),(3,2),(1,3),(3,4)]
-- [(1,3),(3,2),(3,4),(4,1)]
sortBy :: (a -> a -> Ordering) -> Plus a -> Plus a
sortBy cmp = treeFoldMap One (mergelr cmp)

-- | Sort given a selector function. Stable.
-- \(\mathcal{O}(n \log n)\).
--
-- >>> sortOn fst [(4,1),(3,2),(1,3),(3,4)]
-- [(1,3),(3,2),(3,4),(4,1)]
sortOn :: Ord b => (a -> b) -> Plus a -> Plus a
sortOn c = fmap fst . treeFoldMap (\x -> One (x, c x)) (mergelr (comparing snd))

-- | Stable sort.
-- \(\mathcal{O}(n \log n)\).
--
-- >>> sort [4,3,1,3]
-- [1,3,3,4]
sort :: Ord a => Plus a -> Plus a
sort = sortBy compare

-- | Return the last element of a finite list.
--
-- >>> last [1..10]
-- 10
last :: Plus a -> a
last = foldl1 (\_ x -> x)

-- | Unfold a list from a seed.
unfoldr :: (b -> (a, Maybe b)) -> b -> Plus a
unfoldr f b = x :- maybe Nil (Cons . unfoldr f) xs
  where
    (x,xs) = f b

-- | Cycle a list infinitely
--
-- >>> take 10 (cycle [1..3])
-- [1,2,3,1,2,3,1,2,3,1]
cycle :: Plus a -> Plus a
cycle xs = ys
  where
    ys = xs <> ys

-- | Repeatedly apply a function, listing its output.
--
-- >>> take 5 (iterate (+1) 1)
-- [1,2,3,4,5]
iterate :: (a -> a) -> a -> Plus a
iterate f x = x :+ iterate f (f x)

-- | Like 'scanr', but without including the initial element in the output.
--
-- >>> prescanr (+) 0 [1,2,3]
-- [6,5,3]
prescanr :: (a -> b -> b) -> b -> Plus a -> Plus b
prescanr f b (One x) = One (f x b)
prescanr f b (x :+ xs) = f x (head ys) :+ ys
  where
    ys = prescanr f b xs

-- | Like 'scanl', but without including the initial element in the output.
--
-- >>> prescanl (+) 0 [1,2,3]
-- [1,3,6]
prescanl :: (b -> a -> b) -> b -> Plus a -> Plus b
prescanl = prescanlPlus

-- | Reverse a list.
--
-- >>> reverse [1..5]
-- [5,4,3,2,1]
reverse :: Plus a -> Plus a
reverse (x :- xs) = foldl (flip (:+)) (One x) xs

-- | @'take' n xs@ takes the first @n@ elements from @xs@.
--
-- >>> take 5 [1..]
-- [1,2,3,4,5]
--
-- >>> take 5 [1..3]
-- [1,2,3]
take :: Int -> Plus a -> Star a
take = takePlus

-- | Index into a list.
--
-- >>> [0..] !! 3
-- 3
--
-- >>> [0..5] !! 6
-- *** Exception: index: empty list!
(!!) :: Plus a -> Int -> a
(!!) = indexPlus

-- $setup
-- >>> :set -XOverloadedLists