-- 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 CPP #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NoStarIsType #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

-- | Provides 'Vec'-backed tables indexed by 'Finite' types.
--
-- Combined with 'Data.Finite' and its Generics-based derivation, this can
-- effectively provide an array-backed container indexed by finite type.  This
-- is a low-syntactic-overhead way to create 'Representable' functors of any
-- desired shape: just define the index type, tack on the requisite @deriving@
-- clauses, and start using @'Table' MyType@.
--
-- @
--     data PrimaryColor = R | G | B
--       deriving Generic
--       deriving (Finite, Portray) via Wrapped Generic PrimaryColor
--
--     newtype Color = Color { getComponents :: Table PrimaryColor Int8 }
--
--     magenta :: Color
--     magenta = Color (Table $ Vec.fromList [255, 0, 255])
--
--     cyan :: Color
--     cyan = Color $ tabulate (\\case { R -> 0; G -> 255; B -> 255 })
--
--     main = pp $ getComponents magenta
--     -- "mkTable (\\case { R -> 255; G -> 0; B -> 255 })"
-- @

module Data.Finite.Table
         ( -- * Tables
           Table(..), (!), ix, idTable, mkTable, lmapTable, composeTable
           -- * Function Utilities
         , memoize, traverseRep
           -- * Representable Utilities
         , tabulateA, retabulated
         ) where

import Control.Applicative (Applicative(..))
import Data.Foldable (toList, traverse_)
import Data.Maybe (catMaybes, isJust)
import Data.Semigroup (Any(..), All(..))
import Control.DeepSeq (NFData(..))
import GHC.Generics (Generic)

