Simple Calculator

Expr.hs

{-# LANGUAGE DeriveDataTypeable #-}

module Expr where

import Data.Maybe
import Data.Generics
import Control.Monad
import Control.Applicative

import Test.QuickCheck
data Expr = Dbl Double
          | Var String
          | Add Expr Expr
          | Mul Expr Expr
          | Neg Expr
          | Rcp Expr
            deriving (Eq,Show,Data,Typeable)

type Env = [(String,Expr)]

eval :: Expr -> Env -> Double
eval (Dbl a)     = const a
eval (Add e1 e2) = liftA2 (+) (eval e1) (eval e2)
eval (Mul e1 e2) = liftA2 (*) (eval e1) (eval e2)
eval (Neg e)     = negate . eval e
eval (Rcp e)     = recip  . eval e
eval (Var name)  = let res = fromMaybe (error $ "Can't find variable " ++ name) . lookup name
                   in \env -> eval (res env) env

-- We can use number literals as Exprs e.g. 1, 3/2, Var "x" + 4.
instance Num Expr where
  fromInteger = Dbl . fromInteger
  (+) = Add
  (*) = Mul

  a - b = Add a (Neg b)
  negate = Neg

  abs = undefined; signum = undefined

instance Fractional Expr where
  recip = Rcp
  a / b = Mul a (Rcp b)

  fromRational = undefined

optimize :: Expr -> Expr
optimize = reduce . normalize . reduce

reduce :: Expr -> Expr
reduce = everywhere (mkT step)
  where step e = case e of
                   Add (Dbl 0) b -> b
                   Add a (Dbl 0) -> a
                   Mul (Dbl 1) b -> b
                   Mul a (Dbl 1) -> a
--                    Mul (Dbl 0) b -> Dbl 0
--                    Mul a (Dbl 0) -> Dbl 0

                   Add (Dbl a) (Dbl b) -> Dbl (a + b)
                   Mul (Dbl a) (Dbl b) -> Dbl (a * b)

                   Neg (Dbl a) -> Dbl . negate $ a
                   Rcp (Dbl a) -> Dbl . recip  $ a

                   Neg (Neg a) -> a
                   Rcp (Rcp a) -> a

                   _ -> e

normalize :: Expr -> Expr
normalize = everywhere (mkT prepare) . everywhere (mkT leftize)
  where leftize expr = case expr of
                         Add e e'@Add{} -> Add e' e
                         Mul e e'@Mul{} -> Mul e' e

                         Add e@Var{} e'@Dbl{} -> Add e' e
                         Mul e@Var{} e'@Dbl{} -> Mul e' e

                         _ -> expr

        prepare :: Expr -> Expr
        prepare expr = case expr of
                         e@Add{} -> span_apply Add unAdd e
                         e@Mul{} -> span_apply Mul unMul e

                         _ -> expr

        span_apply con dec e = uncurry ($) . split $ e
          where split e = (collect_expr e, collect_const e)
                collect_expr = maybe id r . dec
                  where r (e,d@Dbl{}) = \e' -> collect_expr e e'
                        r (e,f)       = \e' -> con (collect_expr e e') f
                collect_const e = maybe e r . dec $ e
                  where r (e,d@Dbl{}) = con (collect_const e) d
                        r (e,_)       = collect_const e

        unAdd e = case e of Add x y -> Just (x,y); _ -> Nothing
        unMul e = case e of Mul x y -> Just (x,y); _ -> Nothing

satisfy_near :: (Expr -> Expr) -> Property
satisfy_near f = forAll const_expr $ \e -> eval e [] `near` eval (f e) []

near :: Double -> Double -> Bool
near x y = x == y || abs (x - y) < 0.00001

const_expr :: Gen Expr
const_expr = frequency [
              (80, fmap Dbl arbitrary)
             ,(10, fmap Neg const_expr)
             ,(10, fmap Rcp const_expr)
             ,(10, liftM2 Add const_expr const_expr)
             ,(10, liftM2 Mul const_expr const_expr)
             ]

runTests :: IO ()
runTests = mapM_ quickCheck [
            satisfy_near normalize
           ,satisfy_near reduce
           ,satisfy_near optimize
           ]

-- Usage:
--   test (1 + 2 + 3 * (4 + 5)) []
--   test (1 + 2 + 3 * (4 + Var "x" + 5)) [("x",1)]
--   test (Var "x" / 3 / Var "y" * 0) []
--   test (1 + 0 * 1 / 0) []
test :: Expr -> Env -> IO ()
test e env = do
  lbl "Expression" e
  lbl "Normalized" (normalize e)
  lbl "Optimized"  (optimize e)
  lbl "Evaluated"  (eval e env)
  where lbl l e = putStrLn $ l ++ ": " ++ show e