A CLI Calculator with variables, functions and arithmatic

Calculator.hs

{-
King - 1/11/2018
Resources Used:
    - Haskell quick docs: https://www.haskell.org/hoogle/
    - Operator precedence: http://ee.hawaii.edu/~tep/EE160/Book/chap5/_28291_table577.gif
    - Precedence climbing: http://www.engr.mun.ca/~theo/Misc/exp_parsing.htm#climbing
-}

module Main where

import Data.Map (Map)
import Text.Read (readMaybe)
import qualified Data.Map as Map
import Control.Exception (try, evaluate, SomeException)
import System.IO (stdout, hSetBuffering, BufferMode(..))
import Data.Char (isSpace, isNumber, isAlpha, isAlphaNum)

-- Data Types

type EInt = Integer
type EFloat = Double
type EMap = Map String Expr
type ENumber = Either EInt EFloat
type EBinopInt = EInt -> EInt -> EInt
type EBinopFloat = EFloat -> EFloat -> EFloat

data Token =
    Ident String |
    Comma | LParen | RParen |
    Op Char | Number ENumber
    deriving (Eq, Show)

data Expr =
    Var String |
    Const ENumber |
    If Expr Expr Expr |
    Call String [Expr] |
    Binop Char Expr Expr |
    Function [String] Expr
    deriving (Show)

-- Main REPL

main :: IO ()
main = do
    hSetBuffering stdout NoBuffering
    repl (Map.empty :: EMap) (putStrLn "> Calculator Started")

repl :: EMap -> IO () -> IO ()
repl locals output = do
    _ <- output
    _ <- putStr "> "
    expr <- getLine 
    result <- try (evaluate $ calc locals expr) :: IO (Either SomeException (EMap, ENumber))
    case result of
        Left ex -> repl locals (putStrLn $ "[Error] " ++ show ex)
        Right (newLocals, value) -> case value of
            Left n -> repl newLocals (print n)
            Right n -> repl newLocals (print n)
        
calc :: EMap -> String -> (EMap, ENumber)
calc locals str =
    case lexify str of
        [] -> (locals, Left 0)
        tokens -> result
            where (expr, _) = parseExpr tokens 0
                  result = evalExpr locals expr 

-- Lexing

lexify :: String -> [Token]
lexify [] = []
lexify str
    | isSpace char = lexify rest
    | char == ',' = Comma : lexify rest
    | char == '(' = LParen : lexify rest
    | char == ')' = RParen : lexify rest
    | isNumber char = readNumber [char] rest
    | isAlpha char = readIdentifier [char] rest
    | isOperator char = (Op char) : lexify rest
    | otherwise = error $ "Invalid character: " ++ [char]

    where char = head str
          rest = tail str

          isOperator :: Char -> Bool
          isOperator c = char `elem` "+-*/^%=?!<>&|[]"

readIdentifier :: String -> String -> [Token]
readIdentifier str rest =
    readUntil str rest (\char -> isAlphaNum char) (\s -> Ident s)

readNumber :: String -> String -> [Token]
readNumber str rest =
    readUntil str rest (\c -> (isNumber c) || c == '.') (\s -> parseNum s)

readUntil :: String -> String -> (Char -> Bool) -> (String -> Token) -> [Token]
readUntil str [] charTest genTok = [genTok $ reverse $ str]
readUntil str rest charTest genTok =
    if charTest char
        then readUntil (char : str) remaining charTest genTok 
        else (genTok $ reverse $ str) : lexify rest
    where char = head rest
          remaining = tail rest

parseNum :: String -> Token
parseNum str
    | 'e' `elem` str || '.' `elem` str = parseFloat str
    | otherwise = parseInt str
    where 
        parseInt :: String -> Token
        parseInt str = case readMaybe str :: Maybe EInt of
            Just n -> Number $ Left n; Nothing -> parseFloat str

        parseFloat :: String -> Token
        parseFloat str = case readMaybe str :: Maybe EFloat of
            Just n -> Number $ Right n; Nothing -> error "Invalid number literal"

-- Parsing

parseExpr :: [Token] -> Int -> (Expr, [Token])
parseExpr tokens precedence
    | head tokens == Ident "if" = parseIf $ tail tokens
    | otherwise = let (value, rest) = parseOther tokens in
        parseBinop rest precedence value