import Control.Lens (Iso, Lens', from, lens, (&), (.~))
import Data.Default.Class (Default(..))
import Data.Distributive (Distributive(..))
import Data.Foldable.WithIndex (FoldableWithIndex(..))
import Data.Functor.Rep
         ( Representable(..), ifoldMapRep, imapRep, itraverseRep
         , tabulated
         )
import Data.Functor.WithIndex (FunctorWithIndex(..))
import Data.Portray (Portray(..), Portrayal(..))
import Data.Portray.Diff (Diff(..))
import Data.Serialize (Serialize(..))
import Data.Traversable.WithIndex (TraversableWithIndex(..))

import Data.Vec.Short (Vec)
import qualified Data.Vec.Short as V
import qualified Data.Vec.Short.Explicit as VE
import qualified Data.Vec.Short.Lens as V (ix)

import Data.Finite

#if !MIN_VERSION_lens(5,0,0)
import qualified Control.Lens as L
#endif

-- | A compact array of @b@s indexed by @a@, according to @'Finite' a@.
newtype Table a b = Table (Vec (Cardinality a) b)
  deriving (Table a b -> Table a b -> Bool
(Table a b -> Table a b -> Bool)
-> (Table a b -> Table a b -> Bool) -> Eq (Table a b)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall a b. Eq b => Table a b -> Table a b -> Bool
/= :: Table a b -> Table a b -> Bool
$c/= :: forall a b. Eq b => Table a b -> Table a b -> Bool
== :: Table a b -> Table a b -> Bool
$c== :: forall a b. Eq b => Table a b -> Table a b -> Bool
Eq, Eq (Table a b)
Eq (Table a b)
-> (Table a b -> Table a b -> Ordering)
-> (Table a b -> Table a b -> Bool)
-> (Table a b -> Table a b -> Bool)
-> (Table a b -> Table a b -> Bool)
-> (Table a b -> Table a b -> Bool)
-> (Table a b -> Table a b -> Table a b)
-> (Table a b -> Table a b -> Table a b)
-> Ord (Table a b)
Table a b -> Table a b -> Bool
Table a b -> Table a b -> Ordering
Table a b -> Table a b -> Table a b
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a b. Ord b => Eq (Table a b)
forall a b. Ord b => Table a b -> Table a b -> Bool
forall a b. Ord b => Table a b -> Table a b -> Ordering
forall a b. Ord b => Table a b -> Table a b -> Table a b
min :: Table a b -> Table a b -> Table a b
$cmin :: forall a b. Ord b => Table a b -> Table a b -> Table a b
max :: Table a b -> Table a b -> Table a b
$cmax :: forall a b. Ord b => Table a b -> Table a b -> Table a b
>= :: Table a b -> Table a b -> Bool
$c>= :: forall a b. Ord b => Table a b -> Table a b -> Bool
> :: Table a b -> Table a b -> Bool
$c> :: forall a b. Ord b => Table a b -> Table a b -> Bool
<= :: Table a b -> Table a b -> Bool
$c<= :: forall a b. Ord b => Table a b -> Table a b -> Bool
< :: Table a b -> Table a b -> Bool
$c< :: forall a b. Ord b => Table a b -> Table a b -> Bool
compare :: Table a b -> Table a b -> Ordering
$ccompare :: forall a b. Ord b => Table a b -> Table a b -> Ordering
$cp1Ord :: forall a b. Ord b => Eq (Table a b)
Ord, Int -> Table a b -> ShowS
[Table a b] -> ShowS
Table a b -> String
(Int -> Table a b -> ShowS)
-> (Table a b -> String)
-> ([Table a b] -> ShowS)
-> Show (Table a b)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall a b. Show b => Int -> Table a b -> ShowS
forall a b. Show b => [Table a b] -> ShowS
forall a b. Show b => Table a b -> String
showList :: [Table a b] -> ShowS
$cshowList :: forall a b. Show b => [Table a b] -> ShowS
show :: Table a b -> String
$cshow :: forall a b. Show b => Table a b -> String
showsPrec :: Int -> Table a b -> ShowS
$cshowsPrec :: forall a b. Show b => Int -> Table a b -> ShowS
Show, a -> Table a b -> Table a a
(a -> b) -> Table a a -> Table a b
(forall a b. (a -> b) -> Table a a -> Table a b)
-> (forall a b. a -> Table a b -> Table a a) -> Functor (Table a)
forall a b. a -> Table a b -> Table a a
forall a b. (a -> b) -> Table a a -> Table a b
forall a a b. a -> Table a b -> Table a a
forall a a b. (a -> b) -> Table a a -> Table a b
forall (f :: Type -> Type).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> Table a b -> Table a a
$c<$ :: forall a a b. a -> Table a b -> Table a a
fmap :: (a -> b) -> Table a a -> Table a b
$cfmap :: forall a a b. (a -> b) -> Table a a -> Table a b
Functor, a -> Table a a -> Bool
Table a m -> m
Table a a -> [a]
Table a a -> Bool
Table a a -> Int
Table a a -> a
Table a a -> a
Table a a -> a
Table a a -> a
(a -> m) -> Table a a -> m
(a -> m) -> Table a a -> m
(a -> b -> b) -> b -> Table a a -> b
(a -> b -> b) -> b -> Table a a -> b
(b -> a -> b) -> b -> Table a a -> b
(b -> a -> b) -> b -> Table a a -> b
(a -> a -> a) -> Table a a -> a
(a -> a -> a) -> Table a a -> a
(forall m. Monoid m => Table a m -> m)
-> (forall m a. Monoid m => (a -> m) -> Table a a -> m)
-> (forall m a. Monoid m => (a -> m) -> Table a a -> m)
-> (forall a b. (a -> b -> b) -> b -> Table a a -> b)
-> (forall a b. (a -> b -> b) -> b -> Table a a -> b)
-> (forall b a. (b -> a -> b) -> b -> Table a a -> b)
-> (forall b a. (b -> a -> b) -> b -> Table a a -> b)
-> (forall a. (a -> a -> a) -> Table a a -> a)
-> (forall a. (a -> a -> a) -> Table a a -> a)
-> (forall a. Table a a -> [a])
-> (forall a. Table a a -> Bool)
-> (forall a. Table a a -> Int)
-> (forall a. Eq a => a -> Table a a -> Bool)
-> (forall a. Ord a => Table a a -> a)
-> (forall a. Ord a => Table a a -> a)
-> (forall a. Num a => Table a a -> a)
-> (forall a. Num a => Table a a -> a)
-> Foldable (Table a)
forall a. Eq a => a -> Table a a -> Bool
forall a. Num a => Table a a -> a
forall a. Ord a => Table a a -> a
forall m. Monoid m => Table a m -> m
forall a. Table a a -> Bool
forall a. Table a a -> Int
forall a. Table a a -> [a]
forall a. (a -> a -> a) -> Table a a -> a
forall a a. Eq a => a -> Table a a -> Bool
forall a a. Num a => Table a a -> a
forall a a. Ord a => Table a a -> a
forall m a. Monoid m => (a -> m) -> Table a a -> m
forall a m. Monoid m => Table a m -> m
forall a a. Table a a -> Bool
forall a a. Table a a -> Int
forall a a. Table a a -> [a]
forall b a. (b -> a -> b) -> b -> Table a a -> b
forall a b. (a -> b -> b) -> b -> Table a a -> b
forall a a. (a -> a -> a) -> Table a a -> a
forall a m a. Monoid m => (a -> m) -> Table a a -> m
forall a b a. (b -> a -> b) -> b -> Table a a -> b
forall a a b. (a -> b -> b) -> b -> Table a a -> b
forall (t :: Type -> Type).
(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 :: Table a a -> a
$cproduct :: forall a a. Num a => Table a a -> a
sum :: Table a a -> a
$csum :: forall a a. Num a => Table a a -> a
minimum :: Table a a -> a
$cminimum :: forall a a. Ord a => Table a a -> a
maximum :: Table a a -> a
$cmaximum :: forall a a. Ord a => Table a a -> a
elem :: a -> Table a a -> Bool
$celem :: forall a a. Eq a => a -> Table a a -> Bool
length :: Table a a -> Int
$clength :: forall a a. Table a a -> Int
null :: Table a a -> Bool
$cnull :: forall a a. Table a a -> Bool
toList :: Table a a -> [a]
$ctoList :: forall a a. Table a a -> [a]
foldl1 :: (a -> a -> a) -> Table a a -> a
$cfoldl1 :: forall a a. (a -> a -> a) -> Table a a -> a
foldr1 :: (a -> a -> a) -> Table a a -> a
$cfoldr1 :: forall a a. (a -> a -> a) -> Table a a -> a
foldl' :: (b -> a -> b) -> b -> Table a a -> b
$cfoldl' :: forall a b a. (b -> a -> b) -> b -> Table a a -> b
foldl :: (b -> a -> b) -> b -> Table a a -> b
$cfoldl :: forall a b a. (b -> a -> b) -> b -> Table a a -> b
foldr' :: (a -> b -> b) -> b -> Table a a -> b
$cfoldr' :: forall a a b. (a -> b -> b) -> b -> Table a a -> b
foldr :: (a -> b -> b) -> b -> Table a a -> b
$cfoldr :: forall a a b. (a -> b -> b) -> b -> Table a a -> b
foldMap' :: (a -> m) -> Table a a -> m
$cfoldMap' :: forall a m a. Monoid m => (a -> m) -> Table a a -> m
foldMap :: (a -> m) -> Table a a -> m
$cfoldMap :: forall a m a. Monoid m => (a -> m) -> Table a a -> m
fold :: Table a m -> m
$cfold :: forall a m. Monoid m => Table a m -> m
Foldable, (forall x. Table a b -> Rep (Table a b) x)
-> (forall x. Rep (Table a b) x -> Table a b)
-> Generic (Table a b)
forall x. Rep (Table a b) x -> Table a b
forall x. Table a b -> Rep (Table a b) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a b x. Rep (Table a b) x -> Table a b
forall a b x. Table a b -> Rep (Table a b) x
$cto :: forall a b x. Rep (Table a b) x -> Table a b
$cfrom :: forall a b x. Table a b -> Rep (Table a b) x
Generic)

-- | Pretty-print a Table as a 'mkTable' expression.
--
-- @
--     λ> pp $ (tabulate (even . finToInt) :: Table (Fin 3) Bool )
--     mkTable (\\case { 0 -> True; 1 -> False; 2 -> True })
-- @
instance (Finite a, Portray a, Portray b) => Portray (Table a b) where
  portray :: Table a b -> Portrayal
portray (Table Vec (Cardinality a) b
xs) = Portrayal -> [Portrayal] -> Portrayal
Apply (Ident -> Portrayal
Name Ident
"mkTable") ([Portrayal] -> Portrayal) -> [Portrayal] -> Portrayal
forall a b. (a -> b) -> a -> b
$ Portrayal -> [Portrayal]
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Portrayal -> [Portrayal]) -> Portrayal -> [Portrayal]
forall a b. (a -> b) -> a -> b
$ [(Portrayal, Portrayal)] -> Portrayal
LambdaCase ([(Portrayal, Portrayal)] -> Portrayal)
-> [(Portrayal, Portrayal)] -> Portrayal
forall a b. (a -> b) -> a -> b
$
    (a -> b -> (Portrayal, Portrayal))
-> [a] -> [b] -> [(Portrayal, Portrayal)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\a
a b
b -> (a -> Portrayal
forall a. Portray a => a -> Portrayal
portray a
a, b -> Portrayal
forall a. Portray a => a -> Portrayal
portray b
b)) (Finite a => [a]
forall a. Finite a => [a]
enumerate @a) (Vec (Cardinality a) b -> [b]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
toList Vec (Cardinality a) b
xs)

instance (Finite a, Portray a, Diff b) => Diff (Table a b) where
  diff :: Table a b -> Table a b -> Maybe Portrayal
diff (Table Vec (Cardinality a) b
xs) (Table Vec (Cardinality a) b
ys) =
    if Bool
hasDiff
      then Portrayal -> Maybe Portrayal
forall a. a -> Maybe a
Just (Portrayal -> Maybe Portrayal) -> Portrayal -> Maybe Portrayal
forall a b. (a -> b) -> a -> b
$ Portrayal -> [Portrayal] -> Portrayal
Apply (Ident -> Portrayal
Name Ident
"mkTable") ([Portrayal] -> Portrayal) -> [Portrayal] -> Portrayal
forall a b. (a -> b) -> a -> b
$ Portrayal -> [Portrayal]
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Portrayal -> [Portrayal]) -> Portrayal -> [Portrayal]
forall a b. (a -> b) -> a -> b
$ [(Portrayal, Portrayal)] -> Portrayal
LambdaCase ([(Portrayal, Portrayal)] -> Portrayal)
-> [(Portrayal, Portrayal)] -> Portrayal
forall a b. (a -> b) -> a -> b
$
             (if Bool
allDiff then [(Portrayal, Portrayal)] -> [(Portrayal, Portrayal)]
forall a. a -> a
id else ([(Portrayal, Portrayal)]
-> [(Portrayal, Portrayal)] -> [(Portrayal, Portrayal)]
forall a. [a] -> [a] -> [a]
++ [(Text -> Portrayal
Opaque Text
"_", Text -> Portrayal
Opaque Text
"_")])) ([(Portrayal, Portrayal)] -> [(Portrayal, Portrayal)])
-> [(Portrayal, Portrayal)] -> [(Portrayal, Portrayal)]
forall a b. (a -> b) -> a -> b
$
             [Maybe (Portrayal, Portrayal)] -> [(Portrayal, Portrayal)]
