-- Copyright 2019-2021 Google LLC
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--      http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ViewPatterns #-}

-- | A data type of run-length-encoded lists.
--
-- This module is meant to be imported qualified with the exception of the type
-- RLE itself.  It exports names that clash with things in Prelude and many
-- other data structure modules.

module Data.RLE
         ( -- * Run-Length Encoded Lists
           RLE
         , toList, fromList, singleton, empty, cons, uncons
         , reverse, splitAt, take, init, null, length, (++)
         , map, mapInvertible, traverse, zipWith
           -- ** Runs
         , Run(..), toRuns, fromRuns, consRun, unconsRun, runs
         ) where

import Prelude hiding
         ( (++), init, length, map, null, reverse
         , splitAt, take, traverse, zipWith
         )
import qualified Prelude as P

import Control.Applicative (Applicative(..))
import Control.Monad (replicateM)
import Data.Coerce (coerce)
import Data.Functor.Contravariant (Contravariant(..))
import Data.Maybe (fromJust)
import Data.Semigroup (Semigroup(stimes))
import Data.Void (absurd)
import GHC.Exts (IsList, IsString(..))
import qualified GHC.Exts (IsList(..))
import GHC.Generics (Generic)
import GHC.Stack (HasCallStack)

import Control.DeepSeq (NFData)
import Data.Portray (Portray(..), Portrayal(..))
import Data.Portray.Diff (Diff(..))
import Data.Serialize (Serialize)
import Data.Wrapped (Wrapped(..))

infixr 5 :><
-- | @n :>< x@ denotes a sequence of @n@ copies of @x@, as part of an 'RLE'.
data Run a = Int :>< a
  deriving stock (Run a -> Run a -> Bool
(Run a -> Run a -> Bool) -> (Run a -> Run a -> Bool) -> Eq (Run a)
forall a. Eq a => Run a -> Run a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Run a -> Run a -> Bool
$c/= :: forall a. Eq a => Run a -> Run a -> Bool
== :: Run a -> Run a -> Bool
$c== :: forall a. Eq a => Run a -> Run a -> Bool
Eq, Int -> Run a -> ShowS
[Run a] -> ShowS
Run a -> String
(Int -> Run a -> ShowS)
-> (Run a -> String) -> ([Run a] -> ShowS) -> Show (Run a)
forall a. Show a => Int -> Run a -> ShowS
forall a. Show a => [Run a] -> ShowS
forall a. Show a => Run a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Run a] -> ShowS
$cshowList :: forall a. Show a => [Run a] -> ShowS
show :: Run a -> String
$cshow :: forall a. Show a => Run a -> String
showsPrec :: Int -> Run a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Run a -> ShowS
Show, (forall x. Run a -> Rep (Run a) x)
-> (forall x. Rep (Run a) x -> Run a) -> Generic (Run a)
forall x. Rep (Run a) x -> Run a
forall x. Run a -> Rep (Run a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (Run a) x -> Run a
forall a x. Run a -> Rep (Run a) x
$cto :: forall a x. Rep (Run a) x -> Run a
$cfrom :: forall a x. Run a -> Rep (Run a) x
Generic, a -> Run b -> Run a
(a -> b) -> Run a -> Run b
(forall a b. (a -> b) -> Run a -> Run b)
-> (forall a b. a -> Run b -> Run a) -> Functor Run
forall a b. a -> Run b -> Run a
forall a b. (a -> b) -> Run a -> Run b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> Run b -> Run a
$c<$ :: forall a b. a -> Run b -> Run a
fmap :: (a -> b) -> Run a -> Run b
$cfmap :: forall a b. (a -> b) -> Run a -> Run b
Functor)
  deriving anyclass (Run a -> ()
(Run a -> ()) -> NFData (Run a)
forall a. NFData a => Run a -> ()
forall a. (a -> ()) -> NFData a
rnf :: Run a -> ()
$crnf :: forall a. NFData a => Run a -> ()
NFData, Get (Run a)
Putter (Run a)
Putter (Run a) -> Get (Run a) -> Serialize (Run a)
forall a. Serialize a => Get (Run a)
forall a. Serialize a => Putter (Run a)
forall t. Putter t -> Get t -> Serialize t
get :: Get (Run a)
$cget :: forall a. Serialize a => Get (Run a)
put :: Putter (Run a)
$cput :: forall a. Serialize a => Putter (Run a)
Serialize)
  deriving ([Run a] -> Portrayal
Run a -> Portrayal
(Run a -> Portrayal) -> ([Run a] -> Portrayal) -> Portray (Run a)
forall a. Portray a => [Run a] -> Portrayal
forall a. Portray a => Run a -> Portrayal
forall a. (a -> Portrayal) -> ([a] -> Portrayal) -> Portray a
portrayList :: [Run a] -> Portrayal
$cportrayList :: forall a. Portray a => [Run a] -> Portrayal
portray :: Run a -> Portrayal
$cportray :: forall a. Portray a => Run a -> Portrayal
Portray, Run a -> Run a -> Maybe Portrayal
(Run a -> Run a -> Maybe Portrayal) -> Diff (Run a)
forall a. Diff a => Run a -> Run a -> Maybe Portrayal
forall a. (a -> a -> Maybe Portrayal) -> Diff a
diff :: Run a -> Run a -> Maybe Portrayal
$cdiff :: forall a. Diff a => Run a -> Run a -> Maybe Portrayal
Diff) via Wrapped Generic (Run a)

