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

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

module Copilot.Theorem.Misc.SExpr where

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

import Control.Monad

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

data SExpr a = Atom a
             | List [SExpr a]

blank :: SExpr [Char]
blank        = [Char] -> SExpr [Char]
forall a. a -> SExpr a
Atom [Char]
""
atom :: a -> SExpr a
atom         = a -> SExpr a
forall a. a -> SExpr a
Atom                 -- s
unit :: SExpr a
unit         = [SExpr a] -> SExpr a
forall a. [SExpr a] -> SExpr a
List []              -- ()
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)
list :: [SExpr a] -> SExpr a
list         = [SExpr a] -> SExpr a
forall a. [SExpr a] -> SExpr a
List                 -- (ss)
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

instance Show (SExpr String) where
  show :: SExpr [Char] -> [Char]
show = Doc -> [Char]
PP.render (Doc -> [Char]) -> (SExpr [Char] -> Doc) -> SExpr [Char] -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SExpr [Char] -> Doc
show'
    where
      show' :: SExpr [Char] -> Doc
show' (Atom [Char]
s) = [Char] -> Doc
text [Char]
s
      show' (List [SExpr [Char]]
ts) = Doc -> Doc
parens (Doc -> Doc) -> ([SExpr [Char]] -> Doc) -> [SExpr [Char]] -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Doc] -> Doc
hsep ([Doc] -> Doc)
-> ([SExpr [Char]] -> [Doc]) -> [SExpr [Char]] -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SExpr [Char] -> Doc) -> [SExpr [Char]] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map SExpr [Char] -> Doc
show' ([SExpr [Char]] -> Doc) -> [SExpr [Char]] -> Doc
forall a b. (a -> b) -> a -> b
$ [SExpr [Char]]
ts


-- More advanced printing with some basic indentation

indent :: Doc -> Doc
indent = Int -> Doc -> Doc
nest Int
1

toString :: (SExpr a -> Bool) -> (a -> String) -> SExpr a -> String
toString :: (SExpr a -> Bool) -> (a -> [Char]) -> SExpr a -> [Char]
toString SExpr a -> Bool
shouldIndent a -> [Char]
printAtom SExpr a
expr =
  Doc -> [Char]
PP.render ((SExpr a -> Bool) -> (a -> [Char]) -> SExpr a -> Doc
forall a. (SExpr a -> Bool) -> (a -> [Char]) -> SExpr a -> Doc
toDoc SExpr a -> Bool
shouldIndent a -> [Char]
printAtom SExpr a
expr)

toDoc :: (SExpr a -> Bool) -> (a -> String) -> SExpr a -> Doc
toDoc :: (SExpr a -> Bool) -> (a -> [Char]) -> SExpr a -> Doc
toDoc SExpr a -> Bool
shouldIndent a -> [Char]
printAtom SExpr a
expr = case SExpr a
expr of
  Atom a
a  -> [Char] -> Doc
text (a -> [Char]
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 -> [Char]) -> SExpr a -> Doc
forall a. (SExpr a -> Bool) -> (a -> [Char]) -> SExpr a -> Doc
toDoc SExpr a -> Bool
shouldIndent a -> [Char]
printAtom SExpr a
s)
          | Bool
otherwise =
            Doc
doc Doc -> Doc -> Doc
<+> (SExpr a -> Bool) -> (a -> [Char]) -> SExpr a -> Doc
forall a. (SExpr a -> Bool) -> (a -> [Char]) -> SExpr a -> Doc
toDoc SExpr a -> Bool
shouldIndent a -> [Char]
printAtom SExpr a
s

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

parser :: GenParser Char st (SExpr String)
parser :: GenParser Char st (SExpr [Char])
parser =
  [GenParser Char st (SExpr [Char])]