parseIf :: [Token] -> (Expr, [Token])
parseIf tokens
    | null rest = error "Expected 'then'"
    | head rest == Ident "then" = case parseExpr (tail rest) 0 of
        (expr, []) -> (If cond expr (Const $ Left 0), [])
        (expr, other) -> case head other of
            Ident s -> if s == "else" then 
                let (after, remains) = parseExpr (tail other) 0 in
                    (If cond expr after, remains)
                else (If cond expr (Const $ Left 0), other)
            _ -> (If cond expr (Const $ Left 0), other)
    | otherwise = parseIf []
    where (cond, rest) = parseExpr tokens 0

parseOther :: [Token] -> (Expr, [Token])
parseOther [] = error "Unexpected expr"
parseOther tokens = 
    case token of
        Number n ->
            (Const n, rest)
        Ident name ->
            if null rest || (head rest) /= LParen
            then (Var name, rest)
            else parseCall name [] (tail rest)
        Op char ->
            if char == '-'
            then let (value, remaining) = parseExpr rest 0 in
                (Binop '-' (Const $ Left $ 0) value, remaining)
            else error $ "Unexpected unary operator" ++ [char]
        LParen ->
            let (value, remaining) = parseExpr rest 0 in
                if null remaining || (head remaining) /= RParen
                then parseOther [RParen]
                else (value, tail remaining)
        RParen -> error "Unmatched parentheses"
        _ -> parseOther []

    where token = head tokens
          rest  = tail tokens

parseCall :: String -> [Expr] -> [Token] -> (Expr, [Token])
parseCall [] [] [] = error "Unterminated call"
parseCall name args [] = (Call name (reverse args), [])
parseCall name args tokens =
    case head tokens of
        Comma -> error "Unexpected comma"
        RParen -> (Call name args, tail tokens)
        _ -> let (arg, rest) = parseExpr tokens 0 in
            if null rest then parseCall [] [] []
            else case head rest of
                Comma -> parseCall name (arg : args) (tail rest)
                RParen -> parseCall name (arg : args) rest
                _ -> parseCall [] [] []

parseBinop :: [Token] -> Int -> Expr -> (Expr, [Token])
parseBinop [] precedence value = (value, [])
parseBinop tokens precedence lhs
    | getPrecedence token >= precedence =
        let (rhs, rest) = parseExpr (tail tokens) nextPrec in
            parseBinop rest precedence (Binop tokenOp lhs rhs)
    | otherwise = (lhs, tokens)

    where token = head tokens
          tokenOp = getOperator token
          nextPrec = precedence + (if isRightAssociative tokenOp then 1 else 0)

          getOperator :: Token -> Char
          getOperator (Op char) = char
          getOperator _ = ' '

          isRightAssociative :: Char -> Bool
          isRightAssociative char
            | char `elem` "^=" = True
            | otherwise        = False

          getPrecedence :: Token -> Int
          getPrecedence (Op char)
            | char == '='        = 0
            | char == '|'        = 1
            | char == '&'        = 2
            | char `elem` "+-"   = 3
            | char `elem` "*/"   = 4
            | char `elem` "^%"   = 5
            | char `elem` "?!"   = 6
            | char `elem` "<>[]" = 7
            | otherwise          = -1
          getPrecedence _ = getPrecedence $ Op $ ' '
        
-- Evaluation

evalExpr :: EMap -> Expr -> (EMap, ENumber)
evalExpr locals (Const value) = 
    (locals, value)
evalExpr locals (Function args body) =
    (locals, Left $ toInteger $ length $ args)
evalExpr locals (Var name) = 
    case Map.lookup name locals of
        Nothing -> error $ "Undefined variable " ++ name
        Just expr -> evalExpr locals expr
evalExpr locals (Call name args) =
    case Map.lookup name locals of
        Nothing -> error $ "Undefined function " ++ name
        Just func -> case func of
            Function fargs body -> (locals, evalFunc locals args fargs body)
            _ -> error $ name ++ " is a variable and only functions are callable"
evalExpr locals (If cond body after) = 
    case condResult of
        Left result -> if result /= 0
            then evalExpr newLocals body else evalExpr newLocals after
        Right result -> if result /= 0.0 
            then evalExpr newLocals body else evalExpr newLocals after
    where (newLocals, condResult) = evalExpr locals cond