instance Foldable Run where foldMap :: (a -> m) -> Run a -> m
foldMap a -> m
f (Int
n :>< a
x) = Int -> m -> m
forall a b. (Semigroup a, Integral b) => b -> a -> a
stimes Int
n (a -> m
f a
x)

-- | After all, why not?
--
-- This is basically Writer (Product Int).
instance Applicative Run where
  pure :: a -> Run a
pure = (Int
1 Int -> a -> Run a
forall a. Int -> a -> Run a
:><)
  liftA2 :: (a -> b -> c) -> Run a -> Run b -> Run c
liftA2 a -> b -> c
f (Int
m :>< a
x) (Int
n :>< b
y) = Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
n Int -> c -> Run c
forall a. Int -> a -> Run a
:>< a -> b -> c
f a
x b
y
  (Int
m :>< a -> b
f) <*> :: Run (a -> b) -> Run a -> Run b
<*> (Int
n :>< a
x) = Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
n Int -> b -> Run b
forall a. Int -> a -> Run a
:>< a -> b
f a
x

instance Monad Run where (Int
m :>< a
x) >>= :: Run a -> (a -> Run b) -> Run b
>>= a -> Run b
f = case a -> Run b
f a
x of Int
n :>< b
y-> Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
n Int -> b -> Run b
forall a. Int -> a -> Run a
:>< b
y

-- Invariant: 'RLE' never contains two adjacent entries with equal @a@ values.
-- Invariant: 'RLE' never contains zero-length runs.
--
-- These two together ensure we can use generated Eq/Ord instances and can
-- implement certain functions faster by omitting tests for zero or duplicated
-- runs.

