Efficient Free Applicatives/Monads from Okasaki's Functional Data Structures, with an application to fast substitution-based lambda normalising

FastApp.hs

{-# LANGUAGE GADTs, RankNTypes, TypeOperators, ScopedTypeVariables, KindSignatures #-}

module FastApp where

import Prelude hiding (head, tail)
import Control.Applicative

-- Chris Okasaki invented many cool and efficient functional data structures
-- in his book Purely Functional Data Structures.
-- Among them, one is catenable lists supporting amortised O(1)-time |(++)|,
-- |head| and |tail| (by contrast, the ordinary |[]| doesn't support O(1)
-- |(++)|, and the difference list (or Church encoding of |[]|) doesn't support
-- O(1) |head| or |tail|.
--
-- Viewed abstractly, lists are just free monoids on the category of sets. Thus
-- an interesting question is that how far we can generalise Okasaki's data
-- structures to other categories. In this file, I show Okasaki's catenable lists 
-- can be generalised to the monoidal cateogry of endofunctors with a Day tensor,
-- in which monoids are the so-called applicatives functors.


-- We work in the category of endofunctors |* -> *| with morphisms:
--   type (~>) f g = forall a. f a -> g a

-- The monoidal product is the following Day tensor whose unit is the identity
-- functor:
--    type (:*:) f g a = exists b. (f (b -> a), g b)
--       -- or equivalently 
--       -- type (:*:) f g a = exists b c. (f c, g b, c -> b -> a)


-- The free applicative is the free monoids on this category.
-- From the definition we can see that it is roughly a list of operations from |sig|.
data FreeA sig a where
  Pure :: a -> FreeA sig a
  Ap   :: sig b -> FreeA sig c -> (b -> c -> a) -> FreeA sig a
     -- Ap :: f (:*:) FreeA sig ~> FreeA sigo

instance Functor (FreeA sig) where
  fmap f (Pure x)    = Pure (f x)
  fmap f (Ap x xs k) = Ap x xs (\b c -> f (k b c))

instance Applicative (FreeA sig) where
  pure = Pure
  (Pure k)    <*> ys = fmap k ys
  (Ap x xs k) <*> ys = Ap x (fmap (\c a b' -> k b' c a) xs <*> ys) (flip ($))
    -- <*> is list append
    -- TODO: factor out the cryptic stuff manipulating the monoid

reverseA :: forall sig a. FreeA sig a -> FreeA sig a
reverseA xs = go xs (Pure id) where
  go :: forall d. FreeA sig d -> FreeA sig (d -> a) -> FreeA sig a
  go (Pure e)    ys = fmap ($ e) ys
  go (Ap x xs k) ys = go xs (Ap x ys (\b d2a c -> d2a (k b c)))


-- Following Okasaki's catenable lists, first we need a purely functional
-- queue. In this context, this means a representable of the free applicative
-- supporting amortised O(1)-time appending an operation to the end of a
-- computation.
data QueueA sig a where
  Q :: Int -> FreeA sig b -> Int -> FreeA sig c -> (b -> c -> a) -> QueueA sig a

-- Smart constructor maintaining the invariant:
q :: Int -> FreeA sig b -> Int -> FreeA sig c -> (b -> c -> a) -> QueueA sig a
q lenf f lenr r k
  | lenr <= lenf  = Q lenf f lenr r k
  | otherwise     = Q (lenf + lenr) (liftA2 k f (reverseA r)) 0 (Pure ()) (\x _ -> x)


instance Functor (QueueA sig) where
  fmap g (Q lenf f lenr r k) = Q lenf f lenr r (\b c -> g (k b c))

pureQ :: a -> QueueA sig a
pureQ x = q 0 (Pure x) 0 (Pure ()) (\x _ -> x)

isPure :: QueueA sig a -> Maybe a
isPure (Q _ (Pure b) _ (Pure c) k) = Just (k b c)
isPure _                           = Nothing

snoc :: QueueA sig b -> sig c -> (b -> c -> a) -> QueueA sig a
snoc (Q lenf f lenr r k) x k' = q lenf f (lenr + 1) (Ap x r (\c' c b -> k' (k b c) c')) (flip ($))

match :: QueueA sig a 
      -> (a -> d a)
      -> (forall b c. sig b -> QueueA sig c -> (b -> c -> a) -> d a)
      -> d a
match (Q _    (Pure x)     _   (Pure y)    k)  f _ = f (k x y)
match (Q lenf (Ap x xs k) lenr ys          k') _ g = g x (q (lenf - 1) xs lenr ys (\c c' b -> k' (k b c) c')) (flip ($))


-- Now that we have a queue, we can do catenable lists using Okasaki's
-- bootstrapping technique.

data CListA (sig :: * -> *) (a :: *) where
  E :: a -> CListA sig a
  C :: sig b -> QueueA (CListA sig) c -> (b -> c -> a) -> CListA sig a
    -- C :: sig :*: QueueA (CListA sig) ~> sig a

instance Functor (CListA sig) where
  fmap g (E e) = E (g e)
  fmap g (C x xs k) = C x xs (\b c -> g (k b c))

