{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleContexts #-}

module Expr
  ( Expr(Lit, Var, Neg, Add, Sub, Mul, Div)
  , evaluateExpr
  , isValidName
  , parseExpr
  , variablesOf
  ) where

import Control.Applicative ((<|>))
import Data.Char (isAlpha)
import Data.Functor.Identity (Identity)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Text (Text)
import qualified Data.Text as Text
import Text.Parsec
  ( ParseError
  , ParsecT
  , Stream
  , between
  , chainl1
  , digit
  , eof
  , letter
  , many
  , many1
  , option
  , parse
  , try
  )
import Text.Parsec.Char (char, spaces)

data Expr
  = Lit Double
  | Var Text
  | Neg Expr
  | Add Expr Expr
  | Sub Expr Expr
  | Mul Expr Expr
  | Div Expr Expr
  deriving (Int -> Expr -> ShowS
[Expr] -> ShowS
Expr -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Expr] -> ShowS
$cshowList :: [Expr] -> ShowS
show :: Expr -> String
$cshow :: Expr -> String
showsPrec :: Int -> Expr -> ShowS
$cshowsPrec :: Int -> Expr -> ShowS
Show)

isValidName :: Text -> Bool
isValidName :: Text -> Bool
isValidName = (Char -> Bool) -> Text -> Bool
Text.all Char -> Bool
isAlpha

variablesOf :: Expr -> Set Text
variablesOf :: Expr -> Set Text
variablesOf = Set Text -> Expr -> Set Text
go forall a. Set a
Set.empty
  where
    go :: Set Text -> Expr -> Set Text
go Set Text
s (Lit Double
_) = Set Text
s
    go Set Text
s (Var Text
x) = forall a. Ord a => a -> Set a -> Set a
Set.insert Text
x Set Text
s
    go Set Text
s (Neg Expr
a) = Set Text -> Expr -> Set Text
go Set Text
s Expr
a
    go Set Text
s (Add Expr
a Expr
b) = Set Text -> Expr -> Set Text
go (Set Text -> Expr -> Set Text
go Set Text
s Expr
a) Expr
b
    go Set Text
s (Sub Expr
a Expr
b) = Set Text -> Expr -> Set Text
go (Set Text -> Expr -> Set Text
go Set Text
s Expr
a) Expr
b
    go Set Text
s (Mul Expr
a Expr
b) = Set Text -> Expr -> Set Text
go (Set Text -> Expr -> Set Text
go Set Text
s Expr
a) Expr
b
    go Set Text
s (Div Expr
a Expr
b) = Set Text -> Expr -> Set Text
go (Set Text -> Expr -> Set Text
go Set Text
s Expr
a) Expr
b

evaluateExpr :: Map Text Double -> Expr -> Maybe Double
evaluateExpr :: Map Text Double -> Expr -> Maybe Double
evaluateExpr Map Text Double
_ (Lit Double
d) = forall a. a -> Maybe a
Just Double
d
evaluateExpr Map Text Double
env (Var Text
x) = forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
x Map Text Double
env
evaluateExpr Map Text Double
env (Neg Expr
a) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Num a => a -> a
negate (Map Text Double -> Expr -> Maybe Double
evaluateExpr Map Text Double
env Expr
a)
evaluateExpr Map Text Double
env (Add Expr
a Expr
b) = forall a. Num a => a -> a -> a
(+) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Text Double -> Expr -> Maybe Double
evaluateExpr Map Text Double
env Expr
a forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Map Text Double -> Expr -> Maybe Double
evaluateExpr Map Text Double
env Expr
b
evaluateExpr Map Text Double
env (Sub Expr
a Expr
b) = (-) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Text Double -> Expr -> Maybe Double
evaluateExpr Map Text Double
env Expr
a forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Map Text Double -> Expr -> Maybe Double
evaluateExpr Map Text Double
env Expr
b
evaluateExpr Map Text Double
env (Mul Expr
a Expr
b) = forall a. Num a => a -> a -> a
(*) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Text Double -> Expr -> Maybe Double
evaluateExpr Map Text Double
env Expr
a forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Map Text Double -> Expr -> Maybe Double
evaluateExpr Map Text Double
env Expr
b
evaluateExpr Map Text Double
env (Div Expr
a Expr
b) = forall a. Fractional a => a -> a -> a
(/) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Text Double -> Expr -> Maybe Double
evaluateExpr Map Text Double
env Expr
a forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Map Text Double -> Expr -> Maybe Double
evaluateExpr Map Text Double
env Expr
b

parseExpr :: Stream s Identity Char => s -> Either ParseError Expr
parseExpr :: forall s. Stream s Identity Char => s -> Either ParseError Expr
parseExpr = forall s t a.
Stream s Identity t =>
Parsec s () a -> String -> s -> Either ParseError a
parse (forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
expr forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* forall s (m :: * -> *) t u.
(Stream s m t, Show t) =>
ParsecT s u m ()
eof) String
""

expr :: Stream s m Char => ParsecT s u m Expr
expr :: forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
expr = forall s (m :: * -> *) u a.
Stream s m Char =>
ParsecT s u m a -> ParsecT s u m a
pad forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
addOrSub

addOrSub :: Stream s m Char => ParsecT s u m Expr
addOrSub :: forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
addOrSub = forall s (m :: * -> *) t u a.
Stream s m t =>
ParsecT s u m a -> ParsecT s u m (a -> a -> a) -> ParsecT s u m a
chainl1 forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
mulOrDiv (forall {u}. ParsecT s u m (Expr -> Expr -> Expr)
add forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> forall {u}. ParsecT s u m (Expr -> Expr -> Expr)
sub)
  where
    add :: ParsecT s u m (Expr -> Expr -> Expr)