-- | A run-length encoded representation of a @[a]@.
--
-- This doesn't have a 'Functor' or 'Traversable' instance because it would
-- need an 'Eq' constraint on the element type to uphold invariants, but there
-- are 'map' and 'traverse' functions exported.
newtype RLE a = RLE
  { RLE a -> [Run a]
toRuns :: [Run a]
    -- ^ Extract the contents of an 'RLE' as a list of runs.
    --
    -- This is not a retraction of 'fromRuns': @toRuns . fromRuns@ merges
    -- adjacent runs of equal values and eliminates empty runs.
  }
  deriving stock (RLE a -> RLE a -> Bool
(RLE a -> RLE a -> Bool) -> (RLE a -> RLE a -> Bool) -> Eq (RLE a)
forall a. Eq a => RLE a -> RLE a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: RLE a -> RLE a -> Bool
$c/= :: forall a. Eq a => RLE a -> RLE a -> Bool
== :: RLE a -> RLE a -> Bool
$c== :: forall a. Eq a => RLE a -> RLE a -> Bool
Eq, Int -> RLE a -> ShowS
[RLE a] -> ShowS
RLE a -> String
(Int -> RLE a -> ShowS)
-> (RLE a -> String) -> ([RLE a] -> ShowS) -> Show (RLE a)
forall a. Show a => Int -> RLE a -> ShowS
forall a. Show a => [RLE a] -> ShowS
forall a. Show a => RLE a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [RLE a] -> ShowS
$cshowList :: forall a. Show a => [RLE a] -> ShowS
show :: RLE a -> String
$cshow :: forall a. Show a => RLE a -> String
showsPrec :: Int -> RLE a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> RLE a -> ShowS
Show, (forall x. RLE a -> Rep (RLE a) x)
-> (forall x. Rep (RLE a) x -> RLE a) -> Generic (RLE a)
forall x. Rep (RLE a) x -> RLE a
forall x. RLE a -> Rep (RLE a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (RLE a) x -> RLE a
forall a x. RLE a -> Rep (RLE a) x
$cto :: forall a x. Rep (RLE a) x -> RLE a
$cfrom :: forall a x. RLE a -> Rep (RLE a) x
Generic, RLE a -> Bool
(a -> m) -> RLE a -> m
(a -> b -> b) -> b -> RLE a -> b
(forall m. Monoid m => RLE m -> m)
-> (forall m a. Monoid m => (a -> m) -> RLE a -> m)
-> (forall m a. Monoid m => (a -> m) -> RLE a -> m)
-> (forall a b. (a -> b -> b) -> b -> RLE a -> b)
-> (forall a b. (a -> b -> b) -> b -> RLE a -> b)
-> (forall b a. (b -> a -> b) -> b -> RLE a -> b)
-> (forall b a. (b -> a -> b) -> b -> RLE a -> b)
-> (forall a. (a -> a -> a) -> RLE a -> a)
-> (forall a. (a -> a -> a) -> RLE a -> a)
-> (forall a. RLE a -> [a])
-> (forall a. RLE a -> Bool)
-> (forall a. RLE a -> Int)
-> (forall a. Eq a => a -> RLE a -> Bool)
-> (forall a. Ord a => RLE a -> a)
-> (forall a. Ord a => RLE a -> a)
-> (forall a. Num a => RLE a -> a)
-> (forall a. Num a => RLE a -> a)
-> Foldable RLE
forall a. Eq a => a -> RLE a -> Bool
forall a. Num a => RLE a -> a
forall a. Ord a => RLE a -> a
forall m. Monoid m => RLE m -> m
forall a. RLE a -> Bool
forall a. RLE a -> Int
forall a. RLE a -> [a]
forall a. (a -> a -> a) -> RLE a -> a
forall m a. Monoid m => (a -> m) -> RLE a -> m
forall b a. (b -> a -> b) -> b -> RLE a -> b
forall a b. (a -> b -> b) -> b -> RLE a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: RLE a -> a
$cproduct :: forall a. Num a => RLE a -> a
sum :: RLE a -> a
$csum :: forall a. Num a => RLE a -> a
minimum :: RLE a -> a
$cminimum :: forall a. Ord a => RLE a -> a
maximum :: RLE a -> a
$cmaximum :: forall a. Ord a => RLE a -> a
elem :: a -> RLE a -> Bool
$celem :: forall a. Eq a => a -> RLE a -> Bool
length :: RLE a -> Int
$clength :: forall a. RLE a -> Int
null :: RLE a -> Bool
$cnull :: forall a. RLE a -> Bool
toList :: RLE a -> [a]
$ctoList :: forall a. RLE a -> [a]
foldl1 :: (a -> a -> a) -> RLE a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> RLE a -> a
foldr1 :: (a -> a -> a) -> RLE a -> a
$cfoldr1 :: forall a. (a -> a -> a) -> RLE a -> a
foldl' :: (b -> a -> b) -> b -> RLE a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> RLE a -> b
foldl :: (b -> a -> b) -> b -> RLE a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> RLE a -> b
foldr' :: (a -> b -> b) -> b -> RLE a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> RLE a -> b
foldr :: (a -> b -> b) -> b -> RLE a -> b
$cfoldr :: forall a b. (a -> b -> b) -> b -> RLE a -> b
foldMap' :: (a -> m) -> RLE a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> RLE a -> m
foldMap :: (a -> m) -> RLE a -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> RLE a -> m
fold :: RLE m -> m
$cfold :: forall m. Monoid m => RLE m -> m
Foldable)
  deriving anyclass (RLE a -> ()
(RLE a -> ()) -> NFData (RLE a)
forall a. NFData a => RLE a -> ()
forall a. (a -> ()) -> NFData a
rnf :: RLE a -> ()
$crnf :: forall a. NFData a => RLE a -> ()
NFData, Get (RLE a)
Putter (RLE a)
Putter (RLE a) -> Get (RLE a) -> Serialize (RLE a)
forall a. Serialize a => Get (RLE a)
forall a. Serialize a => Putter (RLE a)
forall t. Putter t -> Get t -> Serialize t
get :: Get (RLE a)
$cget :: forall a. Serialize a => Get (RLE a)
put :: Putter (RLE a)
$cput :: forall a. Serialize a => Putter (RLE a)
Serialize)

instance Portray a => Portray (RLE a) where
  portray :: RLE a -> Portrayal
portray RLE a
rle = Portrayal -> [Portrayal] -> Portrayal
Apply (Ident -> Portrayal
Name Ident
"fromRuns") [[Portrayal] -> Portrayal
List ([Portrayal] -> Portrayal) -> [Portrayal] -> Portrayal
forall a b. (a -> b) -> a -> b
$ Run a -> Portrayal
forall a. Portray a => a -> Portrayal
portray (Run a -> Portrayal) -> [Run a] -> [Portrayal]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RLE a -> [Run a]
forall a. RLE a -> [Run a]
toRuns RLE a
rle]

instance (Portray a, Diff a) => Diff (RLE a) where
  diff :: RLE a -> RLE a -> Maybe Portrayal
diff RLE a
x RLE a
y = Portrayal -> [Portrayal] -> Portrayal
Apply (Ident -> Portrayal
Name Ident
"fromRuns") ([Portrayal] -> Portrayal)
-> (Portrayal -> [Portrayal]) -> Portrayal -> Portrayal
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Portrayal -> [Portrayal]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Portrayal -> Portrayal) -> Maybe Portrayal -> Maybe Portrayal
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Run a] -> [Run a] -> Maybe Portrayal
forall a. Diff a => a -> a -> Maybe Portrayal
diff (RLE a -> [Run a]
forall a. RLE a -> [Run a]
toRuns RLE a
x) (RLE a -> [Run a]
forall a. RLE a -> [Run a]
toRuns RLE a
y)

