tricu/src/Parser.hs

module Parser where

import Debug.Trace
import Lexer
import Research hiding (toList)
import Control.Exception (throw)
import Data.List.NonEmpty (toList)
import qualified Data.Set as Set
import Data.Void
import Text.Megaparsec
import Text.Megaparsec.Char
import Text.Megaparsec.Error (errorBundlePretty, ParseErrorBundle)

type Parser = Parsec Void [LToken]

data SaplingAST
  = SVar String
  | SInt Int
  | SStr String
  | SList [SaplingAST]
  | SFunc String [String] SaplingAST
  | SApp SaplingAST SaplingAST
  | TLeaf
  | TStem SaplingAST
  | TFork SaplingAST SaplingAST
  | SLambda [String] SaplingAST
  deriving (Show, Eq, Ord)

parseSapling :: String -> [SaplingAST]
parseSapling input =
  let nonEmptyLines = filter (not . null) (lines input)
  in map parseSingle nonEmptyLines

parseSingle :: String -> SaplingAST
parseSingle input = case runParser parseExpression "" (lexSapling input) of
  Left err -> error $ handleParseError err
  Right ast -> ast

scnParser :: Parser ()
scnParser = skipMany (satisfy isNewline)

parseExpression :: Parser SaplingAST
parseExpression = choice
  [ try parseFunction
  , try parseLambda
  , try parseLambdaExpression
  , try parseListLiteral
  , try parseApplication
  , try parseTreeTerm
  , parseLiteral
  ]

parseFunction :: Parser SaplingAST
parseFunction = do
  LIdentifier name <- satisfy isIdentifier
  args <- many (satisfy isIdentifier)
  satisfy (== LAssign)
  body <- parseExpression
  return (SFunc name (map getIdentifier args) body)

parseAtomicBase :: Parser SaplingAST
parseAtomicBase = choice
    [ try parseVarWithoutAssignment
    , parseTreeLeaf
    , parseGrouped
    ]
parseVarWithoutAssignment :: Parser SaplingAST
parseVarWithoutAssignment = do
    LIdentifier name <- satisfy isIdentifier
    if (name == "t" || name == "__result")
    then fail $ "Reserved keyword: " ++ name ++ " cannot be assigned."
    else notFollowedBy (satisfy (== LAssign)) *> return (SVar name)

parseLambda :: Parser SaplingAST
parseLambda = between (satisfy (== LOpenParen)) (satisfy (== LCloseParen)) $ do
  satisfy (== LBackslash)
  param <- satisfy isIdentifier
  rest <- many (satisfy isIdentifier)
  satisfy (== LColon)
  body <- parseLambdaExpression
  let nestedLambda = foldr (\v acc -> SLambda [v] acc) body (map getIdentifier rest)
  return (SLambda [getIdentifier param] nestedLambda)

parseLambdaExpression :: Parser SaplingAST
parseLambdaExpression = choice
  [ try parseLambdaApplication
  , parseAtomicLambda
  ]

parseAtomicLambda :: Parser SaplingAST
parseAtomicLambda = choice
  [ parseVar
  , parseTreeLeaf
  , parseLiteral
  , parseListLiteral
  , try parseLambda
  , between (satisfy (== LOpenParen)) (satisfy (== LCloseParen)) parseLambdaExpression
  ]

parseApplication :: Parser SaplingAST
parseApplication = do
  func <- parseAtomicBase
  args <- many parseAtomic
  return $ foldl (\acc arg -> SApp acc arg) func args

parseLambdaApplication :: Parser SaplingAST
parseLambdaApplication = do
  func <- parseAtomicLambda
  args <- many parseAtomicLambda
  return $ foldl (\acc arg -> SApp acc arg) func args

isTreeTerm :: SaplingAST -> Bool
isTreeTerm TLeaf = True
isTreeTerm (TStem _) = True
isTreeTerm (TFork _ _) = True
isTreeTerm _ = False

parseTreeLeaf :: Parser SaplingAST
parseTreeLeaf = satisfy isKeywordT *> notFollowedBy (satisfy (== LAssign)) *> pure TLeaf

getIdentifier :: LToken -> String
getIdentifier (LIdentifier name) = name
getIdentifier _ = error "Expected identifier"