forall a. [Maybe a] -> [a]
catMaybes [Maybe (Portrayal, Portrayal)]
labeledDiffs
      else Maybe Portrayal
forall a. Maybe a
Nothing
   where
    (Any Bool
hasDiff, All Bool
allDiff) = (Maybe (Portrayal, Portrayal) -> (Any, All))
-> [Maybe (Portrayal, Portrayal)] -> (Any, All)
forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap
      (\Maybe (Portrayal, Portrayal)
x -> (Bool -> Any
Any (Maybe (Portrayal, Portrayal) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (Portrayal, Portrayal)
x), Bool -> All
All (Maybe (Portrayal, Portrayal) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (Portrayal, Portrayal)
x)))
      [Maybe (Portrayal, Portrayal)]
labeledDiffs
    labeledDiffs :: [Maybe (Portrayal, Portrayal)]
labeledDiffs = (a -> b -> b -> Maybe (Portrayal, Portrayal))
-> [a] -> [b] -> [b] -> [Maybe (Portrayal, Portrayal)]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3
      (\a
a b
x b
y -> (Portrayal, Maybe Portrayal) -> Maybe (Portrayal, Portrayal)
forall (t :: Type -> Type) (f :: Type -> Type) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA (a -> Portrayal
forall a. Portray a => a -> Portrayal
portray a
a, b -> b -> Maybe Portrayal
forall a. Diff a => a -> a -> Maybe Portrayal
diff b
x b
y))
      (Finite a => [a]
forall a. Finite a => [a]
enumerate @a)
      (Vec (Cardinality a) b -> [b]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
toList Vec (Cardinality a) b
xs)
      (Vec (Cardinality a) b -> [b]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
toList Vec (Cardinality a) b
ys)

