Pretty.hs

-- | Generic pretty-printing of expression trees.
--
-- TODO: 
--
--  * make the style configurable (so that we can print the same expression in different formats)
--    (in Haskell98, we have to use the record instead of class trick for this?)
--
--  * corresponding parser?
--

module Data.Generics.Fixplate.Pretty where

--------------------------------------------------------------------------------

import Prelude
import Data.List ( intersperse )

import Data.Generics.Fixplate
import Data.Foldable 

import Text.Show ()

--------------------------------------------------------------------------------

-- | Associativity
data Assoc 
  = NoAssoc 
  | LeftAssoc 
  | RightAssoc
  deriving (Eq,Show)

-- | A pair of matching brackets, eg. @Bracket \"(\" \")\"@ or @Bracket \"[|\" \"|]\"@.
data Bracket = Bracket !String !String deriving (Eq,Show)

-- | A separator, eg. @\",\"@ or @\" | \"@.
type Separator = String

-- | Application style 
data AppStyle
  = Haskell                    -- ^ eg. @(Node arg1 arg2 arg3)@; precedence will be @app_prec == 10@
  | Algol !Bracket !Separator  -- ^ eg. @node[arg1,arg2,arg3]@; precedence will be 11, but child environment precedence will be 0
  deriving (Eq,Show)

-- | Mixfix style. Example: 
--
-- > [ Keyword "if" , Placeholder , keyword "then" , Placeholder , keyword "else" , Placeholder ]
--
data MixWord 
  = Keyword String
  | Placeholder
  deriving (Eq,Show)

mixWords :: [MixWord] -> [ShowS] -> ShowS
mixWords mws args = Prelude.foldr (.) id (intersperse (showChar ' ') (go mws args)) where 
  go :: [MixWord] -> [ShowS] -> [ShowS]
  go (Keyword s   : rest) fs     = showString s : go rest fs
  go (Placeholder : rest) (f:fs) = f : go rest fs
  go (Placeholder : rest) [] = error "mixWords: not enough arguments"
  go [] []     = []
  go [] (f:fs) = f : go [] fs

-- | Fixities. TODO: separate non-fixity stuff like style and words
data Fixity 
  = Atom                   -- ^ eg. @variable@; precedence will be 666
  | Application !AppStyle  -- ^ eg. @(Node arg1 arg2 arg3)@ or @node[arg1,arg2,arg3]@.
  | Prefix  !Int           -- ^ eg. @~arg@; the @Int@ is the precendence
  | Infix !Assoc !Int      -- ^ eg. @x+y@
  | Postfix !Int           -- ^ eg. @arg++@
  | Mixfix [MixWord]       -- ^ eg. @if ... then ... else ... @ or @let ... in ...@. With precedence 0?
  | Custom  !Int           -- ^ for your custom rendering
  deriving (Eq,Show)

fixityPrecedence :: Fixity -> Int
fixityPrecedence f = case f of
  Atom              -> 666
  Application style -> 
    case style of
      Haskell  -> 10
      Algol {} -> 11
  Prefix prec       -> prec
  Infix assoc prec  -> prec
  Postfix prec      -> prec
  Mixfix {}         -> 0
  Custom prec       -> prec

--------------------------------------------------------------------------------

-- | A class encoding fixity and rendering of nodes if the tree.
--
-- Minimum complete definition: 'fixity', and 'showNode' or 'showsPrecNode'.
-- Unless you want some type of rendering not directly supported, you shouldn't specify 'showsPrecNode'.
--
class (Functor f, Foldable f) => Pretty f where
  -- | fixity of the node
  fixity    :: f a -> Fixity                                      

  -- | a string representing the node /without/ the children
  showNode  :: f a -> String                                      

  -- | full rendering of the node. You can redefine this for custom renderings.
  showsPrecNode :: (Int -> a -> ShowS) -> Int -> f a -> ShowS     
  showsPrecNode child d node = showParen (d > prec) $ 
    case fty of       
      Atom -> showString (showNode node)

      Application style -> case style of

        Haskell -> head . args where
          head = showString (showNode node)
          args = Prelude.foldr (.) id [ showChar ' ' . child (prec+1) c | c <- children ]

        Algol (Bracket open close) sep -> head . showString open . args . showString close where
          head = showString (showNode node)
          args = Prelude.foldr (.) id 
               $ intersperse (showString sep) [ child 0 c | c <- children ]

      Prefix prec -> 
        case children of
          [] -> error "showsPrecNode: prefix node with no arguments"
          (c:cs) -> op . arg1 c . args cs
        where
          op = showString (showNode node)
          arg1 c  = child (prec+1) c 
          args cs = Prelude.foldr (.) id [ showChar ' ' . child (prec+1) c | c <- cs ]

      Postfix prec -> 
        case children of
          []  -> error "showsPrecNode: postfix node with no arguments"
          ccs -> let (cs,c) = (Prelude.init ccs, Prelude.last ccs) 
                 in  args cs . arg1 c . op
        where
          op = showString (showNode node)
          arg1 c  = child (prec+1) c 
          args cs = Prelude.foldr (.) id [ child (prec+1) c . showChar ' ' | c <- cs ]

      Infix assoc prec -> 
        case children of
          []  -> error "showsPrecNode: infix node with no arguments"
          [_] -> error "showsPrecNode: infix node with a single argument"
          (c1:c2:cs) -> lhs c1 . op . rhs c2 . rest cs
        where
          lhs  c1 = child lprec c1 
          op      = showString (showNode node) 
          rhs  c2 = child rprec c2
          rest cs = Prelude.foldr (.) id [ showChar ' ' . child (prec+1) c | c <- cs ]
          lprec = case assoc of { LeftAssoc  -> prec ; _ -> prec+1 }
          rprec = case assoc of { RightAssoc -> prec ; _ -> prec+1 }
        
      Mixfix mwords -> mixWords mwords [ child (prec+1) {- ? -} c | c <- children ]

      Custom prec -> error "for custom rendering, you should redefine `showsPrecNode'"

    where
      fty  = fixity node
      prec = fixityPrecedence fty
      children = toList node

--------------------------------------------------------------------------------

-- | Render the expression
pretty :: Pretty f => Mu f -> String 
pretty tree = prettyS tree "" 

prettyS :: Pretty f => Mu f -> ShowS
prettyS = prettyPrec 0 

prettyPrec :: Pretty f => Int -> Mu f -> ShowS
prettyPrec d t = go d t where
  go d (Fix t) = showsPrecNode go d t

--------------------------------------------------------------------------------