{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

module Data.RandomAccessList.Class
  ( RandomAccessList (..)
  , Data.RandomAccessList.Class.head
  , Data.RandomAccessList.Class.tail
  , AsRAL (..)
  ) where

import Data.Kind
import Data.List qualified as List
#if MIN_VERSION_base(4,19,0)
-- Avoid a compiler warning about unused package `extra`.
import Data.List.Extra ()
#else
import Data.List.Extra qualified as List
#endif
import Data.Maybe (fromJust, fromMaybe)
import Data.RAList qualified as RAL
import Data.Vector.NonEmpty qualified as NEV
import Data.Word
import GHC.Exts

{-| Typeclass for various types implementing the "signature" of a random-access list.

A true random-access list should have good asymptotic behaviour for these methods also,
but for convenience we also provide implementations for e.g. '[a]', which has bad
lookup performance. -}
class RandomAccessList e where
  -- | The type of elements in the list.
  type Element e :: Type

  -- | The empty list.
  empty :: e

  -- | Prepend an element to the list.
  cons :: Element e -> e -> e

  -- | Un-prepend an element to the list.
  uncons :: e -> Maybe (Element e, e)

  -- | Get the length of the list. May have linear complexity, but useful.
  length :: e -> Word64

  {-# INLINEABLE consSlab #-}

  {-| Prepend many elements to the list. Has a default implementation, but
  implementations can provide more efficient ones. -}
  consSlab :: NEV.NonEmptyVector (Element e) -> e -> e
  consSlab NonEmptyVector (Element e)
vec e
e = (Element e -> e -> e) -> e -> NonEmptyVector (Element e) -> e
forall a b. (a -> b -> b) -> b -> NonEmptyVector a -> b
NEV.foldr Element e -> e -> e
forall e. RandomAccessList e => Element e -> e -> e
cons e
e NonEmptyVector (Element e)
vec

  {-# INLINEABLE indexZero #-}

  -- | Lookup an element in the list. 0-based index.
  indexZero :: e -> Word64 -> Maybe (Element e)
  indexZero e
e Word64
i = e -> Word64 -> Maybe (Element e)
forall e. RandomAccessList e => e -> Word64 -> Maybe (Element e)
indexOne e
e (Word64
i Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
1)

  {-# INLINEABLE indexOne #-}

  -- | Lookup an element in the list. 1-based index.
  indexOne :: e -> Word64 -> Maybe (Element e)
  indexOne e
_ Word64
0 = Maybe (Element e)
forall a. Maybe a
Nothing
  indexOne e
e Word64
i = e -> Word64 -> Maybe (Element e)
forall e. RandomAccessList e => e -> Word64 -> Maybe (Element e)
indexZero e
e (Word64
i Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1)

  {-# INLINEABLE unsafeIndexZero #-}

  -- | Lookup an element in the list, partially.
  unsafeIndexZero :: e -> Word64 -> Element e
  unsafeIndexZero e
e = Maybe (Element e) -> Element e
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Element e) -> Element e)
-> (Word64 -> Maybe (Element e)) -> Word64 -> Element e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> Word64 -> Maybe (Element e)
forall e. RandomAccessList e => e -> Word64 -> Maybe (Element e)
indexZero e
e

  {-# INLINEABLE unsafeIndexOne #-}

  -- | Lookup an element in the list, partially.
  unsafeIndexOne :: e -> Word64 -> Element e
  unsafeIndexOne e
e = Maybe (Element e) -> Element e
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Element e) -> Element e)
-> (Word64 -> Maybe (Element e)) -> Word64 -> Element e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> Word64 -> Maybe (Element e)
forall e. RandomAccessList e => e -> Word64 -> Maybe (Element e)
indexOne e
e

-- O(1) worst-case
head :: (RandomAccessList e, a ~ Element e) => e -> a
head :: forall e a. (RandomAccessList e, a ~ Element e) => e -> a
head = (a, e) -> a
forall a b. (a, b) -> a
fst ((a, e) -> a) -> (e -> (a, e)) -> e -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, e) -> Maybe (a, e) -> (a, e)
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> (a, e)
forall a. HasCallStack => [Char] -> a
error [Char]
"empty list") (Maybe (a, e) -> (a, e)) -> (e -> Maybe (a, e)) -> e -> (a, e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> Maybe (a, e)
e -> Maybe (Element e, e)
forall e. RandomAccessList e => e -> Maybe (Element e, e)
uncons
{-# INLINEABLE head #-}

-- O(1) worst-case
tail :: RandomAccessList e => e -> e
tail :: forall e. RandomAccessList e => e -> e
tail = (Element e, e) -> e
forall a b. (a, b) -> b
snd ((Element e, e) -> e) -> (e -> (Element e, e)) -> e -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Element e, e) -> Maybe (Element e, e) -> (Element e, e)
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> (Element e, e)
forall a. HasCallStack => [Char] -> a
error [Char]
"empty list") (Maybe (Element e, e) -> (Element e, e))
-> (e -> Maybe (Element e, e)) -> e -> (Element e, e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> Maybe (Element e, e)
forall e. RandomAccessList e => e -> Maybe (Element e, e)
uncons
{-# INLINEABLE tail #-}

instance RandomAccessList [a] where
  type Element [a] = a

  {-# INLINEABLE empty #-}
  empty :: [a]
empty = []
  {-# INLINEABLE cons #-}
  cons :: Element [a] -> [a] -> [a]
cons = (:)
  {-# INLINEABLE uncons #-}
  uncons :: [a] -> Maybe (Element [a], [a])
uncons = [a] -> Maybe (a, [a])
[a] -> Maybe (Element [a], [a])
forall a. [a] -> Maybe (a, [a])
List.uncons
  {-# INLINEABLE length #-}
  length :: [a] -> Word64
length = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word64) -> ([a] -> Int) -> [a] -> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length
  {-# INLINEABLE indexZero #-}
  indexZero :: [a] -> Word64 -> Maybe (Element [a])
indexZero [a]
l Word64
w = [a]
l [a] -> Int -> Maybe a
forall a. [a] -> Int -> Maybe a
List.!? Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
w

instance RandomAccessList (RAL.RAList a) where
  type Element (RAL.RAList a) = a

  {-# INLINEABLE empty #-}
  empty :: RAList a
empty = RAList a
forall a. Monoid a => a
mempty
  {-# INLINEABLE cons #-}
  cons :: Element (RAList a) -> RAList a -> RAList a
cons = a -> RAList a -> RAList a
Element (RAList a) -> RAList a -> RAList a
forall a. a -> RAList a -> RAList a
RAL.cons
  {-# INLINEABLE uncons #-}
  uncons :: RAList a -> Maybe (Element (RAList a), RAList a)
uncons = RAList a -> Maybe (a, RAList a)
RAList a -> Maybe (Element (RAList a), RAList a)
forall a. RAList a -> Maybe (a, RAList a)
RAL.uncons
  {-# INLINEABLE length #-}
  length :: RAList a -> Word64
length = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word64) -> (RAList a -> Int) -> RAList a -> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RAList a -> Int
forall a. RAList a -> Int
RAL.length
  {-# INLINEABLE indexZero #-}
  indexZero :: RAList a -> Word64 -> Maybe (Element (RAList a))
indexZero RAList a
l Word64
w = RAList a
l RAList a -> Int -> Maybe a
forall a. RAList a -> Int -> Maybe a
RAL.!? Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
w

newtype AsRAL a = AsRAL a

instance RandomAccessList e => IsList (AsRAL e) where
  type Item (AsRAL e) = Element e
  fromList :: [Item (AsRAL e)] -> AsRAL e
fromList [Item (AsRAL e)]
l = e -> AsRAL e
forall a. a -> AsRAL a
AsRAL (e -> AsRAL e) -> e -> AsRAL e
forall a b. (a -> b) -> a -> b
$ (Element e -> e -> e) -> e -> [Element e] -> e
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Element e -> e -> e
forall e. RandomAccessList e => Element e -> e -> e
cons e
forall e. RandomAccessList e => e
empty [Item (AsRAL e)]
[Element e]
l
  toList :: AsRAL e -> [Item (AsRAL e)]
toList (AsRAL e
e) = (e -> Maybe (Element e, e)) -> e -> [Element e]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
List.unfoldr e -> Maybe (Element e, e)
forall e. RandomAccessList e => e -> Maybe (Element e, e)
uncons e
e