parseTreeTerm :: Parser SaplingAST
parseTreeTerm = do
  base <- parseTreeLeafOrParenthesized
  rest <- many parseTreeLeafOrParenthesized
  pure $ foldl combine base rest
  where
    combine acc next = case acc of
      TLeaf -> TStem next
      TStem t -> TFork t next
      TFork _ _ -> TFork acc next

parseTreeLeafOrParenthesized :: Parser SaplingAST
parseTreeLeafOrParenthesized = choice
  [ between (satisfy (== LOpenParen)) (satisfy (== LCloseParen)) parseTreeTerm
  , parseTreeLeaf
  ]

foldTree :: [SaplingAST] -> SaplingAST
foldTree [] = TLeaf
foldTree [x] = x
foldTree (x:y:rest) = TFork x (foldTree (y:rest))

parseAtomic :: Parser SaplingAST
parseAtomic = choice
  [ parseVar
  , parseTreeLeaf
  , parseListLiteral
  , parseGrouped
  , parseLiteral
  ]

parseGrouped :: Parser SaplingAST
parseGrouped = between (satisfy (== LOpenParen)) (satisfy (== LCloseParen)) parseExpression

parseLiteral :: Parser SaplingAST
parseLiteral = choice
  [ parseIntLiteral
  , parseStrLiteral
  ]

parens :: Parser SaplingAST -> Parser SaplingAST
parens p = do
  satisfy (== LOpenParen)
  result <- p
  satisfy (== LCloseParen)
  return result

parseListLiteral :: Parser SaplingAST
parseListLiteral = do
  satisfy (== LOpenBracket)
  elements <- many parseListItem
  satisfy (== LCloseBracket)
  return (SList elements)

parseListItem :: Parser SaplingAST
parseListItem = choice
  [ parseGroupedItem
  , parseListLiteral
  , parseSingleItem
  ]

parseGroupedItem :: Parser SaplingAST
parseGroupedItem = do
  satisfy (== LOpenParen)
  inner <- parseExpression
  satisfy (== LCloseParen)
  return inner

parseSingleItem :: Parser SaplingAST
parseSingleItem = do
  token <- satisfy isListItem
  case token of
    LIdentifier name -> return (SVar name)
    LKeywordT -> return TLeaf
    _ -> fail "Unexpected token in list item"

isListItem :: LToken -> Bool
isListItem (LIdentifier _) = True
isListItem LKeywordT = True
isListItem _ = False

parseVar :: Parser SaplingAST
parseVar = do
  LIdentifier name <- satisfy isIdentifier
  if (name == "t" || name == "__result")
    then fail $ "Reserved keyword: " ++ name ++ " cannot be assigned."
    else return (SVar name)

parseIntLiteral :: Parser SaplingAST
parseIntLiteral = do
  LIntegerLiteral value <- satisfy isIntegerLiteral
  return (SInt value)

parseStrLiteral :: Parser SaplingAST
parseStrLiteral = do
  LStringLiteral value <- satisfy isStringLiteral
  return (SStr value)

-- Boolean Helpers
isKeywordT (LKeywordT) = True
isKeywordT _ = False
isIdentifier (LIdentifier _) = True
isIdentifier _ = False
isIntegerLiteral (LIntegerLiteral _) = True
isIntegerLiteral _ = False
isStringLiteral (LStringLiteral _) = True
isStringLiteral _ = False
isLiteral (LIntegerLiteral _) = True
isLiteral (LStringLiteral _) = True
isLiteral _ = False
isNewline (LNewline) = True
isNewline _ = False

-- Error Handling
handleParseError :: ParseErrorBundle [LToken] Void -> String
handleParseError bundle =
  let errors = bundleErrors bundle
      errorList = toList errors
      formattedErrors = map showError errorList
  in unlines ("Parse error(s) encountered:" : formattedErrors)

showError :: ParseError [LToken] Void -> String
showError (TrivialError offset (Just (Tokens tokenStream)) expected) =
  "Parse error at offset " ++ show offset ++ ": unexpected token "
  ++ show tokenStream ++ ", expected one of " ++ show (Set.toList expected)
showError (FancyError offset fancy) =
  "Parse error at offset " ++ show offset ++ ":\n " ++ unlines (map show (Set.toList fancy))
showError (TrivialError offset Nothing expected) =
  "Parse error at offset " ++ show offset ++ ": expected one of "
  ++ show (Set.toList expected)