-- Use like so:
--
-- runghc Desugar.hs < Examples.hs > ExamplesDesugared.hs
--
-- and watch out, the pretty printer sometimes misses (..)'s (wtf!!!)

module Main (main) where

import Data.List

import Language.Haskell.Parser
import Language.Haskell.Syntax
import Language.Haskell.Pretty

main = interact $ \input ->
  let ParseOk (HsModule sl m es is decls) = parseModule input
      decls' = [ decl
               | (exp, stmts, context) <- extractListComps decls
               , let exp' = desugar [0..] exp stmts
               , decl <- [context (HsListComp exp stmts), context exp'] ]
   in prettyPrint (HsModule sl m es is decls') ++ "\n"

extractListComps :: [HsDecl] -> [(HsExp, [HsStmt], HsExp -> HsDecl)]
extractListComps decls =
    [ (exp, stmts,
       \rhs -> HsPatBind sl (HsPVar (HsIdent name))
                            (HsUnGuardedRhs rhs) [])
    | HsPatBind sl (HsPVar (HsIdent name))
                   (HsUnGuardedRhs (HsListComp exp stmts)) [] <- decls ]
 ++ [ (exp, stmts,
       \rhs -> HsFunBind [HsMatch sl name ps (HsUnGuardedRhs rhs) ds])
    | HsFunBind [HsMatch sl name ps
                         (HsUnGuardedRhs (HsListComp exp stmts)) ds] <- decls ]

{-

Streams desugaring for list comprehensions:

<< [ e | Q ] >>


<< [ e | ] >> = return e

<< [ e | b, Q ] >> = gaurd b (<< [ e |  Q ] >>)

<< [ e | p <- l, Q ] >> = bind f g l
  where
    f p = Just (x1, ..., xn)
    f _ = Nothing

    g (x1, ..., xn) = << [ e | Q ] >>

<< [ e | let decls, Q ] >> = declare g (let decls in (x1, .., xn))
  where
    g (x1, .., xn) = << [ e | Q ] >>

We can optimise the common cases:


<< [ e | p <- l, b ] >> = mapFilter f l
  where f p | b = Just e
        f _     = Nothing

<< [ e | p <- l, b, Q ] >> = bind f g l
  where f p | b = Just (x1, ..., xn)
        f _     = Nothing
        
        g (x1, ..., xn) = << [ e | Q ] >>


using the auxillary definitions:

List:
-----

return e = [e]

guard True  xs = xs
guard false xs = []

concatMap' f g xs = foldr h [] xs
  where h x ys = case f x of
                  Just bs -> g bs ++ ys
                  Nothing ->         ys

declare g bs = g bs

mapFilter f xs = foldr h [] xs
  where h x ys = case f x of
                  Just y  -> y : ys
                  Nothing ->     ys


Stream:
-------

return e = Stream next True
  where
    next True  = Yield e False
    next False = Done

guard b (Stream next0 s0) = Stream next (b, s0)
  where
    next (False, s) = Done
    next (True, s)  = case next0 s of
      Done       -> Done
      Skip    s' -> Skip    (True, s')
      Yield x s' -> Yield x (True, s')

bind f g (Stream next0 s0) = Stream next (s, Nothing)
  where
    next (s, Nothing) = case next0 s of
      Done       -> Done
      Skip    s' -> Skip    (s', Nothing)
      Yield x s' -> case f x of
         Just bs -> Skip (s', Just (g bs))
         Nothing -> Skip (s', Nothing)

    next (s, Just (Stream next1 s1)) = case next1 s1 of
      Done        -> Done
      Skip    s1' -> Skip    (s, Just (Stream next1 s1'))
      Yield x s1' -> Yield x (s, Just (Stream next1 s1'))

declare g bs = Stream next (g bs)
  where next (Stream next0 s) = case next0 s of
      Done       -> Done
      Skip    s' -> Skip    (Stream next0 s')
      Yield x s' -> Yield x (Stream next0 s')