-> GenParser Char st (SExpr [Char])
forall s (m :: * -> *) t u a.
Stream s m t =>
[ParsecT s u m a] -> ParsecT s u m a
choice [GenParser Char st (SExpr [Char])
-> GenParser Char st (SExpr [Char])
forall tok st a. GenParser tok st a -> GenParser tok st a
try GenParser Char st (SExpr [Char])
forall u a. ParsecT [Char] u Identity (SExpr a)
unitP, GenParser Char st (SExpr [Char])
forall u. ParsecT [Char] u Identity (SExpr [Char])
nodeP, GenParser Char st (SExpr [Char])
forall u. ParsecT [Char] u Identity (SExpr [Char])
leafP]

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

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

        leafP :: ParsecT [Char] u Identity (SExpr [Char])
leafP      = [Char] -> SExpr [Char]
forall a. a -> SExpr a
atom ([Char] -> SExpr [Char])
-> ParsecT [Char] u Identity [Char]
-> ParsecT [Char] u Identity (SExpr [Char])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ParsecT [Char] u Identity [Char]
forall u. ParsecT [Char] u Identity [Char]
lonelyStr

        nodeP :: ParsecT [Char] u Identity (SExpr [Char])
nodeP      = do ParsecT [Char] u Identity Char -> ParsecT [Char] u Identity ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ParsecT [Char] u Identity Char -> ParsecT [Char] u Identity ())
-> ParsecT [Char] u Identity Char -> ParsecT [Char] u Identity ()
forall a b. (a -> b) -> a -> b
$ Char -> ParsecT [Char] u Identity Char
forall s (m :: * -> *) u.
Stream s m Char =>
Char -> ParsecT s u m Char
char Char
'('
                        ParsecT [Char] u Identity ()
forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m ()
spaces
                        [SExpr [Char]]
st <- ParsecT [Char] u Identity (SExpr [Char])
-> ParsecT [Char] u Identity ()
-> ParsecT [Char] u Identity [SExpr [Char]]
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 [Char] u Identity (SExpr [Char])
forall u. ParsecT [Char] u Identity (SExpr [Char])
parser ParsecT [Char] u Identity ()
forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m ()
spaces
                        ParsecT [Char] u Identity ()
forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m ()
spaces
                        ParsecT [Char] u Identity Char -> ParsecT [Char] u Identity ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ParsecT [Char] u Identity Char -> ParsecT [Char] u Identity ())
-> ParsecT [Char] u Identity Char -> ParsecT [Char] u Identity ()
forall a b. (a -> b) -> a -> b
$ Char -> ParsecT [Char] u Identity Char
forall s (m :: * -> *) u.
Stream s m Char =>
Char -> ParsecT s u m Char
char Char
')'
                        SExpr [Char] -> ParsecT [Char] u Identity (SExpr [Char])
forall (m :: * -> *) a. Monad m => a -> m a
return (SExpr [Char] -> ParsecT [Char] u Identity (SExpr [Char]))
-> SExpr [Char] -> ParsecT [Char] u Identity (SExpr [Char])
forall a b. (a -> b) -> a -> b
$ [SExpr [Char]] -> SExpr [Char]
forall a. [SExpr a] -> SExpr a
List [SExpr [Char]]
st


parseSExpr :: String -> Maybe (SExpr String)
parseSExpr :: [Char] -> Maybe (SExpr [Char])
parseSExpr [Char]
str = case Parsec [Char] () (SExpr [Char])
-> [Char] -> [Char] -> Either ParseError (SExpr [Char])
forall s t a.
Stream s Identity t =>
Parsec s () a -> [Char] -> s -> Either ParseError a
parse Parsec [Char] () (SExpr [Char])
forall u. ParsecT [Char] u Identity (SExpr [Char])
parser [Char]
"" [Char]
str of
  Left ParseError
s -> [Char] -> Maybe (SExpr [Char])
forall a. HasCallStack => [Char] -> a
error (ParseError -> [Char]
forall a. Show a => a -> [Char]
show ParseError
s) -- Nothing
  Right SExpr [Char]
t -> SExpr [Char] -> Maybe (SExpr [Char])
forall a. a -> Maybe a
Just SExpr [Char]
t

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