instance Eq a => IsList (RLE a) where
  type Item (RLE a) = a
  fromList :: [Item (RLE a)] -> RLE a
fromList = [Item (RLE a)] -> RLE a
forall a. Eq a => [a] -> RLE a
fromList
  toList :: RLE a -> [Item (RLE a)]
toList = RLE a -> [Item (RLE a)]
forall a. RLE a -> [a]
toList

instance a ~ Char => IsString (RLE a) where
  fromString :: String -> RLE a
fromString = String -> RLE a
forall a. Eq a => [a] -> RLE a
fromList

instance Eq a => Semigroup (RLE a) where
  <> :: RLE a -> RLE a -> RLE a
(<>) = RLE a -> RLE a -> RLE a
forall a. Eq a => RLE a -> RLE a -> RLE a
(++)

  stimes :: b -> RLE a -> RLE a
stimes b
0  RLE a
_               = RLE a
forall a. RLE a
empty
  stimes b
_  (RLE [])        = RLE a
forall a. RLE a
empty
  stimes b
n  (RLE [Int
nx :>< a
x]) = [Run a] -> RLE a
forall a. [Run a] -> RLE a
RLE [(b -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral b
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nx Int -> a -> Run a
forall a. Int -> a -> Run a
:>< a
x)]
  stimes b
n0 (RLE (Run a
r0:[Run a]
rs0))  = [Run a] -> RLE a
forall a. [Run a] -> RLE a
RLE ([Run a] -> RLE a) -> [Run a] -> RLE a
forall a b. (a -> b) -> a -> b
$ b -> [Run a] -> [Run a]
go (b
n0 b -> b -> b
forall a. Num a => a -> a -> a
- b
1) [Run a]
rs0
   where
    adjustedCycle :: [Run a]
adjustedCycle = RLE a -> [Run a]
forall a. RLE a -> [Run a]
toRuns (RLE a -> [Run a]) -> RLE a -> [Run a]
forall a b. (a -> b) -> a -> b
$ [Run a] -> RLE a
forall a. [Run a] -> RLE a
RLE [Run a]
rs0 RLE a -> RLE a -> RLE a
forall a. Eq a => RLE a -> RLE a -> RLE a
++ [Run a] -> RLE a
forall a. [Run a] -> RLE a
RLE [Run a
r0]

    go :: b -> [Run a] -> [Run a]
go b
0 [Run a]
rs = Run a
r0Run a -> [Run a] -> [Run a]
forall a. a -> [a] -> [a]
:[Run a]
rs
    go b
n [Run a]
rs = b -> [Run a] -> [Run a]
go (b
nb -> b -> b
forall a. Num a => a -> a -> a
-b
1) ([Run a]
adjustedCycle [Run a] -> [Run a] -> [Run a]
forall a. [a] -> [a] -> [a]
P.++ [Run a]
rs)

instance Eq a => Monoid (RLE a) where
  mempty :: RLE a
mempty = RLE a
forall a. RLE a
empty

-- | An empty 'RLE'.
empty :: RLE a
empty :: RLE a
empty = [Run a] -> RLE a
forall a. [Run a] -> RLE a
RLE []

