Work-in-Progress distributive functors a la Aaron Vargo

Distributive.hs

{-# language PolyKinds #-}
{-# language RankNTypes #-}
{-# language TypeOperators #-}
{-# language DefaultSignatures #-}
{-# language TypeFamilies #-}
{-# language FlexibleContexts #-}
{-# language GeneralizedNewtypeDeriving #-}
{-# language DeriveGeneric #-}
{-# language DeriveTraversable #-}
{-# language DeriveAnyClass #-}
{-# language DerivingStrategies #-}
{-# language DerivingVia #-}

import Data.Functor
import Data.Functor.Compose
import Data.Functor.Product
import Data.Functor.Identity
import Data.Kind
import GHC.Generics hiding (Rep)

type f ~> g = forall x. f x -> g x

class FFunctor w where
  ffmap :: (f ~> g) -> w f -> w g

data Person f = Person
  { name :: f String
  , age :: f Int 
  }

instance FFunctor Person where
  ffmap f2g (Person fs fi) = Person (f2g fs) (f2g fi)

class FContravariant w where
  fcontramap :: (f ~> g) -> w g -> w f

newtype Log f = Log { runLog :: forall a. f a -> a }

logId :: Log Identity
logId = Log runIdentity
 
instance FContravariant Log where
  fcontramap f (Log g) = Log (g . f)

indexLog :: f a -> Log f -> a
indexLog fa (Log fa2a) = fa2a fa

newtype Tabulate x f = Tabulate { runTabulate :: Log f -> x }

instance FFunctor (Tabulate x) where
  ffmap f (Tabulate logf2a) = Tabulate $ logf2a . fcontramap f

tabulateLog :: Distributive f => (Log f -> a) -> f a
tabulateLog logf2a = dist (Tabulate logf2a) <&> 
  \(Tabulate logId2a) -> logId2a logId

gdist :: (Generic1 f, Distributive (Rep1 f), FFunctor w) => w f -> f (w Identity)
gdist = to1 . dist . ffmap from1

class Functor f => Distributive f where
  type Rep f
  type Rep f = Log f

  -- distribute :: Functor g => g (f a) -> f (g a)
  dist :: FFunctor w => w f -> f (w Identity)


  default dist :: (Generic1 f, Distributive (Rep1 f), FFunctor w) => w f -> f (w Identity)
  dist = gdist

  index    :: f a -> Rep f -> a
  default index :: (Rep f ~ Log f) => f a -> Rep f -> a
  index = indexLog

  tabulate :: (Rep f -> a) -> f a
  default tabulate :: (Rep f ~ Log f) => (Rep f -> a) -> f a
  tabulate = tabulateLog

dpure :: Distributive f => a -> f a
dpure = tabulate . const 

data Ap a b f = Ap (f (a -> b)) (f a)

instance FFunctor (Ap a b) where
  ffmap f2g (Ap fab fa) = Ap (f2g fab) (f2g fa)

dap :: Distributive f => f (a -> b) -> f a -> f b
dap fab fa = dist (Ap fab fa) <&> \(Ap (Identity ab) (Identity a)) -> ab a

data Bind a b f = Bind (f a) (a -> f b)

instance FFunctor (Bind a b) where
  ffmap f2g (Bind fa a2fb) = Bind (f2g fa) (f2g . a2fb)

dbind :: Distributive f => f a -> (a -> f b) -> f b
dbind fa a2fb = dist (Bind fa a2fb) <&> \(Bind (Identity a) a2ib) ->
  runIdentity (a2ib a)

newtype Dist f a = Dist (f a)
  deriving Functor

instance Distributive f => Applicative (Dist f) where
  pure = Dist . dpure
  Dist fab <*> Dist fa = Dist (dap fab fa)

instance Distributive U1 where
  dist _ = U1

instance (Distributive f, Distributive g) => Distributive (f :*: g) where
  dist wpfg = dist (ffmap (\(f:*:_) -> f) wpfg)
          :*: dist (ffmap (\(_:*:g) -> g) wpfg)

instance Distributive f => Distributive (M1 i c f) where
  dist = M1 . dist . ffmap unM1

instance Distributive ((->) x) where
  dist wx2a x = ffmap (\f -> Identity $ f x) wx2a

-- newtype (f :.: g) a = Comp1 { unComp1 :: f (g a) }

data Post w g f = Post { runPost :: w (f :.: g) }

instance FFunctor w => FFunctor (Post w g) where
  ffmap f2h (Post wcfg) = Post $ ffmap (Comp1 . f2h . unComp1) wcfg

instance (Distributive f, Distributive g) => Distributive (f :.: g) where
  dist wcfg = Comp1 $ dist (Post wcfg) <&>  
    dist . ffmap (runIdentity . unComp1) . runPost

instance Distributive Par1 where
  dist = Par1 . ffmap (Identity . unPar1)

instance Distributive Identity

instance (Distributive f, Distributive g) => Distributive (Compose f g)

instance (Distributive f, Distributive g) => Distributive (Product f g)

instance Distributive f => Distributive (Rec1 f) where
  dist = Rec1 . dist . ffmap unRec1

data Pair a = Pair a a
  deriving (Generic, Generic1, Functor, Distributive)
  deriving (Applicative) via (Dist Pair)

data Stream a = a :- Stream a
  deriving (Generic, Generic1, Functor, Distributive)
  deriving (Applicative) via (Dist Stream)

instance Distributive Stream where
  type Rep f = Natural
  -- tabulate f = ...
  -- index ...

{-
instance Comonad Stream  where
  extract = extractBy 0
  duplicate = duplicateBy (+)
-}

-- Mealy a b ~ NonEmpty a -> b
newtype Mealy a b = Mealy (a -> (b, Mealy a b)
  deriving (Generic, Generic1, Functor, Distributive)
  deriving (Applicative) via (Dist (Moore a))

-- Moore a b ~ [a] -> b
-- type Rep (Moore a) = [a]
data Moore a b = Moore b (a -> Moore a b)
  deriving (Generic, Generic1, Functor, Distributive)
  deriving (Applicative) via (Dist (Moore a))

This is a fun formulation. This is what I came up with from playing with it

collect_ :: Distributive f => FFunctor ho => (ho Identity -> b) -> ho f -> f b
collect_ f = fmap f . dist

class Eval ho where
  eval :: ho a Identity -> a

instance Eval (Bind a) where
  eval :: Bind a b Identity -> b
  eval (Identity as :>>= kont) = runIdentity (kont as)
  
instance Eval (Ap a) where
  eval :: Ap a b Identity -> b
  eval (Identity f :<*> Identity a) = f a

instance Eval Tabulate where
  eval :: Tabulate a Identity -> a
  eval (Tabulate log) = log logId

tabulateLog :: Distributive f => (Log f -> b) -> f b
tabulateLog = collect_ eval . Tabulate

dap :: forall f a b. Distributive f => f (a -> b) -> (f a -> f b)
dap fs as = collect_ eval (fs :<*> as)

dbind :: Distributive f => f a -> (a -> f b) -> f b
dbind as cont = collect_ eval (as :>>= cont)