{-# LANGUAGE CPP #-}

{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

#if MIN_VERSION_GLASGOW_HASKELL(9,8,0,0)
{-# LANGUAGE TypeAbstractions #-}
#endif

module Data.PTraversable.Internal.Day (ptraverseDay) where

import Prelude hiding (Enum)
import Data.Profunctor.Cartesian.Free
import Data.Void (absurd)
import Data.Function ((&))

import Data.Profunctor
import Data.Profunctor.Cartesian
import Data.Finitary.Enum

import Data.Functor.Day ( Day(..) )
import Control.Monad.Trans.State.Lazy (State, state, runState)
import qualified Data.Profunctor.Day as Prof

type Optic p s t a b = p a b -> p s t
type Traversal s t a b = forall p. (Cartesian p, Cocartesian p) => Optic p s t a b

-- An "instance" of (PTraversable t). They are used instead of constraint (PTraversable t),
-- because this module can't depend on Data.PTraversable module to avoid circular import.
type PT t = forall a b. Traversal (t a) (t b) a b

data SomePower a b s t where
  SomePower :: (Enum x) => (s -> (x -> a)) -> ((x -> b) -> t) -> SomePower a b s t

instance Profunctor (SomePower a b) where
  dimap :: forall a b c d.
(a -> b) -> (c -> d) -> SomePower a b b c -> SomePower a b a d
dimap a -> b
f c -> d
g (SomePower b -> x -> a
to (x -> b) -> c
from) = (a -> x -> a) -> ((x -> b) -> d) -> SomePower a b a d
forall x s a b t.
Enum x =>
(s -> x -> a) -> ((x -> b) -> t) -> SomePower a b s t
SomePower (b -> x -> a
to (b -> x -> a) -> (a -> b) -> a -> x -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f) (c -> d
g (c -> d) -> ((x -> b) -> c) -> (x -> b) -> d
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (x -> b) -> c
from)

instance Cartesian (SomePower a b) where
  proUnit :: forall a. SomePower a b a ()
proUnit = (a -> Void -> a) -> ((Void -> b) -> ()) -> SomePower a b a ()
forall x s a b t.
Enum x =>
(s -> x -> a) -> ((x -> b) -> t) -> SomePower a b s t
SomePower ((Void -> a) -> a -> Void -> a
forall a b. a -> b -> a
const Void -> a
forall a. Void -> a
absurd) (() -> (Void -> b) -> ()
forall a b. a -> b -> a
const ())
  SomePower a -> x -> a
to (x -> b) -> b
from *** :: forall a b a' b'.
SomePower a b a b
-> SomePower a b a' b' -> SomePower a b (a, a') (b, b')
*** SomePower a' -> x -> a
to' (x -> b) -> b'
from' =
    ((a, a') -> Either x x -> a)
-> ((Either x x -> b) -> (b, b')) -> SomePower a b (a, a') (b, b')
forall x s a b t.
Enum x =>
(s -> x -> a) -> ((x -> b) -> t) -> SomePower a b s t
SomePower (\(a
s,a'
s') -> (x -> a) -> (x -> a) -> Either x x -> a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (a -> x -> a
to a
s) (a' -> x -> a
to' a'
s')) (\Either x x -> b
k -> ((x -> b) -> b
from (Either x x -> b
k (Either x x -> b) -> (x -> Either x x) -> x -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. x -> Either x x
forall a b. a -> Either a b
Left), (x -> b) -> b'
from' (Either x x -> b
k (Either x x -> b) -> (x -> Either x x) -> x -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. x -> Either x x
forall a b. b -> Either a b
Right)))

type SOP a b s t = FreeCocartesian (SomePower a b) s t

idPow :: SomePower a b a b
idPow :: forall a b. SomePower a b a b
idPow = (a -> () -> a) -> ((() -> b) -> b) -> SomePower a b a b
forall x s a b t.
Enum x =>
(s -> x -> a) -> ((x -> b) -> t) -> SomePower a b s t
SomePower a -> () -> a
forall a b. a -> b -> a
const ((() -> b) -> () -> b
forall a b. (a -> b) -> a -> b
$ ())

idSOP :: SOP a b a b
idSOP :: forall a b. SOP a b a b
idSOP = SomePower a b a b -> FreeCocartesian (SomePower a b) a b
forall (p :: * -> * -> *) a b. p a b -> FreeCocartesian p a b
liftFreeCocartesian SomePower a b a b
forall a b. SomePower a b a b
idPow

ptraverseDay :: forall t u p a b.
     (Cartesian p, Cocartesian p)
  => PT t -> PT u
  -> p a b -> p (Day t u a) (Day t u b)
ptraverseDay :: forall (t :: * -> *) (u :: * -> *) (p :: * -> * -> *) a b.
(Cartesian p, Cocartesian p) =>
PT t -> PT u -> p a b -> p (Day t u a) (Day t u b)
ptraverseDay PT t
travT PT u
travU p a b
p = (Day t u a -> Day' t u a)
-> (Day' t u b -> Day t u b)
-> p (Day' t u a) (Day' t u b)
-> p (Day t u a) (Day t u b)
forall a b c d. (a -> b) -> (c -> d) -> p b c -> p a d
forall (p :: * -> * -> *) a b c d.
Profunctor p =>
(a -> b) -> (c -> d) -> p b c -> p a d
dimap (PT t -> PT u -> Day t u a -> Day' t u a
forall (t :: * -> *) (u :: * -> *) a.
PT t -> PT u -> Day t u a -> Day' t u a
unwrapDay Optic p (t a) (t b) a b
PT t
travT Optic p (u a) (u b) a b
PT u
travU) Day' t u b -> Day t u b
forall (t :: * -> *) (u :: * -> *) a. Day' t u a -> Day t u a
wrapDay (p (Day' t u a) (Day' t u b) -> p (Day t u a) (Day t u b))
-> p (Day' t u a) (Day' t u b) -> p (Day t u a) (Day t u b)
forall a b. (a -> b) -> a -> b
$ PT t -> PT u -> p a b -> p (Day' t u a) (Day' t u b)
forall (t :: * -> *) (u :: * -> *) (p :: * -> * -> *) a b.
(Cartesian p, Cocartesian p) =>
PT t -> PT u -> p a b -> p (Day' t u a) (Day' t u b)
ptraverseDay' Optic p (t a) (t b) a b
PT t
travT Optic p (u a) (u b) a b
PT u
travU p a b
p

ptraverseDay' :: forall t u p a b.
     (Cartesian p, Cocartesian p)
  => PT t -> PT u
  -> p a b -> p (Day' t u a) (Day' t u b)
ptraverseDay' :: forall (t :: * -> *) (u :: * -> *) (p :: * -> * -> *) a b.
(Cartesian p, Cocartesian p) =>
PT t -> PT u -> p a b -> p (Day' t u a) (Day' t u b)
ptraverseDay' PT t
travT PT u
travU p a b
p = SOP Int Int (t Int) (t Int)
-> p (t Int, u (Int -> a)) (t Int, u (Int -> b))
forall s0 t0.
SOP Int Int s0 t0 -> p (s0, u (Int -> a)) (t0, u (Int -> b))
go (Optic
  (FreeMonoidal Either Void (SomePower Int Int))
  (t Int)
  (t Int)
  Int
  Int
PT t
travT SOP Int Int Int Int
forall a b. SOP a b a b
idSOP)
  where
    go :: forall s0 t0. SOP Int Int s0 t0 -> p (s0, u (Int -> a)) (t0, u (Int -> b))
    go :: forall s0 t0.
SOP Int Int s0 t0 -> p (s0, u (Int -> a)) (t0, u (Int -> b))
go (Neutral s0 -> Void
z Void -> t0
_) = ((s0, u (Int -> a)) -> Void)
-> p Void (t0, u (Int -> b))
-> p (s0, u (Int -> a)) (t0, u (Int -> b))
forall a b c. (a -> b) -> p b c -> p a c
forall (p :: * -> * -> *) a b c.
Profunctor p =>
(a -> b) -> p b c -> p a c
lmap (s0 -> Void
z (s0 -> Void)
-> ((s0, u (Int -> a)) -> s0) -> (s0, u (Int -> a)) -> Void
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (s0, u (Int -> a)) -> s0
forall a b. (a, b) -> a
fst) p Void (t0, u (Int -> b))
forall b. p Void b
forall (p :: * -> * -> *) b. Cocartesian p => p Void b
proEmpty
    go (Cons (Prof.Day (SomePower @x a1 -> x -> Int
_ (x -> Int) -> b1
mkT) FreeMonoidal Either Void (SomePower Int Int) a2 b2
rest s0 -> Either a1 a2
opA Either b1 b2 -> t0
opB)) = ((s0, u (Int -> a)) -> Either (u (x -> a)) (a2, u (Int -> a)))
-> (Either (u (x -> b)) (b2, u (Int -> b)) -> (t0, u (Int -> b)))
-> p (Either (u (x -> a)) (a2, u (Int -> a)))
     (Either (u (x -> b)) (b2, u (Int -> b)))
-> p (s0, u (Int -> a)) (t0, u (Int -> b))
forall a b c d. (a -> b) -> (c -> d) -> p b c -> p a d
forall (p :: * -> * -> *) a b c d.
Profunctor p =>
(a -> b) -> (c -> d) -> p b c -> p a d
dimap (s0, u (Int -> a)) -> Either (u (x -> a)) (a2, u (Int -> a))
pre Either (u (x -> b)) (b2, u (Int -> b)) -> (t0, u (Int -> b))
post (p (Either (u (x -> a)) (a2, u (Int -> a)))
   (Either (u (x -> b)) (b2, u (Int -> b)))
 -> p (s0, u (Int -> a)) (t0, u (Int -> b)))
-> p (Either (u (x -> a)) (a2, u (Int -> a)))
     (Either (u (x -> b)) (b2, u (Int -> b)))
-> p (s0, u (Int -> a)) (t0, u (Int -> b))
forall a b. (a -> b) -> a -> b
$ p (u (x -> a)) (u (x -> b))
uFn p (u (x -> a)) (u (x -> b))
-> p (a2, u (Int -> a)) (b2, u (Int -> b))
-> p (Either (u (x -> a)) (a2, u (Int -> a)))
     (Either (u (x -> b)) (b2, u (Int -> b)))
forall a b a' b'. p a b -> p a' b' -> p (Either a a') (Either b b')
forall (p :: * -> * -> *) a b a' b'.
Cocartesian p =>
p a b -> p a' b' -> p (Either a a') (Either b b')
+++ FreeMonoidal Either Void (SomePower Int Int) a2 b2
-> p (a2, u (Int -> a)) (b2, u (Int -> b))
forall s0 t0.
SOP Int Int s0 t0 -> p (s0, u (Int -> a)) (t0, u (Int -> b))
go FreeMonoidal Either Void (SomePower Int Int) a2 b2
rest
      where
        uFn :: p (u (x -> a)) (u (x -> b))
uFn = Optic p (u (x -> a)) (u (x -> b)) (x -> a) (x -> b)
PT u
travU (forall x (p :: * -> * -> *) a b.
(Enum x, Cartesian p) =>
p a b -> p (x -> a) (x -> b)
ptraverseFunction @x p a b
p)
        (x -> Int
x2i, Int -> x
i2x) = forall x. Enum x => (x -> Int, Int -> x)
embeddingToInt @x
        pre :: (s0, u (Int -> a)) -> Either (u (x -> a)) (a2, u (Int -> a))
pre (s0
s0, u (Int -> a)
y) = case s0 -> Either a1 a2
opA s0
s0 of
          Left a1
_ -> u (x -> a) -> Either (u (x -> a)) (a2, u (Int -> a))
forall a b. a -> Either a b
Left (Optic (->) (u (Int -> a)) (u (x -> a)) (Int -> a) (x -> a)
PT u
travU ((Int -> a) -> (x -> Int) -> x -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. x -> Int
x2i) u (Int -> a)
y)
          Right a2
a2 -> (a2, u (Int -> a)) -> Either (u (x -> a)) (a2, u (Int -> a))
forall a b. b -> Either a b
Right (a2
a2, u (Int -> a)
y)
        
        post :: Either (u (x -> b)) (b2, u (Int -> b)) -> (t0, u (Int -> b))
post (Left u (x -> b)
y) = (Either b1 b2 -> t0
opB (b1 -> Either b1 b2
forall a b. a -> Either a b
Left ((x -> Int) -> b1
mkT x -> Int
x2i)), Optic (->) (u (x -> b)) (u (Int -> b)) (x -> b) (Int -> b)
PT u
travU ((x -> b) -> (Int -> x) -> Int -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> x
i2x) u (x -> b)
y)
        post (Right (b2
b2,u (Int -> b)
y)) = (Either b1 b2 -> t0
opB (b2 -> Either b1 b2
forall a b. b -> Either a b
Right b2
b2), u (Int -> b)
y)

embeddingToInt :: forall x. Enum x => (x -> Int, Int -> x)
embeddingToInt :: forall x. Enum x => (x -> Int, Int -> x)
embeddingToInt = (forall (n :: Nat).
 KnownNat n =>
 (x -> Finite n) -> (Finite n -> x) -> (x -> Int, Int -> x))
-> (x -> Int, Int -> x)
forall x r.
Enum x =>
(forall (n :: Nat).
 KnownNat n =>
 (x -> Finite n) -> (Finite n -> x) -> r)
-> r
forall r.
(forall (n :: Nat).
 KnownNat n =>
 (x -> Finite n) -> (Finite n -> x) -> r)
-> r
withEnum (\x -> Finite n
to Finite n -> x
from -> (Finite n -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Finite n -> Int) -> (x -> Finite n) -> x -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. x -> Finite n
to, Finite n -> x
from (Finite n -> x) -> (Int -> Finite n) -> Int -> x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Finite n
forall a b. (Integral a, Num b) => a -> b
fromIntegral))

type Day' t u a = (t Int, u (Int -> a))

unwrapDay :: PT t -> PT u -> Day t u a -> Day' t u a
unwrapDay :: forall (t :: * -> *) (u :: * -> *) a.
PT t -> PT u -> Day t u a -> Day' t u a
unwrapDay PT t
travT PT u
fmapU (Day t b
tb u c
uc b -> c -> a
op) = case PT t -> t b -> (t Int, [b])
forall (t :: * -> *) a. PT t -> t a -> (t Int, [a])
separate Optic p (t a) (t b) a b
PT t
travT t b
tb of
  (t Int
ti, [b]
bs) -> (t Int
ti, Optic (->) (u c) (u (Int -> a)) c (Int -> a)
PT u
fmapU (\c
c Int
i -> b -> c -> a
op ([b]
bs [b] -> Int -> b
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) c
c) u c
uc)

wrapDay :: Day' t u a -> Day t u a
wrapDay :: forall (t :: * -> *) (u :: * -> *) a. Day' t u a -> Day t u a
wrapDay (t Int
ti, u (Int -> a)
uf) = t Int -> u (Int -> a) -> (Int -> (Int -> a) -> a) -> Day t u a
forall (f :: * -> *) (g :: * -> *) a b c.
f b -> g c -> (b -> c -> a) -> Day f g a
Day t Int
ti u (Int -> a)
uf Int -> (Int -> a) -> a
forall a b. a -> (a -> b) -> b
(&)

separate :: PT t -> t a -> (t Int, [a])
separate :: forall (t :: * -> *) a. PT t -> t a -> (t Int, [a])
separate PT t
travT t a
ta =
  let mResult :: State ([a] -> [a], Int) (t Int)
mResult = Star (State ([a] -> [a], Int)) (t a) (t Int)
-> t a -> State ([a] -> [a], Int) (t Int)
forall {k} (f :: k -> *) d (c :: k). Star f d c -> d -> f c
runStar (Optic (Star (State ([a] -> [a], Int))) (t a) (t Int) a Int
PT t
travT Star (State ([a] -> [a], Int)) a Int
forall x. Star (State ([x] -> [x], Int)) x Int
indexStep) t a
ta
      (t Int
ti, ([a] -> [a]
xs, Int
_)) = State ([a] -> [a], Int) (t Int)
-> ([a] -> [a], Int) -> (t Int, ([a] -> [a], Int))
forall s a. State s a -> s -> (a, s)
runState State ([a] -> [a], Int) (t Int)
mResult ([a] -> [a]
forall a. a -> a
id, Int
0)
  in (t Int
ti, [a] -> [a]
xs [])

indexStep :: Star (State ([x] -> [x], Int)) x Int
indexStep :: forall x. Star (State ([x] -> [x], Int)) x Int
indexStep = (x -> State ([x] -> [x], Int) Int)
-> Star (State ([x] -> [x], Int)) x Int
forall {k} (f :: k -> *) d (c :: k). (d -> f c) -> Star f d c
Star ((x -> State ([x] -> [x], Int) Int)
 -> Star (State ([x] -> [x], Int)) x Int)
-> (x -> State ([x] -> [x], Int) Int)
-> Star (State ([x] -> [x], Int)) x Int
forall a b. (a -> b) -> a -> b
$ \x
x -> (([x] -> [x], Int) -> (Int, ([x] -> [x], Int)))
-> State ([x] -> [x], Int) Int
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state ((([x] -> [x], Int) -> (Int, ([x] -> [x], Int)))
 -> State ([x] -> [x], Int) Int)
-> (([x] -> [x], Int) -> (Int, ([x] -> [x], Int)))
-> State ([x] -> [x], Int) Int
forall a b. (a -> b) -> a -> b
$ \([x] -> [x]
xs, Int
n) -> (Int
n, ([x] -> [x]
xs ([x] -> [x]) -> ([x] -> [x]) -> [x] -> [x]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (x
x x -> [x] -> [x]
forall a. a -> [a] -> [a]
:), Int -> Int
forall a. Enum a => a -> a
succ Int
n))