module MProver.Parser where

import MProver.Syntax

import Text.Parsec hiding (parse)
import Text.Parsec.Pos
import Text.Parsec.Expr
import Text.Parsec.Language
import qualified Text.Parsec.Token as T

import Control.Monad

import Data.Char
import Data.Ratio
import Data.Either

import Unbound.LocallyNameless hiding (Infix,Fixity,fixity)

lexer = T.makeTokenParser (haskellDef {
                             T.reservedNames = ["Forall","Foralli","Assume","Assuming","eval","subst","trans","symm","proving","by"] ++ T.reservedNames haskellDef,
                             T.reservedOpNames = [":::","_|_"] ++ T.reservedOpNames haskellDef})

symbol        = T.symbol lexer
reserved      = T.reserved lexer
reservedOp    = T.reservedOp lexer
identifier    = try $ do
                  i <- T.identifier lexer 
                  if isLower (head i)
                     then return i
                     else fail "expecting identifier"
ctor          = try $ do
                  i <- T.identifier lexer
                  if isUpper (head i)
                     then return i
                     else fail "expecting constructor"
lparen        = symbol "("
rparen        = symbol ")"
lcurly        = symbol "{"
rcurly        = symbol "}"
parens        = T.parens lexer
braces        = T.braces lexer
brackets      = T.brackets lexer
comma         = T.comma lexer
operator      = T.operator lexer
semi          = T.semi lexer
natural       = T.natural lexer
float         = T.float lexer
stringLiteral = T.stringLiteral lexer
charLiteral   = T.charLiteral lexer

-- Parser state is a bool indicating whether we are parsing in a context where
-- bottom is okay in expressions/patterns; and a list of fixity declarations
-- currently in scope.
type UserState = (Bool,[FixityDecl])

withBotok b m = do
  (botok,fixitydecls) <- getState
  putState (b,fixitydecls)
  r <- m
  (_,fixitydecls) <- getState
  putState (botok,fixitydecls)
  return r

program = do
  fixitydecls <- lookAhead getFixityDecls
  (botok,_) <- getState
  putState (botok,fixitydecls)
  p <- program'
  eof
  return p

getFixityDecls = do
  optionMaybe (reserved "module" >> ctor >> reserved "where")
  liftM rights (braces (topdecl `sepBy` semi))

program' = do
  reserved "module"
  modname  <- ctor
  reserved "where"
  topdecls <- liftM lefts (braces (topdecl `sepBy` semi))
  return $ Program modname (trec topdecls)
 <|> do
  topdecls <- liftM lefts (braces (topdecl `sepBy` semi))
  return $ Program "Main" (trec topdecls)
 <?> "program"

topdecl = liftM Left typedecl
      <|> liftM Left datadecl
      <|> liftM Left (try defn)
      <|> liftM Left proofdecl
      <|> liftM Right fixitydecl
      <?> "top-level declaration"

typedecl = do
  reserved "type"
  typename <- ctor
  vars     <- many identifier
  reservedOp "="
  t        <- ty
  return $ TypeDecl typename (embed $ bind (map string2Name vars) t)
 <?> "type declaration"

datadecl = do
  reserved "data"
  typename <- ctor
  tyvars   <- many identifier
  reservedOp "="
  constrs  <- constrdecl `sepBy1` reservedOp "|"
  return $ DataDecl typename (embed $ bind (map string2Name tyvars) constrs)
 <?> "datatype declaration"

constrdecl = do
  constrname <- ctor
  tys        <- many aty
  return $ ConstrDecl constrname tys
 <?> "constructor declaration"

defn = do
  (x,t) <- try tysig
  semi
  e <- eqnnamed x
  return (ExprDecl (Just (embed t)) (string2Name x) (embed e))
 <|> do
  (x,e) <- eqn
  return $ (ExprDecl Nothing (string2Name x) (embed e))
 <?> "expression definition"

tysig = do
  x <- identifier
  reservedOp "::"
  t <- ty
  return $ (x,t)
 <|> do
  op <- parens operator
  reservedOp "::"
  t <- ty
  return $ (op,t)
 <?> "type signature declaration"

eqn = do
  x  <- identifier
  reservedOp "="
  e  <- expr
  return (x,e)
 <?> "function binding pattern"

eqnnamed n = do
  symbol n
  reservedOp "="
  e  <- expr
  return e
 <?> "function binding pattern for a function named " ++ n

