Q

q.hs

{-# LANGUAGE ExistentialQuantification, PackageImports #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
import Control.Applicative
import Text.Trifecta
import "mtl" Control.Monad.State
import qualified Data.Map as M
import Data.Monoid
import Data.Ratio

newtype Var = Var String deriving (Eq, Ord, Show)
newtype Numeral = Numeral Rational deriving (Eq, Enum, Ord, Num, Fractional, Show)
data Q = 
  Q Expr
  | Var := Expr
  | Func Var [Var] Expr
  | Q :. Q
  deriving (Eq, Show)
data Expr =
  NumExpr Numeral
  | VarExpr Var
  | Expr :+: Expr | Expr :-: Expr
  | Expr :*: Expr | Expr :/: Expr
  | Expr :^: Expr
  | Lambda [Var] Expr
  | Subst Var [Expr]
  deriving (Eq)

type Env = StateT (M.Map Var Expr) IO

(^%) a (Numeral b)
  | denominator b == 1 = a^(numerator b)
  | otherwise = error "exponential"

instance Show Expr where
  show (NumExpr (Numeral n))
    | denominator n == 1 = show $ numerator n
    | otherwise = show (numerator n) ++ "/" ++ show (denominator n)
  show (VarExpr (Var v)) = v
  show (x :+: y) = "(" ++ show x ++ "+" ++ show y ++ ")"
  show (x :-: y) = "(" ++ show x ++ "-" ++ show y ++ ")"
  show (x :*: y)
    | isFlat x || isFlat y = show x ++ show y
    | otherwise = "(" ++ show x ++ "*" ++ show y ++ ")"
  show (x :/: y)
    | isFlat x && isFlat y = show x ++ "/" ++ show y
    | otherwise = "(" ++ show x ++ "/" ++ show y ++ ")"
  show (x :^: y)
    | isFlat x && isFlat y = show x ++ "^" ++ show y
    | otherwise = "(" ++ show x ++ "^" ++ show y ++ ")"
  show (Lambda vs e) = "\\" ++ foldr1 (\x y -> x ++ " " ++ y) (fmap (\(Var x) -> x) vs) ++ " -> " ++ show e
  show (Subst (Var v) es) = v ++ "(" ++ foldr1 (\x y -> x ++ "," ++ y) (fmap show es) ++ ")"

isFlat (NumExpr _) = True
isFlat (VarExpr _) = True
isFlat _ = False

var :: Parser Var
var = Var <$> some alphaNum <* spaces

parseQ :: Parser Q
parseQ = try line <|> stmt where
  stmt = try def <|> try func <|> try expr
  line = (:.) <$> stmt <*> (symbol ";" >> parseQ)
  expr = Q <$> parseExpr
  def = (:=) <$> var <*> (symbol "=" >> parseExpr)
  func = do
    f <- var
    xs <- parens (var `sepBy` symbol ",")
    Func f xs <$> (symbol "=" *> parseExpr)

parseExpr :: Parser Expr
parseExpr = spaces >> expr where
  -- ^ > *,/ > +,-
  expr = factor `chainl1` expop `chainl1` mulop `chainl1` addop
  factor = try subst <|> try numerals <|> try varexpr <|> try (parens parseExpr)

  addop = (:+:) <$ symbol "+" <|> (:-:) <$ symbol "-"
  mulop = (:*:) <$ symbol "*" <|> (:/:) <$ symbol "/"
  expop = (:^:) <$ symbol "^"
  numerals = NumExpr . Numeral . (% 1) <$> integer
  numvar = (:*:) <$> numerals <*> expr
  varexpr = VarExpr <$> var
  subst = Subst <$> var <*> parens (expr `sepBy` symbol ",")

eval :: Q -> Env ()
eval (q1 :. q2) = eval q1 >> eval q2
eval (v := e) = modify $ M.insert v e
eval (Func f xs e) = modify $ M.insert f (Lambda xs e)
eval (Q q) = do
  lift . putStr $ show q ++ " = "
  lift . print =<< apply q

apply :: Expr -> Env Expr
apply z@(NumExpr _) = return z
apply z@(VarExpr v) = do
  m <- get
  case M.member v m of
    True -> return $ m M.! v
    False -> return z
apply (x :+: y) = apply' (:+:) (+) x y
apply (x :-: y) = apply' (:-:) (-) x y
apply (x :*: y) = apply' (:*:) (*) x y
apply (x :/: y) = apply' (:/:) (/) x y
apply (x :^: (NumExpr 1)) = apply x
apply (x :^: y) = apply' (:^:) (^%) x y
apply (Lambda vs e) = Lambda vs <$> apply e
apply (Subst v es) = do
  m <- get
  let (Lambda us f) = m M.! v
  put $ foldr (\(a,b) -> M.insert a b) m $ zip us es
  x <- apply f
  put $ m
  return x

apply' k f x y = do
  x' <- apply x
  y' <- apply y
  if isNum x' && isNum y'
    then return $ NumExpr $ f (fromNum x') (fromNum y')
    else return $ k x' y'
  where

isNum :: Expr -> Bool
isNum (NumExpr _) = True
isNum _ = False

fromNum :: Expr -> Numeral
fromNum (NumExpr e) = e

main = do
  s <- getContents
  case parseString parseQ mempty s of
    Success a -> execStateT (eval a) M.empty >>= print
    Failure d -> print d