shunting yard

gistfile1.hs

{-# OPTIONS_GHC -Wall #-}
module ShuntingYard where

import qualified Data.Map as M

data Op = Op String Prec Assoc
        | OLParen | ORParen deriving (Show)

type Prec = Int

data Assoc = ARight | ALeft deriving (Show)
     -- `A` prefix added to distinguish from Either data constructors

type OpTable = M.Map String Op

opTable :: OpTable
opTable = M.fromList $ map (\op@(Op id_ _ _) -> (id_, op)) $
    [ (Op "^" 4 ARight)
    , (Op "*" 3 ALeft)
    , (Op "/" 3 ALeft)
    , (Op "+" 2 ALeft)
    , (Op "-" 2 ALeft)
    ]

data Token = TNum Int | TOp String | TLParen | TRParen deriving (Show)

lookupOp :: String -> OpTable -> Op
lookupOp opid tbl =
    case M.lookup opid tbl of
      Nothing -> error ("undefined op: " ++ opid)
      Just op -> op

sy :: OpTable -> [Op] -> [Token] -> [Token]
sy _ opStack [] =
    -- end of input, just add operator stack to output queue
    map (\(Op id_ _ _) -> TOp id_) opStack
sy tbl opStack (TNum i : rest) =
    -- add numbers to output directly
    TNum i : sy tbl opStack rest
sy tbl opStack (TOp newOpid : rest) =
    let opnew@(Op _ prec _) = lookupOp newOpid tbl in
    case opStack of
      (optop@(Op opid prec' assoc') : restOps)
        | prec' > prec ->
            -- operator in the stack has higher prec. so push the new op to
            -- the stack
            TOp opid : sy tbl (opnew : restOps) rest
        | prec' == prec ->
            -- operator in the stack and in the input stream has same
            -- precedences,
            case assoc' of
              ALeft  ->
                -- opeartor on the stack has left assoc, we should add
                -- it to the output queue
                TOp opid : sy tbl (opnew : restOps) rest
              ARight ->
                -- otherwise just push new operator to the stack
                sy tbl (opnew : optop : restOps) rest
        | prec' < prec -> sy tbl (opnew : optop : restOps) rest
      OLParen : _ -> sy tbl (opnew : opStack) rest
      ORParen : _ ->
        -- it's an error to have right paren on the stack
        error "RParen in op stack"
      [] -> sy tbl [opnew] rest
sy tbl opStack (TLParen : rest) = sy tbl (OLParen : opStack) rest
sy tbl opStack (TRParen : rest) =
    -- pop operators from stack until lparen is found and add to queue
    (map (\(Op id_ _ _) -> TOp id_) h ++ sy tbl t rest)
  where
    cond OLParen = False
    cond _       = True
    (h, (_ : t)) = span cond opStack

input1, input2, input3 :: [Token]
input1 = -- from wikipedia
  [ TNum 3, TOp "+", TNum 4, TOp "*", TNum 2, TOp "/"
  , TLParen, TNum 1, TOp "-", TNum 5, TRParen, TOp "^"
  , TNum 2, TOp "^", TNum 3 ]
input2 = [ TNum 1, TOp "*", TNum 2, TOp "*", TNum 3 ]
input3 = [ TNum 1, TOp "^", TNum 2, TOp "^", TNum 4 ]