instance NFData a => NFData (Table k a) where
  rnf :: Table k a -> ()
rnf (Table Vec (Cardinality k) a
vec) = Vec (Cardinality k) a -> ()
forall a. NFData a => a -> ()
rnf Vec (Cardinality k) a
vec

instance (Finite k, Serialize a) => Serialize (Table k a) where
  get :: Get (Table k a)
get = Table k (Get a) -> Get (Table k a)
forall (t :: Type -> Type) (f :: Type -> Type) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA (Table k (Get a) -> Get (Table k a))
-> Table k (Get a) -> Get (Table k a)
forall a b. (a -> b) -> a -> b
$ (k -> Get a) -> Table k (Get a)
forall a b. Finite a => (a -> b) -> Table a b
mkTable (Get a -> k -> Get a
forall a b. a -> b -> a
const Get a
forall t. Serialize t => Get t
get)
  put :: Putter (Table k a)
put = (a -> PutM ()) -> Putter (Table k a)
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ a -> PutM ()
forall t. Serialize t => Putter t
put

instance Finite a => Applicative (Table a) where
  pure :: a -> Table a a
pure = (a -> a) -> Table a a
forall (f :: Type -> Type) a.
Representable f =>
(Rep f -> a) -> f a
tabulate ((a -> a) -> Table a a) -> (a -> a -> a) -> a -> Table a a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a -> a
forall a b. a -> b -> a
const
  liftA2 :: (a -> b -> c) -> Table a a -> Table a b -> Table a c
