{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleContexts #-}
module Data.Profunctor.FinFn(
  FinFn(),
  makeFinFn,
  withFinFn,
  applyFinFn,
  (>>>>),
  (<<<<),

  fromMap,
) where

import Data.Coerce

import Control.Category ((>>>))
import Data.Profunctor
import Data.Profunctor.Cartesian

import qualified Data.Map as Map

import Data.Proxy
import GHC.TypeNats

import Data.Finite (Finite())
import qualified Data.Finite.Internal.Integral as Internal

data FinFn a b = FinFn !Integer (a -> Integer) (Integer -> b)

-- | Safe construction
makeFinFn :: forall n a b.
             KnownNat n
          => (a -> Finite n) -> (Finite n -> b) -> FinFn a b
makeFinFn :: forall (n :: Nat) a b.
KnownNat n =>
(a -> Finite n) -> (Finite n -> b) -> FinFn a b
makeFinFn a -> Finite n
l Finite n -> b
r = Integer -> (a -> Integer) -> (Integer -> b) -> FinFn a b
forall a b.
Integer -> (a -> Integer) -> (Integer -> b) -> FinFn a b
FinFn (Nat -> Integer
forall a. Integral a => a -> Integer
toInteger Nat
nvalue) ((a -> Finite n) -> a -> Integer
forall a b. Coercible a b => a -> b
coerce a -> Finite n
l) ((Finite n -> b) -> Integer -> b
forall a b. Coercible a b => a -> b
coerce Finite n -> b
r)
  where nvalue :: Nat
nvalue = forall (n :: Nat) (proxy :: Nat -> *). KnownNat n => proxy n -> Nat
natVal @n Proxy n
forall {k} (t :: k). Proxy t
Proxy

-- | Safe pattern matching
withFinFn :: FinFn a b
          -> (forall n. KnownNat n => (a -> Finite n) -> (Finite n -> b) -> r)
          -> r
withFinFn :: forall a b r.
FinFn a b
-> (forall (n :: Nat).
    KnownNat n =>
    (a -> Finite n) -> (Finite n -> b) -> r)
-> r
withFinFn (FinFn Integer
n a -> Integer
l Integer -> b
r) forall (n :: Nat).
KnownNat n =>
(a -> Finite n) -> (Finite n -> b) -> r
user =
    case Nat -> SomeNat
someNatVal (Integer -> Nat
forall a. Num a => Integer -> a
fromInteger Integer
n) of
      SomeNat Proxy n
name -> (a -> Finite n) -> (Finite n -> b) -> r
forall (n :: Nat).
KnownNat n =>
(a -> Finite n) -> (Finite n -> b) -> r
user (Proxy n -> Integer -> Finite n
forall (n :: Nat). Proxy n -> Integer -> Finite n
upcast Proxy n
name (Integer -> Finite n) -> (a -> Integer) -> a -> Finite n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Integer
l) (Integer -> b
r (Integer -> b) -> (Finite n -> Integer) -> Finite n -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Proxy n -> Finite n -> Integer
forall (n :: Nat). Proxy n -> Finite n -> Integer
downcast Proxy n
name)
  where
    upcast :: Proxy n -> Integer -> Finite n
    upcast :: forall (n :: Nat). Proxy n -> Integer -> Finite n
upcast Proxy n
_ = Integer -> Finite n
forall a b. Coercible a b => a -> b
coerce
    downcast :: Proxy n -> Finite n -> Integer
    downcast :: forall (n :: Nat). Proxy n -> Finite n -> Integer
downcast Proxy n
_ = Finite n -> Integer
forall a b. Coercible a b => a -> b
coerce

-- | How I would say... 'applyFinFn' is nice, functor-ish.
--   It commutes with 'Profunctor', 'Cartesian', 'Cocartesian',
--   and composition '(>>>>)'.
applyFinFn :: FinFn a b -> a -> b
applyFinFn :: forall a b. FinFn a b -> a -> b
applyFinFn (FinFn Integer
_ a -> Integer
l Integer -> b
r) = a -> Integer
l (a -> Integer) -> (Integer -> b) -> a -> b
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Integer -> b
r

-- | 'FinFn' composes like 'Control.Category.Category', but can't be a Category since
--   it does not have 'Control.Cateogory.id'.
(>>>>) :: FinFn a b -> FinFn b c -> FinFn a c
>>>> :: forall a b c. FinFn a b -> FinFn b c -> FinFn a c
(>>>>) (FinFn Integer
n a -> Integer
l Integer -> b
r) (FinFn Integer
m b -> Integer
s Integer -> c
t)
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
m    = Integer -> (a -> Integer) -> (Integer -> c) -> FinFn a c
forall a b.
Integer -> (a -> Integer) -> (Integer -> b) -> FinFn a b
FinFn Integer
n a -> Integer
l (Integer -> b
r (Integer -> b) -> (b -> c) -> Integer -> c
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> b -> Integer
s (b -> Integer) -> (Integer -> c) -> b -> c
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Integer -> c
t)
  | Bool
otherwise = Integer -> (a -> Integer) -> (Integer -> c) -> FinFn a c
forall a b.
Integer -> (a -> Integer) -> (Integer -> b) -> FinFn a b
FinFn Integer
m (a -> Integer
l (a -> Integer) -> (Integer -> Integer) -> a -> Integer
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Integer -> b
r (Integer -> b) -> (b -> Integer) -> Integer -> Integer
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> b -> Integer
s) Integer -> c
t

infixr 1 >>>>

-- | > (<<<<) = flip (>>>>)
(<<<<) :: FinFn b c -> FinFn a b -> FinFn a c
<<<< :: forall b c a. FinFn b c -> FinFn a b -> FinFn a c
(<<<<) = (FinFn a b -> FinFn b c -> FinFn a c)
-> FinFn b c -> FinFn a b -> FinFn a c
forall a b c. (a -> b -> c) -> b -> a -> c
flip FinFn a b -> FinFn b c -> FinFn a c
forall a b c. FinFn a b -> FinFn b c -> FinFn a c
(>>>>)

infixr 1 <<<<

-- * Special constructions

-- | > applyFinFn . fromMap = flip Map.lookup
fromMap :: (Ord a) => Map.Map a b -> FinFn a (Maybe b)
fromMap :: forall a b. Ord a => Map a b -> FinFn a (Maybe b)
fromMap Map a b
m = Integer
-> (a -> Integer) -> (Integer -> Maybe b) -> FinFn a (Maybe b)
forall a b.
Integer -> (a -> Integer) -> (Integer -> b) -> FinFn a b
FinFn (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n) (Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Integer) -> (a -> Int) -> a -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Int
l) (Int -> Maybe b
r (Int -> Maybe b) -> (Integer -> Int) -> Integer -> Maybe b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Int
forall a. Num a => Integer -> a
fromInteger)
  where
    n :: Int
n = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Map a b -> Int
forall k a. Map k a -> Int
Map.size Map a b
m
    l :: a -> Int
l a
a = Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+) (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ a -> Map a b -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe Int
Map.lookupIndex a
a Map a b
m
    r :: Int -> Maybe b
r Int
0 = Maybe b
forall a. Maybe a
Nothing
    r Int
i = b -> Maybe b
forall a. a -> Maybe a
Just (b -> Maybe b) -> ((a, b) -> b) -> (a, b) -> Maybe b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, b) -> b
forall a b. (a, b) -> b
snd ((a, b) -> Maybe b) -> (a, b) -> Maybe b
forall a b. (a -> b) -> a -> b
$ Int -> Map a b -> (a, b)
forall k a. Int -> Map k a -> (k, a)
Map.elemAt (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Map a b
m

instance Profunctor FinFn where
  dimap :: forall a b c d. (a -> b) -> (c -> d) -> FinFn b c -> FinFn a d
dimap a -> b
l c -> d
r (FinFn Integer
n b -> Integer
l' Integer -> c
r') = Integer -> (a -> Integer) -> (Integer -> d) -> FinFn a d
forall a b.
Integer -> (a -> Integer) -> (Integer -> b) -> FinFn a b
FinFn Integer
n (a -> b
l (a -> b) -> (b -> Integer) -> a -> Integer
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> b -> Integer
l') (Integer -> c
r' (Integer -> c) -> (c -> d) -> Integer -> d
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> c -> d
r)

instance Cartesian FinFn where
  proUnit :: forall a. FinFn a ()
proUnit = Integer -> (a -> Integer) -> (Integer -> ()) -> FinFn a ()
forall a b.
Integer -> (a -> Integer) -> (Integer -> b) -> FinFn a b
FinFn Integer
1 (Integer -> a -> Integer
forall a b. a -> b -> a
const Integer
0) (() -> Integer -> ()
forall a b. a -> b -> a
const ())
  FinFn Integer
n a -> Integer
l Integer -> b
r *** :: forall a b a' b'. FinFn a b -> FinFn a' b' -> FinFn (a, a') (b, b')
*** FinFn Integer
m a' -> Integer
s Integer -> b'
t = Integer
-> ((a, a') -> Integer)
-> (Integer -> (b, b'))
-> FinFn (a, a') (b, b')
forall a b.
Integer -> (a -> Integer) -> (Integer -> b) -> FinFn a b
FinFn (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
m) (a, a') -> Integer
ls Integer -> (b, b')
rt
    where
      ls :: (a, a') -> Integer
ls (a
a,a'
a') = a -> Integer
l a
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
m Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ a' -> Integer
s a'
a'
      rt :: Integer -> (b, b')
rt Integer
i = let (Integer
j,Integer
k) = Integer
i Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Integer
m
             in (Integer -> b
r Integer
j, Integer -> b'
t Integer
k)

instance Cocartesian FinFn where
  proEmpty :: forall b. FinFn Void b
proEmpty = Integer -> (Void -> Integer) -> (Integer -> b) -> FinFn Void b
forall a b.
Integer -> (a -> Integer) -> (Integer -> b) -> FinFn a b
FinFn Integer
0 ([Char] -> Void -> Integer
forall a. HasCallStack => [Char] -> a
error [Char]
"empty function") ([Char] -> Integer -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"empty function")
  FinFn Integer
n a -> Integer
l Integer -> b
r +++ :: forall a b a' b'.
FinFn a b -> FinFn a' b' -> FinFn (Either a a') (Either b b')
+++ FinFn Integer
m a' -> Integer
s Integer -> b'
t = Integer
-> (Either a a' -> Integer)
-> (Integer -> Either b b')
-> FinFn (Either a a') (Either b b')
forall a b.
Integer -> (a -> Integer) -> (Integer -> b) -> FinFn a b
FinFn (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
m) Either a a' -> Integer
ls Integer -> Either b b'
rt
    where
      ls :: Either a a' -> Integer
ls = (a -> Integer) -> (a' -> Integer) -> Either a a' -> Integer
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either a -> Integer
l ((Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+) (Integer -> Integer) -> (a' -> Integer) -> a' -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a' -> Integer
s)
      rt :: Integer -> Either b b'
rt Integer
i | Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
n     = b -> Either b b'
forall a b. a -> Either a b
Left (Integer -> b
r Integer
i)
           | Bool
otherwise = b' -> Either b b'
forall a b. b -> Either a b
Right (Integer -> b'
t (Integer
iInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
n))