-- | Returns 'True' iff the argument contains no elements.
null :: RLE a -> Bool
null :: RLE a -> Bool
null = [Run a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
P.null ([Run a] -> Bool) -> (RLE a -> [Run a]) -> RLE a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RLE a -> [Run a]
forall a. RLE a -> [Run a]
toRuns

-- | 'Data.Foldable.length' specialized to 'RLE'.
length :: RLE a -> Int
length :: RLE a -> Int
length (RLE [Run a]
rs0) = [Run a] -> Int
forall a. [Run a] -> Int
go [Run a]
rs0
 where
  go :: [Run a] -> Int
go [] = Int
0
  go ((Int
n :>< a
_) : [Run a]
rs) = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Run a] -> Int
go [Run a]
rs

-- | Run-length-encode a list by testing adjacent elements for equality.
fromList :: Eq a => [a] -> RLE a
fromList :: [a] -> RLE a
fromList = (a -> RLE a -> RLE a) -> RLE a -> [a] -> RLE a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> RLE a -> RLE a
forall a. Eq a => a -> RLE a -> RLE a
cons RLE a
forall a. RLE a
empty

-- | 'Data.Foldable.toList' specialized to 'RLE'.
toList :: RLE a -> [a]
toList :: RLE a -> [a]
toList (RLE [])          = []
toList (RLE ((Int
n :>< a
x):[Run a]
xs)) = Int -> a -> [a]
forall a. Int -> a -> [a]
replicate Int
n a
x [a] -> [a] -> [a]
forall a. Semigroup a => a -> a -> a
<> RLE a -> [a]
forall a. RLE a -> [a]
toList ([Run a] -> RLE a
forall a. [Run a] -> RLE a
RLE [Run a]
xs)

-- | Add an element onto the beginning of the sequence.
cons :: Eq a => a -> RLE a -> RLE a
cons :: a -> RLE a -> RLE a
cons = Run a -> RLE a -> RLE a
forall a. Eq a => Run a -> RLE a -> RLE a
consRun (Run a -> RLE a -> RLE a) -> (a -> Run a) -> a -> RLE a -> RLE a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int
1 Int -> a -> Run a
forall a. Int -> a -> Run a
:><)

consRun_ :: Eq a => Run a -> [Run a] -> [Run a]
consRun_ :: Run a -> [Run a] -> [Run a]
consRun_ (Int
nx :>< a
x) ((Int
ny :>< a
y) : [Run a]
rs) | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y = (Int
nxInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
ny Int -> a -> Run a
forall a. Int -> a -> Run a
:>< a
x) Run a -> [Run a] -> [Run a]
forall a. a -> [a] -> [a]
: [Run a]
rs
consRun_ (Int
0 :>< a
_)  [Run a]
rs                         =                 [Run a]
rs
consRun_ Run a
r          [Run a]
rs                         = Run a
r             Run a -> [Run a] -> [Run a]
forall a. a -> [a] -> [a]
: [Run a]
rs

-- | Add a run of equal elements onto the beginning of the sequence.
consRun :: forall a. Eq a => Run a -> RLE a -> RLE a
consRun :: Run a -> RLE a -> RLE a
consRun = (Run a -> [Run a] -> [Run a]) -> Run a -> RLE a -> RLE a
coerce (Eq a => Run a -> [Run a] -> [Run a]
forall a. Eq a => Run a -> [Run a] -> [Run a]
consRun_ @a)

-- | Split the first element from the rest of the sequence.
uncons :: Eq a => RLE a -> Maybe (a, RLE a)
uncons :: RLE a -> Maybe (a, RLE a)
uncons (RLE a -> Maybe (Run a, RLE a)
forall a. RLE a -> Maybe (Run a, RLE a)
unconsRun -> Just (Int
n :>< a
a, RLE a
rest)) = (a, RLE a) -> Maybe (a, RLE a)
forall a. a -> Maybe a
Just (a
a, Run a -> RLE a -> RLE a
forall a. Eq a => Run a -> RLE a -> RLE a
consRun (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1 Int -> a -> Run a
forall a. Int -> a -> Run a
:>< a
a) RLE a
rest)
uncons RLE a
_                                   = Maybe (a, RLE a)
forall a. Maybe a
Nothing

-- | Split the first run of equal elements from the rest of the sequence.
unconsRun :: RLE a -> Maybe (Run a, RLE a)
unconsRun :: RLE a -> Maybe (Run a, RLE a)
unconsRun (RLE (Run a
r:[Run a]
rs)) = (Run a, RLE a) -> Maybe (Run a, RLE a)
forall a. a -> Maybe a
Just (Run a
r, [Run a] -> RLE a
forall a. [Run a] -> RLE a
RLE [Run a]
rs)
unconsRun RLE a
_            = Maybe (Run a, RLE a)
forall a. Maybe a
Nothing

