dyck.hs

{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE KindSignatures       #-}
{-# LANGUAGE PolyKinds            #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE StandaloneDeriving   #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE UnicodeSyntax        #-}

{-# OPTIONS_GHC -Wall -fno-warn-unticked-promoted-constructors #-}

import Data.Kind     (Type)
import Data.Foldable (foldlM,foldrM)
import Data.List     (inits,tails)

data Nat = Z | S Nat

data DyckSuff (n :: Nat) :: Type where
  Done :: DyckSuff Z
  Open :: DyckSuff (S n) -> DyckSuff n
  Clos :: DyckSuff n     -> DyckSuff (S n)

type Dyck = DyckSuff Z

deriving instance Eq (DyckSuff n)
deriving instance Ord (DyckSuff n)

instance Show (DyckSuff n) where
  showsPrec _ xs' = showChar '\"' . toString xs'
    where
      toString :: DyckSuff n -> String -> String
      toString Done      = showChar '\"'
      toString (Open xs) = showChar '(' . toString xs
      toString (Clos xs) = showChar ')' . toString xs

data Natty (n :: Nat) :: Type where
  Zy :: Natty Z
  Sy :: Natty n -> Natty (S n)

enumDyck :: Int -> [Dyck]
enumDyck sz = go Zy sz Done []
  where
    go, zero, left, right :: Natty n -> Int -> DyckSuff n -> [Dyck] -> [Dyck]
    
    go n m k = zero n m k . left n m k . right n m k

    zero Zy 0 k = (k:)
    zero _  _ _ = id
    
    left (Sy n) m k = go n m (Open k)
    left Zy     _ _ = id
    
    right _ 0 _ = id
    right n m k = go (Sy n) (m-1) (Clos k)

infixr 5 :-
data Stack (a :: Type) (n :: Nat)  :: Type where
  Nil  :: Stack a Z
  (:-) :: a -> Stack a n -> Stack a (S n)

data Tree = Leaf | Tree :*: Tree deriving (Eq, Ord, Show)

dyckToTree :: Dyck -> Tree
dyckToTree dy = go dy (Leaf :- Nil)
  where
    go :: DyckSuff n -> Stack Tree (S n) -> Tree
    go (Open d) ts               = go d (Leaf :- ts)
    go (Clos d) (t1 :- t2 :- ts) = go d (t2 :*: t1 :- ts)
    go Done     (t  :- Nil)      = t
    
treeToDyck :: Tree -> Dyck
treeToDyck t = go t Done
  where
    go :: Tree -> DyckSuff n -> DyckSuff n
    go Leaf        = id
    go (xs :*: ys) = go xs . Open . go ys . Clos
    
data Expr (a :: Type) where
  Val   :: a -> Expr a
  (:+:) :: Expr a -> Expr a -> Expr a
  
data Code (n :: Nat) (a :: Type) :: Type where
  HALT :: Code (S Z) a
  PUSH :: a -> Code (S n) a -> Code n a
  ADD  :: Code (S n) a -> Code (S (S n)) a
  
type family (+) (n :: Nat) (m :: Nat) :: Nat where
  Z   + m = m
  S n + m = S (n + m)

exec :: Code n Int -> Stack Int (n + m) -> Stack Int (S m)
exec HALT         st              = st
exec (PUSH v is)  st              = exec is (v :- st)
exec (ADD    is) (t1 :- t2 :- st) = exec is (t2 + t1 :- st)

comp :: Expr a -> Code Z a
comp e = comp' e HALT
  where
    comp' :: Expr a -> Code (S n) a -> Code n a
    comp' (Val     x) = PUSH x
    comp' (xs :+: ys) = comp' xs . comp' ys . ADD

foldlCode :: (∀ n. a -> b n -> b (S n))
          -> (∀ n. b (S (S n)) -> b (S n))
          -> b m
          -> Code m a -> b (S Z)
foldlCode _ _ h  HALT       = h
foldlCode p a h (PUSH x xs) = foldlCode p a (p x h) xs
foldlCode p a h (ADD    xs) = foldlCode p a (a   h) xs

shift :: Int -> Stack Int n -> Stack Int (S n)
shift x xs = x :- xs

reduce :: Stack Int (S (S n)) -> Stack Int (S n)
reduce (t1 :- t2 :- st) = t2 + t1 :- st

execFold :: Code Z Int -> Int
execFold = pop . foldlCode shift reduce Nil

pop :: Stack a (S n) -> a
pop (x :- _) = x

enumTrees :: [a] -> [Expr a]
enumTrees = fmap (foldl1 (flip (:+:))) . foldlM f []
  where    
    f []         v = [[Val v]]
    f [t1]       v = [[Val v, t1]]
    f (t1:t2:st) v = (Val v : t1 : t2 : st) : f ((t2 :+: t1) : st) v
    
data Rose a = a :& Forest a
type Forest a = [Rose a]

dyckToForest :: Dyck -> Forest ()
dyckToForest dy = go dy ([] :- Nil)
  where
    go :: DyckSuff n -> Stack (Forest ()) (S n) -> Forest ()
    go (Open d) ts               = go d ([] :- ts)
    go (Clos d) (t1 :- t2 :- ts) = go d ((() :& t2 : t1) :- ts)
    go Done     (t  :- Nil)      = t

forestToDyck :: Forest () -> Dyck
forestToDyck t = go t Done
  where
    go :: Forest () -> DyckSuff n -> DyckSuff n
    go []          = id
    go ((() :& x):xs) = go x . Open . go xs . Clos
    
enumForests :: [a] -> [Forest a]
enumForests = foldrM f []
  where
    f x xs = zipWith ((:) . (:&) x) (inits xs) (tails xs)