liftA2 a -> b -> c
f Table a a
x Table a b
y = (Rep (Table a) -> c) -> Table a c
forall (f :: Type -> Type) a.
Representable f =>
(Rep f -> a) -> f a
tabulate ((a -> b -> c) -> (a -> a) -> (a -> b) -> a -> c
forall (f :: Type -> Type) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> b -> c
f (Table a a -> Rep (Table a) -> a
forall (f :: Type -> Type) a. Representable f => f a -> Rep f -> a
index Table a a
x) (Table a b -> Rep (Table a) -> b
forall (f :: Type -> Type) a. Representable f => f a -> Rep f -> a
index Table a b
y))
  Table a (a -> b)
f <*> :: Table a (a -> b) -> Table a a -> Table a b
<*> Table a a
x = (Rep (Table a) -> b) -> Table a b
forall (f :: Type -> Type) a.
Representable f =>
(Rep f -> a) -> f a
tabulate (Table a (a -> b) -> Rep (Table a) -> a -> b
forall (f :: Type -> Type) a. Representable f => f a -> Rep f -> a
index Table a (a -> b)
f (a -> a -> b) -> (a -> a) -> a -> b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Table a a -> Rep (Table a) -> a
forall (f :: Type -> Type) a. Representable f => f a -> Rep f -> a
index Table a a
x)

instance (Finite a, Default b) => Default (Table a b) where
  def :: Table a b
def = b -> Table a b
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure b
forall a. Default a => a
def

-- | 'Data.Profunctor.lmap' for a constrained 'Data.Profunctor.Profunctor'.
lmapTable :: (Finite b, Finite c) => (b -> c) -> Table c a -> Table b a
lmapTable :: (b -> c) -> Table c a -> Table b a
lmapTable b -> c
f Table c a
t = (Rep (Table b) -> a) -> Table b a
forall (f :: Type -> Type) a.
Representable f =>
(Rep f -> a) -> f a
tabulate ((Rep (Table b) -> a) -> Table b a)
-> (Rep (Table b) -> a) -> Table b a
forall a b. (a -> b) -> a -> b
$ \Rep (Table b)
x -> Table c a
t Table c a -> Rep (Table c) -> a
forall (f :: Type -> Type) a. Representable f => f a -> Rep f -> a
`index` b -> c
f b
Rep (Table b)
x

