{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE Safe              #-}

-- | A representation for structured expression trees, with support for pretty
-- printing.
module Copilot.Theorem.Misc.SExpr where

import Text.PrettyPrint.HughesPJ as PP hiding (char, Str)

import Control.Monad

-- | A structured expression is either an atom, or a sequence of expressions,
-- where the first in the sequence denotes the tag or label of the tree.
data SExpr a = Atom a
             | List [SExpr a]

-- | Empty string expression.
blank :: SExpr String
blank = String -> SExpr String
forall a. a -> SExpr a
Atom String
""

-- | Atomic expression constructor.
atom :: a -> SExpr a
atom = a -> SExpr a
forall a. a -> SExpr a
Atom                 -- s

-- | Empty expression (empty list).
unit :: SExpr a
unit = [SExpr a] -> SExpr a
forall a. [SExpr a] -> SExpr a
List []              -- ()

-- | Single expression.
singleton :: a -> SExpr a
singleton a
a  = [SExpr a] -> SExpr a
forall a. [SExpr a] -> SExpr a
List [a -> SExpr a
forall a. a -> SExpr a
Atom a
a]        -- (s)

-- | Sequence of expressions.
list :: [SExpr a] -> SExpr a
list = [SExpr a] -> SExpr a
forall a. [SExpr a] -> SExpr a
List                 -- (ss)

-- | Sequence of expressions with a root or main note, and a series of
-- additional expressions or arguments.
node :: a -> [SExpr a] -> SExpr a
node a
a [SExpr a]
l = [SExpr a] -> SExpr a
forall a. [SExpr a] -> SExpr a
List (a -> SExpr a
forall a. a -> SExpr a
Atom a
a SExpr a -> [SExpr a] -> [SExpr a]
forall a. a -> [a] -> [a]
: [SExpr a]
l)    -- (s ss)

-- A straightforward string representation for 'SExpr's of Strings that
-- parenthesizes lists of expressions.
instance Show (SExpr String) where
  show :: SExpr String -> String
show = Doc -> String
PP.render (Doc -> String) -> (SExpr String -> Doc) -> SExpr String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SExpr String -> Doc
show'
    where
      show' :: SExpr String -> Doc
show' (Atom String
s) = String -> Doc
text String
s
      show' (List [SExpr String]
ts) = Doc -> Doc
parens (Doc -> Doc) -> ([SExpr String] -> Doc) -> [SExpr String] -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Doc] -> Doc
hsep ([Doc] -> Doc)
-> ([SExpr String] -> [Doc]) -> [SExpr String] -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SExpr String -> Doc) -> [SExpr String] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map SExpr String -> Doc
show' ([SExpr String] -> Doc) -> [SExpr String] -> Doc
forall a b. (a -> b) -> a -> b
$ [SExpr String]
ts

-- More advanced printing with some basic indentation

-- | Indent by a given number.
indent :: Doc -> Doc
indent = Int -> Doc -> Doc
nest Int
1

-- | Pretty print a structured expression as a String.
toString :: (SExpr a -> Bool)  -- ^ True if an expression should be indented.
         -> (a -> String)      -- ^ Pretty print the value inside as 'SExpr'.
         -> SExpr a            -- ^ Root of 'SExpr' tree.
         -> String
toString :: forall a. (SExpr a -> Bool) -> (a -> String) -> SExpr a -> String
toString SExpr a -> Bool
shouldIndent a -> String
printAtom SExpr a
expr =
  Doc -> String
PP.render ((SExpr a -> Bool) -> (a -> String) -> SExpr a -> Doc
forall a. (SExpr a -> Bool) -> (a -> String) -> SExpr a -> Doc
toDoc SExpr a -> Bool
shouldIndent a -> String
printAtom SExpr a
expr)

-- | Pretty print a structured expression as a 'Doc', or set of layouts.
toDoc :: (SExpr a -> Bool)  -- ^ True if an expression should be indented.
      -> (a -> String)      -- ^ Pretty print the value inside as 'SExpr'.
      -> SExpr a            -- ^ Root of 'SExpr' tree.
      -> Doc
toDoc :: forall a. (SExpr a -> Bool) -> (a -> String) -> SExpr a -> Doc
toDoc SExpr a -> Bool
shouldIndent a -> String
printAtom SExpr a
expr = case SExpr a
expr of
  Atom a
a  -> String -> Doc
text (a -> String
printAtom a
a)
  List [SExpr a]
l  -> Doc -> Doc
parens ((Doc -> SExpr a -> Doc) -> Doc -> [SExpr a] -> Doc
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Doc -> SExpr a -> Doc
renderItem Doc
empty [SExpr a]
l)

  where
    renderItem :: Doc -> SExpr a -> Doc
renderItem Doc
doc SExpr a
s
      | SExpr a -> Bool
shouldIndent SExpr a
s =
        Doc
doc Doc -> Doc -> Doc
$$ Doc -> Doc
indent ((SExpr a -> Bool) -> (a -> String) -> SExpr a -> Doc
forall a. (SExpr a -> Bool) -> (a -> String) -> SExpr a -> Doc
toDoc SExpr a -> Bool
shouldIndent a -> String
printAtom SExpr a
s)
      | Bool
otherwise =
        Doc
doc Doc -> Doc -> Doc
<+> (SExpr a -> Bool) -> (a -> String) -> SExpr a -> Doc
forall a. (SExpr a -> Bool) -> (a -> String) -> SExpr a -> Doc
toDoc SExpr a -> Bool
shouldIndent a -> String
printAtom SExpr a
s