-- | Return an 'RLE' containing the first @n@ elements of the input.
take :: Int -> RLE a -> RLE a
take :: Int -> RLE a -> RLE a
take Int
n (RLE [Run a]
xs) = [Run a] -> RLE a
forall a. [Run a] -> RLE a
RLE (Int -> [Run a] -> [Run a]
forall a. Int -> [Run a] -> [Run a]
go Int
n [Run a]
xs)
  where
    go :: Int -> [Run a] -> [Run a]
go Int
0 [Run a]
_ = []
    go Int
_ [] = []
    go Int
i ((Int
l :>< a
x):[Run a]
rs) = (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
i Int
l Int -> a -> Run a
forall a. Int -> a -> Run a
:>< a
x) Run a -> [Run a] -> [Run a]
forall a. a -> [a] -> [a]
: Int -> [Run a] -> [Run a]
go (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l)) [Run a]
rs

-- | Returns a tuple where the first element contains the first n elements of
-- the sequence, and the second element is the remainder of the sequence.
splitAt :: (HasCallStack, Eq a) => Int -> RLE a -> (RLE a, RLE a)
splitAt :: Int -> RLE a -> (RLE a, RLE a)
splitAt Int
n RLE a
rle = RLE a -> Int -> RLE a -> (RLE a, RLE a)
forall a. Eq a => RLE a -> Int -> RLE a -> (RLE a, RLE a)
go RLE a
rle Int
n RLE a
forall a. RLE a
empty
  where
    go :: RLE a -> Int -> RLE a -> (RLE a, RLE a)
go RLE a
r Int
i RLE a
prev
      | RLE a -> Bool
forall a. RLE a -> Bool
null RLE a
r Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = (RLE a -> RLE a
forall a. RLE a -> RLE a
reverse RLE a
prev, RLE a
r)
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len = ( RLE a -> RLE a
forall a. RLE a -> RLE a
reverse ((Int
i Int -> a -> Run a
forall a. Int -> a -> Run a
:>< a
a) Run a -> RLE a -> RLE a
forall a. Eq a => Run a -> RLE a -> RLE a
`consRun` RLE a
prev)
                  , Run a -> RLE a -> RLE a
forall a. Eq a => Run a -> RLE a -> RLE a
consRun (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> a -> Run a
forall a. Int -> a -> Run a
:>< a
a) RLE a
r')
      | Bool
otherwise = RLE a -> Int -> RLE a -> (RLE a, RLE a)
go RLE a
r' (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len) ((Int
len Int -> a -> Run a
forall a. Int -> a -> Run a
:>< a
a) Run a -> RLE a -> RLE a
forall a. Eq a => Run a -> RLE a -> RLE a
`consRun` RLE a
prev)
      where
        -- Safe since we check for null above
        ((Int
len :>< a
a), RLE a
r') = Maybe (Run a, RLE a) -> (Run a, RLE a)
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Run a, RLE a) -> (Run a, RLE a))
-> Maybe (Run a, RLE a) -> (Run a, RLE a)
forall a b. (a -> b) -> a -> b
$ RLE a -> Maybe (Run a, RLE a)
forall a. RLE a -> Maybe (Run a, RLE a)
unconsRun RLE a
r

-- | Reverse the order of the elements in the sequence.
reverse :: RLE a -> RLE a
reverse :: RLE a -> RLE a
reverse (RLE [Run a]
r) = [Run a] -> RLE a
forall a. [Run a] -> RLE a
RLE ([Run a] -> [Run a]
forall a. [a] -> [a]
P.reverse [Run a]
r)

-- | Creates an RLE with a single element.
singleton :: a -> RLE a
singleton :: a -> RLE a
singleton a
a = [Run a] -> RLE a
forall a. [Run a] -> RLE a
RLE [Int
1 Int -> a -> Run a
forall a. Int -> a -> Run a
:>< a
a]

