{-# LANGUAGE CPP #-}
module PmExpr (
        PmExpr(..), PmLit(..), SimpleEq, ComplexEq, toComplex, eqPmLit,
        truePmExpr, falsePmExpr, isTruePmExpr, isFalsePmExpr, isNotPmExprOther,
        lhsExprToPmExpr, hsExprToPmExpr, substComplexEq, filterComplex,
        pprPmExprWithParens, runPmPprM
    ) where
#include "HsVersions.h"
import GhcPrelude
import HsSyn
import Id
import Name
import NameSet
import DataCon
import ConLike
import TysWiredIn
import Outputable
import Util
import SrcLoc
import Data.Maybe (mapMaybe)
import Data.List (groupBy, sortBy, nubBy)
import Control.Monad.Trans.State.Lazy
data PmExpr = PmExprVar   Name
            | PmExprCon   ConLike [PmExpr]
            | PmExprLit   PmLit
            | PmExprEq    PmExpr PmExpr  
            | PmExprOther (HsExpr GhcTc)  
mkPmExprData :: DataCon -> [PmExpr] -> PmExpr
mkPmExprData dc args = PmExprCon (RealDataCon dc) args
data PmLit = PmSLit (HsLit GhcTc)                               
           | PmOLit Bool  (HsOverLit GhcTc) 
eqPmLit :: PmLit -> PmLit -> Bool
eqPmLit (PmSLit    l1) (PmSLit    l2) = l1 == l2
eqPmLit (PmOLit b1 l1) (PmOLit b2 l2) = b1 == b2 && l1 == l2
  
eqPmLit _              _              = False
nubPmLit :: [PmLit] -> [PmLit]
nubPmLit = nubBy eqPmLit
type SimpleEq  = (Id, PmExpr) 
type ComplexEq = (PmExpr, PmExpr)
toComplex :: SimpleEq -> ComplexEq
toComplex (x,e) = (PmExprVar (idName x), e)
truePmExpr :: PmExpr
truePmExpr = mkPmExprData trueDataCon []
falsePmExpr :: PmExpr
falsePmExpr = mkPmExprData falseDataCon []
isNotPmExprOther :: PmExpr -> Bool
isNotPmExprOther (PmExprOther _) = False
isNotPmExprOther _expr           = True
isNegatedPmLit :: PmLit -> Bool
isNegatedPmLit (PmOLit b _) = b
isNegatedPmLit _other_lit   = False
isTruePmExpr :: PmExpr -> Bool
isTruePmExpr (PmExprCon c []) = c == RealDataCon trueDataCon
isTruePmExpr _other_expr      = False
isFalsePmExpr :: PmExpr -> Bool
isFalsePmExpr (PmExprCon c []) = c == RealDataCon falseDataCon
isFalsePmExpr _other_expr      = False
isNilPmExpr :: PmExpr -> Bool
isNilPmExpr (PmExprCon c _) = c == RealDataCon nilDataCon
isNilPmExpr _other_expr     = False
isPmExprEq :: PmExpr -> Maybe (PmExpr, PmExpr)
isPmExprEq (PmExprEq e1 e2) = Just (e1,e2)
isPmExprEq _other_expr      = Nothing
isConsDataCon :: DataCon -> Bool
isConsDataCon con = consDataCon == con
substPmExpr :: Name -> PmExpr -> PmExpr -> (PmExpr, Bool)
substPmExpr x e1 e =
  case e of
    PmExprVar z | x == z    -> (e1, True)
                | otherwise -> (e, False)
    PmExprCon c ps -> let (ps', bs) = mapAndUnzip (substPmExpr x e1) ps
                      in  (PmExprCon c ps', or bs)
    PmExprEq ex ey -> let (ex', bx) = substPmExpr x e1 ex
                          (ey', by) = substPmExpr x e1 ey
                      in  (PmExprEq ex' ey', bx || by)
    _other_expr    -> (e, False) 
                                 
substComplexEq :: Name -> PmExpr -> ComplexEq -> Either ComplexEq ComplexEq
substComplexEq x e (ex, ey)
  | bx || by  = Left  (ex', ey')
  | otherwise = Right (ex', ey')
  where
    (ex', bx) = substPmExpr x e ex
    (ey', by) = substPmExpr x e ey
lhsExprToPmExpr :: LHsExpr GhcTc -> PmExpr
lhsExprToPmExpr (L _ e) = hsExprToPmExpr e
hsExprToPmExpr :: HsExpr GhcTc -> PmExpr
hsExprToPmExpr (HsVar         x) = PmExprVar (idName (unLoc x))
hsExprToPmExpr (HsConLikeOut  c) = PmExprVar (conLikeName c)
hsExprToPmExpr (HsOverLit  olit) = PmExprLit (PmOLit False olit)
hsExprToPmExpr (HsLit       lit) = PmExprLit (PmSLit lit)
hsExprToPmExpr e@(NegApp _ neg_e)
  | PmExprLit (PmOLit False ol) <- synExprToPmExpr neg_e
  = PmExprLit (PmOLit True ol)
  | otherwise = PmExprOther e
hsExprToPmExpr (HsPar (L _ e)) = hsExprToPmExpr e
hsExprToPmExpr e@(ExplicitTuple ps boxity)
  | all tupArgPresent ps = mkPmExprData tuple_con tuple_args
  | otherwise            = PmExprOther e
  where
    tuple_con  = tupleDataCon boxity (length ps)
    tuple_args = [ lhsExprToPmExpr e | L _ (Present e) <- ps ]
hsExprToPmExpr e@(ExplicitList _elem_ty mb_ol elems)
  | Nothing <- mb_ol = foldr cons nil (map lhsExprToPmExpr elems)
  | otherwise        = PmExprOther e 
  where
    cons x xs = mkPmExprData consDataCon [x,xs]
    nil       = mkPmExprData nilDataCon  []
hsExprToPmExpr (ExplicitPArr _elem_ty elems)
  = mkPmExprData (parrFakeCon (length elems)) (map lhsExprToPmExpr elems)
hsExprToPmExpr e@(RecordCon   _ _ _ _) = PmExprOther e
hsExprToPmExpr (HsTick            _ e) = lhsExprToPmExpr e
hsExprToPmExpr (HsBinTick       _ _ e) = lhsExprToPmExpr e
hsExprToPmExpr (HsTickPragma  _ _ _ e) = lhsExprToPmExpr e
hsExprToPmExpr (HsSCC           _ _ e) = lhsExprToPmExpr e
hsExprToPmExpr (HsCoreAnn       _ _ e) = lhsExprToPmExpr e
hsExprToPmExpr (ExprWithTySig     e _) = lhsExprToPmExpr e
hsExprToPmExpr (ExprWithTySigOut  e _) = lhsExprToPmExpr e
hsExprToPmExpr (HsWrap            _ e) =  hsExprToPmExpr e
hsExprToPmExpr e = PmExprOther e 
synExprToPmExpr :: SyntaxExpr GhcTc -> PmExpr
synExprToPmExpr = hsExprToPmExpr . syn_expr  
type PmNegLitCt = (Name, (SDoc, [PmLit]))
filterComplex :: [ComplexEq] -> [PmNegLitCt]
filterComplex = zipWith rename nameList . map mkGroup
              . groupBy name . sortBy order . mapMaybe isNegLitCs
  where
    order x y = compare (fst x) (fst y)
    name  x y = fst x == fst y
    mkGroup l = (fst (head l), nubPmLit $ map snd l)
    rename new (old, lits) = (old, (new, lits))
    isNegLitCs (e1,e2)
      | isFalsePmExpr e1, Just (x,y) <- isPmExprEq e2 = isNegLitCs' x y
      | isFalsePmExpr e2, Just (x,y) <- isPmExprEq e1 = isNegLitCs' x y
      | otherwise = Nothing
    isNegLitCs' (PmExprVar x) (PmExprLit l) = Just (x, l)
    isNegLitCs' (PmExprLit l) (PmExprVar x) = Just (x, l)
    isNegLitCs' _ _             = Nothing
    
    nameList :: [SDoc]
    nameList = map text ["p","q","r","s","t"] ++
                 [ text ('t':show u) | u <- [(0 :: Int)..] ]
runPmPprM :: PmPprM a -> [PmNegLitCt] -> (a, [(SDoc,[PmLit])])
runPmPprM m lit_env = (result, mapMaybe is_used lit_env)
  where
    (result, (_lit_env, used)) = runState m (lit_env, emptyNameSet)
    is_used (x,(name, lits))
      | elemNameSet x used = Just (name, lits)
      | otherwise         = Nothing
type PmPprM a = State ([PmNegLitCt], NameSet) a
addUsed :: Name -> PmPprM ()
addUsed x = modify (\(negated, used) -> (negated, extendNameSet used x))
checkNegation :: Name -> PmPprM (Maybe SDoc) 
checkNegation x = do
  negated <- gets fst
  return $ case lookup x negated of
    Just (new, _) -> Just new
    Nothing       -> Nothing
pprPmExpr :: PmExpr -> PmPprM SDoc
pprPmExpr (PmExprVar x) = do
  mb_name <- checkNegation x
  case mb_name of
    Just name -> addUsed x >> return name
    Nothing   -> return underscore
pprPmExpr (PmExprCon con args) = pprPmExprCon con args
pprPmExpr (PmExprLit l)        = return (ppr l)
pprPmExpr (PmExprEq _ _)       = return underscore 
pprPmExpr (PmExprOther _)      = return underscore 
needsParens :: PmExpr -> Bool
needsParens (PmExprVar   {}) = False
needsParens (PmExprLit    l) = isNegatedPmLit l
needsParens (PmExprEq    {}) = False 
needsParens (PmExprOther {}) = False 
needsParens (PmExprCon (RealDataCon c) es)
  | isTupleDataCon c || isPArrFakeCon c
  || isConsDataCon c || null es = False
  | otherwise                   = True
needsParens (PmExprCon (PatSynCon _) es) = not (null es)
pprPmExprWithParens :: PmExpr -> PmPprM SDoc
pprPmExprWithParens expr
  | needsParens expr = parens <$> pprPmExpr expr
  | otherwise        =            pprPmExpr expr
pprPmExprCon :: ConLike -> [PmExpr] -> PmPprM SDoc
pprPmExprCon (RealDataCon con) args
  | isTupleDataCon con = mkTuple <$> mapM pprPmExpr args
  |  isPArrFakeCon con = mkPArr  <$> mapM pprPmExpr args
  |  isConsDataCon con = pretty_list
  where
    mkTuple, mkPArr :: [SDoc] -> SDoc
    mkTuple = parens     . fsep . punctuate comma
    mkPArr  = paBrackets . fsep . punctuate comma
    
    pretty_list :: PmPprM SDoc
    pretty_list = case isNilPmExpr (last list) of
      True  -> brackets . fsep . punctuate comma <$> mapM pprPmExpr (init list)
      False -> parens   . hcat . punctuate colon <$> mapM pprPmExpr list
    list = list_elements args
    list_elements [x,y]
      | PmExprCon c es <- y,  RealDataCon nilDataCon == c
          = ASSERT(null es) [x,y]
      | PmExprCon c es <- y, RealDataCon consDataCon == c
          = x : list_elements es
      | otherwise = [x,y]
    list_elements list  = pprPanic "list_elements:" (ppr list)
pprPmExprCon cl args
  | conLikeIsInfix cl = case args of
      [x, y] -> do x' <- pprPmExprWithParens x
                   y' <- pprPmExprWithParens y
                   return (x' <+> ppr cl <+> y')
      
      list   -> pprPanic "pprPmExprCon:" (ppr list)
  | null args = return (ppr cl)
  | otherwise = do args' <- mapM pprPmExprWithParens args
                   return (fsep (ppr cl : args'))
instance Outputable PmLit where
  ppr (PmSLit     l) = pmPprHsLit l
  ppr (PmOLit neg l) = (if neg then char '-' else empty) <> ppr l
instance Outputable PmExpr where
  ppr e = fst $ runPmPprM (pprPmExpr e) []