link :: CListA sig b -> CListA sig c -> (b -> c -> a) -> CListA sig a
link (C x q k) s k' = C x (snoc q s (\c c' b -> k' (k b c) c')) (flip ($))

pureC :: a -> CListA sig a
pureC e = E e

isPureC :: CListA sig a -> Maybe a
isPureC (E a) = Just a
isPureC _     = Nothing

append :: CListA sig b -> CListA sig c -> (b -> c -> a) -> CListA sig a
append xs (E c) k = fmap (\b -> k b c) xs
append (E b) ys k = fmap (\c -> k b c) ys
append xs ys    k = link xs ys k

inj :: sig a -> CListA sig a
inj x = C x (pureQ ()) (\x () -> x)

consC :: sig b -> CListA sig c -> (b -> c -> a) -> CListA sig a
consC x xs k = append (inj x) xs k

snocC :: CListA sig b -> sig c -> (b -> c -> a) -> CListA sig a
snocC xs x k = append xs (inj x) k

matchC :: CListA sig a 
       -> (a -> d a)
       -> (forall b c. sig b -> CListA sig c -> (b -> c -> a) -> d a)
       -> d a
matchC (E e) f _     = f e
matchC (C x q k) _ g = g x tail k where
  tail = case isPure q of
    Just c -> E c
    _      -> linkAll q