mapFilter f (Stream next0 s0) = Stream next s0
  where
    next s = case next0 s of
      Done       -> Done
      Skip    s' -> Skip    s'
      Yield x s' -> case f x of
                  Just y  -> Yield y s'
                  Nothing -> Skip    s'

-}

hsJust       = HsVar (UnQual (HsIdent "Just"))
hsNothing    = HsVar (UnQual (HsIdent "Nothing"))

hsReturn     = HsVar (Qual (Module "L") (HsIdent "return"))
hsGuard      = HsVar (Qual (Module "L") (HsIdent "guard"))
hsConcatMap  = HsVar (Qual (Module "L") (HsIdent "concatMap"))
hsConcatMap' = HsVar (Qual (Module "L") (HsIdent "concatMap'"))
hsDeclare    = HsVar (Qual (Module "L") (HsIdent "declare"))
hsMapFilter  = HsVar (Qual (Module "L") (HsIdent "mapFilter"))
hsMap        = HsVar (Qual (Module "L") (HsIdent "map"))

desugar :: [Int] -> HsExp -> [HsStmt] -> HsExp

-- [ e | ] = return e
desugar ns e [] = HsApp hsReturn e

-- [ e | b, Q ] = gaurd b [ e | Q ]
desugar ns e (HsQualifier b : qs) = HsApp (HsApp hsGuard b)
                                           (HsParen (desugar ns e qs))

-- special case: final generator for simple irrefutable pattern p
-- [ e | p <- l ] = let f p = e
--                   in map f l
desugar (n:ns) e (HsGenerator sloc p@(HsPVar _) l : []) =
  HsLet [HsFunBind [HsMatch sloc (HsIdent (f n)) [p]
                   {- = -} (HsUnGuardedRhs e) []
                   ]
        ]
        (HsApp (HsApp hsMap (HsVar (UnQual (HsIdent (f n)))))
               (HsParen l))

-- special case: final generator
-- [ e | p <- l ] = let f p = Just e
--                      f _ = Nothing
--                   in mapFilter f l
desugar (n:ns) e (HsGenerator sloc p l : []) =
  HsLet [HsFunBind [HsMatch sloc (HsIdent (f n)) [p]
                   {- = -} (HsUnGuardedRhs (HsApp hsJust e)) []
                   ,HsMatch sloc (HsIdent (f n)) [HsPWildCard]
                   {- = -} (HsUnGuardedRhs hsNothing)       []
                   ]
        ]
        (HsApp (HsApp hsMapFilter (HsVar (UnQual (HsIdent (f n)))))
               (HsParen l))

-- special case: final generator with trailing condition
-- [ e | p <- l, b ] = let f p | b = Just e
--                         f _     = Nothing
--                      in mapFilter f l
desugar (n:ns) e (HsGenerator sloc p l : HsQualifier b : []) =
  HsLet [HsFunBind [HsMatch sloc (HsIdent (f n)) [p]
                   {- = -} (HsGuardedRhss [HsGuardedRhs sloc b
                                           (HsApp hsJust e)]) []
                   ,HsMatch sloc (HsIdent (f n)) [HsPWildCard]
                   {- = -} (HsUnGuardedRhs hsNothing) []
                   ]
        ]
        (HsApp (HsApp hsMapFilter (HsVar (UnQual (HsIdent (f n)))))
               (HsParen l))

-- special case: generator with trailing condition
-- [ e | p <- l, b, Q ] = let f p | b = Just (x1, ..., xn)
--                            f _ = Nothing
--                            g (x1, ..., xn) = [ e | Q ]
--                         in concatMap' f g l
desugar (n:ns) e (HsGenerator sloc p l : HsQualifier b : qs) =
  let xs = patBinders p
           `intersect` stmtsFV e qs -- optimisation: only bind used vars
   in HsLet [HsFunBind [HsMatch sloc (HsIdent (f n)) [p]
                       {- = -} (HsGuardedRhss
                                 [HsGuardedRhs sloc b
                                 (HsApp hsJust (mkTupleExpr xs))]) []
                       ,HsMatch sloc (HsIdent (f n)) [HsPWildCard]
                       {- = -} (HsUnGuardedRhs hsNothing) []
                       ]
            ,HsFunBind [HsMatch sloc (HsIdent (g n)) [mkTuplePat xs]
                       {- = -} (HsUnGuardedRhs (desugar ns e qs)) []
                       ]
            ]
            (HsApp (HsApp (HsApp hsConcatMap'
                                 (HsVar (UnQual (HsIdent (f n)))))
                          (HsVar (UnQual (HsIdent (g n)))))
                   (HsParen l))


