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

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

import Text.ParserCombinators.Parsec
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 (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

-- | Parser for strings of characters separated by spaces into a structured
-- tree.
--
-- Parentheses are interpreted as grouping elements, that is, defining a
-- 'List', which may be empty.
parser :: GenParser Char st (SExpr String)
parser :: forall st. GenParser Char st (SExpr String)
parser =
  [ParsecT String st Identity (SExpr String)]
-> ParsecT String st Identity (SExpr String)
forall s (m :: * -> *) t u a.
Stream s m t =>
[ParsecT s u m a] -> ParsecT s u m a
choice [ParsecT String st Identity (SExpr String)
-> ParsecT String st Identity (SExpr String)
forall tok st a. GenParser tok st a -> GenParser tok st a
try ParsecT String st Identity (SExpr String)
forall {u} {a}. ParsecT String u Identity (SExpr a)
unitP, ParsecT String st Identity (SExpr String)
forall st. GenParser Char st (SExpr String)
nodeP, ParsecT String st Identity (SExpr String)
forall st. GenParser Char st (SExpr String)
leafP]

  where
    symbol :: ParsecT String u Identity Char
symbol     = String -> ParsecT String u Identity Char
forall s (m :: * -> *) u.
Stream s m Char =>
String -> ParsecT s u m Char
oneOf String
"!#$%&|*+-/:<=>?@^_~."
    lonelyStr :: ParsecT String u Identity String
lonelyStr  = ParsecT String u Identity Char -> ParsecT String u Identity String
forall s (m :: * -> *) t u a.
Stream s m t =>
ParsecT s u m a -> ParsecT s u m [a]
many1 (ParsecT String u Identity Char
forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Char
alphaNum ParsecT String u Identity Char
-> ParsecT String u Identity Char -> ParsecT String u Identity Char
forall s u (m :: * -> *) a.
ParsecT s u m a -> ParsecT s u m a -> ParsecT s u m a
<|> ParsecT String u Identity Char
forall {u}. ParsecT String u Identity Char
symbol)

    unitP :: ParsecT String u Identity (SExpr a)
unitP      = String -> ParsecT String u Identity String
forall s (m :: * -> *) u.
Stream s m Char =>
String -> ParsecT s u m String
string String
"()" ParsecT String u Identity String
-> ParsecT String u Identity (SExpr a)
-> ParsecT String u Identity (SExpr a)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> SExpr a -> ParsecT String u Identity (SExpr a)
forall (m :: * -> *) a. Monad m => a -> m a
return SExpr a
forall {a}. SExpr a
unit

    leafP :: ParsecT String u Identity (SExpr String)
leafP      = String -> SExpr String
forall a. a -> SExpr a
atom (String -> SExpr String)
-> ParsecT String u Identity String
-> ParsecT String u Identity (SExpr String)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ParsecT String u Identity String
forall {u}. ParsecT String u Identity String
lonelyStr

    nodeP :: ParsecT String u Identity (SExpr String)
nodeP      = do ParsecT String u Identity Char -> ParsecT String u Identity ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ParsecT String u Identity Char -> ParsecT String u Identity ())
-> ParsecT String u Identity Char -> ParsecT String u Identity ()
forall a b. (a -> b) -> a -> b
$ Char -> ParsecT String u Identity Char
forall s (m :: * -> *) u.
Stream s m Char =>
Char -> ParsecT s u m Char
char Char
'('
                    ParsecT String u Identity ()
forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m ()
spaces
                    [SExpr String]
st <- ParsecT String u Identity (SExpr String)
-> ParsecT String u Identity ()
-> ParsecT String u Identity [SExpr String]
forall s (m :: * -> *) t u a sep.
Stream s m t =>
ParsecT s u m a -> ParsecT s u m sep -> ParsecT s u m [a]
sepBy ParsecT String u Identity (SExpr String)
forall st. GenParser Char st (SExpr String)
parser ParsecT String u Identity ()
forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m ()
spaces
                    ParsecT String u Identity ()
forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m ()
spaces
                    ParsecT String u Identity Char -> ParsecT String u Identity ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ParsecT String u Identity Char -> ParsecT String u Identity ())
-> ParsecT String u Identity Char -> ParsecT String u Identity ()
forall a b. (a -> b) -> a -> b
$ Char -> ParsecT String u Identity Char
forall s (m :: * -> *) u.
Stream s m Char =>
Char -> ParsecT s u m Char
char Char
')'
                    SExpr String -> ParsecT String u Identity (SExpr String)
forall (m :: * -> *) a. Monad m => a -> m a
return (SExpr String -> ParsecT String u Identity (SExpr String))
-> SExpr String -> ParsecT String u Identity (SExpr String)
forall a b. (a -> b) -> a -> b
$ [SExpr String] -> SExpr String
forall a. [SExpr a] -> SExpr a
List [SExpr String]
st

-- | Parser for strings of characters separated by spaces into a structured
-- tree.
--
-- Parentheses are interpreted as grouping elements, that is, defining a
-- 'List', which may be empty.
parseSExpr :: String -> Maybe (SExpr String)
parseSExpr :: String -> Maybe (SExpr String)
parseSExpr String
str = case Parsec String () (SExpr String)
-> String -> String -> Either ParseError (SExpr String)
forall s t a.
Stream s Identity t =>
Parsec s () a -> String -> s -> Either ParseError a
parse Parsec String () (SExpr String)
forall st. GenParser Char st (SExpr String)
parser String
"" String
str of
  Left ParseError
s -> String -> Maybe (SExpr String)
forall a. HasCallStack => String -> a
error (ParseError -> String
forall a. Show a => a -> String
show ParseError
s) -- Nothing
  Right SExpr String
t -> SExpr String -> Maybe (SExpr String)
forall a. a -> Maybe a
Just SExpr String
t