--     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 #-}  -- only needed in FreeSectAnnotated.hs
{- # LANGUAGE MultiParamTypeClasses #-}
{- # LANGUAGE RankNTypes #-}  -- needed for the path accumulators
{- # LANGUAGE ExistentialQuantification #-}
{- # LANGUAGE GADTs #-}
{- # LANGUAGE ScopedTypeVariables #-}  -- needed for a pattern type sig

{- # NOINLINE fs_warn_flaw #-}

-- CPP definitions are set using compiler options; see ./z and ./ile.

  module FreeSect(fs_module) 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 Data.IORef(IORef,newIORef)

  import HSE

  import Util

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

  -- No type signatures are needed anywhere in this file (confirmed),
  -- but in FreeSectAnnotated some /are/ needed.

  -- 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

--fs_warn_flaw :: Data a => a -> a  -- unnec.
  fs_warn_flaw m = m'
   where
    m' = everywhere (mkT step) m
--  step :: Exp -> Exp  -- unnec.
    step x@(App p@(Paren (App FSWildcard _)) _) = warning True p x
    step x@(App p@(Paren (App _ FSWildcard)) _) = warning False p x
    step x = x
--  warning :: Data a => Bool -> a -> a -> a  -- unnec.
    warning b p x =   unsafePerformIO
                $ do hPutStr stderr $ warning_message b p x
                     hFlush stderr
                     return x
    warning_message b p x
     =    "Warning:\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"
       ++ "  (Compile freesect with ANNOTATED=1 to get location info.)\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 -> Rhs  -- unnec.
    step1 (UnGuardedRhs e) = UnGuardedRhs e'
     where e' = fs_rhs_exp fresh e
    step1 x = x
--  step2 :: Rhs -> Rhs  -- unnec.
    step2 x@(UnGuardedRhs e)
     | still_fsss = UnGuardedRhs e''
     | otherwise = x
     where
       still_fsss = 0 < gcount (False `mkQ` p) x
--     p :: Exp -> Bool
       p FSWildcard = True
       p _ = False
       e'' = fs_rhs_exp fresh e'
       e' = Paren 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 -> GuardedRhs  -- unnec.
    step1 (GuardedRhs srcLoc slst e) = GuardedRhs srcLoc slst e'
     where e' = fs_rhs_exp fresh e
    step1 x = x
--  step2 :: GuardedRhs -> GuardedRhs  -- unnec.
    step2 x@(GuardedRhs srcLoc slst e)
     | still_fsss = GuardedRhs srcLoc slst e''
     | otherwise = x
     where
       still_fsss = 0 < gcount (False `mkQ` p) x
--     p :: Exp -> Bool
       p FSWildcard = True
       p _ = False
       e'' = fs_rhs_exp fresh e'
       e' = Paren 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 -> Bool
    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 -> Rhs  -- unnec.
    step1 (UnGuardedRhs (FSContext e)) = UnGuardedRhs e
#if CLEAN_EXTRANEOUS_GROUPINGS
    step1 (UnGuardedRhs
           (InfixApp
             (FSContext e1)
             (QVarOp (UnQual (Symbol "$")))
             e2))
             = UnGuardedRhs
                (App
                  (Paren e1)
                  e2)
#endif
    step1 x = x
--  step2 :: GuardedRhs -> GuardedRhs  -- unnec.
    step2 x@(GuardedRhs srcLoc slst (FSContext e)) = GuardedRhs srcLoc slst e
#if CLEAN_EXTRANEOUS_GROUPINGS
    step2 (GuardedRhs srcLoc slst
           (InfixApp
             (FSContext e1)
             (QVarOp (UnQual (Symbol "$")))
             e2))
             = GuardedRhs srcLoc slst
                (App
                  (Paren e1)
                  e2)
#endif
    step2 x = x
--  step3 :: Exp -> Exp  -- once was nec. but not now?
    step3 (FSContext e) = Paren 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 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 srcLoc ts) = fs_lambda [] x  -- quick test
#endif
    step x = x
    num_fss_remaining = gcount (False `mkQ` p) rhs_top_exp'
--  p :: Exp -> Bool
    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.
--  step2 :: Exp -> Exp
    step2 x@(Paren e)
      | num_fss_remaining == gcount (False `mkQ` p) e
         = x'_
      | otherwise
         = x
      where
        -- (We safely discarded the Paren from the AST since FSContext will
        -- give Paren grouping behaviour in addition to freesect contexting.)
        x_ = FSContext e
--      x_ = x
        (x',(ps,_)) = fs_name_slots fresh x_
        ps' = reverse ps
        x'_ = fs_lambda_old ps' x'
    step2 x@(InfixApp e1 qop e2)
      |                    not good_op  =  x
      |  num_fss_x < num_fss_remaining  =  x
      |                num_fss_e2 == 0  =  InfixApp e1'_ qop e2
      |                num_fss_e1 == 0  =  InfixApp e1   qop e2'_
      |                      otherwise  =  x'_
      where
        e1_ = FSContext e1
        e2_ = FSContext e2
        x_ = FSContext x
        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.
    step2 x@(List _) = process x
    step2 x@(EnumFrom _) = process x
    step2 x@(EnumFromTo _ _) = process x
    step2 x@(EnumFromThen _ _) = process x
    step2 x@(EnumFromThenTo _ _ _) = process x
    step2 x@(ListComp e slst) = process x
    step2 x = x
    process x
      | num_fss_remaining == gcount (False `mkQ` p) x
         = x'_
      | otherwise
         = x
      where
        x_ = FSContext 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!
--fs_lambda :: [String] ->  Exp -> Exp  -- unnec.
  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"
                ++ "(Source location unavailable; try compiling freesect "
                ++ "with ANNOTATED set to 1.)\n"
   | otherwise = lambda
   where
         lambda = Paren $ Lambda srcloc ps_lambda' e_lambda''
         ps_lambda' = map (\x->(PVar (Ident x))) ps_lambda
         e_lambda'' = e_lambda
--       e_lambda'@(FSContext e) = e_lambda
--       e_lambda'' = e
         srcloc = SrcLoc "" 0 0

  -- :: [String] -> Exp -> Exp always, at present
--fs_lambda_old :: Data a => [String] -> a -> a  -- must NOT give this one!
--fs_lambda_old :: [String] ->  Exp -> Exp  -- unnec.
  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 e_lambda
#else
      = error $    "Error: Free section contains no wildcards.\n"
                ++ "(Source location unavailable; try compiling freesect "
                ++ "with ANNOTATED set to 1.)\n"
#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 $ Lambda srcloc ps_lambda' e_lambda''
--       lambda = Paren $ Lambda srcloc ps_lambda' e_lambda''
         ps_lambda' = map (\x->(PVar (Ident x))) ps_lambda
         e_lambda'@(FSContext e) = e_lambda
         e_lambda'' = e
         srcloc = SrcLoc "" 0 0

  showSLorSSI (SrcLoc n l c)
   = n ++ ": line=" ++ show l ++ " col=" ++ show c

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

  -- 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
--  step :: MonadState ([String],Int) m => Exp -> m Exp  -- unnec.
    step FSWildcard
     = do (ss,n) <- get
          let s = fresh ++ show n
          put ((s:ss),(1+n))
          return $ Var $ UnQual $ Ident $ s
    step x = return x