-- special case: for simple irrefutable pattern p
-- [ e | p <- l, Q ] = let g p = [ e | Q ]
--                      in concatMap g l
desugar (n:ns) e (HsGenerator sloc p@(HsPVar _) l : qs) =
  HsLet [HsFunBind [HsMatch sloc (HsIdent (g n)) [p]
                   {- = -} (HsUnGuardedRhs (desugar ns e qs)) []
                   ]
        ]
        (HsApp (HsApp hsConcatMap (HsVar (UnQual (HsIdent (g n)))))
               (HsParen l))

-- general case of a generator
-- [ e | p <- l, Q ] = let f p = Just (x1, ..., xn)
--                         f _ = Nothing
--                         g (x1, ..., xn) = [ e | Q ]
--                       in concatMap' f g l
desugar (n:ns) e (HsGenerator sloc p l : qs) =
  let xs = patBinders p 
           `intersect` stmtsFV e qs -- optimisation: only bind used vars
   in HsLet [HsFunBind [HsMatch sloc (HsIdent (f n)) [p]
                       {- = -} (HsUnGuardedRhs
                                 (HsApp hsJust (mkTupleExpr xs))) []
                       ,HsMatch sloc (HsIdent (f n)) [HsPWildCard]
                       {- = -} (HsUnGuardedRhs hsNothing) []
                       ]
            ,HsFunBind [HsMatch sloc (HsIdent (g n)) [mkTuplePat xs]
                       {- = -} (HsUnGuardedRhs (desugar ns e qs)) []
                       ]
            ]
            (HsApp (HsApp hsConcatMap (HsVar (UnQual (HsIdent (g n)))))
                   (HsParen l))

{-
-- TODO: need to get binders in decls for this one:

-- [ e | let decls, Q ] = let g (x1, .., xn) = [ e | Q ]
--                         in declare g (let decls in (x1, .., xn))
desugar ns e (HsLetStmt decls : qs) =
  HsLet decls (desugar ns e qs)
-}

patBinders :: HsPat -> [HsName]
patBinders (HsPVar name)         = [name]
patBinders (HsPLit _)            = []
patBinders (HsPNeg p)            = patBinders p
patBinders (HsPInfixApp p1 _ p2) = patBinders p1 ++ patBinders p2
patBinders (HsPApp _ ps)         = concatMap patBinders ps
patBinders (HsPTuple ps)         = concatMap patBinders ps
patBinders (HsPList ps)          = concatMap patBinders ps
patBinders (HsPParen p)          = patBinders p
patBinders (HsPRec _ fields)     = concat [ patBinders p
                                          | HsPFieldPat _ p <- fields ]
patBinders (HsPAsPat _ p)        = patBinders p
patBinders (HsPWildCard)         = []
patBinders (HsPIrrPat p)         = patBinders p


-- The following is for an optimisation, perhaps it's unnecessary.

nameFV (Qual _ _)  = []
nameFV (UnQual n)  = [n]
nameFV (Special _) = []

qopFV (HsQVarOp n) = nameFV n
qopFV _            = []