-- | Append two sequences.
(++) :: Eq a => RLE a -> RLE a -> RLE a
++ :: RLE a -> RLE a -> RLE a
(++) (RLE (Run a
x0:xs :: [Run a]
xs@(Run a
_:[Run a]
_))) = \RLE a
ys -> [Run a] -> RLE a
forall a. [Run a] -> RLE a
RLE ([Run a] -> RLE a) -> [Run a] -> RLE a
forall a b. (a -> b) -> a -> b
$ Run a
x0 Run a -> [Run a] -> [Run a]
forall a. a -> [a] -> [a]
: RLE a -> [Run a]
forall a. RLE a -> [Run a]
toRuns ([Run a] -> RLE a
forall a. [Run a] -> RLE a
RLE [Run a]
xs RLE a -> RLE a -> RLE a
forall a. Eq a => RLE a -> RLE a -> RLE a
++ RLE a
ys)
(++) (RLE [Run a
r])           = Run a -> RLE a -> RLE a
forall a. Eq a => Run a -> RLE a -> RLE a
consRun Run a
r
(++) (RLE [])            = RLE a -> RLE a
forall a. a -> a
id

-- | Map the given function over each element of the sequence.
map :: Eq b => (a -> b) -> RLE a -> RLE b
map :: (a -> b) -> RLE a -> RLE b
map a -> b
f (RLE [Run a]
xs) = [Run b] -> RLE b
forall a. Eq a => [Run a] -> RLE a
fromRuns ((Run a -> Run b) -> [Run a] -> [Run b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> b) -> Run a -> Run b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) [Run a]
xs)

-- | Map the given invertible function over each element of the sequence. This
-- is only safe when the function is invertible.
--
-- This is slightly faster than @map@ and does not require an Eq constraint on
-- the result type.
mapInvertible :: (a -> b) -> RLE a -> RLE b
mapInvertible :: (a -> b) -> RLE a -> RLE b
mapInvertible a -> b
f (RLE [Run a]
xs) = [Run b] -> RLE b
forall a. [Run a] -> RLE a
RLE ((Run a -> Run b) -> [Run a] -> [Run b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> b) -> Run a -> Run b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) [Run a]
xs)

-- | Visit each element of the sequence in an 'Applicative'.
--
-- @
-- traverse :: Eq b => Traversal (RLE a) (RLE b) a b
-- @
traverse :: (Eq b, Applicative f) => (a -> f b) -> RLE a -> f (RLE b)
traverse :: (a -> f b) -> RLE a -> f (RLE b)
traverse a -> f b
f RLE a
rle = case RLE a -> Maybe (Run a, RLE a)
forall a. RLE a -> Maybe (Run a, RLE a)
unconsRun RLE a
rle of
  Maybe (Run a, RLE a)