evalExpr locals (Binop op lhs rhs) = case op of
    '+' -> (locals, evalCalc (e lhs) (e rhs) (\a b -> a + b) (\a b -> a + b))
    '-' -> (locals, evalCalc (e lhs) (e rhs) (\a b -> a - b) (\a b -> a - b))
    '*' -> (locals, evalCalc (e lhs) (e rhs) (\a b -> a * b) (\a b -> a * b))
    '^' -> (locals, evalCalc (e lhs) (e rhs) (\a b -> a ^ b) (\a b -> a ** b))
    '/' -> (locals, evalCalc (e lhs) (e rhs) (\a b -> a `div` b) (\a b -> a / b))
    '%' -> (locals, evalCalc (e lhs) (e rhs) (\a b -> a `mod` b) (\a b -> error "Cannot modulo decimals"))
    '<' -> (locals, evalCalc (e lhs) (e rhs) (eCmp 1 0 $ \a b -> a < b) (eCmp 1.0 0.0 $ \a b -> a < b))
    '>' -> (locals, evalCalc (e lhs) (e rhs) (eCmp 1 0 $ \a b -> a > b) (eCmp 1.0 0.0 $ \a b -> a > b))
    '[' -> (locals, evalCalc (e lhs) (e rhs) (eCmp 1 0 $ \a b -> a <= b) (eCmp 1.0 0.0 $ \a b -> a <= b))
    ']' -> (locals, evalCalc (e lhs) (e rhs) (eCmp 1 0 $ \a b -> a >= b) (eCmp 1.0 0.0 $ \a b -> a >= b))
    '?' -> (locals, evalCalc (e lhs) (e rhs) (eCmp 1 0 $ \a b -> a == b) (eCmp 1.0 0.0 $ \a b -> a == b))
    '!' -> (locals, evalCalc (e lhs) (e rhs) (eCmp 1 0 $ \a b -> a /= b) (eCmp 1.0 0.0 $ \a b -> a /= b))
    '&' -> (locals, evalCalc (e lhs) (e rhs) (eCmp 1 0 $ \a b -> a /=0 && b /=0) (eCmp 1.0 0.0 $ \a b -> a /=0.0 && b /=0.0))
    '|' -> (locals, evalCalc (e lhs) (e rhs) (eCmp 1 0 $ \a b -> a /=0 || b /=0) (eCmp 1.0 0.0 $ \a b -> a /=0.0 || b/=0.0))
    '=' -> case lhs of
        Var name -> evalInsert locals name rhs rhs
        Call name args -> evalLambda locals name rhs [] args
        expr -> error $ (show expr) ++ " is not assignable"
    _ -> error $ "Invalid operator " ++ [op] 

    where 
        e :: Expr -> ENumber
        e expr = snd $ evalExpr locals expr

        eCmp :: a -> a -> (a -> a -> Bool) -> (a -> a -> a)
        eCmp tru fal check = \x y -> if check x y then tru else fal

        evalLambda :: EMap -> String -> Expr -> [String] -> [Expr] -> (EMap, ENumber)
        evalLambda locals name body args [] = 
            evalInsert locals name (Function (reverse args) body) (Const $ Left $ 0)
        evalLambda locals name body args (expr:rest) = case expr of
            Var argName -> evalLambda locals name body (argName : args) rest
            _ -> error "Function can only have names as argument identifiers"

        evalInsert :: EMap -> String -> Expr -> Expr -> (EMap, ENumber)
        evalInsert locals name value retValue =
            evalExpr newLocals retValue
            where newLocals = Map.insert name value locals

evalCalc :: ENumber -> ENumber -> EBinopInt -> EBinopFloat -> ENumber
evalCalc lhsNumber rhsNumber calcInt calcFloat =
    case lhsNumber of
        Left lhs -> case rhsNumber of
            Left rhs -> Left $ calcInt lhs rhs
            Right rhs -> Left $ calcInt lhs (round rhs)
        Right lhs -> case rhsNumber of
            Right rhs -> Right $ calcFloat lhs rhs
            Left rhs -> Right $ calcFloat lhs (fromInteger rhs)     

evalFunc :: EMap -> [Expr] -> [String] -> Expr -> ENumber
evalFunc locals [] [] body = value
    where (_, value) = evalExpr locals body
evalFunc locals args fargs body
    | argLen == fargLen =
        let (_, argValue) = evalExpr locals (head args) in
            let newLocals = Map.insert (head fargs) (Const argValue) locals in
                evalFunc newLocals (tail args) (tail fargs) body
    | otherwise = 
        error $ "Expected " ++ (show fargLen) ++ " arguments. Got " ++ (show argLen)
    where argLen = length args
          fargLen = length fargs