instance Finite a => Traversable (Table a) where
  traverse :: (a -> f b) -> Table a a -> f (Table a b)
traverse a -> f b
f (Table Vec (Cardinality a) a
vec) = Vec (Cardinality a) b -> Table a b
forall a b. Vec (Cardinality a) b -> Table a b
Table (Vec (Cardinality a) b -> Table a b)
-> f (Vec (Cardinality a) b) -> f (Table a b)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Vec (Cardinality a) a -> f (Vec (Cardinality a) b)
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Vec (Cardinality a) a
vec

instance Finite a => Distributive (Table a) where
  collect :: (a -> Table a b) -> f a -> Table a (f b)
collect a -> Table a b
f f a
fa =
    let fgb :: f (Table a b)
fgb = a -> Table a b
f (a -> Table a b) -> f a -> f (Table a b)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
fa
    in  Vec (Cardinality a) (f b) -> Table a (f b)
forall a b. Vec (Cardinality a) b -> Table a b
Table (Vec (Cardinality a) (f b) -> Table a (f b))
-> Vec (Cardinality a) (f b) -> Table a (f b)
forall a b. (a -> b) -> a -> b
$ SInt (Cardinality a)
-> (Fin (Cardinality a) -> f b) -> Vec (Cardinality a) (f b)
forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
VE.mkVec (Finite a => SInt (Cardinality a)
forall a. Finite a => SInt (Cardinality a)
cardinality @a) (\Fin (Cardinality a)
i -> (Table a b -> a -> b) -> a -> Table a b -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip Table a b -> a -> b
forall (f :: Type -> Type) a. Representable f => f a -> Rep f -> a
index (Fin (Cardinality a) -> a
forall a. Finite a => Fin (Cardinality a) -> a
fromFin Fin (Cardinality a)
i) (Table a b -> b) -> f (Table a b) -> f b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Table a b)
fgb)

instance Finite a => Representable (Table a) where
  type Rep (Table a) = a
  tabulate :: (Rep (Table a) -> a) -> Table a a
tabulate Rep (Table a) -> a
f = Vec (Cardinality a) a -> Table a a
forall a b. Vec (Cardinality a) b -> Table a b
Table (Vec (Cardinality a) a -> Table a a)
-> Vec (Cardinality a) a -> Table a a
forall a b. (a -> b) -> a -> b
$ SInt (Cardinality a)
-> (Fin (Cardinality a) -> a) -> Vec (Cardinality a) a
forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
VE.mkVec (Finite a => SInt (Cardinality a)
forall a. Finite a => SInt (Cardinality a)
cardinality @a) (a -> a
Rep (Table a) -> a
f (a -> a) -> (Fin (Cardinality a) -> a) -> Fin (Cardinality a) -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Fin (Cardinality a) -> a
forall a. Finite a => Fin (Cardinality a) -> a
fromFin)
  index :: Table a a -> Rep (Table a) -> a
index (Table Vec (Cardinality a) a
vec) Rep (Table a)
i = Vec (Cardinality a) a
vec Vec (Cardinality a) a -> Fin (Cardinality a) -> a
forall (n :: Nat) a. Vec n a -> Fin n -> a
V.! a -> Fin (Cardinality a)
forall a. Finite a => a -> Fin (Cardinality a)
toFin a
Rep (Table a)
i

