{-# LANGUAGE BangPatterns       #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFoldable     #-}
{-# LANGUAGE DeriveFunctor      #-}
{-# LANGUAGE DeriveGeneric      #-}
{-# LANGUAGE DeriveTraversable  #-}
{-# LANGUAGE PatternSynonyms    #-}
{-# LANGUAGE TypeFamilies       #-}

-- |
-- Module      : Data.List.Kleene.Internal
-- Description : Common utility functions and definitions for the kleene-list package.
-- Copyright   : (c) Donnacha Oisín Kidney, 2020
-- License     : Apache
-- Maintainer  : mail@doisinkidney.com
-- Stability   : experimental
-- Portability : ghc
--
-- = WARNING
--
-- This module is considered __internal__.
--
-- The Package Versioning Policy __does not apply__.
--
-- This contents of this module may change __in any way whatsoever__
-- and __without any warning__ between minor versions of this package.
--
-- Authors importing this module are expected to track development
-- closely.

{-# OPTIONS_HADDOCK not-home #-}
module Data.List.Kleene.Internal where

import           Control.DeepSeq      (NFData (rnf))
import           Data.Data            (Data, Typeable)
import           Data.Functor.Classes
import           GHC.Generics         (Generic)

import           GHC.Exts             (IsList)
import qualified GHC.Exts

import           Control.Applicative
import           Control.Monad
import           Control.Monad.Fix
import           Control.Monad.Zip
import           Data.Foldable

import           Prelude              hiding (filter, head, scanl, scanr, tail)

-- | A list, based on the Kleene star.
-- This type is isomorphic to Haskell's standard @[]@ type, so it can be used
-- in the same way.
data Star a
  = Nil
  | Cons (Plus a)
  deriving (Eq, Ord, Generic, Data, Typeable, Functor, Traversable)

infixr 5 :-
-- | A non-empty list type, based on the Kleene plus.
-- This type is isomorphic to 'Data.List.NonEmpty.NonEmpty' type, so it
-- can be used in the same way.
data Plus a
  = (:-)
  { head :: a
  , tail :: Star a
  } deriving (Eq, Ord, Generic, Data, Typeable, Functor, Traversable)

instance Foldable Star where
    foldr _ b Nil       = b
    foldr f b (Cons xs) = foldr f b xs

    foldl _ b Nil       = b
    foldl f b (Cons xs) = foldl f b xs

    foldl' _ !b Nil       = b
    foldl' f !b (Cons xs) = foldl' f b xs

    foldl1 _ Nil       = errorWithoutStackTrace "foldl1: empty list"
    foldl1 f (Cons xs) = foldl1 f xs

    foldr1 _ Nil       = errorWithoutStackTrace "foldr1: empty list"
    foldr1 f (Cons xs) = foldr1 f xs

    foldMap _ Nil       = mempty
    foldMap f (Cons xs) = foldMap f xs

    minimum Nil       = errorWithoutStackTrace "minimum: empty list"
    minimum (Cons xs) = minimum xs

    maximum Nil       = errorWithoutStackTrace "maximum: empty list"
    maximum (Cons xs) = maximum xs

instance Foldable Plus where
    foldr f b ~(x :- xs) = f x (foldr f b xs)

    foldl f b ~(x :- xs) = foldl f (f b x) xs

    foldl' f !b ~(x :- xs) = foldl' f (f b x) xs

    foldl1 f ~(x :- xs) = foldl f x xs

    foldr1 f = go
      where
        go (x :- xs) = case xs of
          Nil     -> x
          Cons ys -> f x (go ys)

    foldMap f ~(x :- xs) = f x <> foldMap f xs

    null _ = False

    minimum = foldr1 min
    maximum = foldr1 max

instance Eq1 Star where
  liftEq _ Nil Nil              = True
  liftEq eq (Cons xs) (Cons ys) = liftEq eq xs ys
  liftEq _ _ _                  = False

instance Eq1 Plus where
  liftEq eq ~(x :- xs) (y :- ys) = eq x y && liftEq eq xs ys

instance Ord1 Star where
  liftCompare _ Nil Nil             = EQ
  liftCompare _ Nil (Cons _)        = LT
  liftCompare _ (Cons _) Nil        = GT
  liftCompare c (Cons xs) (Cons ys) = liftCompare c xs ys

instance Ord1 Plus where
  liftCompare c ~(x :- xs) ~(y :- ys) = c x y <> liftCompare c xs ys

instance Show1 Plus where
  liftShowsPrec _ sp _ = sp . foldr (:) []

instance Show1 Star where
  liftShowsPrec _ sp _ = sp . foldr (:) []

-- | A pattern for building up star lists as cons-lists.
--
-- >>> 1 :* 2 :* 3 :* Nil
-- [1,2,3]
infixr 5 :*
pattern (:*) :: a -> Star a -> Star a
pattern (:*) x xs = Cons (x :- xs)
{-# COMPLETE (:*), Nil #-}

-- | A pattern for building up plus lists as cons-lists.
--
-- >>> 1 :+ 2 :+ One 3
-- [1,2,3]
infixr 5 :+
pattern (:+) :: a -> Plus a -> Plus a
pattern (:+) x xs = x :- Cons xs

-- | A pattern for a singleton plus list.
pattern One :: a -> Plus a
pattern One x = x :- Nil
{-# COMPLETE (:+), One #-}

instance IsList (Star a) where
  type Item (Star a) = a
  fromList = foldr (:*) Nil
  toList = foldr (:) []

instance IsList (Plus a) where
  type Item (Plus a) = a
  fromList []     = errorWithoutStackTrace "Cannot make plus from empty list"
  fromList (x:xs) = x :- GHC.Exts.fromList xs
  toList = foldr (:) []

instance Show a => Show (Star a) where
  showsPrec n = showsPrec n . toList

instance Show a => Show (Plus a) where
  showsPrec n = showsPrec n . toList

instance NFData a => NFData (Star a) where
  rnf Nil       = ()
  rnf (Cons xs) = rnf xs

instance NFData a => NFData (Plus a) where
  rnf (x :- xs) = rnf x `seq` rnf xs

instance Semigroup (Plus a) where
  ~(x :- xs) <> ys = x :+ (xs *<>+ ys)

(*<>+) :: Star a -> Plus a -> Plus a
Nil     *<>+ ys = ys
Cons xs *<>+ ys = xs <> ys

instance Semigroup (Star a) where
  Nil     <> ys = ys
  Cons xs <> ys = Cons (xs +<>* ys)

(+<>*) :: Plus a -> Star a -> Plus a
~(x :- xs) +<>* ys = x :- (xs <> ys)

instance Monoid (Star a) where
  mempty = Nil

instance Applicative Star where
  pure = Cons . pure

  Nil     <*> _  = Nil
  f :* fs <*> xs = foldr ((:*) . f) (fs <*> xs) xs

  liftA2 _ Nil       _  = Nil
  liftA2 f (x :* xs) ys = foldr ((:*) . f x) (liftA2 f xs ys) ys

-- |
-- >>> (,) <$> (1 :+ 2 :+ One 3) <*> ('a' :+ 'b' :+ One 'c')
-- [(1,'a'),(1,'b'),(1,'c'),(2,'a'),(2,'b'),(2,'c'),(3,'a'),(3,'b'),(3,'c')]
--
-- >>> liftA2 (,) (1 :+ 2 :+ One 3) ('a' :+ 'b' :+ One 'c')
-- [(1,'a'),(1,'b'),(1,'c'),(2,'a'),(2,'b'),(2,'c'),(3,'a'),(3,'b'),(3,'c')]
instance Applicative Plus where
  pure = One

  ~(f' :- fs') <*> xs = f' (head xs) :- foldr ((:*) . f') (go fs') (tail xs)
    where
      go Nil       = Nil
      go (f :* fs) = foldr ((:*) . f) (go fs) xs

  liftA2 f ~(x' :- xs') ys = f x' (head ys) :- foldr ((:*) . f x') (go xs') (tail ys)
    where
      go Nil       = Nil
      go (x :* xs) = foldr ((:*) . f x) (go xs) ys

instance Monad Star where
    xs >>= f = foldr ((<>) . f) Nil xs

instance Monad Plus where
    ~(x :- xs) >>= f = f x +<>* go xs
      where
        go Nil       = Nil
        go (Cons ys) = Cons (ys >>= f)

instance Alternative Star where
    (<|>) = (<>)
    empty = Nil

instance MonadPlus Star

instance MonadFix Plus where
  mfix f = case fix (f . head) of
             ~(x :- _) -> x :- mfix (tail . f)

instance MonadFix Star where
  mfix f = case fix (f . head . unStar) of
             Nil      -> Nil
             (x :* _) -> x :* mfix (tail . unStar . f)
    where
      unStar ~(Cons xs) = xs

instance MonadZip Plus where
    mzip ~(x :- xs) ~(y :- ys) = (x, y) :- mzip xs ys

    mzipWith f ~(x :- xs) ~(y :- ys) = f x y :- mzipWith f xs ys

    munzip ~(~(y,z) :- xs) = (y :- ys, z :- zs)
      where
        ~(ys,zs) = munzip xs

instance MonadZip Star where
    mzip Nil _               = Nil
    mzip _ Nil               = Nil
    mzip (Cons xs) (Cons ys) = Cons (mzip xs ys)

    mzipWith _ Nil _               = Nil
    mzipWith _ _ Nil               = Nil
    mzipWith f (Cons xs) (Cons ys) = Cons (mzipWith f xs ys)

    munzip Nil = (Nil, Nil)
    munzip (Cons xs) = (Cons ys, Cons zs)
      where
        ~(ys,zs) = munzip xs

merge :: (a -> a -> Ordering) -> Star a -> Star a -> Star a
merge _ Nil       ys   = ys
merge cmp (Cons xs) ys = Cons (mergel cmp xs ys)

mergel :: (a -> a -> Ordering) -> Plus a -> Star a -> Plus a
mergel _ xs Nil         = xs
mergel cmp xs (Cons ys) = mergelr cmp xs ys

merger :: (a -> a -> Ordering) -> Star a -> Plus a -> Plus a
merger _ Nil ys         = ys
merger cmp (Cons xs) ys = mergelr cmp xs ys

mergelr :: (a -> a -> Ordering) -> Plus a -> Plus a -> Plus a
mergelr cmp xss@ ~(x :- xs) yss@ ~(y :- ys) = case cmp x y of
  LT -> x :+ merger cmp xs yss
  EQ -> x :+ y :- merge cmp xs ys
  GT -> y :+ mergel cmp xss ys

treeFoldMap :: (a -> b) -> (b -> b -> b) -> Plus a -> b
treeFoldMap c f = go
  where
    go (One x)        = c x
    go (x :+ y :- xs) = go' (f (c x) (c y) :- pairMap xs)

    pairMap (x1 :* x2 :* xs) = f (c x1) (c x2) :* pairMap xs
    pairMap (x1 :* Nil)      = c x1 :* Nil
    pairMap Nil              = Nil

    go' (One x)        = x
    go' (x :+ y :- xs) = go' (f x y :- pairMap' xs)

    pairMap' (x1 :* x2 :* xs) = f x1 x2 :* pairMap' xs
    pairMap' xs               = xs

-- |
-- >>> prescanlPlus (+) 0 [1,2,3]
-- [1,3,6]
prescanlPlus :: (b -> a -> b) -> b -> Plus a -> Plus b
prescanlPlus f b (x :- xs) = scanl f (f b x) xs

-- |
-- >>> prescanlStar (+) 0 [1,2,3]
-- [1,3,6]
prescanlStar :: (b -> a -> b) -> b -> Star a -> Star b
prescanlStar _ _ Nil       = Nil
prescanlStar f b (Cons xs) = Cons (prescanlPlus f b xs)

-- | Functions the same as 'Data.List.scanl' in "Data.List".
--
-- >>> scanl (+) 0 [1,2,3]
-- [0,1,3,6]
scanl :: (b -> a -> b) -> b -> Star a -> Plus b
scanl f b xs = b :- prescanlStar f b xs

-- | Functions the same as 'Data.List.scanr' in "Data.List".
--
-- >>> scanr (+) 0 ([1,2,3] :: Star Int)
-- [6,5,3,0]
scanr :: Foldable f => (a -> b -> b) -> b -> f a -> Plus b
scanr f b = foldr (\x xs -> f x (head xs) :+ xs) (One b)

-- | Functions the same as 'Data.List.filter' in "Data.List".
--
-- >>> filter even ([1..5] :: Star Int)
-- [2,4]
filter :: Foldable f => (a -> Bool) -> f a -> Star a
filter p = foldr f Nil
  where
    f x xs
      | p x = x :* xs
      | otherwise = xs

takeStar :: Int -> Star a -> Star a
takeStar _ Nil       = Nil
takeStar i (Cons xs) = takePlus i xs

takePlus :: Int -> Plus a -> Star a
takePlus 0 _          = Nil
takePlus i ~(x :- xs) = x :* takeStar (i-1) xs

indexPlus :: Plus a -> Int -> a
indexPlus xs 0 = head xs
indexPlus xs i = indexStar (tail xs) (i-1)

indexStar :: Star a -> Int -> a
indexStar Nil _       = errorWithoutStackTrace "index: empty list!"
indexStar (Cons xs) i = indexPlus xs i


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