{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

{- |
Module      : Primus.Rep
Description : representable methods for use with fixed containers
Copyright   : (c) Grant Weyburne, 2022
License     : BSD-3
-}
module Primus.Rep (
  buildRepL,
  buildRepR,
  fillRep,
  toEnumRep,
  izipWithR,
  izipWithRF,
  ipostscanr,
  ipostscanl,
  unfoldlRep,
  unfoldrRep,
) where

import Data.Bool
import Data.Distributive
import Data.Foldable
import Data.Functor.Rep
import qualified Data.List as L
import qualified Data.List.NonEmpty as N
import Primus.Enum
import Primus.Error
import Primus.Fold

-- | builds a representable from the left using past and future inputs
buildRepL ::
  forall f a b.
  (Traversable f, Representable f) =>
  ([Rep f] -> [Rep f] -> b -> Rep f -> (b, a)) ->
  b ->
  (b, f a)
buildRepL :: ([Rep f] -> [Rep f] -> b -> Rep f -> (b, a)) -> b -> (b, f a)
buildRepL [Rep f] -> [Rep f] -> b -> Rep f -> (b, a)
f b
b0 = ([Rep f] -> [Rep f] -> b -> Rep f -> (b, a))
-> b -> f (Rep f) -> (b, f a)
forall (t :: * -> *) a z b.
Traversable t =>
([a] -> [a] -> z -> a -> (z, b)) -> z -> t a -> (z, t b)
histMapL [Rep f] -> [Rep f] -> b -> Rep f -> (b, a)
f b
b0 ((Rep f -> Rep f) -> f (Rep f)
forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate Rep f -> Rep f
forall a. a -> a
id)

-- | same as 'buildRepL' but associates to the right
buildRepR ::
  forall f a b.
  (Traversable f, Representable f) =>
  ([Rep f] -> [Rep f] -> b -> Rep f -> (b, a)) ->
  b ->
  (b, f a)
buildRepR :: ([Rep f] -> [Rep f] -> b -> Rep f -> (b, a)) -> b -> (b, f a)
buildRepR [Rep f] -> [Rep f] -> b -> Rep f -> (b, a)
f b
b0 = ([Rep f] -> [Rep f] -> b -> Rep f -> (b, a))
-> b -> f (Rep f) -> (b, f a)
forall (t :: * -> *) a z b.
Traversable t =>
([a] -> [a] -> z -> a -> (z, b)) -> z -> t a -> (z, t b)
histMapR [Rep f] -> [Rep f] -> b -> Rep f -> (b, a)
f b
b0 ((Rep f -> Rep f) -> f (Rep f)
forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate Rep f -> Rep f
forall a. a -> a
id)

-- | fill a representable container with a foldable
fillRep ::
  forall f a.
  (Representable f, Traversable f) =>
  [a] ->
  Either String ([a], f a)
fillRep :: [a] -> Either String ([a], f a)
fillRep = f (Rep f) -> [a] -> Either String ([a], f a)
forall (t :: * -> *) a z.
Traversable t =>
t z -> [a] -> Either String ([a], t a)
fillTraversable ((Rep f -> Rep f) -> f (Rep f)
forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate Rep f -> Rep f
forall a. a -> a
id)

-- | load a fixed container with "a"s using the relative position "i"
toEnumRep ::
  forall f a.
  (Traversable f, Representable f, Enum a, Bounded a) =>
  Integer ->
  Either String (f a)
toEnumRep :: Integer -> Either String (f a)
toEnumRep = f (Rep f) -> Integer -> Either String (f a)
forall a (f :: * -> *) z.
(Traversable f, Enum a, Bounded a) =>
f z -> Integer -> Either String (f a)
toEnumTraversable ((Rep f -> Rep f) -> f (Rep f)
forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate Rep f -> Rep f
forall a. a -> a
id)

-- | 'Data.List.zipWith' with rep index
izipWithR ::
  Representable f =>
  (Rep f -> a -> b -> c) ->
  f a ->
  f b ->
  f c
izipWithR :: (Rep f -> a -> b -> c) -> f a -> f b -> f c
izipWithR Rep f -> a -> b -> c
f f a
as f b
bs = (Rep f -> c) -> f c
forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate ((Rep f -> c) -> f c) -> (Rep f -> c) -> f c
forall a b. (a -> b) -> a -> b
$ \Rep f
k -> Rep f -> a -> b -> c
f Rep f
k (f a -> Rep f -> a
forall (f :: * -> *) a. Representable f => f a -> Rep f -> a
index f a
as Rep f
k) (f b -> Rep f -> b
forall (f :: * -> *) a. Representable f => f a -> Rep f -> a
index f b
bs Rep f
k)

-- | 'Control.Monad.zipWithM' with rep index
izipWithRF ::
  (Representable f, Distributive g) =>
  (Rep f -> a -> b -> g c) ->
  f a ->
  f b ->
  g (f c)
izipWithRF :: (Rep f -> a -> b -> g c) -> f a -> f b -> g (f c)
izipWithRF Rep f -> a -> b -> g c
f = (g c -> g c) -> f (g c) -> g (f c)
forall (g :: * -> *) (f :: * -> *) a b.
(Distributive g, Functor f) =>
(a -> g b) -> f a -> g (f b)
collect g c -> g c
forall a. a -> a
id (f (g c) -> g (f c))
-> (f a -> f b -> f (g c)) -> f a -> f b -> g (f c)
forall c d a b. (c -> d) -> (a -> b -> c) -> a -> b -> d
.@ (Rep f -> a -> b -> g c) -> f a -> f b -> f (g c)
forall (f :: * -> *) a b c.
Representable f =>
(Rep f -> a -> b -> c) -> f a -> f b -> f c
izipWithR Rep f -> a -> b -> g c
f

{- | like 'Data.List.scanr'
 passes in the 'Rep' index and removes the first element
-}
ipostscanr :: (Representable f, Traversable f) => (Rep f -> a -> b -> b) -> b -> f a -> f b
ipostscanr :: (Rep f -> a -> b -> b) -> b -> f a -> f b
ipostscanr Rep f -> a -> b -> b
f b
c f a
ta =
  Either String (f b) -> f b
forall a. HasCallStack => Either String a -> a
frp (Either String (f b) -> f b) -> Either String (f b) -> f b
forall a b. (a -> b) -> a -> b
$ f a -> [b] -> Either String (f b)
forall (f :: * -> *) a z.
Traversable f =>
f z -> [a] -> Either String (f a)
fillTraversableExact f a
ta ([b] -> Either String (f b)) -> [b] -> Either String (f b)
forall a b. (a -> b) -> a -> b
$ NonEmpty b -> [b]
forall a. NonEmpty a -> [a]
N.init (NonEmpty b -> [b]) -> NonEmpty b -> [b]
forall a b. (a -> b) -> a -> b
$ ((Rep f, a) -> b -> b) -> b -> [(Rep f, a)] -> NonEmpty b
forall (f :: * -> *) a b.
Foldable f =>
(a -> b -> b) -> b -> f a -> NonEmpty b
N.scanr ((Rep f -> a -> b -> b) -> (Rep f, a) -> b -> b
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Rep f -> a -> b -> b
f) b
c [(Rep f, a)]
xs
 where
  xs :: [(Rep f, a)]
xs = f (Rep f, a) -> [(Rep f, a)]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (f (Rep f, a) -> [(Rep f, a)]) -> f (Rep f, a) -> [(Rep f, a)]
forall a b. (a -> b) -> a -> b
$ (Rep f -> a -> (Rep f, a)) -> f a -> f (Rep f, a)
forall (r :: * -> *) a a'.
Representable r =>
(Rep r -> a -> a') -> r a -> r a'
imapRep (,) f a
ta

{- | like 'Data.List.scanl'
 passes in the 'Rep' index and removes the last element
-}
ipostscanl :: (Representable f, Traversable f) => (Rep f -> b -> a -> b) -> b -> f a -> f b
ipostscanl :: (Rep f -> b -> a -> b) -> b -> f a -> f b
ipostscanl Rep f -> b -> a -> b
f b
c f a
ta =
  Either String (f b) -> f b
forall a. HasCallStack => Either String a -> a
frp (Either String (f b) -> f b) -> Either String (f b) -> f b
forall a b. (a -> b) -> a -> b
$ f a -> [b] -> Either String (f b)
forall (f :: * -> *) a z.
Traversable f =>
f z -> [a] -> Either String (f a)
fillTraversableExact f a
ta ([b] -> Either String (f b)) -> [b] -> Either String (f b)
forall a b. (a -> b) -> a -> b
$ NonEmpty b -> [b]
forall a. NonEmpty a -> [a]
N.tail (NonEmpty b -> [b]) -> NonEmpty b -> [b]
forall a b. (a -> b) -> a -> b
$ (b -> (Rep f, a) -> b) -> b -> f (Rep f, a) -> NonEmpty b
forall (f :: * -> *) b a.
Foldable f =>
(b -> a -> b) -> b -> f a -> NonEmpty b
N.scanl b -> (Rep f, a) -> b
g b
c f (Rep f, a)
xs
 where
  xs :: f (Rep f, a)
xs = (Rep f -> a -> (Rep f, a)) -> f a -> f (Rep f, a)
forall (r :: * -> *) a a'.
Representable r =>
(Rep r -> a -> a') -> r a -> r a'
imapRep (,) f a
ta
  g :: b -> (Rep f, a) -> b
g b
b (Rep f
i, a
a) = Rep f -> b -> a -> b
f Rep f
i b
b a
a

-- | left/right unfold from the right into a Representable
unfoldlRep
  , unfoldrRep ::
    (Representable f, Traversable f) =>
    (Rep f -> s -> (s, a)) ->
    s ->
    (s, f a)
unfoldlRep :: (Rep f -> s -> (s, a)) -> s -> (s, f a)
unfoldlRep = Bool -> (Rep f -> s -> (s, a)) -> s -> (s, f a)
forall (f :: * -> *) s a.
(Representable f, Traversable f) =>
Bool -> (Rep f -> s -> (s, a)) -> s -> (s, f a)
unfoldRepImpl Bool
False
unfoldrRep :: (Rep f -> s -> (s, a)) -> s -> (s, f a)
unfoldrRep = Bool -> (Rep f -> s -> (s, a)) -> s -> (s, f a)
forall (f :: * -> *) s a.
(Representable f, Traversable f) =>
Bool -> (Rep f -> s -> (s, a)) -> s -> (s, f a)
unfoldRepImpl Bool
True

unfoldRepImpl :: (Representable f, Traversable f) => Bool -> (Rep f -> s -> (s, a)) -> s -> (s, f a)
unfoldRepImpl :: Bool -> (Rep f -> s -> (s, a)) -> s -> (s, f a)
unfoldRepImpl Bool
isright Rep f -> s -> (s, a)
f s
s = ((s -> Rep f -> (s, a)) -> s -> f (Rep f) -> (s, f a))
-> ((s -> Rep f -> (s, a)) -> s -> f (Rep f) -> (s, f a))
-> Bool
-> (s -> Rep f -> (s, a))
-> s
-> f (Rep f)
-> (s, f a)
forall a. a -> a -> Bool -> a
bool (s -> Rep f -> (s, a)) -> s -> f (Rep f) -> (s, f a)
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
L.mapAccumL (s -> Rep f -> (s, a)) -> s -> f (Rep f) -> (s, f a)
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
L.mapAccumR Bool
isright ((Rep f -> s -> (s, a)) -> s -> Rep f -> (s, a)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Rep f -> s -> (s, a)
f) s
s ((Rep f -> Rep f) -> f (Rep f)
forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate Rep f -> Rep f
forall a. a -> a
id)