Trees That Grow

ttg.hs

{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE StandaloneDeriving #-}

import Control.Lens
import Data.Void(Void, absurd)

-- page 6
-- https://www.microsoft.com/en-us/research/uploads/prod/2016/11/trees-that-grow.pdf

data Typ =
  IntTyp
  | Fun Typ Typ
  deriving (Eq, Ord, Show)

data Exp =
  Lit Int
  | Var String
  | Ann Exp Typ
  | Abs String Exp
  | App Exp Exp
  deriving (Eq, Ord, Show)
  
-- Expₓ ξ
data Exp' e =
  Lit' !(XLit e) Int
  | Var' !(XVar e) String
  | Ann' !(XAnn e) (Exp' e) Typ
  | Abs' !(XAbs e) String (Exp' e)
  | App' !(XApp e) (Exp' e) (Exp' e)
  | Exp' !(XExp e)

deriving instance (Eq (XLit e), Eq (XVar e), Eq (XAnn e), Eq (XAbs e), Eq (XApp e), Eq (XExp e)) =>
    Eq (Exp' e)
deriving instance (Ord (XLit e), Ord (XVar e), Ord (XAnn e), Ord (XAbs e), Ord (XApp e), Ord (XExp e)) =>
    Ord (Exp' e)
deriving instance (Show (XLit e), Show (XVar e), Show (XAnn e), Show (XAbs e), Show (XApp e), Show (XExp e)) => Show (Exp' e)

class HasExp' a e | a -> e where
  exp ::
    Lens' a (Exp' e)
  xexp ::
    (
      XLit e ~ x
    , XVar e ~ x
    , XAnn e ~ x
    , XAbs e ~ x
    , XApp e ~ x
    , XExp e ~ Void
    ) =>
    Lens' a x

instance HasExp' (Exp' e) e where
  exp =
    id
  xexp f (Lit' x i) =
    fmap (\x' -> Lit' x' i) (f x)
  xexp f (Var' x s) =
    fmap (\x' -> Var' x' s) (f x)
  xexp f (Ann' x e t) =
    fmap (\x' -> Ann' x' e t) (f x)
  xexp f (Abs' x s e) =
    fmap (\x' -> Abs' x' s e) (f x)
  xexp f (App' x e1 e2) =
    fmap (\x' -> App' x' e1 e2) (f x)
  xexp _ (Exp' x) =
    absurd x

class AsExp' a e | a -> e where
  _Exp' ::
    Prism' a (Exp' e)
  _Lit ::
    Prism' a (XLit e, Int)
  _Lit' ::
    XLit e ~ () =>
    Prism' a Int
  _Lit' =
    _Lit . unproduct
  _Var ::
    Prism' a (XVar e, String)
  _Var' ::
    XVar e ~ () =>
    Prism' a String
  _Var' =
    _Var . unproduct
  _Ann ::
    Prism' a (XAnn e, (Exp' e, Typ))
  _Ann' ::
    XAnn e ~ () =>
    Prism' a (Exp' e, Typ)
  _Ann' =
    _Ann . unproduct
  _Abs ::
    Prism' a (XAbs e, (String, Exp' e))
  _Abs' ::
    XAbs e ~ () =>
    Prism' a (String, Exp' e)
  _Abs' =
    _Abs . unproduct
  _App ::
    Prism' a (XApp e, (Exp' e, Exp' e))
  _App' ::
    XApp e ~ () =>
    Prism' a (Exp' e, Exp' e)
  _App' =
    _App . unproduct
  _XExp ::
    Prism' a (XExp e)

instance AsExp' (Exp' e) e where
  _Exp' =
    id
  _Lit =
    prism'
      (\(x, i) -> Lit' x i)
      (
        \case
          Lit' x i ->
            Just (x, i)
          _ ->
            Nothing
      )
  _Var =
    prism'
      (\(x, s) -> Var' x s)
      (
        \case
          Var' x s ->
            Just (x, s)
          _ ->
            Nothing
      )
  _Ann =
    prism'
      (\(x, (e, t)) -> Ann' x e t)
      (
        \case
          Ann' x e t ->
            Just (x, (e, t))
          _ ->
            Nothing
      )
  _Abs =
    prism'
      (\(x, (s, e)) -> Abs' x s e)
      (
        \case
          Abs' x s e ->
            Just (x, (s, e))
          _ ->
            Nothing
      )
  _App =
    prism'
      (\(x, (e1, e2)) -> App' x e1 e2)
      (
        \case
          App' x e1 e2 ->
            Just (x, (e1, e2))
          _ ->
            Nothing
      )
  _XExp =
    prism'
      Exp'
      (
        \case
          Exp' x ->
            Just x
          _ ->
            Nothing
      )

type family XLit e
type family XVar e
type family XAnn e
type family XAbs e 
type family XApp e
type family XExp e

-- Expᵁᴰ
type Exp'' =
  Exp' ()

-- use () not Void for undecorated type

type instance XLit () =
  ()
type instance XVar () =
  ()
type instance XAnn () =
  ()
type instance XAbs () =
  ()
type instance XApp () =
  ()
type instance XExp () =
  Void

-- pattern Litᵁᴰ
pattern Lit'' ::
  Int
  -> Exp''
pattern Lit'' i <- Lit' _ i
  where Lit'' i = Lit' () i

-- pattern Varᵁᴰ
pattern Var'' ::
  String
  -> Exp''
pattern Var'' s <- Var' _ s
  where Var'' s = Var' () s

-- pattern Annᵁᴰ
pattern Ann'' ::
  Exp''
  -> Typ
  -> Exp''
pattern Ann'' e t <- Ann' _ e t
  where Ann'' e t = Ann' () e t

-- pattern Absᵁᴰ
pattern Abs'' ::
  String
  -> Exp''
  -> Exp''
pattern Abs'' s e <- Abs' _ s e
  where Abs'' s e = Abs' () s e

-- pattern Appᵁᴰ
pattern App'' ::
  Exp''
  -> Exp''
  -> Exp''
pattern App'' e1 e2 <- App' _ e1 e2
  where App'' e1 e2 = App' () e1 e2

pattern Exp'' ::
  Void
  -> Exp''
pattern Exp'' v <- Exp' v
  where Exp'' v = Exp' v

incLit ::
  Exp' e
  -> Exp' e
incLit (Lit' x i) =
  Lit' x (i + 1)
incLit e =
  e

----

data AddField =
  AddField Int
  deriving (Eq, Show)

wrapAddField ::
  Iso' AddField Int
wrapAddField =
  iso
    (\(AddField n) -> n)
    AddField

type ExpAddField =
  Exp' AddField

type instance XLit AddField =
  AddField
type instance XVar AddField =
  AddField
type instance XAnn AddField =
  AddField
type instance XAbs AddField =
  AddField
type instance XApp AddField =
  AddField
type instance XExp AddField =
  Void

testExpAddFieldPattern ::
  ExpAddField
  -> Int
testExpAddFieldPattern (Lit' (AddField n) _) =
  n
testExpAddFieldPattern (Var' (AddField n) _) =
  n
testExpAddFieldPattern (Ann' (AddField n) _ _) =
  n
testExpAddFieldPattern (Abs' (AddField n) _ _) =
  n
testExpAddFieldPattern (App' (AddField n) _ _) =
  n
testExpAddFieldPattern (Exp' v) =
  absurd v

incExpAddFieldPattern ::
  ExpAddField
  -> ExpAddField
incExpAddFieldPattern (Lit' (AddField n) i) =
  incExpAddFieldPattern (Lit' (AddField (n + 1)) i)
incExpAddFieldPattern (Var' (AddField n) s) =
  incExpAddFieldPattern (Var' (AddField (n + 1)) s)
incExpAddFieldPattern (Ann' (AddField n) t u) =
  incExpAddFieldPattern (Ann' (AddField (n + 1)) t u)
incExpAddFieldPattern (Abs' (AddField n) s t) =
  incExpAddFieldPattern (Abs' (AddField (n + 1)) s t)
incExpAddFieldPattern (App' (AddField n) e1 e2) =
  incExpAddFieldPattern (App' (AddField (n + 1)) e1 e2)
incExpAddFieldPattern (Exp' v) =
  absurd v

testExpAddFieldLens ::
  ExpAddField
  -> Int
testExpAddFieldLens =
  view (xexp . wrapAddField)

incExpAddFieldLens ::
  ExpAddField
  -> ExpAddField
incExpAddFieldLens =
  over (xexp . wrapAddField) (+1)

----

data AddCtor =
  AddCtor Int
  deriving (Eq, Show)

wrapAddCtor ::
  Iso' AddCtor Int
wrapAddCtor =
  iso
    (\(AddCtor n) -> n)
    AddCtor

type ExpAddCtor =
  Exp' AddCtor

type instance XLit AddCtor =
  ()
type instance XVar AddCtor =
  ()
type instance XAnn AddCtor =
  ()
type instance XAbs AddCtor =
  ()
type instance XApp AddCtor =
  ()
type instance XExp AddCtor =
  AddCtor

testExpAddCtorPattern ::
  ExpAddCtor
  -> Maybe Int
testExpAddCtorPattern (Exp' (AddCtor x)) =
  Just x
testExpAddCtorPattern _ =
  Nothing

incExpAddCtorPattern ::
  ExpAddCtor
  -> ExpAddCtor
incExpAddCtorPattern (Exp' (AddCtor n)) =
  (Exp' (AddCtor (n + 1)))
incExpAddCtorPattern e =
  e

testExpAddCtorPrism ::
  ExpAddCtor
  -> Maybe Int
testExpAddCtorPrism =
  preview (_XExp . wrapAddCtor)

incExpAddCtorPrism ::
  ExpAddCtor
  -> ExpAddCtor
incExpAddCtorPrism =
  over (_XExp . wrapAddCtor) (+1)

---- do these exist somewhere?

unproduct ::
  Iso ((), a) ((), b) a b
unproduct =
  iso
    (\((), a) -> a)
    (\a -> ((), a))

uncoproduct ::
  Iso (Either Void a) (Either Void b) a b
uncoproduct =
  iso
    (either absurd id) 
    Right