EADT with profunctor lenses and prisms

LensEADT.purs

module Main where

import Prelude

import Control.Lazy (fix)
import Control.MonadZero (guard, (<|>))
import Data.Foldable (oneOfMap)
import Data.Functor.Mu (Mu, roll, unroll)
import Data.Functor.Variant (VariantF)
import Data.Functor.Variant as VF
import Data.Maybe (Maybe(..), maybe)
import Data.Profunctor (dimap)
import Data.Symbol (class IsSymbol)
import Data.Traversable (class Foldable, class Traversable, foldlDefault, foldrDefault, oneOf, sequenceDefault, traverse)
import Data.Lens (AnIso, Iso, Prism', Traversal', iso, over, prism', re, wander, withIso, (^.), (^?))
import Data.Tuple (Tuple(..), uncurry)
import Effect (Effect)
import Effect.Console (log)
import Matryoshka (Algebra, CoalgebraM, cata, traverseR)
import Prim.Row as Row
import Type.Equality (class TypeEquals)
import Type.Equality as TE

transformOf ∷ forall a b. ((a -> b) -> a -> b) -> (b -> b) -> a -> b
transformOf = fix (\r l f x -> f (over l (r l f) x))

rewriteOf ∷ forall a b. ((a -> b) -> a -> b) -> (b -> Maybe a) -> a -> b
rewriteOf = fix (\r l f -> transformOf l (\v -> maybe v (r l f) (f v)))

from ∷ forall s t a b. AnIso s t a b -> Iso b a t s
from l = withIso l $ \ sa bt -> iso bt sa

type RowApply (f ∷ # Type -> # Type) (a ∷ # Type) = f a
infixr 0 type RowApply as +

type EADT t = Mu (VariantF t)

injEADT 
  ∷ forall f s a b
  . Row.Cons s (VF.FProxy f) a b 
  => IsSymbol s 
  => Functor f 
  => VF.SProxy s 
  -> Algebra f (EADT b)
injEADT label = roll <<< VF.inj label

prjEADT 
  :: forall f s a b
  . Row.Cons s (VF.FProxy f) a b 
  => IsSymbol s
  => Functor f
  => VF.SProxy s 
  -> CoalgebraM Maybe f (EADT b)
prjEADT label = VF.prj label <<< unroll

_VariantF
  ∷ forall l f v a
   . IsSymbol l
  => Functor f
  => Row.Cons l (VF.FProxy f) _ v
  => VF.SProxy l
  -> Prism' (VF.VariantF v a) (f a)
_VariantF l = prism' (VF.inj l) (VF.prj l)

_EADT 
  :: forall l f v
  . Row.Cons l (VF.FProxy f) _ v 
  => IsSymbol l 
  => Functor f 
  => VF.SProxy l 
  -> Prism' (EADT v) (f (EADT v))
_EADT l = prism' (injEADT l) (prjEADT l) 

plateMu ∷ forall f. Traversable f => Traversal' (Mu f) (Mu f)
plateMu = wander go where
  go ∷ forall g. Applicative g => (Mu f -> g (Mu f)) -> Mu f -> g (Mu f)
  go = traverseR <<< traverse

data ValF a = ValF Int
derive instance functorValF ∷ Functor ValF

instance foldableValF :: Foldable ValF where
  foldl f z (ValF a) = z
  foldr f z (ValF a) = z
  foldMap f _ = mempty

instance traversableValF :: Traversable ValF where
  sequence = sequenceDefault
  traverse f (ValF a) = pure (ValF a)

data AddF a = AddF a a
derive instance functorAddF ∷ Functor AddF

instance foldableAddF :: Foldable AddF where
  foldl f z (AddF a b) = f (f z a) b
  foldr f z (AddF a b) = f a (f b z)
  foldMap f (AddF a b) = f a <> f b

instance traversableAddF :: Traversable AddF where
  sequence = sequenceDefault
  traverse f (AddF a b) = AddF <$> f a <*> f b

data MulF a = MulF a a
derive instance functorMulF ∷ Functor MulF

instance foldableMulF :: Foldable MulF where
  foldl f = foldlDefault f
  foldr f = foldrDefault f
  foldMap f (MulF a b) = f a <> f b

instance traversableMulF :: Traversable MulF where
  sequence = sequenceDefault
  traverse f (MulF a b) = MulF <$> f a <*> f b

data AnnF a e = AnnF a e
derive instance functorAnnF ∷ Functor (AnnF a)

type Val r = (val ∷ VF.FProxy ValF | r)

type Add r = (add ∷ VF.FProxy AddF | r)

type Mul r = (mul ∷ VF.FProxy MulF | r)

type Ann a r = (ann ∷ VF.FProxy (AnnF a) | r)

type BaseExpr r = Val + Add + r

_val = VF.SProxy ∷ _ "val"
_add = VF.SProxy ∷ _ "add"
_mul = VF.SProxy ∷ _ "mul"
_ann = VF.SProxy ∷ _ "ann"

_Mu ∷ forall f g. Iso (f (Mu f)) (g (Mu g)) (Mu f) (Mu g)
_Mu = iso roll unroll

class AsValF s a | s -> a where
  _ValF ∷ Prism' s Int

instance asValFValF ∷ AsValF (ValF a) a where
  _ValF = prism' ValF (\(ValF a) -> Just a)

else instance asValFVariant :: (Functor f, AsValF (f a) a, TypeEquals (VariantF ( val :: VF.FProxy f | tail ) a) (VariantF row a)) => AsValF (VariantF row a) a where
  _ValF = dimap TE.from TE.to <<< _VariantF _val <<< _ValF

else instance asValFFMu :: (Functor f, AsValF (f (Mu f)) a) => AsValF (Mu f) a where
  _ValF = re _Mu <<< _ValF

----------------------------------------------------------------

class AsAddF s a | s -> a where
  _AddF ∷ Prism' s (Tuple a a)

instance asAddFAddF ∷ AsAddF (AddF a) a where
  _AddF = prism' (uncurry AddF) (\(AddF a b) -> Just (Tuple a b))

else instance asAddFVariant :: (Functor f, AsAddF (f a) a, TypeEquals (VariantF ( add :: VF.FProxy f | tail ) a) (VariantF row a)) => AsAddF (VariantF row a) a where
  _AddF = dimap TE.from TE.to <<< _VariantF _add <<< _AddF

else instance asAddFMu :: (Functor f, AsAddF (f (Mu f)) a) => AsAddF (Mu f) a where
  _AddF = re _Mu <<< _AddF

----------------------------------------------------------------

class AsMulF s a | s -> a where
  _MulF ∷ Prism' s (Tuple a a)

instance asMulFMulF ∷ AsMulF (MulF a) a where
  _MulF = prism' (uncurry MulF) (\(MulF a b) -> Just (Tuple a b))

else instance asMulFVariant :: (Functor f, AsMulF (f a) a, TypeEquals (VariantF ( mul :: VF.FProxy f | tail ) a) (VariantF row a)) => AsMulF (VariantF row a) a where
  _MulF = dimap TE.from TE.to <<< _VariantF _mul <<< _MulF

else instance asMulFMu :: (Functor f, AsMulF (f (Mu f)) a) => AsMulF (Mu f) a where
  _MulF = re _Mu <<< _MulF

val ∷ forall r. Int -> EADT (Val r)
val v = v ^. re _ValF 

add ∷ forall r. EADT (Add + r) -> EADT (Add + r) -> EADT (Add + r)
add x y = Tuple x y ^. re _AddF

mul ∷ forall r. EADT (Mul + r) -> EADT (Mul + r) -> EADT (Mul + r)
mul x y = Tuple x y ^. re _MulF

optimize :: forall m. Traversable m => Array (Mu m -> Maybe (Mu m)) -> Mu m -> Mu m
optimize = rewriteOf plateMu <<< flip (oneOfMap <<< (#))

---- Optimizations

elimPlusZero :: forall r. EADT (Add + Val + r) -> Maybe (EADT (Add + Val + r))
elimPlusZero m = do
  Tuple x y <- m ^? _AddF
  y <$ is0 x <|> x <$ is0 y
    where
      is0 v = guard <<< (_ == 0) =<< v ^? _ValF

elimMulZero :: forall r. EADT (BaseExpr + Mul + r) -> Maybe (EADT (BaseExpr + Mul + r))
elimMulZero m = do
  Tuple x y <- m ^? _MulF
  val 0 <$ is0 x <|> val 0 <$ is0 y
    where
      is0 v = guard <<< (_ == 0) =<< v ^? _ValF

distr ∷ forall r. EADT (Add + Mul + r) -> Maybe (EADT (Add + Mul + r))
distr m = do
  Tuple a b <- m ^? _MulF
  oneOf [ do
            Tuple c d <- b ^? _AddF
            pure $ add (mul a c) (mul a d)
        , do
            Tuple c d <- a ^? _AddF
            pure $ add (mul b c) (mul b d)
        ]

---- Algebras

exprAlg ∷ forall r. Algebra (VariantF r) Int -> Algebra (VariantF (BaseExpr + r)) Int
exprAlg = VF.onMatch
  { val: case _ of ValF x -> x
  , add: case _ of AddF x y -> x + y }

exprAlg2 ∷ forall r. Algebra (VariantF r) Int -> Algebra (VariantF (BaseExpr + Mul + r)) Int
exprAlg2 = exprAlg
  >>> VF.on _mul case _ of MulF x y -> x * y

exprShowAlg ∷ forall r. Algebra (VariantF r) String -> Algebra (VariantF (BaseExpr + r)) String
exprShowAlg = VF.onMatch
  { val: case _ of ValF x -> show x
  , add: case _ of AddF x y -> "(" <> x <> " + " <> y <> ")" }

exprShowAlg2 ∷ forall r. Algebra (VariantF r) String -> Algebra (VariantF (BaseExpr + Mul + r)) String
exprShowAlg2 = exprShowAlg
  >>> VF.on _mul case _ of MulF x y -> "(" <> x <> " * " <> y <> ")"

expr3 ∷ EADT (Val + Add + Mul + ())
expr3 = mul (add (val 10) (val 20)) (add (val 30) (val 40))

expr4 ∷ EADT (Val + Add + Mul + ())
expr4 = add (mul (add (val 10) (val 0)) (add (val 30) (mul (val 40) (val 0)))) (val 10)

main :: Effect Unit
main = do
  log $ cata (VF.case_ # exprShowAlg2) expr4
  log $ cata (VF.case_ # exprShowAlg2) $ optimize [elimMulZero, elimPlusZero] expr4

  log "----------------------------------------------------------------"

  log $ cata (VF.case_ # exprShowAlg2) expr3
  log $ cata (VF.case_ # exprShowAlg2) $ optimize [distr] expr3

  log "----------------------------------------------------------------"

  log $ cata (VF.case_ # exprShowAlg2) expr3
  log $ cata (VF.case_ # exprShowAlg2) $ optimize [distr] expr3