-- Taken from https://github.com/raichoo/TheFreeAndTheFurious
-- Modified to use `free` and `kan-extensions` library instead of writing the
-- required functions from those libraries by hand.
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE UndecidableInstances #-}

module Main where

import Control.Monad (when)
import Control.Monad.Codensity
import Control.Monad.Free
import Data.DList qualified as DL
import Prelude hiding (getChar, putChar)
import Prelude qualified as P (getChar, putChar)

newtype DList a = DList {runDList :: [a] -> [a]}

type List a = Codensity (Free ((,) a)) ()

empty :: List a
empty = pure ()

singleton :: a -> List a
singleton x = liftF (x, ())

append :: List a -> List a -> List a
append xs ys = xs >> ys

exec :: Free ((,) a) () -> [a]
exec (Pure _) = []
exec (Free (x, xs)) = x : exec xs

data TeletypeF k
  = PutChar Char k
  | GetChar (Char -> k)
  deriving (Functor)

type Teletype a = forall m. (MonadFree TeletypeF m) => m a

getChar :: Teletype Char
getChar = liftF (GetChar id)

putChar :: Char -> Teletype ()
putChar c = liftF (PutChar c ())

revEcho :: Teletype ()
revEcho = do
  c <- getChar
  when (c /= ' ') $ do
    revEcho
    putChar c

runIO :: Free TeletypeF a -> IO a
runIO (Pure x) = return x
runIO (Free (GetChar k)) = P.getChar >>= runIO . k
runIO (Free (PutChar c k)) = P.putChar c >> runIO k

data Output a = Read (Output a) | Print Char (Output a) | Finish a
  deriving (Show, Eq)

input :: [Char]
input = replicate 10000 'c' ++ ' ' : repeat 'c'

runPure :: Free TeletypeF a -> [Char] -> Output a
runPure (Pure x) _ = Finish x
runPure (Free (GetChar k)) [] = Read (runPure (k ' ') [])
runPure (Free (GetChar k)) (i : is) = Read (runPure (k i) is)
runPure (Free (PutChar c k)) [] = Print c (runPure k [])
runPure (Free (PutChar c k)) (_ : is) = Print c (runPure k is)

type Log = List String

type FastLog = DL.DList String

prog1 :: String
prog1 = last . exec . lowerCodensity $ prog' empty 20000
  where
    prog' :: Log -> Int -> Log
    prog' logMsg 0 = logMsg
    prog' logMsg n = prog' (logMsg `append` singleton (show n)) (n - 1)

prog2 :: String
prog2 = last (DL.toList (prog' [] 20000))
  where
    prog' :: FastLog -> Int -> FastLog
    prog' logMsg 0 = logMsg
    prog' logMsg n = prog' (logMsg <> [show n]) (n - 1)

slow :: IO ()
slow = print $ runPure revEcho input

fast :: IO ()
fast = print $ runPure (improve revEcho) input

main :: IO ()
main = fast