My solution to the 2025 Haskell January Test

module Parser where

import Lexer
import Types

import Data.Bifunctor (Bifunctor (second))

------------------------------------------------------------------------------
-- Given...

showToken :: Token -> String
showToken (Ident v) = v
showToken (Nat v) = show v
showToken WhileTok = "while"
showToken t = [head [c | (c, t') <- tokenTable, t == t']]

printParse :: String -> IO ()
printParse input = either printError printOK (parse input)
 where
  printOK prog = putStrLn "Parse successful..." >> print prog
  printError err = putStr "Parse error: " >> printError' err
  printError'' t s =
    putStrLn
      ( s
          ++ " expected, but "
          ++ maybe "nothing" showToken t
          ++ " found"
      )
  printError' (BadChar c) = do
    putStr "Unrecognised character: "
    putStrLn [c]
  printError' (Unexpected t t') = printError'' t (showToken t')
  printError' (StmtNotFound t) = printError'' t "Statement"
  printError' (ExprNotFound t) = printError'' t "Expression"
  printError' (IntNotFound t) = printError'' t "Integer literal"
  printError' (UnparsedInput toks) =
    putStrLn
      ( "Unparsed input: "
          ++ unwords (map showToken toks)
      )

------------------------------------------------------------------------------

-- Given...
mHead :: [a] -> Maybe a
mHead (x : _) = Just x
mHead _ = Nothing

checkTok :: Token -> [Token] -> Either Error [Token]
checkTok tok toks
  | h == Just tok = Right (tail toks)
  | otherwise = Left (Unexpected h tok)
 where
  h = mHead toks

parseAtom :: Parser Expr
parseAtom [] = Left (ExprNotFound Nothing)
parseAtom (Ident x : toks) = Right (toks, Var x)
parseAtom (Nat x : toks) = Right (toks, Val x)
parseAtom (Minus : Nat x : toks) = Right (toks, Val (negate x))
parseAtom (Minus : toks) = Left (IntNotFound (mHead toks))
parseAtom (LParen : toks) = do
  (toks', expr) <- parseExpr toks
  toks'' <- checkTok RParen toks'
  return (toks'', expr)
parseAtom toks = Left (ExprNotFound (mHead toks))

parseTerm :: Parser Expr
parseTerm = parseE parseAtom Times Mul

parseExpr :: Parser Expr
parseExpr = parseE parseTerm Plus Add

parseE :: Parser Expr -> Token -> (Expr -> Expr -> Expr) -> Parser Expr
parseE p tok constructor toks = do
  (toks', term) <- p toks
  parseE' p tok constructor term toks'

parseE' :: Parser Expr -> Token -> (Expr -> Expr -> Expr) -> Expr -> Parser Expr
parseE' p tok constructor t ts@(tok' : toks)
  | tok == tok' = do
      (toks', term) <- p toks
      parseE' p tok constructor (constructor t term) toks'
  | otherwise = return (ts, t)
parseE' _ _ _ t toks = Right (toks, t)

parseStmt :: Parser Stmt
parseStmt ((Ident v) : Eq : toks) = do
  (toks', expr) <- parseExpr toks
  return (toks', Asgn v expr)
parseStmt (WhileTok : toks) = do
  (toks', expr) <- parseExpr toks
  toks'' <- checkTok LBrace toks'
  (toks''', block) <- parseBlock toks''
  -- it's obvious that you want a State monad at this point
  toks'''' <- checkTok RBrace toks'''
  return (toks'''', While expr block)
parseStmt toks = Left (StmtNotFound (mHead toks))

parseBlock :: Parser Block
parseBlock toks = fmap (second reverse) (parseBlock' [] toks)

parseBlock' :: [Stmt] -> Parser Block
parseBlock' [] toks = do
  (toks', stmt) <- parseStmt toks
  parseBlock' [stmt] toks'
parseBlock' stmts (Semi : toks) = do
  (toks', stmt) <- parseStmt toks
  parseBlock' (stmt : stmts) toks'
parseBlock' stmt toks = return (toks, stmt)

parse :: String -> Either Error Program
parse input = do
  toks <- tokenise input
  (toks', program) <- parseBlock toks
  case toks' of
    [] -> Right program
    ts -> Left (UnparsedInput ts)