linkAll :: QueueA (CListA sig) a -> CListA sig a
linkAll q = 
  match q (error "This should be non-empty") 
          (\t q' k -> case isPure q' of 
                        (Just c) -> fmap (\b -> k b c) t
                        _        -> link t (linkAll q') k)

-- Finally! We have an efficient free applicative supporting both
--   amortised O(1)-time <*> and pattern matching!
instance Applicative (CListA sig) where
  pure = pureC
  xs <*> ys = append xs ys (\f a -> f a)


runAp :: Applicative g => (forall x. f x -> g x) -> CListA f a -> g a
runAp h xxs = matchC xxs pure (\x xs k -> pure k <*> h x <*> runAp h xs)

runAp_ :: Monoid m => (forall a. f a -> m) -> CListA f b -> m
runAp_ f = getConst . runAp (Const . f)

liftAp :: sig x -> CListA sig x
liftAp = inj

FastFree.hs

{-# LANGUAGE GADTs, RankNTypes, TypeOperators, ScopedTypeVariables, KindSignatures, DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances, FlexibleContexts, MultiParamTypeClasses, FunctionalDependencies, InstanceSigs #-}

module FastFree where

import Prelude hiding (head, tail)
import Control.Monad
import Data.Bool (bool)
import Data.Functor.Const

data Free sig a where
  Ret  :: a -> Free sig a
  Cons :: sig a -> (a -> Free sig b) -> Free sig b

instance Functor (Free sig) where
  fmap = liftM

instance Applicative (Free sig) where
  pure = return
  (<*>) = ap

instance Monad (Free sig) where
  return = Ret

  (Ret a)     >>= k = k a
  (Cons x xs) >>= k = Cons x (fmap (>>= k) xs)

-- Snoc lists
---------------------------

data FreeSnoc sig a where
  Ret' :: a -> FreeSnoc sig a
  Snoc :: FreeSnoc sig a -> (a -> sig b) -> FreeSnoc sig b

reverseF :: FreeSnoc sig a -> Free sig a
reverseF xs = go xs Ret where
  go :: FreeSnoc sig a -> (a -> Free sig b) -> Free sig b
  go (Ret' x) k    = k x
  go (Snoc sx x) k = go sx (\a -> Cons (x a) k)


-- Bankers' queue
--------------------------

data QueueF sig a where
  Q :: Int -> Free sig a -> Int -> (a -> FreeSnoc sig b) -> QueueF sig b
  -- Invariants : for all |Q lenf f lenr r|, 
  --   |lenf| is the depth for |f|, 
  --   |lenr| is the depth for |r a| for all |a|,
  --   and |lenr <= lenf|.

q :: Int -> Free sig a -> Int -> (a -> FreeSnoc sig b) -> QueueF sig b
q lenf f lenr r 
  | lenr <= lenf  = Q lenf f lenr r
  | otherwise     = Q (lenf + lenr) (f >>= (reverseF . r)) 0 Ret'

retQ :: a -> QueueF sig a
retQ a = q 0 (Ret a) 0 Ret' 

snoc :: QueueF sig a -> (a -> sig b) -> QueueF sig b
snoc (Q lenf f lenr r) k = q lenf f (lenr + 1) (\x -> Snoc (r x) k)


match :: QueueF sig a 
      -> (a -> d a)
      -> (forall b . sig b -> (b -> QueueF sig a) -> d a)
      -> d a
match (Q lenf (Ret b) lenr r) ret _  = case r b of Ret' a -> ret a
                                                   _      -> error "Invariant violation"

match (Q lenf (Cons x xs) lenr r) _ cons = cons x (\b -> q (lenf - 1) (xs b) lenr r)


-- Catenable lists for free monads
-------------------------------

data CListF (sig :: * -> *) (a :: *) where
  E :: a -> CListF sig a
  C :: sig b -> (b -> QueueF (CListF sig) a) -> CListF sig a

-- Technical note:
-- In the original algorithm from Okasaki in CatList.hs, the QueueF in CListF only stores
-- non-E CListF, but we can't decide if a |a -> CListF sig a| is E or not, so here
-- we do allow QueueF to store E. Thus the |link| function and |++| function in CatList.hs
-- become the same one here.

-- Concatenation is essentially snoc
append :: CListF sig a -> (a -> CListF sig b) -> CListF sig b
append (E a)    ys = ys a
append (C x xs) ys = C x (\b -> snoc (xs b) ys)

inj :: sig a -> CListF sig a
inj x = C x retQ

consC :: sig a -> (a -> CListF sig b) -> CListF sig b
consC x xs = append (inj x) xs

snocC :: CListF sig a -> (a -> sig b) -> CListF sig b
snocC xs x = append xs (\a -> inj (x a))


matchC :: CListF sig a 
       -> (a -> d)
       -> (forall b . sig b -> (b -> CListF sig a) -> d)
       -> d
matchC (E a) ret _     = ret a
matchC (C x xs) _ cons = cons x (fmap tail xs) 
  where tail xs' = linkAll xs'

-- This is a kind of concat
linkAll :: QueueF (CListF sig) a -> CListF sig a
linkAll q = match q E (\xs xss -> xs `append` (fmap linkAll xss))


instance Functor (CListF sig) where
  fmap = liftM

instance Applicative (CListF sig) where
  pure = return
  (<*>) = ap

-- CListF is a representation of Free supporting amortised O(1)-time
-- pattern matching and >>=.
instance Monad (CListF sig) where
  return = E
  (>>=) = append

-- Some simple tests
--------------------------

data State s a = Get (s -> a) | Put s a
  deriving Functor

class Monad m => MonadState m s | m -> s where
  get :: m s
  put :: s -> m ()

instance MonadState (CListF (State s)) s where
  get :: CListF (State s) s
  get = inj (Get id)
  
  put :: s -> CListF (State s) ()
  put s = inj (Put s ())

instance MonadState (Free (State s)) s where
  get = Cons (Get id) Ret
  put s = Cons (Put s ()) Ret


-- C++'s ++ operation
incr :: MonadState m Int => m Int
incr = do i <- get
          put (i+1)
          return i

handleSt :: CListF (State s) a -> s -> a
handleSt p = matchC p (\a s -> a) hdl 
  where hdl (Get x) k s    = handleSt (k (x s)) s
        hdl (Put s' x) k s = handleSt (k x) s'

handleSt' :: Free (State s) a -> s -> a
handleSt' p = case p of
  (Ret a) -> \s -> a
  (Cons f k) -> hdl f k
  where hdl (Get x) k s    = handleSt' (k (x s)) s
        hdl (Put s' x) k s = handleSt' (k x) s'


-- Basic correctness check
test0 = handleSt (incr >> incr >> (incr >> incr) >> get) 0 
  -- test0 == 4

-- CListF should be much faster than Free for left associated computations
incrN :: MonadState m Int => Int -> m Int
incrN 0 = return 0
incrN n = incrN (n-1) >> incr -- This is deliberately associated in the way

test1 = handleSt (incrN 10000 >> get) 0  -- This should be 10000 computed in several seconds

test1' = handleSt' (incrN 10000 >> get) 0  -- This should take a lot of time


-- CListF supports pattern matching 
test2 = getConst (go 10 (incrN 10000)) == "GPGPGPGPGP"
  where
    go 0 p = Const []
    go n p = matchC p (\_ -> Const []) 
                      (\op k -> case op of
                                 (Put _ a) -> Const ('P' : getConst (go (n - 1) (k a)))
                                 (Get a)   -> Const ('G' : getConst (go (n - 1) (k (a 0)))))

FastFreeMem.hs

{-# LANGUAGE GADTs, RankNTypes, TypeOperators, ScopedTypeVariables, KindSignatures, DeriveFunctor, BangPatterns #-}
{-# LANGUAGE FlexibleInstances, FlexibleContexts, MultiParamTypeClasses, FunctionalDependencies, InstanceSigs #-}

-- One possible reason that FastFree.hs is slow is that the monoidal product
--   (f . g) x = exists y. (f y, y -> g x)
-- has a function in the second binder, and functions are not memoized by
-- Haskell, which may destroys the lazy structure in Okasaki's algorithm.
-- Therefore this file (together with LamPHOASMem.hs) tries to memoize functions
-- in the data structure. However, it seems only the memoization in linkAll matters
-- in the test cases.

module FastFreeMem where

import Prelude hiding (head, tail)
import Data.Function.Memoize
import Control.Monad

type M a = (Memoizable a)
type M2 a b = (Memoizable a, Memoizable b)

class MonadM m where
  returnM :: M a => a -> m a
  bindM :: M2 a b => m a -> (a -> m b) -> m b

data Free sig a where
  Ret  :: M a => a -> Free sig a
  Cons :: M2 a b => sig a -> (a -> Free sig b) -> Free sig b

instance MonadM (Free sig) where
  returnM = Ret

  bindM (Ret a)     k = k a
  -- bindM (Cons x xs) k = Cons x (memoize (\a -> xs a `bindM` k))
  bindM (Cons x xs) k = Cons x (\a -> xs a `bindM` k)

matchF :: M a
       => Free sig a 
       -> (a -> d)
       -> (forall b . M b => sig b -> (b -> Free sig a) -> d)
       -> d
matchF (Ret a) ret _      = ret a
matchF (Cons x xs) _ cons = cons x xs

-- Snoc lists
---------------------------

data FreeSnoc sig a where
  Ret' :: M a => a -> FreeSnoc sig a
  Snoc :: M2 a b => FreeSnoc sig a -> (a -> sig b) -> FreeSnoc sig b

reverseF :: M a => FreeSnoc sig a -> Free sig a
reverseF xs = go xs Ret where
  go :: M2 a b => FreeSnoc sig a -> (a -> Free sig b) -> Free sig b
  go (Ret' x) k    = k x
  --go (Snoc sx x) k = go sx (memoize (\a -> Cons (x a) k))
  go (Snoc sx x) k = go sx (\a -> Cons (x a) k)


-- Bankers' queue
--------------------------

data QueueF sig a where
  Q :: M2 a b => !Int -> Free sig a -> !Int -> (a -> FreeSnoc sig b) -> QueueF sig b
  -- Invariants : for all |Q lenf f lenr r|, 
  --   |lenf| is the depth for |f|, 
  --   |lenr| is the depth for |r a| for all |a|,
  --   and |lenr <= lenf|.

q :: M2 a b => Int -> Free sig a -> Int -> (a -> FreeSnoc sig b) -> QueueF sig b
q lenf f lenr r 
  | lenr <= lenf  = Q lenf f lenr r
  -- | otherwise     = Q (lenf + lenr) (f `bindM` memoize (\b -> reverseF (r b))) 0 Ret'
  | otherwise     = Q (lenf + lenr) (f `bindM` (\b -> reverseF (r b))) 0 Ret'


retQ :: M a => a -> QueueF sig a
retQ a = q 0 (Ret a) 0 Ret' 

snoc :: M2 a b => QueueF sig a -> (a -> sig b) -> QueueF sig b
--snoc (Q lenf f lenr r) k = q lenf f (lenr + 1) (memoize (\x -> Snoc (r x) k))
snoc (Q lenf f lenr r) k = q lenf f (lenr + 1) (\x -> Snoc (r x) k)

match :: M a => QueueF sig a 
      -> (a -> d a)
      -> (forall b . M b => sig b -> (b -> QueueF sig a) -> d a)
      -> d a
match (Q lenf (Ret b) lenr r) ret _  = case r b of Ret' a -> ret a
                                                   _      -> error "Invariant violation"

--match (Q lenf (Cons x xs) lenr r) _ cons = cons x (memoize (\b -> q (lenf - 1) (xs b) lenr r))
match (Q lenf (Cons x xs) lenr r) _ cons = cons x (\b -> q (lenf - 1) (xs b) lenr r)


-- Catenable lists for free monads
-------------------------------

data CListF (sig :: * -> *) (a :: *) where
  E :: M a => a -> CListF sig a
  C :: M2 a b => sig b -> (b -> QueueF (CListF sig) a) -> CListF sig a

-- Technical note:
-- In the original algorithm from Okasaki in CatList.hs, the QueueF in CListF only stores
-- non-E CListF, but we can't decide if a |a -> CListF sig a| is E or not, so here
-- we do allow QueueF to store E. Thus the |link| function and |++| function in CatList.hs
-- become the same one here.

-- Concatenation is essentially snoc
append :: M2 a b => CListF sig a -> (a -> CListF sig b) -> CListF sig b
append (E a)    ys = ys a
--append (C x xs) ys = C x (memoize (\b -> snoc (xs b) ys))
append (C x xs) ys = C x (\b -> snoc (xs b) ys)

inj :: M a => sig a -> CListF sig a
inj x = C x retQ

consC :: M2 a b => sig a -> (a -> CListF sig b) -> CListF sig b
consC x xs = append (inj x) xs

snocC :: M2 a b => CListF sig a -> (a -> sig b) -> CListF sig b
--snocC xs x = append xs (memoize (\a -> inj (x a)))
snocC xs x = append xs (\a -> inj (x a))

matchC :: M a 
       => CListF sig a 
       -> (a -> d)
       -> (forall b . M b => sig b -> (b -> CListF sig a) -> d)
       -> d
matchC (E a) ret _     = ret a
--matchC (C x xs) _ cons = cons x (memoize (\b -> linkAll (xs b))) 
matchC (C x xs) _ cons = cons x (\b -> linkAll (xs b)) 

-- This is a kind of concat
linkAll :: M a => QueueF (CListF sig) a -> CListF sig a
linkAll q = match q E (\xs xss -> xs `append` memoize (\b -> linkAll (xss b)))
                                           -- ^^^^^^^
                                           --    |
                                           --  This seems the only useful memoize


-- CListF is a representation of Free supporting amortised O(1)-time
-- pattern matching and >>=.
instance MonadM (CListF sig) where
  returnM = E
  bindM  = append

FastFreeMixed.hs

{-# LANGUAGE GADTs, RankNTypes, TypeApplications, TypeOperators, ScopedTypeVariables, KindSignatures, DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances, FlexibleContexts, MultiParamTypeClasses, FunctionalDependencies, InstanceSigs, Strict #-}

-- Using FastFree to implement a normaliser of lambda calculus based on
-- parametric higher-order abstract syntax

import FastFreeMixed
import Data.Bool (bool)

-- for names we use an arbitrary type `a` that can be compared 
-- and has a stream of distinct elements
class Eq a => Atom a where
  ns :: [a]

  -- and we require a name can be printed as an identifier
  showName :: a -> String

instance Atom Int where
  ns = [1..]
  showName n = "x" ++ show n

data LamS a x = App x x | Abs (a -> x)

type LamP a = CListF (LamS a) a

type Lam = forall a. Atom a => LamP a

instance Atom a => Show (LamP a) where
  show t = fst (go ns t) where
    go :: [a] -> LamP a -> (String, [a])
    go ns t = matchC t (\a -> (showName a, ns)) f where
      f :: LamS a b -> (b -> LamP a) -> (String, [a])
      f (App fun opr) k = let (s1, ns') = go ns (k fun)
                              (s2, ns'') = go ns' (k opr)
                          in (s1 ++ " " ++ s2, ns'')
      f (Abs body) k = let (n:ns') = ns 
                           (s, ns'') = go ns' (k (body n))
                       in ("(λ " ++ showName n ++ ". " ++  s ++ ")", ns'')

printLam :: Lam -> IO ()
printLam t = print (t :: LamP Int)

app :: LamP a -> LamP a -> LamP a
app t1 t2 = consC (App False True) f where
  l = t1
  r = t2
  f False = l
  f True =  r

lam :: Atom a => (a -> LamP a) -> LamP a
lam t = consC (Abs id) t

var :: a -> LamP a
var = return

t1 :: Lam
t1 = app (lam (\x -> var x `app` var x)) (lam (\x -> var x))

-- replaces every occurrence of x in t with m
subst :: Atom a => LamP a -> a -> LamP a -> LamP a
subst t x m = t >>= (\y -> if y == x then m else var y)

-- evaluate a closed lambda term (in the lazy evaluation order)
-- no evaluation under binders
eval :: Lam -> Lam
eval t = fst (go ns t) where 
  go :: forall a. Atom a => [a] -> LamP a -> (LamP a, [a])
  go ns t = matchC t (\a -> (return a, ns)) f where
    f :: forall b. LamS a b -> (b -> LamP a) -> (LamP a, [a])
    f (Abs body)  k = (consC (Abs body) k, ns)
    f (App t1 t2) k = let (t1', ns') = go ns (k t1)

                          g :: forall b. LamS a b -> (b -> LamP a) -> (LamP a, [a])
                          g (Abs b) k2 = let (n:ns'') = ns'
                                         in go ns'' (subst (k2 (b n)) n (k t2))
                          g _       k2 = (app t1' (k t2), ns')

                      in  matchC t1' (\a -> (app t1' (k t2), ns')) g 

-- full normalisation in the applicative order
norm :: Lam -> Lam
norm t = fst (go ns t) where 
  go :: forall a. Atom a => [a] -> LamP a -> (LamP a, [a])
  go ns t = matchC t (\a -> (return a, ns)) f where
    f :: forall b. LamS a b -> (b -> LamP a) -> (LamP a, [a])
    f (Abs body)  k = let (n:ns') = ns                                   -- this looks odd! I am not sure we should
                          (body',ns'') = go ns' (k (body n))             -- concretise, normalise, and abstract. 
                                                                         -- Can we delay the normalisation until the binder
                                                                         -- actually gets a concrete atom? (That would 
                                                                         -- make fresh name management more difficult though!)
                      in (lam (\a -> subst body' n (var a)), ns'')
    f (App t1 t2) k = let (t1', ns') = go ns (k t1)
                          (t2', ns'') = go ns' (k t2)

                          g :: forall b. LamS a b -> (b -> LamP a) -> (LamP a, [a])
                          g (Abs b) k2 = let (n:ns''') = ns''
                                         in go ns''' (subst (k2 (b n)) n t2')
                          g _       k2 = (app t1' t2', ns'')

                      in  matchC t1' (\a -> (app t1' t2', ns'')) g 

t2 :: Lam
t2 = lam (\f -> var f `app` var f)
       `app`
       lam (\x -> lam (\y -> var x `app` var y))

-- Church numeral of 3
c3 :: Lam
c3 = lam (\f -> lam (\x -> var f `app` (var f `app` (var f `app` var x))))

infixl 4 `app` 
cadd :: Lam
cadd = lam (\n -> lam (\m -> lam (\f -> lam (\x -> 
          var n `app` var f 
            `app` (var m `app` var f `app` var x)))))

-- 6
t3 :: Lam
t3 = cadd `app` c3 `app` c3

-- printLam (norm t3)
-- (λ x1. (λ x2. x1 x1 x1 x1 x1 x1 x2))

cmul :: Lam
cmul = lam (\n -> lam (\m -> lam (\f -> lam (\x -> 
         var n `app` (var m `app` var f) `app` var x))))

-- 36
t4 :: Lam
t4 = cmul `app` t3 `app` t3

-- 36 * 36
t5 :: Lam
t5 = cmul `app` t4 `app` t4

-- 36 * 36 * 36
t6 :: Lam
t6 = cmul `app` t4 `app` t5

t7 :: Lam
t7 = cmul `app` t5 `app` t4

cnum :: Int -> Lam
cnum n = lam (\f -> lam (\x -> go f x n)) where
  go f x 0 = var x
  go f x n = var f `app` (go f x (n - 1))

t8 :: Lam
t8 = (cmul `app` t6) `app` t3   -- slow

t9 :: Lam
t9 = cmul `app` t3 `app` t6

type Tm = Lam

-- \x y . x
tt :: Tm
tt = lam (\x -> lam (\y -> var x))

-- \x y . y
ff :: Tm
ff = lam (\x -> lam (\y -> var y))

-- \b x y . b y x
neg :: Tm
neg = lam (\b -> lam (\x -> lam (\y -> (var b `app` var y) `app` var x)))

-- \b c x y . b (c x y) y
cand :: Tm
cand = lam (\b -> lam (\c -> lam (\x -> lam (\y ->
  (var b `app` (var c `app` var x `app` var y)) `app` var y))))

-- \n . n neg tt 
ceven :: Tm 
ceven = lam (\n -> (var n `app` neg) `app` tt)

t10 :: Tm
t10 = cnum 50000 `app` (lam (\x -> cand `app` (f `app` var x) `app` var x)) `app` tt where
  expensive = ceven `app` (cnum 50001)
  f = lam (\_ -> expensive)

--main = print (length (show (norm t10 :: LamP Int)))
main = printLam (norm t10)

FastFreeNested.hs

{-# LANGUAGE GADTs, RankNTypes, TypeOperators, ScopedTypeVariables, KindSignatures, DeriveFunctor, BangPatterns #-}
{-# LANGUAGE FlexibleInstances, FlexibleContexts, MultiParamTypeClasses, FunctionalDependencies, InstanceSigs #-}

module FastFreeNested where

-- A version of FastFree that use the composition monoidal product
-- rather than the coYoneda version (f . g) x = exists y. f y , y -> g x.
-- Some special care about fmaps must be taken in this file compared to FastFree.hs

import Prelude hiding (head, tail)
import Control.Monad

data Free sig a where
  Ret  :: a -> Free sig a
  Fmap :: (a -> b) -> Free sig a -> Free sig b
  Cons :: sig (Free sig b) -> Free sig b

instance Functor sig => Functor (Free sig) where
  fmap f (Ret x)    = Ret (f x)
  fmap f (Fmap g x) = Fmap (f . g) x
  fmap f x          = Fmap f x

pushFmap :: Functor sig => (a -> b) -> Free sig a -> Free sig b
pushFmap f (Ret x) = Ret (f x)
pushFmap f (Fmap g x) = Fmap (f . g) x
pushFmap f (Cons xs) = Cons (fmap (fmap f) xs) 

instance Functor sig => Applicative (Free sig) where
  pure = return
  (<*>) = ap

instance Functor sig => Monad (Free sig) where
  return = Ret

  (Ret a)     >>= k = k a
  (Fmap f xs) >>= k = xs >>= (k . f)
  (Cons xxs)  >>= k = Cons (fmap (>>= k) xxs)

-- Snoc lists
---------------------------

data FreeSnoc sig a where
  Ret' :: a -> FreeSnoc sig a
  Fmap' :: (a -> b) -> FreeSnoc sig a -> FreeSnoc sig b
  Snoc :: FreeSnoc sig (sig b) -> FreeSnoc sig b

instance Functor sig => Functor (FreeSnoc sig) where
  fmap f (Ret' x) = Ret' (f x)
  fmap f (Fmap' g xs) = Fmap' (f . g) xs 
  fmap f xs = Fmap' f xs

reverseF :: forall sig a. Functor sig => FreeSnoc sig a -> Free sig a
reverseF xs = go xs Ret where
  go :: forall a b. FreeSnoc sig a -> (a -> Free sig b) -> Free sig b
  go (Ret' x) k     = k x
  go (Fmap' f xs) k = go xs (k . f)
  go (Snoc sxx) k   = go sxx (\sa -> Cons (fmap k sa))

-- Bankers' queue
--------------------------

data QueueF sig a where
  Q :: !Int -> !Int -> Free sig (FreeSnoc sig b) -> QueueF sig b
  -- Invariants : for all |Q lenf f lenr r|, 
  --   |lenf| is the depth for |f|, 
  --   |lenr| is the depth for |r a| for all |a|,
  --   and |lenr <= lenf|.

q :: Functor sig => Int -> Int -> Free sig (FreeSnoc sig b) -> QueueF sig b
q !lenf !lenr fr 
  | lenr <= lenf  = Q lenf lenr fr
  | otherwise     = Q (lenf + lenr) 0 (do r <- fr; x <- reverseF r; return (Ret' x))

retQ :: Functor sig => a -> QueueF sig a
retQ a = q 0 0 (Ret (Ret' a))

snoc :: Functor sig => QueueF sig (sig b) -> QueueF sig b
snoc (Q lenf lenr fr) = q lenf (lenr + 1) (fmap Snoc fr)

instance Functor sig => Functor (QueueF sig) where
  fmap f (Q lenf lenr fr) = q lenf lenr (fmap (fmap f) fr)

match :: Functor sig 
      => QueueF sig a 
      -> (a -> d)
      -> (sig (QueueF sig a) -> d)
      -> d
match (Q lenf lenr (Ret b)) ret _  = case b of Ret' a -> ret a
                                               _      -> error "Invariant violation"
match (Q lenf lenr (Fmap f xs)) ret cons  = match (Q lenf lenr (pushFmap f xs)) ret cons
match (Q lenf lenr (Cons xxs)) _ cons = cons (fmap (\xs -> q (lenf - 1) lenr xs) xxs)


-- Catenable lists for free monads
-------------------------------

data CListF (sig :: * -> *) (a :: *) where
  E :: a -> CListF sig a
  C :: sig (QueueF (CListF sig) a) -> CListF sig a

-- Technical note:
-- In the original algorithm from Okasaki in CatList.hs, the QueueF in CListF only stores
-- non-E CListF, but we can't decide if a |a -> CListF sig a| is E or not, so here
-- we do allow QueueF to store E. Thus the |link| function and |++| function in CatList.hs
-- become the same one here.


-- Concatenation is essentially snoc
append :: Functor sig => CListF sig (CListF sig b) -> CListF sig b
append (E ys) = ys
append (C xxsys) = C (fmap (\xsys -> snoc xsys) xxsys)

inj :: Functor sig => sig a -> CListF sig a
inj x = C (fmap retQ x) 

consC :: Functor sig => sig (CListF sig b) -> CListF sig b
consC xxs = append (inj xxs)

snocC :: Functor sig => CListF sig (sig b) -> CListF sig b
snocC sxx = append (fmap inj sxx) 

matchC :: Functor sig 
       => CListF sig a 
       -> (a -> d)
       -> (sig (CListF sig a) -> d)
       -> d
matchC (E a) ret _    = ret a
matchC (C xxs) _ cons = cons (fmap linkAll xxs)

-- This is a kind of concat
linkAll :: Functor sig => QueueF (CListF sig) a -> CListF sig a
linkAll q = match q E (\xsxss -> append (fmap linkAll xsxss))

instance Functor sig => Functor (CListF sig) where
  fmap f (E x) = E (f x)
  fmap f (C xxs) = C (fmap (fmap f) xxs)

instance Functor sig => Applicative (CListF sig) where
  pure = return
  (<*>) = ap

-- CListF is a representation of Free supporting amortised O(1)-time
-- pattern matching and >>=.
instance Functor sig => Monad (CListF sig) where
  return  = E
  x >>= k = append (fmap k x)

-- Some simple tests
--------------------------

data BState a = Get a a | Put Bool a
  deriving Functor

class Monad m => MonadBState m where
  get :: m Bool
  put :: Bool -> m ()

instance MonadBState (CListF BState) where
  get = inj (Get False True)
  
  put s = inj (Put s ())

instance MonadBState (Free BState) where
  get = Cons (Get (return False) (return True))
  put s = Cons (Put s (return ()))


-- C++'s ++ operation
incr :: MonadBState m => m Bool
incr = do i <- get
          put (not i)
          return i

handleSt :: CListF BState a -> Bool -> a
handleSt p = matchC p (\a s -> a) hdl 
  where hdl (Get x y) s  = if s then handleSt y s else handleSt x s
        hdl (Put s' x) s = handleSt x s'

handleSt' :: Free BState a -> Bool -> a
handleSt' p = case p of
  (Ret a) -> \s -> a
  (Cons f) -> hdl f
  where hdl (Get x y) s  = if s then handleSt' y s else handleSt' x s
        hdl (Put s' x) s = handleSt' x s'

-- Basic correctness check
test0 = handleSt (incr >> incr >> (incr >> incr) >> get) False
  -- test0 == 4

-- CListF should be much faster than Free for left associated computations
incrN :: MonadBState m => Int -> m Bool
incrN 0 = return False
incrN n = incrN (n-1) >> incr -- This is deliberately associated in the way

test1 = handleSt (incrN 10000 >> get) False -- This should be 10000 computed in several seconds

test1' = handleSt' (incrN 10000 >> get) False  -- This should take a lot of time


-- CListF supports pattern matching 
test2 = go 10 (incrN 10000) == "GPGPGPGPGP"
  where
    go :: Int -> CListF BState a -> String
    go 0 p = []
    go n p = matchC p (\_ -> []) 
                      (\op -> case op of
                                (Put _ a) -> 'P' : go (n - 1) a
                                (Get a b) -> 'G' : go (n - 1) a)


-- A test case of lambda terms with higher-order abstract syntax
-- can be found in LamHOAS.hs

LamPHOAS.hs

{-# LANGUAGE GADTs, RankNTypes, TypeApplications, TypeOperators, ScopedTypeVariables, KindSignatures, DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances, FlexibleContexts, MultiParamTypeClasses, FunctionalDependencies, InstanceSigs #-}

module LamPHOAS where

-- Using FastFree to implement a normaliser of lambda calculus based on
-- parametric higher-order abstract syntax

import FastFree
import Data.Bool (bool)

-- for names we use an arbitrary type `a` that can be compared 
-- and has a stream of distinct elements
class Eq a => Atom a where
  ns :: [a]

  -- and we require a name can be printed as an identifier
  showName :: a -> String

instance Atom Int where
  ns = [1..]
  showName n = "x" ++ show n

data LamS a x = App x x | Abs (a -> x)

type LamP a = CListF (LamS a) a

type Lam = forall a. Atom a => LamP a

instance Atom a => Show (LamP a) where
  show t = fst (go ns t) where
    go :: [a] -> LamP a -> (String, [a])
    go ns t = matchC t (\a -> (showName a, ns)) f where
      f :: LamS a b -> (b -> LamP a) -> (String, [a])
      f (App fun opr) k = let (s1, ns') = go ns (k fun)
                              (s2, ns'') = go ns' (k opr)
                          in (s1 ++ " " ++ s2, ns'')
      f (Abs body) k = let (n:ns') = ns 
                           (s, ns'') = go ns' (k (body n))
                       in ("(λ " ++ showName n ++ ". " ++  s ++ ")", ns'')

printLam :: Lam -> IO ()
printLam t = print (t :: LamP Int)

app :: LamP a -> LamP a -> LamP a
app t1 t2 = consC (App False True) (bool t1 t2) 

lam :: (a -> LamP a) -> LamP a
lam t = consC (Abs id) t

var :: a -> LamP a
var = return

t1 :: Lam
t1 = app (lam (\x -> var x `app` var x)) (lam (\x -> var x))

-- replaces every occurrence of x in t with m
subst :: Atom a => LamP a -> a -> LamP a -> LamP a
subst t x m = t >>= (\y -> if y == x then m else var y)

-- evaluate a closed lambda term (in the lazy evaluation order)
-- no evaluation under binders
eval :: Lam -> Lam
eval t = fst (go ns t) where 
  go :: forall a. Atom a => [a] -> LamP a -> (LamP a, [a])
  go ns t = matchC t (\a -> (return a, ns)) f where
    f :: forall b. LamS a b -> (b -> LamP a) -> (LamP a, [a])
    f (Abs body)  k = (consC (Abs body) k, ns)
    f (App t1 t2) k = let (t1', ns') = go ns (k t1)

                          g :: forall b. LamS a b -> (b -> LamP a) -> (LamP a, [a])
                          g (Abs b) k2 = let (n:ns'') = ns'
                                         in go ns'' (subst (k2 (b n)) n (k t2))
                          g _       k2 = (app t1' (k t2), ns')

                      in  matchC t1' (\a -> (app t1' (k t2), ns')) g 

-- full normalisation in the applicative order
norm :: Lam -> Lam
norm t = fst (go ns t) where 
  go :: forall a. Atom a => [a] -> LamP a -> (LamP a, [a])
  go ns t = matchC t (\a -> (return a, ns)) f where
    f :: forall b. LamS a b -> (b -> LamP a) -> (LamP a, [a])
    f (Abs body)  k = let (n:ns') = ns                                   -- this looks odd! I am not sure we should
                          (body',ns'') = go ns' (k (body n))             -- concretise, normalise, and abstract. 
                                                                         -- Can we delay the normalisation until the binder
                                                                         -- actually gets a concrete atom? (That would 
                                                                         -- make fresh name management more difficult though!)
                      in (lam (\a -> subst body' n (var a)), ns'')
    f (App t1 t2) k = let (t1', ns') = go ns (k t1)
                          (t2', ns'') = go ns' (k t2)

                          g :: forall b. LamS a b -> (b -> LamP a) -> (LamP a, [a])
                          g (Abs b) k2 = let (n:ns''') = ns''
                                         in go ns''' (subst (k2 (b n)) n t2')
                          g _       k2 = (app t1' t2', ns'')

                      in  matchC t1' (\a -> (app t1' t2', ns'')) g 

t2 :: Lam
t2 = lam (\f -> var f `app` var f)
       `app`
       lam (\x -> lam (\y -> var x `app` var y))

-- Church numeral of 3
c3 :: Lam
c3 = lam (\f -> lam (\x -> var f `app` (var f `app` (var f `app` var x))))

infixl 4 `app` 
cadd :: Lam
cadd = lam (\n -> lam (\m -> lam (\f -> lam (\x -> 
          var n `app` var f 
            `app` (var m `app` var f `app` var x)))))

t3 :: Lam
t3 = cadd `app` c3 `app` c3

-- printLam (norm t3)
-- (λ x1. (λ x2. x1 x1 x1 x1 x1 x1 x2))