s-cargot/Data/SCargot/Pretty.hs

{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Data.SCargot.Pretty
         ( LayoutOptions(..)
         , Indent(..)
         , basicPrint
         , flatPrint
         , prettyPrintSExpr
         ) where

import           Data.Monoid ((<>))
import           Data.Text (Text)
import qualified Data.Text as T

import           Data.SCargot.Repr

data Indent
  = Swing
  | SwingAfter Int
  | Align
    deriving (Eq, Show)

-- | A 'LayoutOptions' value describes how to pretty-print a 'SExpr'.
--   It describes how to print atoms, what horizontal space to fit
--   it into, and other related options.
--
--   The 'swingIndent' value might require a big of explanation: in
--   pretty-printing s-expressions, you have the option of whether
--   to 'swing' expressions which get pushed to subsequent lines
--   to the left, or to align them along the right. e.g. the
--   s-expression @(foo a b)@ could use a non-swing indent as
--
--   > (foo arg-one
--   >      arg-two)
--
--   or a swing indent as
--
--   > (foo arg-one
--   >   arg-two)
--
--   often, in formatting Lisp code, control structures will
--   swing subsequent expressions, as in
--
--   > (define (factorial n)
--   >   (if (= n 0)
--   >     1
--   >     (* n (fact (- n 1)))))
--
--   but most functions will _not_ swing:
--
--   > (call-my-func arg-number-one
--   >               arg-number-two
--   >               arg-number-three)
--
--   The 'swingIndent' field lets you choose whether or not to
--   swing subsequent s-expressions based on the atom in the car
--   position of a list. You can default to always swinging subsequent
--   expressions with @const True@ and never with @const False@, or
--   choose based on some more advanced criteria. _If_ a swing happens,
--   subsequent lines are indented based on the 'indentAmount' variable;
--   otherwise, subsequent lines are indented based on the size of the
--   @car@ of the list.
data LayoutOptions a = LayoutOptions
  { atomPrinter  :: a -> Text         -- ^ How to serialize a given atom to 'Text'.
  , swingIndent  :: SExpr a -> Indent -- ^ Whether or not to swing
  , indentAmount :: Int               -- ^ How much to indent after a swing
  , maxWidth     :: Maybe Int         -- ^ The maximum width (if any)
  }

flatPrint :: (a -> Text) -> LayoutOptions a
flatPrint printer = LayoutOptions
  { atomPrinter  = printer
  , swingIndent  = const Swing
  , indentAmount = 2
  , maxWidth     = Nothing
  }

basicPrint :: (a -> Text) -> LayoutOptions a
basicPrint printer = LayoutOptions
  { atomPrinter  = printer
  , swingIndent  = const Swing
  , indentAmount = 2
  , maxWidth     = Just 80
  }

-- Sort of like 'unlines' but without the trailing newline
joinLines :: [Text] -> Text
joinLines = T.intercalate "\n"

-- Indents a line by n spaces
indent :: Int -> Text -> Text
indent n ts = T.replicate n " " <> ts

-- Indents every line n spaces, and adds a newline to the beginning
indentAll :: Int -> [Text] -> Text
indentAll n = ("\n" <>) . joinLines . map (indent n)

-- Indents every line but the first by some amount
indentSubsequent :: Int -> [Text] -> Text
indentSubsequent _ [] = ""
indentSubsequent _ [t] = t
indentSubsequent n (t:ts) = joinLines (t : go ts)
  where go = map (indent n)

-- oh god this code is so disgusting
-- i'm sorry to everyone i let down by writing this
-- i swear i'll do better in the future i promise i have to
-- for my sake and for everyone's
prettyPrintSExpr :: LayoutOptions a -> SExpr a -> Text
prettyPrintSExpr LayoutOptions { .. } = pHead 0
  where pHead _   SNil         = "()"
        pHead _   (SAtom a)    = atomPrinter a
        pHead ind (SCons x xs) = gather ind x xs id
        gather _   _ (SAtom _)    _ = error "no dotted pretty printing yet!"
        gather ind h (SCons x xs) k = gather ind h xs (k . (x:))
        gather ind h SNil         k = "(" <> hd <> body <> ")"
          where hd   = indentSubsequent ind [pHead (ind+1) h]
                lst  = k []
                flat = T.unwords (map (pHead (ind+1)) lst)
                headWidth = T.length hd + 1
                indented =
                  case swingIndent h of
                    SwingAfter n ->
                      let (l, ls) = splitAt n lst
                          t  = T.unwords (map (pHead (ind+1)) l)
                          ts = indentAll (ind + indentAmount)
                                 (map (pHead (ind + indentAmount)) ls)
                      in t <> ts
                    Swing ->
                      indentAll (ind + indentAmount)
                        (map (pHead (ind + indentAmount)) lst)
                    Align ->
                      indentSubsequent (ind + headWidth + 1)
                        (map (pHead (ind + headWidth + 1)) lst)
                body
                  | length lst == 0                = ""
                  | Just maxAmt <- maxWidth
                  , (T.length flat + ind) > maxAmt = " " <> indented
                  | otherwise                      = " " <> flat