expFV :: HsExp -> [HsName]
expFV (HsVar n) = nameFV n
expFV (HsCon _) = []
expFV (HsLit _) = []
expFV (HsInfixApp e1 op e2) = qopFV op ++ expFV e1 ++ expFV e2
expFV (HsApp e1 e2) = expFV e1 ++ expFV e2
expFV (HsNegApp e) = expFV e
expFV (HsLambda _ ps e) = nub (expFV e) \\ concatMap patBinders ps
--expFV (HsLet [HsDecl] HsExp)
expFV (HsIf e1 e2 e3) = expFV e1 ++ expFV e2 ++ expFV e3
--expFV (HsCase HsExp [HsAlt])
expFV (HsDo stmts) = stmtsFV unit_con stmts
expFV (HsTuple es) = concatMap expFV es
expFV (HsList es) = concatMap expFV es
expFV (HsParen e) = expFV e
expFV (HsLeftSection e op)  = qopFV op ++ expFV e
expFV (HsRightSection op e) = qopFV op ++ expFV e
--expFV (HsRecConstr HsQName [HsFieldUpdate])	record construction expression
--expFV (HsRecUpdate HsExp [HsFieldUpdate])	record update expression
expFV (HsEnumFrom e) = expFV e
expFV (HsEnumFromTo e1 e2) = expFV e1 ++ expFV e2
expFV (HsEnumFromThen e1 e2) = expFV e1 ++ expFV e2
expFV (HsEnumFromThenTo e1 e2 e3) = expFV e1 ++ expFV e2 ++ expFV e3
expFV (HsListComp e stmts) = stmtsFV e stmts
expFV (HsExpTypeSig _ e _) = expFV e

stmtsFV :: HsExp -> [HsStmt] -> [HsName]
stmtsFV e [] = expFV e
stmtsFV e (HsGenerator _ p e' : stmts) = nub (expFV e' ++ stmtsFV e stmts) \\ patBinders p
stmtsFV e (HsQualifier     e' : stmts) = expFV e' ++ stmtsFV e stmts
--stmtsFV e (HsLetStmt decls) = 

mkTupleExpr []  = HsCon (Special HsUnitCon)
mkTupleExpr [n] = HsVar (UnQual n)
mkTupleExpr ns  = HsTuple [ HsVar (UnQual n) | n <- ns ]

mkTuplePat []  = HsPWildCard
mkTuplePat [n] = HsPVar n
mkTuplePat ns  = HsPTuple [ HsPVar n | n <- ns ]

f n = "f_" ++ show n
g n = "g_" ++ show n












{-

Alternative streams desugaring for list comprehensions:

<< [ e | Q ] P D >> :: Exp -> [Qual] -> (Pat -> Pat) -> Exp -> ([Clause], Exp)

Translation scheme takes the list comp's expression and list of qualifiers
and as extra args, a pattern (with a hole) that gives the state pattern and an
expression giving the state to return to when done. It returns a list of
clauses for a 'next' stepper function and an initialiser expression for the
stream's initial state.

To construct the top level we use:

[ e | Q ] = Stream (FunDef "next" clauses) init
  where (clauses, init) = << [ e | Q ] [.] Done >>


<< [ e | ] P D >> = (next:[], True)
  where
     next P[True]  = Yield e P[False]
     next P[False] = D


<< [ e | b, Q ] P D >> = (next:nexts, Nothing)
  where
     (nexts, init) = << [ e | Q ] P[Just[.]] D >>

     next P[Nothing] | b         = Skip (P[Just init])
                     | otherwise = D


<< [ e | let decls, Q ] P D >> = (next:nexts, Nothing)
  where
     (nexts, init) = << [ e | Q ] P[Just((x_n_0, , x_n_m) :!: [.])] D >>

     next P[Nothing] = Skip P[Just(let decls in (x_n_0, , x_n_m) :!: init)]
     
     {x_n_0, , x_n_m} = DV (decls) -- Defined variables


<< [ e | p <- l, Q ] P >> = (next:nexts, stream l :!: Nothing)
  where
     (nexts, init) = << [ e | Q ] P[(Stream next_n s_n :!: Just (p :!: [.]))]
                                  P[(Stream next_n s_n :!: Nothing)] >>

     next P[(Stream next_n s_n :!: Nothing)] = case next_n s_n of
       Done         -> D
       Skip    s_n' -> Skip P[(Stream next_n s_n' :!: Nothing)]
       Yield x s_n' -> case x of
                  p -> Skip P[(Stream next_n s_n' :!: Just (x :!: init))]
                  _ -> Skip P[(Stream next_n s_n' :!: Nothing)]

-}