data Fixity = InfixL | InfixR | InfixNone deriving (Eq,Show)
data FixityDecl = FixityDecl Fixity (Maybe Integer) [Identifier]

fixity = (reserved "infixl" >> return InfixL)
     <|> (reserved "infixr" >> return InfixR)
     <|> (reserved "infix" >> return InfixNone)
     <?> "fixity"

fixitydecl = do
  f <- fixity
  p <- optionMaybe natural
  ops <- many operator
  return $ FixityDecl f p ops
 <?> "fixity declaration"

ty = do
  btys <- bty `sepBy1` (reservedOp "->")
  return (foldr1 TyArrow btys)
 <?> "type"

bty = do
  atys <- many1 aty
  return (foldl1 TyApp atys)
 <?> "type application"

aty = liftM (TyVar . string2Name) identifier
 <|>  liftM TyCon ctor
 <|> do
  parens (do
           tys <- ty `sepBy` comma
           case tys of
             []  -> return (TyCon "()")
             [t] -> return t
             _   -> return $ foldl TyApp (TyCon $ "(" ++ replicate (length tys - 1) ',' ++ ")") tys
         )
 <?> "atomic type"

pat = do
  (botok,_) <- getState
  (if botok
     then (reservedOp "_|_" >> return PatBottom) <|> pat'
     else pat') <?> "pattern"
    
pat' = do
  v <- identifier
  return (PatVar (string2Name v))
 <|> do
  c <- ctor
  return (PatCtor c)
 <|> do
  l <- literal
  return (PatLiteral l)
 <|> do
  reserved "_"
  return PatWildcard
 <|> try (do
    parens (do
              c  <- ctor
              ps <- many1 pat
              return (PatApp c ps)
           ))
 <|> do
  parens (do
            pats <- pat `sepBy` comma
            case pats of
              []  -> return (PatCtor "()")
              [p] -> return p
              _   -> return (PatApp ("(" ++ replicate (length pats - 1) ',' ++ ")") pats)
         )

lamexpr = do
  reservedOp "\\"
  boundvar <- identifier
  reservedOp "->"
  body     <- expr
  return $ Lambda (bind (string2Name boundvar) body)
 <?> "lambda expression"

letexpr = do
  reserved "let"
  lbs <- braces (lb `sepBy1` semi)
  reserved "in"
  body  <- expr
  return $ Let (bind (rec lbs) body)
 <?> "let-expression"

lb = do
  x <- identifier
  reservedOp "="
  e <- expr
  return (string2Name x,embed e)

nontopdecl = defn

caseexpr = do
  reserved "case"
  scrut <- expr
  reserved "of"
  alts <- braces (alt `sepBy` semi)
  return (Case scrut alts)
 <?> "case expression"

literal = liftM LitInteger natural
      <|> liftM LitFrac float 
      <|> liftM LitChar charLiteral
      <?> "literal"

alt = do
  p <- pat
  reservedOp "->"
  e <- expr
  return (bind p e)

term = do
  (botok,_) <- getState
  (if botok
     then (term' <|> (reservedOp "_|_" >> return Bottom))
     else term') <?> "expression term"

term' = lamexpr
   <|> try (liftM Literal literal)
   <|> letexpr
   <|> caseexpr
   <|> liftM (Var . string2Name) identifier
   <|> liftM Ctor ctor
   <|> (try $ parens (return ()) >> return (Ctor "()"))
   <|> (try $ parens expr)
   <|> (parens $ expr `sepBy1` comma >>= \ es -> return $ foldl App (Ctor $ "(" ++ replicate (length es - 1) ',' ++ ")") es)

mkOperatorTable = do
  (_,fixitydecls) <- getState
  let fromFixity InfixL    = AssocLeft
      fromFixity InfixR    = AssocRight
      fromFixity InfixNone = AssocNone

      fromPrec (Just x) = x
      fromPrec Nothing  = 9

      mkCols (FixityDecl fx prec ids) =
        map (\ i ->
              Infix (do { symbol i
                        ; return (\ e1 e2 -> (App (App (Var $ string2Name i) e1) e2)) })
                    (fromFixity fx))
            ids

      getPrec (FixityDecl _ prec _) = fromPrec prec

      getOpers (FixityDecl _ _ os) = os

      mkRow n = concat [mkCols fd | fd <- fixitydecls, getPrec fd == n]

      allBound = concatMap getOpers fixitydecls

      approw = [Infix (do { return (\ e1 e2 -> (App e1 e2)) }) AssocLeft]
      row9 = [Infix (try $ do { oper <- operator ; if oper `elem` allBound then fail "<<shouldn't see this>>" else return (\ e1 e2 -> (App (App (Var $ string2Name oper) e1) e2)) }) AssocLeft] ++ mkRow 9
      otherRows = map mkRow [8,7,6,5,4,3,2,1,0]

  return (approw:row9:otherRows)

expr = do
         table <- mkOperatorTable
         buildExpressionParser table term <?> "expression"

proofdecl = do
  (x,fo) <- try proofsig
  semi
  pr     <- proofnamed x
  return (ProofDecl (embed fo) (string2Name x) (embed pr))
  
proofsig = do
  x <- identifier
  reservedOp ":::"
  fo <- formula
  return $ (x,fo)

proofnamed x = do
  symbol x
  reservedOp "="
  p  <- proof
  return p
 <?> "proof named " ++ x

data ProofTerm = PTProof Proof | PTExpr Expr | PTCtor String

proofterm = do
  reserved "Foralli"
  lparen
  x <- identifier
  reservedOp "::"
  t <- ty
  rparen
  comma
  pr <- proof
  return (PTProof $ ForallIExpr t (bind (string2Name x) pr))
 <|> do
  reserved "Assume"
  lparen
  x  <- identifier
  reservedOp ":::"
  fo <- formula
  rparen
  comma
  pr <- proof
  return (PTProof $ ProofImpl fo (bind (string2Name x) pr))
 <|> (liftM PTProof $ caseproof)
 <|> (reserved "eval" >> return (PTProof Eval))
 <|> do
  reserved "trans"
  (PTProof p1) <- proofterm
  (PTProof p2) <- proofterm
  return (PTProof $ Trans p1 p2)
 <|> do
  reserved "subst"
  x <- identifier
  reserved "by"
  p <- proof
  reserved "in"
  e <- withBotok True expr
  return (PTProof $ Subst (bind (string2Name x) (p,e)))
 <|> do
  reserved "symm"
  p <- proof
  return (PTProof $ Symm p)
 <|> (liftM (PTProof . ProofVar . string2Name) $ identifier)
 <|> (liftM PTCtor ctor)
 <|> (liftM PTProof $ parens proof)
 <|> (liftM PTExpr $ brackets (withBotok True expr))
 <?> "proof term"

caseproof = do
  reserved "case"
  scrut <- withBotok True expr
  reserved "proving"
  fo <- formula
  reserved "by"
  x <- identifier
  reserved "of"
  proofalts <- braces (proofalt `sepBy` semi)
  return (ProofCase scrut fo (bind (string2Name x) proofalts))
 <?> "case expression"

proofalt = do
  p <- withBotok True pat
  reservedOp "->"
  e <- proof
  return (bind p e)

formula = do
  reserved "Forall"
  lparen
  x <- identifier
  reservedOp "::"
  t <- ty
  rparen
  comma
  fo <- formula
  return (ForallExpr t (bind (string2Name x) fo))
 <|> do
  reserved "Assuming"
  fo  <- formula
  comma
  fo' <- formula
  return (ForallProof fo fo')
 <|> do
  lcurly
  e  <- withBotok True expr
  reservedOp "="
  e' <- withBotok True expr
  rcurly
  return (FormulaEq e e')
 <?> "formula"

proof = do
  pts <- many1 proofterm
  mr  <- optionMaybe ( do
           reservedOp ":::"
           formula
         )
  let mkProofApp (PTProof p) (PTExpr e)   = PTProof (ProofAppExpr p e)
      mkProofApp (PTProof p) (PTProof p') = PTProof (ProofAppProof p p')
      mkProofApp _ _                      = error "Zabonga FIXME"
      (PTProof pr) = case pts of
        (PTCtor ctor:pts') -> PTProof $ ProofBisim ctor (map (\(PTProof p) -> p) pts')
        _                  -> foldl1 mkProofApp pts
  case mr of
    (Just fo) -> return (ProofAnno pr fo)
    Nothing   -> return pr

parse :: FilePath -> String -> Either ParseError Program
parse filename stream = runParser program (False,[]) filename stream

parseExpr :: FilePath -> String -> Either ParseError Expr
parseExpr filename stream = runParser (expr >>= \ e -> eof >> return e) (False,[]) filename stream