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

{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE FlexibleContexts #-}  -- needed for some of the type sigs.
{-# LANGUAGE CPP #-}
{- # LANGUAGE MultiParamTypeClasses #-}
{- # LANGUAGE RankNTypes #-}  -- needed for the path accumulators.
{- # LANGUAGE ExistentialQuantification #-}
{- # LANGUAGE GADTs #-}
{- # LANGUAGE ScopedTypeVariables #-}  -- needed for a pattern type sig.

-- CPP definitions are now set using compiler options; see ./z and ./ile.
-- #define ANNOTATED 0
-- #define PARALLEL 0
-- #define GHC_F 1

-- Most helpful sources:
--  - http://hpaste.org/steps/10722  -- use of everywhereM with State
--  - #haskell (thanks dreixel, quintessence, eyebloom, ...)

  module Main(main) where

  import System.Environment(getArgs)
  import System.IO(writeFile,hFlush,stdout)

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

#if PARALLEL
  import Control.Parallel.Strategies(rpar,parTraversable,runEval)
--import Control.Parallel  -- Not imported with Strategies, but
                           -- contains only par and pseq.
--import Control.Concurrent
---import Control.Exception
---import System.IO.Unsafe
---import Foreign
#endif

  import Control.Monad.State

  import System.Random(StdGen,mkStdGen,next)

  import Data.List(isPrefixOf)

#if ANNOTATED
  import HSE.Annotated
  import HSE.SrcLoc
#else
  import HSE
#endif

  import HSE.Extension

  import FilesAndParsing

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

#if ANNOTATED
  fs_module :: Module SrcSpanInfo -> Module SrcSpanInfo
#else
  fs_module :: Module -> Module
#endif
  fs_module m0 = m4
   where
    m1 = everywhere (mkT step) m0
    m2 = everywhere (mkT step2) m1
    m3 = fs_guarded_rhss m2
    m4 = fs_error_if_any_remain m3
#if ANNOTATED
    step :: Rhs SrcSpanInfo -> Rhs SrcSpanInfo  -- seems nec.
#else
--  step :: Rhs -> Rhs  -- unnec.
#endif
#if ANNOTATED
    step (UnGuardedRhs srcSpanInfo e) = UnGuardedRhs srcSpanInfo e'
#else
    step (UnGuardedRhs e) = UnGuardedRhs e'
#endif
     where e' = fs_rhs_exp fresh e
    step x = x
#if ANNOTATED
    step2 :: Rhs SrcSpanInfo -> Rhs SrcSpanInfo  -- seems nec.
#else
--  step2 :: Rhs -> Rhs  -- unnec.
#endif
#if ANNOTATED
    step2 x@(UnGuardedRhs srcSpanInfo e)
     | still_fsss = UnGuardedRhs srcSpanInfo e''
#else
    step2 x@(UnGuardedRhs e)
     | still_fsss = UnGuardedRhs e''
#endif
     | otherwise = x
     where
       still_fsss = 0 < ( ( gcount (False `mkQ` p) x ) :: Int )
#if ANNOTATED
       p :: Exp SrcSpanInfo -> Bool
       p x@(FreeSectSlot _) = True
#else
--     p :: Exp -> Bool
       p x@FreeSectSlot = True
#endif
       p _              = False
       e'' = fs_rhs_exp fresh e'
#if ANNOTATED
       e' = Paren srcSpanInfo e
#else
       e' = Paren e
#endif
    step2 x = x
    fresh = fs_fresh_name m0

#if ANNOTATED
  fs_guarded_rhss :: Module SrcSpanInfo -> Module SrcSpanInfo
#else
  fs_guarded_rhss :: Module -> Module
#endif
  fs_guarded_rhss m = m''
   where
    m' = everywhere (mkT step) m
    m'' = everywhere (mkT step2) m'
#if ANNOTATED
    step :: GuardedRhs SrcSpanInfo -> GuardedRhs SrcSpanInfo  -- seems nec.
#else
--  step :: GuardedRhs -> GuardedRhs  -- unnec.
#endif
#if ANNOTATED
    step (GuardedRhs srcSpanInfo slst e) = GuardedRhs srcSpanInfo slst e'
#else
    step (GuardedRhs srcLoc slst e) = GuardedRhs srcLoc slst e'
#endif
     where e' = fs_rhs_exp fresh e
    step x = x
#if ANNOTATED
    step2 :: GuardedRhs SrcSpanInfo -> GuardedRhs SrcSpanInfo  -- seems nec.
#else
--  step2 :: GuardedRhs -> GuardedRhs  -- unnec.
#endif
#if ANNOTATED
    step2 x@(GuardedRhs srcSpanInfo slst e)
     | still_fsss = GuardedRhs srcSpanInfo slst e''
#else
    step2 x@(GuardedRhs srcLoc slst e)
     | still_fsss = GuardedRhs srcLoc slst e''
#endif
     | otherwise = x
     where
       still_fsss = 0 < ( ( gcount (False `mkQ` p) x ) :: Int )
#if ANNOTATED
       p :: Exp SrcSpanInfo -> Bool
       p x@(FreeSectSlot _) = True
#else
--     p :: Exp -> Bool
       p x@FreeSectSlot = True
#endif
       p _              = False
       e'' = fs_rhs_exp fresh e'
#if ANNOTATED
       e' = Paren srcSpanInfo e
#else
       e' = Paren e
#endif
    step2 x = x
    fresh = fs_fresh_name m

#if ANNOTATED
  fs_error_if_any_remain :: Module SrcSpanInfo -> Module SrcSpanInfo
#else
  fs_error_if_any_remain :: Module -> Module
#endif
  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 ) :: Int )
#if ANNOTATED
    p :: Exp SrcSpanInfo -> Bool
    p x@(FreeSectSlot _) = True
    p x@(FSContext _ _) = True
#else
--  p :: Exp -> Bool
    p x@FreeSectSlot = True
    p x@(FSContext _) = True
#endif
    p _              = False

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

  -- v.0.0.5:
  fs_rhs_exp :: Data a => String -> a -> a
--fs_rhs_exp fresh rhs = rhs'
  fs_rhs_exp fresh rhs = rhs''
   where
    rhs' = everywhere (mkT step) rhs
    rhs'' | num_fss_remaining > 0  = everywhere (mkT step2) rhs'
          | otherwise              = rhs'
#if ANNOTATED
    step x@(FSContext srcSpanInfo e) = fs_lambda ps' x'
#else
    step x@(FSContext e) = fs_lambda ps' x'
#endif
     where (x',(ps,_)) = fs_name_slots fresh x
           ps' = reverse ps
    step x = x
    num_fss_remaining = ( gcount (False `mkQ` p) rhs' ) :: Int
--  num_fss_remaining = ( error $ show $ ( gcount (False `mkQ` p) rhs' ) :: Int ) :: Int
#if ANNOTATED
    p :: Exp SrcSpanInfo -> Bool
    p x@(FreeSectSlot _) = True
#else
--  p :: Exp -> Bool
    p x@FreeSectSlot = True
#endif
    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.
    --
    -- Would prefer to use 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.
#if ANNOTATED
    step2 :: Exp SrcSpanInfo -> Exp SrcSpanInfo
#else
--  step2 :: Exp -> Exp
#endif
#if ANNOTATED
    -- | 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
#else
    step2 x@(Paren e)
      | num_fss_remaining == gcount (False `mkQ` p) e
         = x'_
      | otherwise
         = x
      where
        x_ = FSContext e
#endif
        -- (We safely discarded the Paren from the AST since FSContext will
        -- give Paren grouping behaviour in addition to freesect contexting.)
        (x',(ps,_)) = fs_name_slots fresh x_
        ps' = reverse ps
        x'_ = fs_lambda ps' x'
#if ANNOTATED
    -- | 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
#else
    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
#endif
        (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 ps1' e1'
        e2'_ = fs_lambda ps2' e2'
        x'_ = fs_lambda ps' x'
        num_fss_e1 = ( gcount (False `mkQ` p) e1 ) :: Int
        num_fss_e2 = ( gcount (False `mkQ` p) e2 ) :: Int
        num_fss_x = gcount (False `mkQ` p) x
    step2 x = x

#if ANNOTATED
  fs_lambda :: [String] -> Exp SrcSpanInfo -> Exp SrcSpanInfo  -- seems nec.
#else
--fs_lambda :: [String] ->  Exp -> Exp  -- unnec.
#endif
  fs_lambda ps_lambda e_lambda
#if ANNOTATED
   | null ps_lambda  = error $    "Error: Free section contains no wildcards.\n"
                               ++ showSLorSSI ssi
#else
   | null ps_lambda  = error $    "Error: Free section contains no wildcards.\n"
                               ++ "(Source location not available, try -annotated.)\n"
#endif
   | otherwise  = lambda
   where
#if ANNOTATED
         lambda = 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
#else
         lambda = 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
#endif

#if ANNOTATED
  showSLorSSI :: SrcSpanInfo -> String
  showSLorSSI (SrcSpanInfo si _)
   = fileName si ++ ": line=" ++ show (startLine si)
                 ++ " col=" ++ show (startColumn si)
#else
  showSLorSSI :: SrcLoc -> String
  showSLorSSI sl@(SrcLoc n l c)
   = n ++ ": line=" ++ show l ++ " col=" ++ show c
#endif

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

  -- 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!
  -- 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.
  fs_name_slots :: Data a => String -> a -> (a,([String],Int))
  fs_name_slots fresh
   = flip runState ([],0) . everywhereM (mkM step)
   where
#if ANNOTATED
    step :: MonadState ([String],Int) m => Exp SrcSpanInfo -> m (Exp SrcSpanInfo)  -- seems nec.
#else
--  step :: MonadState ([String],Int) m => Exp -> m Exp  -- unnec.
#endif
#if ANNOTATED
    step (FreeSectSlot srcSpanInfo)
#else
    step FreeSectSlot
#endif
     = do (ss,n) <- get
          let s = fresh ++ show n
--        let s = "freeSect_" ++ show n
          put ((s:ss),(1+n))
#if ANNOTATED
          return $ ( Var srcSpanInfo ( UnQual srcSpanInfo ( Ident srcSpanInfo s ) ) )
#else
          return $ Var $ UnQual $ Ident $ s
#endif
    step x = return x

  fs_all_identifiers :: Data a => a -> [String]
  fs_all_identifiers = everything (++) ([] `mkQ` f)
   where
#if ANNOTATED
         f :: (Name SrcSpanInfo) -> [String]  -- seems nec.
#else
--       f :: Name -> [String]  -- unnec.
#endif
#if ANNOTATED
         f (Ident _ x) = [x]
#else
         f (Ident x) = [x]
#endif
         f _ = []

  -- The names which FreeSect inserts will never conflict with each other.
  -- We only need to assure they don't conflict with any existing names.
  -- Actually, we need to make sure the name created here is not a prefix
  -- of any existing name, because we add _XY to freesect slot names.
#if ANNOTATED
  fs_fresh_name :: Module SrcSpanInfo -> String
#else
  fs_fresh_name :: Module -> String
#endif
  fs_fresh_name m = f g
   where
         ss = fs_all_identifiers m
         g = mkStdGen 123  -- arbitrary seed
         -- The following was much simpler when accept whole of r
         -- as the random part of the name -- however, that made for
         -- ugly long names, and so we try for the shortest possible
         -- first.  (If you never inspect the intermediate code, you
         -- wouldn't care if the var names were ugly...)
         f :: StdGen -> String
         f g | b          = s
             | otherwise  = f g'  -- unlikely
             where (r,g') = next g
                   (b,s) = f' rs (0,length rs)
                   rs = show r
         f' :: String -> (Int,Int) -> (Bool,String)
         f' s (n,ntop) | n > ntop   = (False,"")
                       | not fail   = (True,s'')
                       | otherwise  = f' s (1+n,ntop)
                       where s' = take n s
                             s'' =         s' ++ "_"
--                           s'' = "fs" ++ s' ++ "_"
                             fail = or $ map (isPrefixOf s'') ss

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

#if ANNOTATED
  stripFSPragma :: Module SrcSpanInfo -> Module SrcSpanInfo
#else
  stripFSPragma :: Module -> Module
#endif
#if ANNOTATED
  stripFSPragma (Module x1 x2 prags x4 x5)
   = Module x1 x2 prags' x4 x5
#else
  stripFSPragma (Module x1 x2 prags x4 x5 x6 x7)
   = Module x1 x2 prags' x4 x5 x6 x7
#endif
   where prags' = map f prags
         f (LanguagePragma sl_or_ssi ns)
          = LanguagePragma sl_or_ssi $ filter p ns
         f x = x
#if ANNOTATED
         p n@(Ident sl_or_ssi "FreeSections") = False
#else
         p n@(Ident           "FreeSections") = False
#endif
         p _ = True

#if ANNOTATED
  stripEmptyPragmaList :: Module SrcSpanInfo -> Module SrcSpanInfo
#else
  stripEmptyPragmaList :: Module -> Module
#endif
#if ANNOTATED
  stripEmptyPragmaList (Module x1 x2 prags x4 x5)
   = Module x1 x2 prags' x4 x5
#else
  stripEmptyPragmaList (Module x1 x2 prags x4 x5 x6 x7)
   = Module x1 x2 prags' x4 x5 x6 x7
#endif
   where prags' = filter p prags
         p (LanguagePragma ssi []) = False
         p _ = True

#if ANNOTATED
  fixModuleName :: String -> Module SrcSpanInfo -> Module SrcSpanInfo
#else
  fixModuleName :: String -> Module -> Module
#endif
#if ANNOTATED
  fixModuleName name (Module x1 x2 x3 x4 x5)
   = Module x1 name' x3 x4 x5
#else
  fixModuleName name (Module x1 x2 x3 x4 x5 x6 x7)
   = Module x1 name' x3 x4 x5 x6 x7
#endif
   where
#if ANNOTATED
         (Just (ModuleHead ssi _ mwt mesl)) = x2
         name' = Just (ModuleHead ssi (ModuleName ssi name) mwt mesl)
#else
         (ModuleName _) = x2
         name' = ModuleName name
#endif
         p (LanguagePragma ssi []) = False
         p _ = True

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

  main:: IO ()
  main
   = do (
          outfile :
          lexsrc_pathnames
          ) <- getArgs
        lexsrc_serials_ <- mapM readSourcesFromFileOrDir lexsrc_pathnames
        let
          (pnames,lexsrc_serials) = unzip $ concat $ reverse lexsrc_serials_
#if ANNOTATED
#if PARALLEL
          parsedsrc_maybes = (runEval $ parTraversable rpar $ doParsing pnames lexsrc_serials) :: [ParseResult (Module SrcSpanInfo)]
#else
          parsedsrc_maybes = (doParsing pnames lexsrc_serials) :: [ParseResult (Module SrcSpanInfo)]
#endif
#else
          parsedsrc_maybes = (doParsing pnames lexsrc_serials) :: [ParseResult Module]
#endif
        let
--        parsed_srcs = error $ ( ( concatMap prettyPrint $ ( ( testParses parsedsrc_maybes ) :: [Module] ) ) :: String )
#if ANNOTATED
          parsed_srcs = ( testParses parsedsrc_maybes ) :: [Module SrcSpanInfo]
#else
          parsed_srcs = ( testParses parsedsrc_maybes ) :: [Module]
#endif
{--
        let
          test = ( error $ show $ map fs_FSS_lineal_chains parsed_srcs ) :: String
        print test
--}
        let
#if ANNOTATED
#if PARALLEL
          transformed_srcs = ( runEval $ parTraversable rpar $ 
                               map fs_module parsed_srcs
                               ) :: [Module SrcSpanInfo]
#else
          transformed_srcs = ( map fs_module parsed_srcs
                               ) :: [Module SrcSpanInfo]
#endif
#else
#if PARALLEL
          transformed_srcs = ( runEval $ parTraversable rpar $ 
                               map fs_module parsed_srcs
                               ) :: [Module]
#else
          transformed_srcs = ( map fs_module parsed_srcs
                               ) :: [Module]
#endif
#endif
--        transformed_srcs = parsed_srcs

        let
          transformed_srcs' = map stripFSPragma transformed_srcs
          transformed_srcs'' = map stripEmptyPragmaList transformed_srcs'
#if GHC_F
#else
          transformed_srcs''' = map (fixModuleName outfile) transformed_srcs''
#endif

#if GHC_F
        debug parsed_srcs transformed_srcs''
        writeFile outfile $  -- with ghc -F
                        concatMap prettyPrint transformed_srcs''
#else
        debug parsed_srcs transformed_srcs'''
        writeFile (outfile++".hs") $
                        concatMap prettyPrint transformed_srcs'''
#endif

        hFlush stdout

#if ANNOTATED
  debug :: [Module SrcSpanInfo] -> [Module SrcSpanInfo] -> IO ()
#else
  debug :: [Module] -> [Module] -> IO ()
#endif
  debug ms ms'
   = do
#if 0
        putStrLn $ show ms
        putStrLn $ show ms'
#endif
#if 1
        putStrLn $ concatMap prettyPrint ms
        putStrLn $ concatMap prettyPrint ms'
#endif
        return ()