add = forall s u (m :: * -> *) a. ParsecT s u m a -> ParsecT s u m a
try (forall s (m :: * -> *) u a.
Stream s m Char =>
ParsecT s u m a -> ParsecT s u m a
pad (Expr -> Expr -> Expr
Add forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall s (m :: * -> *) u.
Stream s m Char =>
Char -> ParsecT s u m Char
char Char
'+'))
    sub :: ParsecT s u m (Expr -> Expr -> Expr)
sub = forall s u (m :: * -> *) a. ParsecT s u m a -> ParsecT s u m a
try (forall s (m :: * -> *) u a.
Stream s m Char =>
ParsecT s u m a -> ParsecT s u m a
pad (Expr -> Expr -> Expr
Sub forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall s (m :: * -> *) u.
Stream s m Char =>
Char -> ParsecT s u m Char
char Char
'-'))

mulOrDiv :: Stream s m Char => ParsecT s u m Expr
mulOrDiv :: forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
mulOrDiv = forall s (m :: * -> *) t u a.
Stream s m t =>
ParsecT s u m a -> ParsecT s u m (a -> a -> a) -> ParsecT s u m a
chainl1 forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
negative (forall {u}. ParsecT s u m (Expr -> Expr -> Expr)
mul forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> forall {u}. ParsecT s u m (Expr -> Expr -> Expr)
div)
  where
    mul :: ParsecT s u m (Expr -> Expr -> Expr)
mul = forall s u (m :: * -> *) a. ParsecT s u m a -> ParsecT s u m a
try (forall s (m :: * -> *) u a.
Stream s m Char =>
ParsecT s u m a -> ParsecT s u m a
pad (Expr -> Expr -> Expr
Mul forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall s (m :: * -> *) u.
Stream s m Char =>
Char -> ParsecT s u m Char
char Char
'*'))
    div :: ParsecT s u m (Expr -> Expr -> Expr)
div = forall s u (m :: * -> *) a. ParsecT s u m a -> ParsecT s u m a
try (forall s (m :: * -> *) u a.
Stream s m Char =>
ParsecT s u m a -> ParsecT s u m a
pad (Expr -> Expr -> Expr
Div forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall s (m :: * -> *) u.
Stream s m Char =>
Char -> ParsecT s u m Char
char Char
'/'))

negative :: Stream s m Char => ParsecT s u m Expr
negative :: forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
negative = forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m ()
spaces forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall b c a. (b -> c) -> (a -> b) -> a -> c
(.) forall a. a -> a
id forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s u (m :: * -> *) a. ParsecT s u m a -> ParsecT s u m [a]
many (Expr -> Expr
Neg forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall s (m :: * -> *) u.
Stream s m Char =>
Char -> ParsecT s u m Char
char Char
'-') forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
parensOrAtom)

parensOrAtom :: Stream s m Char => ParsecT s u m Expr
parensOrAtom :: forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
parensOrAtom = forall s u (m :: * -> *) a. ParsecT s u m a -> ParsecT s u m a
try (forall s (m :: * -> *) u a.
Stream s m Char =>
ParsecT s u m a -> ParsecT s u m a
parens forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
expr) forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> forall s u (m :: * -> *) a. ParsecT s u m a -> ParsecT s u m a
try forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
literal forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
variable

literal :: Stream s m Char => ParsecT s u m Expr
literal :: forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
literal = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Double -> Expr
Lit forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Read a => String -> a
read) forall {u}. ParsecT s u m String
double
  where
    double :: ParsecT s u m String
double = forall a. [a] -> [a] -> [a]
(++) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {u}. ParsecT s u m String
natural forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {u}. ParsecT s u m String
decimal
    natural :: ParsecT s u m String
natural = forall s (m :: * -> *) t u a.
Stream s m t =>
ParsecT s u m a -> ParsecT s u m [a]
many1 forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Char
digit
    decimal :: ParsecT s u m String
decimal = forall s (m :: * -> *) t a u.
Stream s m t =>
a -> ParsecT s u m a -> ParsecT s u m a
option String
"" ((:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s (m :: * -> *) u.
Stream s m Char =>
Char -> ParsecT s u m Char
char Char
'.' forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {u}. ParsecT s u m String
natural)

variable :: Stream s m Char => ParsecT s u m Expr
variable :: forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Expr
variable = Text -> Expr
Var forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
Text.pack forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s (m :: * -> *) t u a.
Stream s m t =>
ParsecT s u m a -> ParsecT s u m [a]
many1 forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m Char
letter

pad :: Stream s m Char => ParsecT s u m a -> ParsecT s u m a
pad :: forall s (m :: * -> *) u a.
Stream s m Char =>
ParsecT s u m a -> ParsecT s u m a
pad = forall s (m :: * -> *) t u open close a.
Stream s m t =>
ParsecT s u m open
-> ParsecT s u m close -> ParsecT s u m a -> ParsecT s u m a
between forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m ()
spaces forall s (m :: * -> *) u. Stream s m Char => ParsecT s u m ()
spaces

parens :: Stream s m Char => ParsecT s u m a -> ParsecT s u m a
parens :: forall s (m :: * -> *) u a.
Stream s m Char =>
ParsecT s u m a -> ParsecT s u m a
parens = forall s (m :: * -> *) t u open close a.
Stream s m t =>
ParsecT s u m open
-> ParsecT s u m close -> ParsecT s u m a -> ParsecT s u m a
between (forall s (m :: * -> *) u.
Stream s m Char =>
Char -> ParsecT s u m Char
char Char
'(') (forall s (m :: * -> *) u.
Stream s m Char =>
Char -> ParsecT s u m Char
char Char
')')