instance Finite a => FunctorWithIndex a (Table a) where imap :: (a -> a -> b) -> Table a a -> Table a b
imap = (a -> a -> b) -> Table a a -> Table a b
forall (r :: Type -> Type) a a'.
Representable r =>
(Rep r -> a -> a') -> r a -> r a'
imapRep
instance Finite a => FoldableWithIndex a (Table a) where ifoldMap :: (a -> a -> m) -> Table a a -> m
ifoldMap = (a -> a -> m) -> Table a a -> m
forall (r :: Type -> Type) m a.
(Representable r, Foldable r, Monoid m) =>
(Rep r -> a -> m) -> r a -> m
ifoldMapRep
instance Finite a => TraversableWithIndex a (Table a) where
  itraverse :: (a -> a -> f b) -> Table a a -> f (Table a b)
itraverse = (a -> a -> f b) -> Table a a -> f (Table a b)
forall (r :: Type -> Type) (f :: Type -> Type) a a'.
(Representable r, Traversable r, Applicative f) =>
(Rep r -> a -> f a') -> r a -> f (r a')
itraverseRep

#if !MIN_VERSION_lens(5,0,0)
instance Finite a => L.FunctorWithIndex a (Table a) where imap = imapRep
instance Finite a => L.FoldableWithIndex a (Table a) where ifoldMap = ifoldMapRep
instance Finite a => L.TraversableWithIndex a (Table a) where
  itraverse = itraverseRep
#endif

-- | The identity morphism of a constrained category of 'Table's.
idTable :: Finite a => Table a a
idTable :: Table a a
idTable = (Rep (Table a) -> a) -> Table a a
forall (f :: Type -> Type) a.
Representable f =>
(Rep f -> a) -> f a
tabulate Rep (Table a) -> a
forall a. a -> a
id

-- | The composition of a constrained category of 'Table's.
composeTable :: (Finite a, Finite b) => Table b c -> Table a b -> Table a c
composeTable :: Table b c -> Table a b -> Table a c
composeTable Table b c
tbc Table a b
tab = (Rep (Table a) -> c) -> Table a c
forall (f :: Type -> Type) a.
Representable f =>
(Rep f -> a) -> f a
tabulate ((Rep (Table a) -> c) -> Table a c)
-> (Rep (Table a) -> c) -> Table a c
forall a b. (a -> b) -> a -> b
$ Table b c -> Rep (Table b) -> c
forall (f :: Type -> Type) a. Representable f => f a -> Rep f -> a
index Table b c
tbc (b -> c) -> (a -> b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Table a b -> Rep (Table a) -> b
forall (f :: Type -> Type) a. Representable f => f a -> Rep f -> a
index Table a b
tab

-- | 'traverse' a function whose argument is a finite enumerable type.
traverseRep
  :: forall x a b f
   . (Finite x, Applicative f)
  => (a -> f b) -> (x -> a) -> f (x -> b)
traverseRep :: (a -> f b) -> (x -> a) -> f (x -> b)
traverseRep a -> f b
f = (Table x b -> x -> b) -> f (Table x b) -> f (x -> b)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Table x b -> x -> b
forall (f :: Type -> Type) a. Representable f => f a -> Rep f -> a
index (f (Table x b) -> f (x -> b))
-> ((x -> a) -> f (Table x b)) -> (x -> a) -> f (x -> b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> f b) -> Table x a -> f (Table x b)
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f (Table x a -> f (Table x b))
-> ((x -> a) -> Table x a) -> (x -> a) -> f (Table x b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a.
Representable (Table x) =>
(Rep (Table x) -> a) -> Table x a
forall (f :: Type -> Type) a.
Representable f =>
(Rep f -> a) -> f a
tabulate @(Table _)

-- | Memoize a function by using a 'Vec' as a lazy lookup table.
--
-- Given a function whose argument is a 'Finite' type, return a new function
-- that looks up the argument in a table constructed by applying the original
-- function to every possible value.  Since 'Vec' stores its elements boxed,
-- none of the applications of @f@ in the table are forced until they're forced
-- by calling the memoized function and forcing the result.
memoize :: Finite a => (a -> b) -> a -> b
memoize :: (a -> b) -> a -> b
memoize = Table a b -> a -> b
forall (f :: Type -> Type) a. Representable f => f a -> Rep f -> a
index (Table a b -> a -> b)
-> ((a -> b) -> Table a b) -> (a -> b) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a.
Representable (Table a) =>
(Rep (Table a) -> a) -> Table a a
forall (f :: Type -> Type) a.
Representable f =>
(Rep f -> a) -> f a
tabulate @(Table _)

-- | An 'Iso' between two 'Representable' Functors with the same 'Rep' type.
retabulated
  :: (Representable f, Representable g, Rep f ~ Rep g)
  => Iso (f a) (f b) (g a) (g b)
retabulated :: Iso (f a) (f b) (g a) (g b)
retabulated = AnIso (Rep g -> b) (Rep g -> a) (f b) (f a)
-> Iso (f a) (f b) (Rep g -> a) (Rep g -> b)
forall s t a b. AnIso s t a b -> Iso b a t s
from AnIso (Rep g -> b) (Rep g -> a) (f b) (f a)
forall (f :: Type -> Type) (g :: Type -> Type)
       (p :: Type -> Type -> Type) (h :: Type -> Type) a b.
(Representable f, Representable g, Profunctor p, Functor h) =>
p (f a) (h (g b)) -> p (Rep f -> a) (h (Rep g -> b))
tabulated (p (Rep g -> a) (f (Rep g -> b)) -> p (f a) (f (f b)))
-> (p (g a) (f (g b)) -> p (Rep g -> a) (f (Rep g -> b)))
-> p (g a) (f (g b))
-> p (f a) (f (f b))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. p (g a) (f (g b)) -> p (Rep g -> a) (f (Rep g -> b))
forall (f :: Type -> Type) (g :: Type -> Type)
       (p :: Type -> Type -> Type) (h :: Type -> Type) a b.
(Representable f, Representable g, Profunctor p, Functor h) =>
p (f a) (h (g b)) -> p (Rep f -> a) (h (Rep g -> b))
tabulated

-- | Infix 'index', monomorphized.
(!) :: Finite a => Table a b -> a -> b
(!) = Table a b -> a -> b
forall (f :: Type -> Type) a. Representable f => f a -> Rep f -> a
index

-- | Lens on a single element.
ix :: Finite a => a -> Lens' (Table a b) b
ix :: a -> Lens' (Table a b) b
ix a
a = a
a a
-> ((b -> f b) -> Table a b -> f (Table a b))
-> (b -> f b)
-> Table a b
-> f (Table a b)
`seq` (Table a b -> b)
-> (Table a b -> b -> Table a b) -> Lens' (Table a b) b
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens (Table a b -> a -> b
forall a b. Finite a => Table a b -> a -> b
! a
a) (\(Table Vec (Cardinality a) b
vec) b
b -> Vec (Cardinality a) b -> Table a b
forall a b. Vec (Cardinality a) b -> Table a b
Table (Vec (Cardinality a) b
vec Vec (Cardinality a) b
-> (Vec (Cardinality a) b -> Vec (Cardinality a) b)
-> Vec (Cardinality a) b
forall a b. a -> (a -> b) -> b
& Fin (Cardinality a) -> Lens' (Vec (Cardinality a) b) b
forall (n :: Nat) a. Fin n -> Lens' (Vec n a) a
V.ix (a -> Fin (Cardinality a)
forall a. Finite a => a -> Fin (Cardinality a)
toFin a
a) ((b -> Identity b)
 -> Vec (Cardinality a) b -> Identity (Vec (Cardinality a) b))
-> b -> Vec (Cardinality a) b -> Vec (Cardinality a) b
forall s t a b. ASetter s t a b -> b -> s -> t
.~ b
b))

-- | Monomorphized 'tabulate'.  Can be useful for type ambiguity reasons.
mkTable :: Finite a => (a -> b) -> Table a b
mkTable :: (a -> b) -> Table a b
mkTable = (a -> b) -> Table a b
forall (f :: Type -> Type) a.
Representable f =>
(Rep f -> a) -> f a
tabulate

-- | Convenience function for building any 'Representable' as if by 'traverse'.
--
-- > tabulateA f = sequenceA (tabulate f) = traverse f (tabulate id)
tabulateA
  :: (Traversable t, Representable t, Applicative f)
  => (Rep t -> f b) -> f (t b)
tabulateA :: (Rep t -> f b) -> f (t b)
tabulateA = t (f b) -> f (t b)
forall (t :: Type -> Type) (f :: Type -> Type) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA (t (f b) -> f (t b))
-> ((Rep t -> f b) -> t (f b)) -> (Rep t -> f b) -> f (t b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Rep t -> f b) -> t (f b)
forall (f :: Type -> Type) a.
Representable f =>
(Rep f -> a) -> f a
tabulate