Van Laarhoven versions of Jules Hedges' stochastic lenses

jules.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}

import Data.Coerce (coerce)
import Data.Functor.Compose (Compose (..))

data Vector
data Stochastic a
instance Functor Stochastic
instance Applicative Stochastic
instance Monad Stochastic


type JulesLens s t a b = s -> Stochastic (a, (Vector, b) -> Stochastic (Vector, t))


class (Functor f, Functor g) => HasStochWithV f g | f -> g, g -> f where
  stochastic :: f a -> Stochastic (g (Stochastic (Vector, a)))
  unstochastic :: Stochastic (g (Stochastic (Vector, a))) -> f a

type JulesLensVL s t a b = forall f g. HasStochWithV f g => (a -> f b) -> s -> f t


toVL :: JulesLens s t a b -> JulesLensVL s t a b
toVL l afb s = unstochastic $ do
  (a, bt) <- l s
  fmap (>>= bt) <$> stochastic (afb a)


-- This instance is used by fromVL
instance
  HasStochWithV
    (Compose Stochastic (Compose (Compose ((,) a) ((->) b)) (Compose Stochastic ((,) Vector))))
    (Compose ((,) a) ((->) b))
  where
  stochastic = coerce
  unstochastic = coerce

fromVL :: JulesLensVL s t a b -> JulesLens s t a b
fromVL l s = coerce . stochastic $ l (\a -> unstochastic . pure . Compose $ (a, pure)) s


-- for comparison the conversion for regular lenses:

toVLLens :: Functor f => (s -> (a, b -> t)) -> (a -> f b) -> (s -> f t)
toVLLens l afb s = case l s of (a, bt) -> fmap bt (afb a)

fromVLLens :: (forall f. Functor f => (a -> f b) -> (s -> f t)) -> s -> (a, b -> t)
fromVLLens l s = getCompose $ l (\a -> Compose $ (a, id)) s