Nothing           -> RLE b -> f (RLE b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure RLE b
forall a. RLE a
empty
  Just (Int
n :>< a
x, RLE a
rs) -> (RLE b -> [b] -> RLE b) -> [b] -> RLE b -> RLE b
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((b -> RLE b -> RLE b) -> RLE b -> [b] -> RLE b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr b -> RLE b -> RLE b
forall a. Eq a => a -> RLE a -> RLE a
cons)
    ([b] -> RLE b -> RLE b) -> f [b] -> f (RLE b -> RLE b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f b -> f [b]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n (a -> f b
f a
x)
    f (RLE b -> RLE b) -> f (RLE b) -> f (RLE b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> f b) -> RLE a -> f (RLE b)
forall b (f :: * -> *) a.
(Eq b, Applicative f) =>
(a -> f b) -> RLE a -> f (RLE b)
traverse a -> f b
f RLE a
rs

-- | @Fold@ over the contained runs in order.
--
-- This is as strong a type as this can have without breaking any laws, due to
-- the invariants that no empty or mergeable runs exist: if we make it a
-- @Traversal@, it can end up changing the number of targets, and if we make it
-- an @Iso@ to @[(Int, a)]@, the reverse direction is not an isomorphism.
--
-- If you want to use a law-breaking @Iso@ or @Traversal@ for this anyway, use
-- @iso 'fromRuns' 'toRuns'@ to inline the problematic @Iso@.
--
-- @
-- runs :: Fold (RLE a) (Int, a)
-- @
runs
  :: (Contravariant f, Applicative f)
  => (Run a -> f (Run a))
  -> RLE a -> f (RLE a)
runs :: (Run a -> f (Run a)) -> RLE a -> f (RLE a)
runs Run a -> f (Run a)
f RLE a
rle = (Void -> RLE a) -> f Void -> f (RLE a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Void -> RLE a
forall a. Void -> a
absurd (f Void -> f (RLE a)) -> f Void -> f (RLE a)
forall a b. (a -> b) -> a -> b
$ (Void -> [Run a]) -> f [Run a] -> f Void
forall (f :: * -> *) a b. Contravariant f => (a -> b) -> f b -> f a
contramap Void -> [Run a]
forall a. Void -> a
absurd (f [Run a] -> f Void) -> f [Run a] -> f Void
forall a b. (a -> b) -> a -> b
$ (Run a -> f (Run a)) -> [Run a] -> f [Run a]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
P.traverse Run a -> f (Run a)
f ([Run a] -> f [Run a]) -> [Run a] -> f [Run a]
forall a b. (a -> b) -> a -> b
$ RLE a -> [Run a]
forall a. RLE a -> [Run a]
toRuns RLE a
rle

-- | Construct an 'RLE' from a list of runs.
--
-- This is a retraction of 'toRuns'.
fromRuns :: Eq a => [Run a] -> RLE a
fromRuns :: [Run a] -> RLE a
fromRuns = (Run a -> RLE a -> RLE a) -> RLE a -> [Run a] -> RLE a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Run a -> RLE a -> RLE a
forall a. Eq a => Run a -> RLE a -> RLE a
consRun RLE a
forall a. RLE a
empty

-- | Zip two sequences together.
zipWith :: Eq c => (a -> b -> c) -> RLE a -> RLE b -> RLE c
zipWith :: (a -> b -> c) -> RLE a -> RLE b -> RLE c
zipWith a -> b -> c
f (RLE [Run a]
xs0) (RLE [Run b]
ys0) = [Run c] -> RLE c
forall a. [Run a] -> RLE a
RLE ([Run c] -> RLE c) -> [Run c] -> RLE c
forall a b. (a -> b) -> a -> b
$ [Run a] -> [Run b] -> [Run c]
go [Run a]
xs0 [Run b]
ys0
 where
  go :: [Run a] -> [Run b] -> [Run c]
go [] [Run b]
_ = []
  go [Run a]
_ [] = []
  go ((Int
nx :>< a
x) : [Run a]
xs) ((Int
ny :>< b
y) : [Run b]
ys) = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
nx Int
ny of
    Ordering
LT -> (Int
nx Int -> c -> Run c
forall a. Int -> a -> Run a
:>< a -> b -> c
f a
x b
y) Run c -> [Run c] -> [Run c]
forall a. Eq a => Run a -> [Run a] -> [Run a]
`consRun_` [Run a] -> [Run b] -> [Run c]
go [Run a]
xs ((Int
nyInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
nx Int -> b -> Run b
forall a. Int -> a -> Run a
:>< b
y) Run b -> [Run b] -> [Run b]
forall a. a -> [a] -> [a]
: [Run b]
ys)
    Ordering
GT -> (Int
ny Int -> c -> Run c
forall a. Int -> a -> Run a
:>< a -> b -> c
f a
x b
y) Run c -> [Run c] -> [Run c]
forall a. Eq a => Run a -> [Run a] -> [Run a]
`consRun_` [Run a] -> [Run b] -> [Run c]
go ((Int
nxInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
ny Int -> a -> Run a
forall a. Int -> a -> Run a
:>< a
x) Run a -> [Run a] -> [Run a]
forall a. a -> [a] -> [a]
: [Run a]
xs) [Run b]
ys
    Ordering
EQ -> (Int
nx Int -> c -> Run c
forall a. Int -> a -> Run a
:>< a -> b -> c
f a
x b
y) Run c -> [Run c] -> [Run c]
forall a. Eq a => Run a -> [Run a] -> [Run a]
`consRun_` [Run a] -> [Run b] -> [Run c]
go [Run a]
xs [Run b]
ys

-- | Return an 'RLE' containing all but the last element of the input.
init :: HasCallStack => RLE a -> RLE a
init :: RLE a -> RLE a
init (RLE [Run a]
rs0) = [Run a] -> RLE a
forall a. [Run a] -> RLE a
RLE ([Run a] -> RLE a) -> [Run a] -> RLE a
forall a b. (a -> b) -> a -> b
$ [Run a] -> [Run a]
forall a. [Run a] -> [Run a]
go [Run a]
rs0
 where
  go :: [Run a] -> [Run a]
go []        = String -> [Run a]
forall a. HasCallStack => String -> a
error String
"RLE.init: empty RLE"
  go (Run a
r0:Run a
r:[Run a]
rs) = Run a
r0 Run a -> [Run a] -> [Run a]
forall a. a -> [a] -> [a]
: [Run a] -> [Run a]
go (Run a
rRun a -> [Run a] -> [Run a]
forall a. a -> [a] -> [a]
:[Run a]
rs)
  go [Int
n :>< a
x] = [Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1 Int -> a -> Run a
forall a. Int -> a -> Run a
:>< a
x | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1]