Traversal.hs

{-# language RankNTypes #-}
{-# language GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
module Data.Optics where

import Data.Vec.Lazy as Vec
import Data.Profunctor
import Data.Type.Nat
import Data.Maybe (fromJust)


data Cofree f where
  Run :: Functor f => (a, f (Cofree f)) -> Cofree f -- \int^{f a} Φ_f(a) = (a, f (Φ_f)
  Cons :: a -> Cofree f -> Cofree f


data Analytic (ts :: Cofree f) a where
  Simple :: t -> Analytic ('Cons t ts) a
  Complex :: t -> Analytic ts a -> Analytic ('Cons t ts) a

data ExTraversal (n :: Nat) s t a b where
  ExT :: (s -> (Vec n a, c)) -> ((Vec n b, c) -> t) -> ExTraversal n s t a b

type ConcTraversal s t a b = s -> ([a], [b] -> t)

mkConcTraversal :: SNatI n => ExTraversal n s t a b -> ConcTraversal s t a b
mkConcTraversal (ExT f g) s = (Vec.toList $ fst (f s), \bs -> g (fromJust $ Vec.fromList bs, snd (f s)))

mkExTraversal :: SNatI n => ConcTraversal s t a b -> ExTraversal n s t a b
mkExTraversal f = ExT (\s -> (fromJust $ Vec.fromList $ fst (f s), s)) (\(bs, c) -> snd (f c) $ Vec.toList bs)

class Profunctor p => FinPoly p where
  natl :: SNatI n => p a b -> p (Vec n a, c) (Vec n b, c)
  natr :: SNatI n => p a b -> p (c, Vec n a) (c, Vec n b)

class FinPoly p => Polynomial p where
  stretchl :: p a b -> p ([a], c) ([b], c)
  stretchr :: p a b -> p (c, [a]) (c, [b])

type FinTraversal s t a b = forall p. FinPoly p => p a b -> p s t
type Traversal s t a b = forall p. Polynomial p => p a b -> p s t


newtype Theta p a b = Theta (forall c. p ([a], c) ([b], c))

data Phi p a b where
  Phi :: (a -> ([u], c)) -> p u v -> (b -> ([v], c)) -> Phi p a b

finTraversal :: SNatI n => (s -> (Vec n a, c)) -> ((Vec n b, c) -> t) -> FinTraversal s t a b
finTraversal f g pab = dimap f g (natl pab)

traversal :: (s -> ([a], c), ([b], c) -> t) -> Traversal s t a b
traversal (f, g) pab = dimap f g (stretchl pab)

-- example

data Tree a = Leaf a | Node (Tree a) (Tree a)
  deriving (Eq, Show, Functor, Foldable, Traversable)

toShape_ :: Tree a -> Tree ()
toShape_ = fmap (const ())

fromList_ :: Tree () -> [a] -> Tree a
fromList_ t as = fst (go t as)
  where
    go :: Tree () -> [a] -> (Tree a, [a])
    go (Leaf ()) (a : as) = (Leaf a, as)
    go (Node l r) as =
      let (l', as1) = go l as
          (r', as') = go r as1
      in (Node l' r', as')


t :: Tree Int
t = Node (Node (Leaf 1) (Leaf 2)) (Node (Leaf 3) (Node (Leaf 4) (Leaf 5)))

sh :: Tree ()
sh = toShape_ t