--     Package:  freesect-0.8
-- Description:  Extend Haskell to support free sections
--     Example:  zipWith (f __ b __ d) as cs
--      Author:  Andrew Seniuk <rasfar@gmail.com>
--        Date:  March 11, 2012
--     License:  BSD3 (./LICENSE)
--  Executable:  freesect
--       Usage:  See accompanying files 000-readme and z

{-# LANGUAGE CPP #-}  -- just a couple #if 0/1's
{- # LANGUAGE DeriveDataTypeable #-}  -- not needed!
{-# LANGUAGE FlexibleContexts #-}  -- needed for one of the type sigs
{- # LANGUAGE MultiParamTypeClasses #-}
{- # LANGUAGE RankNTypes #-}  -- needed for the path accumulators
{- # LANGUAGE ExistentialQuantification #-}
{- # LANGUAGE GADTs #-}
{- # LANGUAGE ScopedTypeVariables #-}  -- needed for a pattern type sig

{- # NOINLINE showSSI #-}
{- # NOINLINE FreeSectAnnotated.showSSI #-}

  module FreeSectAnnotated where

  import Data.Data(Data,gmapQi)
  import Data.Generics.Aliases(mkQ,mkT,mkM)
  import Data.Generics.Schemes(everywhere,everywhereM,gcount)
--import Data.Generics  -- this suffices to import all the above

  import Control.Monad.State

  import System.IO.Unsafe(unsafePerformIO)  -- warning message to stderr
  import System.IO(hFlush,stderr,hPutStr)
  import System.IO(putStrLn)
--import Data.IORef(IORef,newIORef)

  import HSE.Annotated

  import Util

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

  -- Why does GHC complain of pattern matches overlapping in some of
  -- the generic transformers, but not all?  The code structure is
  -- completely analogous so far as I can see.

  -- :: Module -> Module always, at present
--fs_module :: Data a => a -> a  -- typesig not nec.
  fs_module m0 = m5
   where
    -- It's a bit annoying, but GuardedRhs and UnGuardedRhs are
    -- not constructors of the same data type, so we cannot use
    -- a single generic traversal to handle both.  Could the
    -- duplication be avoided?
    m1 = fs_warn_flaw           m0  -- check/warn re. <ski>'s flaw
    m2 = fs_unguarded_rhss      m1  -- translate UnGuardedRhs's
    m3 = fs_guarded_rhss        m2  -- translate GuardedRhs's
    m4 = fs_error_if_any_remain m3  -- error if any freesects remain
    m5 = fs_cleanup             m4  -- remove some redundant Paren's

  -- :: Module -> Module always, at present
--fs_warn_flaw :: Data a => a -> a  -- unnec.
  fs_warn_flaw m = m'
   where
    m' = everywhere (mkT step) m
--  step :: Exp SrcSpanInfo -> Exp SrcSpanInfo  -- unnec.
    step x@(App _ p@(Paren ssi (App _ (FSWildcard _) _)) _) = warning True ssi p x
    step x@(App _ p@(Paren ssi (App _ _ (FSWildcard _))) _) = warning False ssi p x
    step x = x
--  warning :: Data a => Bool -> SrcSpanInfo -> a -> a -> a  -- unnec.
    warning b ssi p x =   unsafePerformIO
                          $ do
                               hPutStr stderr $ warning_message b ssi p x
                               hFlush stderr
                               return x
    warning_message b ssi p x
     =    showSSI ssi ++ " Warning:\n"
--   =    (error $ showSSI ssi) ++ "\n"
       ++ "  Inferring free section context of loose wildcard(s) occurring\n"
       ++ "  in redundantly-parenthesised application\n"
       ++ "    " ++ prettyPrint p ++ "\n"
       ++ "  in the expression\n"
       ++ "    " ++ prettyPrint x ++ "\n"
       ++ ( if b then "  This means for e.g. that (f __) y is rewritten to (\\x->f x) y.\n" else "  This means for e.g. that (__ x) y is rewritten to (\\f->f x) y.\n" )  -- parentheses are really key here...
       ++ "  If this is not what you want, remove the redundant parentheses\n"
       ++ "  or use explicit _[...]_ free section context syntax.\n"

  -- :: Module -> Module always, at present
--fs_unguarded_rhss :: Data a => a -> a  -- typesig not nec.
  fs_unguarded_rhss m = m''
   where
    m' = everywhere (mkT step1) m  -- explicitly _[...]_ grouped freesects
    m'' = everywhere (mkT step2) m'  -- remaining __'s get inferred context
    step1 :: Rhs SrcSpanInfo -> Rhs SrcSpanInfo  -- seems nec.
    step1 (UnGuardedRhs srcSpanInfo e) = UnGuardedRhs srcSpanInfo e'
     where e' = fs_rhs_exp fresh e
    step1 x = x
    step2 :: Rhs SrcSpanInfo -> Rhs SrcSpanInfo  -- seems nec.
    step2 x@(UnGuardedRhs srcSpanInfo e)
     | still_fsss = UnGuardedRhs srcSpanInfo e''
     | otherwise = x
     where
       still_fsss = 0 < gcount (False `mkQ` p) x
       p :: Exp SrcSpanInfo -> Bool  -- nec.
       p (FSWildcard _) = True
       p _ = False
       e'' = fs_rhs_exp fresh e'
       e' = Paren srcSpanInfo e
    step2 x = x
    fresh = fs_fresh_name m

-- Unfortunate about the cloning here (see comment heading fs_module above).
  -- :: Module -> Module always, at present
--fs_guarded_rhss :: Data a => a -> a  -- typesig not nec.
  fs_guarded_rhss m = m''
   where
    m' = everywhere (mkT step1) m     -- explicitly _[...]_ grouped freesects
    m'' = everywhere (mkT step2) m'  -- remaining __'s get inferred context
    step1 :: GuardedRhs SrcSpanInfo -> GuardedRhs SrcSpanInfo  -- seems nec.
    step1 (GuardedRhs srcSpanInfo slst e) = GuardedRhs srcSpanInfo slst e'
     where e' = fs_rhs_exp fresh e
    step1 x = x
    step2 :: GuardedRhs SrcSpanInfo -> GuardedRhs SrcSpanInfo  -- seems nec.
    step2 x@(GuardedRhs srcSpanInfo slst e)
     | still_fsss = GuardedRhs srcSpanInfo slst e''
     | otherwise = x
     where
       still_fsss = 0 < gcount (False `mkQ` p) x
       p :: Exp SrcSpanInfo -> Bool  -- nec.
       p (FSWildcard _) = True
       p _ = False
       e'' = fs_rhs_exp fresh e'
       e' = Paren srcSpanInfo e
    step2 x = x
    fresh = fs_fresh_name m

  -- :: Module -> Module always, at present
--fs_error_if_any_remain :: Data a => a -> a  -- typesig not nec.
  fs_error_if_any_remain m = m'
   where
    m' | still_fsss = error "Free sections can only occur in RHS Exp contexts."
       | otherwise  = m
    still_fsss = 0 < gcount (False `mkQ` p) m
    p :: Exp SrcSpanInfo -> Bool  -- nec.
    p (FSWildcard _) = True
--  p (FSContext _ _) = True  -- dealt with subsequently in fs_cleanup
    p _ = False

  -- :: Module -> Module always, at present
--fs_cleanup :: Data a => a -> a  -- typesig not nec.
  fs_cleanup m0 = m3
   where
    m1 = everywhere (mkT step1) m0  -- for the Rhs's (un-guarded)
    m2 = everywhere (mkT step2) m1  -- for the GuardedRhs's
    m3 = everywhere (mkT step3) m2  -- for remaining FSContext -> Paren
    step1 :: Rhs SrcSpanInfo -> Rhs SrcSpanInfo  -- seems nec.
    step1 (UnGuardedRhs srcSpanInfo (FSContext _ e)) = UnGuardedRhs srcSpanInfo e
#if CLEAN_EXTRANEOUS_GROUPINGS
    step1 (UnGuardedRhs ssi
           (InfixApp _
             (FSContext _ e1)
             (QVarOp _ (UnQual _ (Symbol _ "$")))
             e2))
             = UnGuardedRhs ssi
                (App ssi
                  (Paren ssi e1)
                  e2)
#endif
    step1 x = x
    step2 :: GuardedRhs SrcSpanInfo -> GuardedRhs SrcSpanInfo  -- seems nec.
    step2 x@(GuardedRhs srcSpanInfo slst (FSContext _ e)) = GuardedRhs srcSpanInfo slst e
#if CLEAN_EXTRANEOUS_GROUPINGS
    step2 (GuardedRhs ssi slst
           (InfixApp _
             (FSContext _ e1)
             (QVarOp _ (UnQual _ (Symbol _ "$")))
             e2))
             = GuardedRhs ssi slst
                (App ssi
                  (Paren ssi e1)
                  e2)
#endif
    step2 x = x
    step3 :: Exp SrcSpanInfo -> Exp SrcSpanInfo  -- seems nec.
    step3 (FSContext ssi e) = Paren ssi e
    step3 x = x

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

  -- Actually perform freesect translations in the immediate subexpression
  -- of a given RHS in the AST.  Since the caller is itself a bottom-up
  -- generic traversal, nested freesects will get rewritten before
  -- enclosing freesects are processed.
  -- :: String -> Exp -> Exp
--fs_rhs_exp :: Data a => String -> a -> a  -- typesig not nec.
  fs_rhs_exp fresh rhs_top_exp = rhs_top_exp''
   where
    rhs_top_exp' = everywhere (mkT step) rhs_top_exp
    rhs_top_exp'' | num_fss_remaining > 0  = everywhere (mkT step2) rhs_top_exp'
          | otherwise              = rhs_top_exp'
-- FSContext is the grouping node in the AST produced by freesect _[ ]_ syntax.
-- The default context inferencing cases follow this explicit FSContext case.
-- The Exp -> Exp type sig for step (though it works) is not needed here...
--  step :: Data a => a -> a  -- ...although this one won't work.
    step x@(FSContext srcSpanInfo e) = fs_lambda_old ps' x'
     where (x',(ps,_)) = fs_name_slots fresh x
           ps' = reverse ps
#if 0
-- Just a test of generic power of SYB.  A single traversal is generic, but
-- only permits transformation of nodes of a single specific type.  The above
-- case is Exp -> Exp, while this is Decl -> Decl.
    step x@(DefaultDecl srcSpanInfo ts) = fs_lambda [] x  -- quick test
#endif
    step x = x
    num_fss_remaining = gcount (False `mkQ` p) rhs_top_exp'
    p :: Exp SrcSpanInfo -> Bool  -- nec.
    p (FSWildcard _) = True
    p _ = False
    -- Default context inference works as follows:
    -- The (semilattice) join of all unbracketed __'s in a RHS is found.
    -- Then, the innermost enclosing Paren or infix $ determines the context,
    -- or -- if neither exists -- the whole RHS is taken as context.
    -- (Later: Added list braces (list enumerations and comprehensions) to
    -- the set of delimiters.  This was motivated by consideration of
    -- primitives.html, but may need reconsideration when see more
    -- real-world examples.)
    --
    -- Would prefer to use SYB "everywhereBut" or "something", to stop
    -- searching farther, but ... would need an "everywhereButM" I think,
    -- since need to pass on the info that an amenable Paren
    -- has already been found.
-- Strangely the type sig not needed here...
--  step2 :: Exp SrcSpanInfo -> Exp SrcSpanInfo
    -- | Paren l (Exp l)
    step2 x@(Paren srcSpanInfo e)
      | num_fss_remaining == gcount (False `mkQ` p) e
         = x'_
      | otherwise
         = x
      where
        x_ = FSContext srcSpanInfo e
--      x_ = x
        (x',(ps,_)) = fs_name_slots fresh x_
        ps' = reverse ps
        x'_ = fs_lambda_old ps' x'
    -- | InfixApp l (Exp l) (QOp l) (Exp l)
    step2 x@(InfixApp srcSpanInfo e1 qop e2)
      |                    not good_op  =  x
      |  num_fss_x < num_fss_remaining  =  x
      |                num_fss_e2 == 0  =  InfixApp srcSpanInfo e1'_ qop e2
      |                num_fss_e1 == 0  =  InfixApp srcSpanInfo e1   qop e2'_
      |                      otherwise  =  x'_
      where
        e1_ = FSContext srcSpanInfo e1
        e2_ = FSContext srcSpanInfo e2
        x_ = FSContext srcSpanInfo x
        -- May want to broaden this category? remember, it's a trade off,
        -- if you use an op for a freesect context delimiter, it can't
        -- be used inside a freesect with defaulting context.
        -- To see why $ was chosen, check out the S23.hs test file.
        good_op = case qop of
          QVarOp _ (UnQual _ (Symbol _ "$")) -> True
          _ -> False
        (e1',(ps1,_)) = fs_name_slots fresh e1_
        ps1' = reverse ps1
        (e2',(ps2,_)) = fs_name_slots fresh e2_
        ps2' = reverse ps2
        (x',(ps,_)) = fs_name_slots fresh x_
        ps' = reverse ps
        e1'_ = fs_lambda_old ps1' e1'
        e2'_ = fs_lambda_old ps2' e2'
        x'_ = fs_lambda_old ps' x'
        num_fss_e1 = gcount (False `mkQ` p) e1
        num_fss_e2 = gcount (False `mkQ` p) e2
        num_fss_x = num_fss_e1 + num_fss_e2
--      num_fss_x = gcount (False `mkQ` p) x
    -- These are simpler, since we are just wrapping the node in FSContext
    -- and calling the fs_lambda transformer.
    -- | List l [Exp l]
    step2 x@(List srcSpanInfo _) = process srcSpanInfo x
    -- | EnumFrom l (Exp l)
    step2 x@(EnumFrom srcSpanInfo _) = process srcSpanInfo x
    -- | EnumFromTo l (Exp l) (Exp l)
    step2 x@(EnumFromTo srcSpanInfo _ _) = process srcSpanInfo x
    -- | EnumFromThen l (Exp l) (Exp l)
    step2 x@(EnumFromThen srcSpanInfo _ _) = process srcSpanInfo x
    -- | EnumFromThenTo l (Exp l) (Exp l) (Exp l)
    step2 x@(EnumFromThenTo srcSpanInfo _ _ _) = process srcSpanInfo x
    -- | ListComp l (Exp l) [(QualStmt l)]
    step2 x@(ListComp srcSpanInfo e slst) = process srcSpanInfo x
    -- | ParComp l (Exp l) [[(QualStmt l)]]
    -- HSE generates syntax errors when try to use this extension.
    step2 x = x

    process ssi x
      | num_fss_remaining == gcount (False `mkQ` p) x
         = x'_
      | otherwise
         = x
      where
        x_ = FSContext ssi x
        x'_ = fs_lambda_old ps' x'
        (x',(ps,_)) = fs_name_slots fresh x_
        ps' = reverse ps

  -- Actually rewrite the passed Exp branch as a Lambda.  The argument is,
  -- at least at present, always an FSContext, but any Exp branch would be
  -- treated analogously without changing fs_lambda.
  -- Note that the Lambda itself is wrapped in a Paren; this does not
  -- change the semantics of the AST, but is necessary in general to
  -- preserve the semantics when pretty-printing as lexical sourcecode.
  -- :: [String] -> Exp -> Exp always, at present
--fs_lambda :: Data a => [String] -> a -> a  -- must NOT give this one!
-- Strangely the type sig not needed here...
--fs_lambda :: [String] -> Exp SrcSpanInfo -> Exp SrcSpanInfo
  fs_lambda ps_lambda e_lambda
-- XXX See fs_lambda_old for what we should do here now...
   | null ps_lambda
      = error $    "Error: Free section contains no wildcards.\n"
                ++ showSSI ssi
   | otherwise = lambda
   where
         lambda = Paren ssi $ Lambda ssi ps_lambda' e_lambda''
         ps_lambda' = map (\x->(PVar ssi (Ident ssi x))) ps_lambda
         e_lambda'' = e_lambda
         ssi = head $ gmapQi 0 ([] `mkQ` ((\x->[x])::SrcSpanInfo->[SrcSpanInfo])) e_lambda
--       e_lambda'@(FSContext ssi e) = e_lambda
--       e_lambda'' = e

  -- :: [String] -> Exp -> Exp always, at present
--fs_lambda_old :: Data a => [String] -> a -> a  -- must NOT give this one!
-- Strangely the type sig not needed here...
-- ...But it /is/ when changed the "null ps_lambda" case.
  fs_lambda_old :: [String] -> Exp SrcSpanInfo -> Exp SrcSpanInfo
  fs_lambda_old ps_lambda e_lambda
-- Now, rather than report the error, we silently convert them
-- to Paren's.  No harm is done with this interpretation (it
-- is natural), and it allows us to keep the FSContext nodes
-- around until a post-translation cleanup where they are made use of.
   | null ps_lambda
#if 1
      = FSContext ssi e_lambda
#else
      = error $    "Error: Free section contains no wildcards.\n"
                ++ showSSI ssi
#endif
   | otherwise = lambda
   where
         -- The idea with leaving the FSContext's is, we can use
         -- them as markers to indicate where the rewrites happened
         -- (i.e. which Lambda's are due to freesect rewrites)
         -- and, in fs_clean, can use this to make the rewritten
         -- code a little bit cleaner (removing superfluous groupings
         -- or $ opertators).
         lambda = FSContext ssi $ Lambda ssi ps_lambda' e_lambda''
--       lambda = Paren ssi $ Lambda ssi ps_lambda' e_lambda''
         ps_lambda' = map (\x->(PVar ssi (Ident ssi x))) ps_lambda
         e_lambda'@(FSContext ssi e) = e_lambda
         e_lambda'' = e

  showSSI (SrcSpanInfo si _)
#if 0
#elif 0
   -- Prints, for example, "S28.hs:6:0: 7" (sans quotes), which is wrong.
   -- We expect "S28.hs:6:7" (sans quotes).  This wierdness is somehow
   -- connected to the use from unsafePerformIO, since "error $ showSSI ssi"
   -- prints the expected output.
   = fileName si ++ ":" ++ show (startLine si)
                 ++ ":" ++ show (startColumn si)
#elif 1
   -- prints correctly!
   = fileName si ++ ":\0" ++ show (startLine si)
                 ++ ":" ++ show (startColumn si)
#elif 0
   -- prints correctly (except for the extra space...)
   = fileName si ++ ": " ++ show (startLine si)
                 ++ ":" ++ show (startColumn si)
#elif 0
   -- (prints fine, but we prefer the terser, standard GHC location designator)
   = fileName si ++ ": line=" ++ show (startLine si)
                 ++ " col=" ++ show (startColumn si)
#endif

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

  -- We need to construct the fresh names in this recursion anyway, so
  -- may as well collect them rather than recompute them in the caller,
  -- although we could because they are canonically constructable from
  -- fresh and n, the Int part of the state.
  -- Perhaps ironically, I don't like using partially-point-free function
  -- declarations like this, but I couldn't figure out what to do with
  -- the second parameter if I made it explicit!
--fs_name_slots :: Data a => String -> a -> (a,([String],Int))  -- not needed
  fs_name_slots fresh = flip runState ([],0) . everywhereM (mkM step)
   where
-- Type signature necessary, and it seems that -XFlexibleContexts is
-- needed for it to be written?
--  step :: MonadState m => Exp SrcSpanInfo -> m (Exp SrcSpanInfo)
--  step :: MonadState s m => Exp SrcSpanInfo -> m (Exp SrcSpanInfo)
--  step :: MonadState (a,b) m => Exp SrcSpanInfo -> m (Exp SrcSpanInfo)
--  step :: MonadState ([a],b) m => Exp SrcSpanInfo -> m (Exp SrcSpanInfo)
    step :: MonadState ([String],Int) m => Exp SrcSpanInfo -> m (Exp SrcSpanInfo)
    step (FSWildcard srcSpanInfo)
     = do (ss,n) <- get
          let s = fresh ++ show n
          put ((s:ss),(1+n))
          return $ Var srcSpanInfo $ UnQual srcSpanInfo $ Ident srcSpanInfo s
    step x = return x