"A Prettier Printer" in continuation passing style

PrettierCPS.hs

{- cabal:
build-depends: base, lens, generic-lens
-}

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE RankNTypes #-}

import Prelude hiding ((<>))
import Control.Lens
import Data.Generics.Labels
import GHC.Generics


data State = State {
    nestingLevel :: Int,
    docStream :: [Either String Int],
    written :: Int
} deriving (Generic)

type Cont r = OnFailure r -> State -> r
newtype OnFailure r = OnFailure { runFailure :: Cont r -> Cont r }

type Doc = forall r. Cont r -> Cont r

pprint :: Doc -> String
pprint doc =
    let stream = doc (\_ s -> s ^. #docStream) (OnFailure id) (State 0 [] 0)
    in go stream ""
    where
        go [] s = s
        go (Left t : xs) s = go xs (t ++ s)
        go (Right n : xs) s = go xs (replicate n ' ' ++ s)

empty :: Doc
empty = id

line :: Doc
line k q s = k (OnFailure id) state
    where
        state = s
            & #written .~ n
            & #docStream %~ (Right n :)
        n = s ^. #nestingLevel

text :: String -> Doc
text t k q s =
    let cont = if len > 80 then runFailure q k else k
    in cont q state
    where
        len = s ^. #written + length t
        state = s & #docStream %~ (Left t :) & #written .~ len

nest :: Int -> Doc -> Doc
nest m d k q s = d (\q' s' -> k q' (s' & #nestingLevel -~ m)) q (s & #nestingLevel +~ m)

(<>) :: Doc -> Doc -> Doc
(<>) = (.)

(<|>) :: Doc -> Doc -> Doc
(x <|> y) k q s = x k (OnFailure $ \_ _ _ -> y k q s) s

main = return ()