haskell-awk-1.1.1: tests/Everything.hs

{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE OverloadedStrings, PackageImports, RecordWildCards #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE OverloadedStrings, PackageImports #-}
{-# LANGUAGE PackageImports, RankNTypes #-}
module Everything where


import Control.Monad.Trans.Class
import qualified Data.ByteString.Char8 as B

import Data.HaskellSource

-- Most of the API is re-exported from those submodules
import Data.HaskellModule.Base

import "mtl" Control.Monad.Trans
import qualified Data.ByteString.Char8 as B
import Data.List
import Language.Haskell.Exts
import Text.Printf

import Data.HaskellModule.Base
import Data.HaskellSource
import Language.Haskell.Exts.Location


import Language.Haskell.Interpreter hiding (Option, languageExtensions)
import Text.Printf (printf)

import System.Console.Hawk.Args
import System.Console.Hawk.Args.Spec
import System.Console.Hawk.Help
import System.Console.Hawk.Interpreter
import System.Console.Hawk.Version

import System.FilePath

import System.Console.Hawk

import qualified Data.ByteString.Char8 as B

import Data.HaskellModule.Base

import Data.ByteString (ByteString)
import Data.ByteString.Char8 (pack)
import Text.Printf

import Control.Monad.Trans.OptionParser

import Control.Applicative
import Data.Char                                 (isSpace)
import "mtl" Control.Monad.Trans

import Control.Monad.Trans.OptionParser
import qualified System.Console.Hawk.Args.Option as Option
import           System.Console.Hawk.Args.Option (HawkOption, options)
import           System.Console.Hawk.Args.Spec
import           System.Console.Hawk.Context.Dir

import Control.Applicative
import Control.Monad.Trans.Class
import Data.ByteString as B
import Text.Printf

import System.Console.Hawk.Sandbox


import Control.Applicative
import Data.Maybe

import System.Console.Hawk.UserPrelude.Defaults


import Control.Applicative
import "mtl" Control.Monad.Trans
import "mtl" Control.Monad.Identity
import "transformers" Control.Monad.Trans.Error hiding (Error)
import "transformers" Control.Monad.Trans.Writer
import System.Exit
import System.IO
import Text.Printf


-- |
-- >>> B.putStr $ showModule "orig.hs" $ emptyModule
-- 
-- >>> B.putStr $ showModule "orig.hs" $ addExtension "OverloadedStrings" $ addImport ("Data.ByteString.Char8", Just "B") $ addExtension "RecordWildCards" $ addImport ("Prelude", Nothing) $ emptyModule
-- {-# LANGUAGE OverloadedStrings #-}
-- {-# LANGUAGE RecordWildCards #-}
-- import qualified Data.ByteString.Char8 as B
-- import Prelude
showModule :: FilePath -- ^ the original's filename,
                       --   used for fixing up line numbers
           -> HaskellModule -> B.ByteString
showModule orig (HaskellModule {..}) = showSource orig fullSource
  where
    fullSource = Data.List.concat [ pragmaSource
                        , moduleSource
                        , importSource
                        , codeSource
                        ]

writeModule :: FilePath -- ^ the original's filename,
                        --   used for fixing up line numbers
            -> FilePath
            -> HaskellModule
            -> IO ()
writeModule orig f = B.writeFile f . showModule orig


compileModule :: FilePath -- ^ the original's filename,
                          --   used for fixing up line numbers
              -> FilePath -- ^ new filename, because ghc compiles from disk.
                          --   the compiled output will be in the same folder.
              -> HaskellModule
              -> UncertainT IO ()
compileModule = compileModuleWithArgs []

compileModuleWithArgs :: [String] -- ^ extra ghc args
                      -> FilePath -- ^ the original's filename,
                                  --   used for fixing up line numbers
                      -> FilePath -- ^ new filename, because ghc compiles from disk.
                                  --   the compiled output will be in the same folder.
                      -> HaskellModule
                      -> UncertainT IO ()
compileModuleWithArgs args orig f m = do
    lift $ writeModule orig f m
    compileFileWithArgs args f
-- | In which a Haskell module is deconstructed into extensions and imports.

locatedExtensions :: [ModulePragma] -> Located [ExtensionName]
locatedExtensions = fmap go . located
  where
    go :: [ModulePragma] -> [ExtensionName]
    go = concatMap extNames
    
    extNames :: ModulePragma -> [ExtensionName]
    extNames (LanguagePragma _ exts) = Data.List.map prettyPrint exts
    extNames (OptionsPragma _ _ _) = []  -- TODO: accept "-XExtName"
    extNames _ = []

locatedImports :: [ImportDecl] -> Located [QualifiedModule]
locatedImports = fmap go . located
  where
    go :: [ImportDecl] -> [QualifiedModule]
    go = Data.List.map qualify
    
    qualify :: ImportDecl -> QualifiedModule
    qualify decl = (fullName decl, qualifiedName decl)
    
    fullName :: ImportDecl -> String
    fullName = prettyPrint . importModule
    
    qualifiedName :: ImportDecl -> Maybe String
    qualifiedName = fmap prettyPrint . importAs

locatedModule :: SrcLoc -> HaskellSource -> ModuleName -> Located (Maybe String)
locatedModule srcLoc source (ModuleName mName) = case moduleLine of
    Nothing -> return Nothing
    Just line -> located (srcLoc {srcLine = line}) >> return (Just mName)
  where
    isModuleDecl :: Either B.ByteString String -> Bool
    isModuleDecl (Left xs) = "module " `B.isPrefixOf` xs
    isModuleDecl (Right xs) = "module " `isPrefixOf` xs
    
    moduleLine :: Maybe Int
    moduleLine = fmap index2line $ findIndex isModuleDecl source


-- line numbers start at 1, list indices start at 0.
line2index, index2line :: Int -> Int
line2index = subtract 1
index2line = (+ 1)


-- | A variant of `splitAt` which makes it easy to make `snd` empty.
-- 
-- >>> maybeSplitAt Nothing "abc"
-- ("abc","")
-- 
-- >>> maybeSplitAt (Just 0) "abc"
-- ("","abc")
maybeSplitAt :: Maybe Int -> [a] -> ([a], [a])
maybeSplitAt Nothing  ys = (ys, [])
maybeSplitAt (Just i) ys = splitAt i ys

-- | Given n ordered indices before which to split, split the list into n+1 pieces.
--   Omitted indices will produce empty pieces.
-- 
-- >>> multiSplit [] "foo"
-- ["foo"]
-- 
-- >>> multiSplit [Just 0, Just 1, Just 2] "foo"
-- ["","f","o","o"]
-- 
-- >>> multiSplit [Just 0, Just 1, Nothing] "foo"
-- ["","f","oo",""]
-- 
-- >>> multiSplit [Just 0, Nothing, Just 2] "foo"
-- ["","fo","","o"]
-- 
-- >>> multiSplit [Just 0, Nothing, Nothing] "foo"
-- ["","foo","",""]
-- 
-- >>> multiSplit [Nothing, Just 1, Just 2] "foo"
-- ["f","","o","o"]
-- 
-- >>> multiSplit [Nothing, Just 1, Nothing] "foo"
-- ["f","","oo",""]
-- 
-- >>> multiSplit [Nothing, Nothing, Just 2] "foo"
-- ["fo","","","o"]
-- 
-- >>> multiSplit [Nothing, Nothing, Nothing] "foo"
-- ["foo","","",""]
multiSplit :: [Maybe Int] -> [a] -> [[a]]
multiSplit []           xs = [xs]
multiSplit (j:js) xs = ys1 : ys2 : yss
  where
    (ys:yss) = multiSplit js xs
    (ys1, ys2) = maybeSplitAt j ys

-- | Given n ordered source locations, split the source into n+1 pieces.
--   Omitted source locations will produce empty pieces.
splitSource :: [Maybe SrcLoc] -> HaskellSource -> [HaskellSource]
splitSource = multiSplit . (fmap . fmap) (line2index . srcLine)


-- Due to a limitation of haskell-parse-exts, there is no `parseModule`
-- variant of `readModule` which would parse from a String instead of a file.
-- 
-- According to the documentation [1], only `parseFile` honors language
-- pragmas, without which PackageImport-style imports will fail to parse.
-- 
-- [1] http://hackage.haskell.org/package/haskell-src-exts-1.14.0.1/docs/Language-Haskell-Exts-Parser.html#t:ParseMode

readModule :: FilePath -> UncertainT IO HaskellModule
readModule f = do
    s <- lift $ readSource f
    r <- lift $ parseFile f
    case r of
      ParseOk (Module srcLoc moduleDecl pragmas _ _ imports decls)
        -> return $ go s srcLoc pragmas moduleDecl imports decls
      ParseFailed loc err -> multilineFail msg
        where
          -- we start with a newline to match ghc's errors
          msg = printf "\n%s:%d:%d: %s" (srcFilename loc) (srcLine loc) (srcColumn loc) err
  where
    go source srcLoc pragmas moduleDecl imports decls = HaskellModule {..}
      where
        (languageExtensions,      _) = runLocated (locatedExtensions pragmas)
        (moduleName,      moduleLoc) = runLocated (locatedModule srcLoc source moduleDecl)
        (importedModules, importLoc) = runLocated (locatedImports imports)
        (_,                 declLoc) = runLocated (located decls)
        
        sourceParts = splitSource [moduleLoc, importLoc, declLoc] source
        [pragmaSource, moduleSource, importSource, codeSource] = sourceParts
--   Copyright 2013 Mario Pastorelli (pastorelli.mario@gmail.com) Samuel Gélineau (gelisam@gmail.com)
--
--   Licensed under the Apache License, Version 2.0 (the "License");
--   you may not use this file except in compliance with the License.
--   You may obtain a copy of the License at
--
--       http://www.apache.org/licenses/LICENSE-2.0
--
--   Unless required by applicable law or agreed to in writing, software
--   distributed under the License is distributed on an "AS IS" BASIS,
--   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
--   See the License for the specific language governing permissions and
--   limitations under the License.

-- | Hawk as seen from the outside world: parsing command-line arguments,
--   evaluating user expressions.

-- | Same as if the given arguments were passed to Hawk on the command-line.
processArgs :: [String] -> IO ()
processArgs args = do
    r <- runWarningsIO $ parseArgs args
    case r of
      Left err -> failHelp err
      Right spec -> processSpec spec

-- | A variant of `processArgs` which accepts a structured specification
--   instead of a sequence of strings.
processSpec :: HawkSpec -> IO ()
processSpec Help          = help
processSpec Version       = putStrLn versionString
processSpec (Eval  e   o) = applyExpr (wrapExpr "const" e) noInput o
processSpec (Apply e i o) = applyExpr e                    i       o
processSpec (Map   e i o) = applyExpr (wrapExpr "map"   e) i       o

wrapExpr :: String -> ExprSpec -> ExprSpec
wrapExpr f e = e'
  where
    u = userExpression e
    u' = printf "%s (%s)" (prel f) u
    e' = e { userExpression = u' }

applyExpr :: ExprSpec -> InputSpec -> OutputSpec -> IO ()
applyExpr e i o = do
    let contextDir = userContextDirectory e
    let expr = userExpression e
    
    processRuntime <- runUncertainIO $ runHawkInterpreter $ do
      applyContext contextDir
      interpret' $ processTable' $ tableExpr expr
    runHawkIO $ processRuntime hawkRuntime
  where
    interpret' expr = do
      interpret expr (as :: HawkRuntime -> HawkIO ())
    
    hawkRuntime = HawkRuntime i o
    
    processTable' :: String -> String
    processTable' = printf "(%s) (%s) (%s)" (prel "flip")
                                            (runtime "processTable")
    
    -- turn the user expr into an expression manipulating [[B.ByteString]]
    tableExpr :: String -> String
    tableExpr = (`compose` fromTable)
      where
        fromTable = case inputFormat i of
            RawStream            -> head' `compose` head'
            Records _ RawRecord  -> map' head'
            Records _ (Fields _) -> prel "id"
    
    compose :: String -> String -> String
    compose f g = printf "(%s) %s (%s)" f (prel ".") g
    
    head' :: String
    head' = prel "head"
    
    map' :: String -> String
    map' = printf "(%s) (%s)" (prel "map")

-- we cannot use any unqualified symbols in the user expression,
-- because we don't know which modules the user prelude will import.
qualify :: String -> String -> String
qualify moduleName = printf "%s.%s" moduleName

prel, runtime :: String -> String
prel = qualify "Prelude"
runtime = qualify "System.Console.Hawk.Runtime"
-- | Tests which require particular prelude files.

-- | A helper for specifying which arguments to pass to Hawk.
--   Simplifies testing.
testBuilder :: FilePath -> FilePath -> [String] -> String -> FilePath -> IO ()
testBuilder preludeBase preludeBasename flags expr inputBasename
  = processArgs args
  where
    args = preludeArgs preludeBasename
        ++ flags
        ++ [expr]
        ++ inputArgs inputBasename
    
    preludePath f = preludeBase </> f
    inputPath f = "tests" </> "inputs" </> f
    
    preludeArgs f = ["-c", preludePath f]
    
    inputArgs "" = []
    inputArgs f = [inputPath f]

-- | A version of `testBuilder` without a custom prelude.
-- 
-- We still need to specify a prelude file, because the user running these
-- tests might have installed a custom prelude.
test :: [String] -> String -> FilePath -> IO ()
test = testBuilder ("tests" </> "preludes") "default"

-- | A version of `test` without a custom input file either.
testEval :: [String] -> String -> IO ()
testEval flags expr = test flags expr ""


-- | A version of `testBuilder` using the preludes from "tests/preludes".
-- 
-- The first example from the README:
-- 
-- >>> test ["-d:", "-m"] "head" "passwd"
-- root
-- 
-- 
-- The second example, which adds `takeLast` to the user prelude:
-- 
-- >>> testPrelude "readme" ["-a"] "takeLast 3" "0-100"
-- 98
-- 99
-- 100
-- 
-- 
-- The last example from the README, a quick test to validate that Hawk was
-- properly installed:
-- 
-- >>> testEval [] "[1..3]"
-- 1
-- 2
-- 3
-- 
-- 
-- Making sure that we don't assume the user prelude exports "map":
-- 
-- >>> testPrelude "set" ["-m"] "const $ \"hello\"" "1-3"
-- hello
-- hello
-- hello
-- 
-- Making sure that we can find "map" even with NoImplicitPrelude:
-- 
-- >>> testPrelude "noImplicitPrelude" ["-m"] "\\_ -> hello" "1-3"
-- hello
-- hello
-- hello
-- 
-- Making sure sequences of whitespace count as one delimiter:
-- 
-- >>> testPrelude "default" ["-a"] "L.transpose" "1-12"
-- 1 4 7 10
-- 2 5 8 11
-- 3 6 9 12
testPrelude :: FilePath -> [String] -> String -> FilePath -> IO ()
testPrelude = testBuilder ("tests" </> "preludes")

-- | A version of `testBuilder` using the preludes from the documentation.
-- 
-- All the examples from the documentation:
-- 
-- >>> test ["-a"] "L.reverse" "1-3"
-- 3
-- 2
-- 1
-- 
-- >>> test ["-ad"] "L.takeWhile (/=\"7\") . L.dropWhile (/=\"3\")" "1-10"
-- 3
-- 4
-- 5
-- 6
-- 
-- >>> testDoc "between" ["-ad"] "between \"3\" \"7\"" "1-10"
-- 3
-- 4
-- 5
-- 6
-- 
-- >>> test ["-a"] "L.take 3" "1-10"
-- 1
-- 2
-- 3
-- 
-- >>> test ["-m"] "L.reverse" "1-9"
-- 3 2 1
-- 6 5 4
-- 9 8 7
-- 
-- >>> testDoc "postorder" ["-ad"] "postorder (\\x -> printf \"(%s)\" . L.intercalate \" \" . (unpack x:))" "example.in"
-- (foo (bar1) (bar2 (baz)) (bar3))
-- 
-- >>> test ["-ad"] "sum . L.map (read . B.unpack)" "1-3"
-- 6
-- 
-- >>> testDoc "conversions" ["-ad"] "sum . L.map toInt" "1-3"
-- 6
-- 
-- >>> testEval [] "2 ^ 100"
-- 1267650600228229401496703205376
-- 
-- >>> test ["-a"] "L.take 2" "1-9"
-- 1 2 3
-- 4 5 6
-- 
-- >>> test ["-m"] "L.take 2" "1-9"
-- 1 2
-- 4 5
-- 7 8
-- 
-- >>> test ["-a"] "show" "1-9"
-- [["1","2","3"],["4","5","6"],["7","8","9"]]
-- 
-- >>> test ["-a"] "id :: [[B.ByteString]] -> [[B.ByteString]]" "1-9"
-- 1 2 3
-- 4 5 6
-- 7 8 9
-- 
-- >>> test ["-a", "-d\\t"] "id" "1-9tabs"
-- 1	2	3
-- 4	5	6
-- 7	8	9
-- 
-- >>> test ["-ad,"] "id" "1-9commas"
-- 1,2,3
-- 4,5,6
-- 7,8,9
-- 
-- >>> test ["-D + ", "-d*", "-a"] "L.transpose" "equation"
-- x1*x2 + y1*y2 + z1*z2
-- 
-- >>> test ["-d", "-a"] "show :: [B.ByteString] -> String" "1-3"
-- ["1","2","3"]
-- 
-- >>> test ["-d", "-D", "-a"] "show :: B.ByteString -> String" "1-3"
-- "1\n2\n3\n"
-- 
-- >>> testEval [] "[[B.pack \"1\",B.pack \"2\"], [B.pack \"3\",B.pack \"4\"]]"
-- 1 2
-- 3 4
-- 
-- >>> testEval [] "[[\"1\",\"2\"], [\"3\",\"4\"]]"
-- 1 2
-- 3 4
-- 
-- >>> testEval [] "[[1,2], [3,4]] :: [[Float]]"
-- 1.0 2.0
-- 3.0 4.0
-- 
-- >>> testEval [] "1 :: Double"
-- 1.0
-- 
-- >>> testEval [] "(True,False)"
-- True
-- False
-- 
-- >>> testEval [] "[(1,2),(3,4)] :: [(Int,Float)]"
-- 1 2.0
-- 3 4.0
-- 
-- >>> testEval ["-O or "] "(True,False)"
-- True or False
-- 
-- >>> testEval ["-o\\t"] "[(1,2),(3,4.0)] :: [(Int,Float)]"
-- 1	2.0
-- 3	4.0
-- 
-- >>> test ["-m", "-d ", "-o*", "-D\\n", "-O+"] "id" "1-6"
-- 1*2*3+4*5*6
-- 
-- >> testEval ["-a"] "L.length"
-- 3
testDoc :: String -> [String] -> String -> FilePath -> IO ()
testDoc = testBuilder "doc"

-- | Test that `readModule` splits prelude files into correct sections.
-- 
-- >>> testM "tests/preludes/default/prelude.hs"
-- "{-# LANGUAGE ExtendedDefaultRules, OverloadedStrings #-}"
-- ["ExtendedDefaultRules","OverloadedStrings"]
-- ===
-- Nothing
-- ===
-- "import Prelude"
-- "import qualified Data.ByteString.Lazy.Char8 as B"
-- "import qualified Data.List as L"
-- [("Prelude",Nothing),("Data.ByteString.Lazy.Char8",Just "B"),("Data.List",Just "L")]
-- ===
-- 
-- >>> testM "tests/preludes/readme/prelude.hs"
-- "{-# LANGUAGE ExtendedDefaultRules, OverloadedStrings #-}"
-- ["ExtendedDefaultRules","OverloadedStrings"]
-- ===
-- Nothing
-- ===
-- "import Prelude"
-- "import qualified Data.ByteString.Lazy.Char8 as B"
-- "import qualified Data.List as L"
-- [("Prelude",Nothing),("Data.ByteString.Lazy.Char8",Just "B"),("Data.List",Just "L")]
-- ===
-- "takeLast n = reverse . take n . reverse"
-- 
-- >>> testM "tests/preludes/moduleName/prelude.hs"
-- []
-- ===
-- "module MyPrelude where"
-- Just "MyPrelude"
-- ===
-- []
-- ===
-- "t = take"
-- 
-- >>> testM "tests/preludes/moduleNamedMain/prelude.hs"
-- []
-- ===
-- "module Main where"
-- Just "Main"
-- ===
-- []
-- ===
-- "t = take"
testM :: FilePath -> IO ()
testM f = do
    m <- runUncertainIO $ readModule f
    putSource (pragmaSource m)
    print (languageExtensions m)
    putStrLn "==="
    putSource (moduleSource m)
    print (moduleName m)
    putStrLn "==="
    putSource (importSource m)
    print (importedModules m)
    putStrLn "==="
    putSource (codeSource m)
  where
    putSource = mapM_ (print . either id B.pack)
-- | The string-typed version of Hawk's command-line arguments.

data HawkOption
    = Apply
    | Map
    | FieldDelimiter
    | RecordDelimiter
    | OutputFieldDelimiter
    | OutputRecordDelimiter
    | Version
    | Help
    | ContextDirectory
  deriving (Show, Eq, Enum, Bounded)

-- | In the order listed by --help.
options :: [HawkOption]
options = enumFrom minBound


delimiter :: OptionType
delimiter = nullable (Setting "delim")

-- | Interpret escape sequences, but don't worry if they're invalid.
-- 
-- >>> parseDelimiter ","
-- ","
-- 
-- >>> parseDelimiter "\\n"
-- "\n"
-- 
-- >>> parseDelimiter "\\t"
-- "\t"
-- 
-- >>> parseDelimiter "\\"
-- "\\"
parseDelimiter :: String -> ByteString
parseDelimiter s = pack $ case reads (printf "\"%s\"" s) of
    [(s', "")] -> s'
    _          -> s

-- | Almost like a string, except escape sequences are interpreted.
consumeDelimiter :: (Functor m, Monad m) => OptionConsumer m ByteString
consumeDelimiter = fmap parseDelimiter . consumeNullable "" consumeString

instance Option HawkOption where
  shortName Apply                 = 'a'
  shortName Map                   = 'm'
  shortName FieldDelimiter        = 'd'
  shortName RecordDelimiter       = 'D'
  shortName OutputFieldDelimiter  = 'o'
  shortName OutputRecordDelimiter = 'O'
  shortName Version               = 'v'
  shortName Help                  = 'h'
  shortName ContextDirectory      = 'c'
  
  longName Apply                 = "apply"
  longName Map                   = "map"
  longName FieldDelimiter        = "field-delimiter"
  longName RecordDelimiter       = "record-delimiter"
  longName OutputFieldDelimiter  = "output-field-delim"
  longName OutputRecordDelimiter = "output-record-delim"
  longName Version               = "version"
  longName Help                  = "help"
  longName ContextDirectory      = "context-directory"
  
  helpMsg Apply                      = ["apply <expr> to the entire table"]
  helpMsg Map                        = ["apply <expr> to each row"]
  helpMsg FieldDelimiter             = ["default whitespace"]
  helpMsg RecordDelimiter            = ["default '\\n'"]
  helpMsg OutputFieldDelimiter       = ["default <field-delim>"]
  helpMsg OutputRecordDelimiter      = ["default <record-delim>"]
  helpMsg Version                    = ["print version and exit"]
  helpMsg Help                       = ["this help"]
  helpMsg ContextDirectory           = ["<ctx-dir> directory, default is"
                                       ,"'~/.hawk'"]
  
  optionType Apply                 = flag
  optionType Map                   = flag
  optionType FieldDelimiter        = delimiter
  optionType RecordDelimiter       = delimiter
  optionType OutputFieldDelimiter  = delimiter
  optionType OutputRecordDelimiter = delimiter
  optionType Version               = flag
  optionType Help                  = flag
  optionType ContextDirectory      = filePath
-- | In which Hawk's command-line arguments are structured into a `HawkSpec`.

-- $setup
-- >>> let testP parser = runUncertainIO . runOptionParserT options parser


-- | (record separator, field separator)
type CommonSeparators = (Separator, Separator)

-- | Extract '-D' and '-d'. We perform this step separately because those two
--   delimiters are used by both the input and output specs.
-- 
-- >>> let test = testP commonSeparators
-- 
-- >>> test []
-- (Delimiter "\n",Whitespace)
-- 
-- >>> test ["-D\\n", "-d\\t"]
-- (Delimiter "\n",Delimiter "\t")
-- 
-- >>> test ["-D|", "-d,"]
-- (Delimiter "|",Delimiter ",")
commonSeparators :: (Functor m, Monad m)
                 => OptionParserT HawkOption m CommonSeparators
commonSeparators = do
    r <- lastSep Option.RecordDelimiter defaultRecordSeparator
    f <- lastSep Option.FieldDelimiter defaultFieldSeparator
    return (r, f)
  where
    lastSep opt def = consumeLast opt def consumeSep
    consumeSep = fmap Delimiter . Option.consumeDelimiter


-- | The input delimiters have already been parsed, but we still need to
--   interpret them and to determine the input source.
-- 
-- >>> :{
-- let test = testP $ do { c <- commonSeparators
--                       ; _ <- consumeExtra consumeString  -- skip expr
--                       ; i <- inputSpec c
--                       ; lift $ print $ inputSource i
--                       ; lift $ print $ inputFormat i
--                       }
-- :}
-- 
-- >>> test []
-- UseStdin
-- Records (Delimiter "\n") (Fields Whitespace)
-- 
-- >>> test ["-d", "-a", "L.reverse"]
-- UseStdin
-- Records (Delimiter "\n") RawRecord
-- 
-- >>> test ["-D", "-a", "B.reverse"]
-- UseStdin
-- RawStream
-- 
-- >>> test ["-d:", "-m", "L.head", "/etc/passwd"]
-- InputFile "/etc/passwd"
-- Records (Delimiter "\n") (Fields (Delimiter ":"))
inputSpec :: (Functor m, Monad m)
          => CommonSeparators -> OptionParserT HawkOption m InputSpec
inputSpec (r, f) = InputSpec <$> source <*> format
  where
    source = do
        r <- consumeExtra consumeString
        return $ case r of
          Nothing -> UseStdin
          Just f  -> InputFile f
    format = return streamFormat
    streamFormat | r == Delimiter "" = RawStream
                 | otherwise         = Records r recordFormat
    recordFormat | f == Delimiter ""   = RawRecord
                 | otherwise           = Fields f

-- | The output delimiters take priority over the input delimiters, regardless
--   of the order in which they appear.
-- 
-- >>> :{
-- let test = testP $ do { c <- commonSeparators
--                       ; o <- outputSpec c
--                       ; let OutputFormat r f = outputFormat o
--                       ; lift $ print $ outputSink o
--                       ; lift $ print (r, f)
--                       }
-- :}
-- 
-- >>> test []
-- UseStdout
-- ("\n"," ")
-- 
-- >>> test ["-D;", "-d", "-a", "L.reverse"]
-- UseStdout
-- (";","")
-- 
-- >>> test ["-o\t", "-d,", "-O|"]
-- UseStdout
-- ("|","\t")
outputSpec :: (Functor m, Monad m)
           => CommonSeparators -> OptionParserT HawkOption m OutputSpec
outputSpec (r, f) = OutputSpec <$> sink <*> format
  where
    sink = return UseStdout
    format = OutputFormat <$> record <*> field
    record = consumeLast Option.OutputRecordDelimiter r' Option.consumeDelimiter
    field = consumeLast Option.OutputFieldDelimiter f' Option.consumeDelimiter
    r' = fromSeparator r
    f' = fromSeparator f


-- | The information we need in order to evaluate a user expression:
--   the expression itself, and the context in which it should be evaluated.
--   In Hawk, that context is the user prelude.
-- 
-- >>> :{
-- let test = testP $ do { e <- exprSpec
--                       ; lift $ print $ userExpression e
--                       ; lift $ print $ userContextDirectory e
--                       }
-- :}
-- 
-- >>> test []
-- error: missing user expression
-- *** Exception: ExitFailure 1
-- 
-- >>> test [""]
-- error: user expression cannot be empty
-- *** Exception: ExitFailure 1
--
-- >>> test ["-D;", "-d", "-a", "L.reverse","-c","somedir"]
-- "L.reverse"
-- "somedir"
exprSpec :: (Functor m, MonadIO m)
         => OptionParserT HawkOption m ExprSpec
exprSpec = ExprSpec <$> contextDir <*> expr
  where
    contextDir = do
      dir <- consumeLast Option.ContextDirectory "" consumeString
      if null dir
        then liftIO findContextFromCurrDirOrDefault
        else return dir
    expr = do
        r <- consumeExtra consumeString
        case r of
          Just e  -> if all isSpace e
                      then fail "user expression cannot be empty"
                      else return e
          Nothing -> fail "missing user expression"


-- | Parse command-line arguments to construct a `HawkSpec`.
-- 
-- TODO: complain if some arguments are unused (except perhaps "-d" and "-D").
-- 
-- >>> :{
-- let test args = do { spec <- runUncertainIO $ parseArgs args
--                    ; case spec of
--                        Help        -> putStrLn "Help"
--                        Version     -> putStrLn "Version"
--                        Eval  e   o -> putStrLn "Eval"  >> print (userExpression e)                                         >> print (recordDelimiter (outputFormat o), fieldDelimiter (outputFormat o))
--                        Apply e i o -> putStrLn "Apply" >> print (userExpression e, inputSource i) >> print (inputFormat i) >> print (recordDelimiter (outputFormat o), fieldDelimiter (outputFormat o))
--                        Map   e i o -> putStrLn "Map"   >> print (userExpression e, inputSource i) >> print (inputFormat i) >> print (recordDelimiter (outputFormat o), fieldDelimiter (outputFormat o))
--                    }
-- :}
-- 
-- >>> test []
-- Help
-- 
-- >>> test ["--help"]
-- Help
-- 
-- >>> test ["--version"]
-- Version
-- 
-- >>> test ["-d\\t", "L.head"]
-- Eval
-- "L.head"
-- ("\n","\t")
-- 
-- >>> test ["-D\r\n", "-d\\t", "-m", "L.head"]
-- Map
-- ("L.head",UseStdin)
-- Records (Delimiter "\r\n") (Fields (Delimiter "\t"))
-- ("\r\n","\t")
-- 
-- >>> test ["-D", "-O\n", "-m", "L.head", "file.in"]
-- Map
-- ("L.head",InputFile "file.in")
-- RawStream
-- ("\n"," ")
parseArgs :: (Functor m,MonadIO m) => [String] -> UncertainT m HawkSpec
parseArgs [] = return Help
parseArgs args = runOptionParserT options parser args
  where
    parser = do
        lift $ return ()  -- silence a warning
        cmd <- consumeExclusive assoc eval
        c <- commonSeparators
        cmd c
    assoc = [ (Option.Help,    help)
            , (Option.Version, version)
            , (Option.Apply,   apply)
            , (Option.Map,     map')
            ]
    
    help, version, eval, apply, map' :: (Functor m,MonadIO m) => CommonSeparators
                                     -> OptionParserT HawkOption m HawkSpec
    help    _ = return Help
    version _ = return Version
    eval    c = Eval  <$> exprSpec <*>                 outputSpec c
    apply   c = Apply <$> exprSpec <*> inputSpec c <*> outputSpec c
    map'    c = Map   <$> exprSpec <*> inputSpec c <*> outputSpec c
-- | In which the user prelude is massaged into the form hint needs.



type UserPrelude = HaskellModule


testC :: FilePath -> IO ()
testC f = do
    let orig = printf "tests/preludes/%s/prelude.hs" f
    m <- runUncertainIO $ readModule orig
    B.putStr $ showModule orig (canonicalizeUserPrelude m)

-- |
-- >>> testC "default"
-- {-# LANGUAGE ExtendedDefaultRules, OverloadedStrings #-}
-- module System.Console.Hawk.CachedPrelude where
-- {-# LINE 2 "tests/preludes/default/prelude.hs" #-}
-- import Prelude
-- import qualified Data.ByteString.Lazy.Char8 as B
-- import qualified Data.List as L
-- 
-- >>> testC "moduleName"
-- module MyPrelude where
-- import Prelude
-- {-# LINE 2 "tests/preludes/moduleName/prelude.hs" #-}
-- t = take
canonicalizeUserPrelude :: HaskellModule -> UserPrelude
canonicalizeUserPrelude = extendModuleName . extendImports

readUserPrelude :: FilePath -> UncertainT IO UserPrelude
readUserPrelude f = canonicalizeUserPrelude <$> readModule f


compileUserPrelude :: FilePath -- ^ the original's filename,
                               --   used for fixing up line numbers
                   -> FilePath -- ^ new filename, because ghc compiles from disk.
                               --   the compiled output will be in the same folder.
                   -> UserPrelude
                   -> UncertainT IO ()
compileUserPrelude = compileUserPreludeWithArgs []

compileUserPreludeWithArgs :: [String] -- ^ extra ghc args
                           -> FilePath -- ^ the original's filename,
                                       --   used for fixing up line numbers
                           -> FilePath -- ^ new filename, because ghc compiles from disk.
                                       --   the compiled output will be in the same folder.
                           -> UserPrelude
                           -> UncertainT IO ()
compileUserPreludeWithArgs args orig f m = do
    extraArgs <- lift $ extraGhcArgs
    let args' = (extraArgs ++ args)
    compileModuleWithArgs args' orig f m
-- | In which the implicit defaults are explicitly added.


-- | We cannot import a module unless it has a name.
extendModuleName :: HaskellModule -> HaskellModule
extendModuleName = until hasModuleName
                       $ addDefaultModuleName defaultModuleName
  where
    hasModuleName = isJust . moduleName


moduleNames :: HaskellModule -> [String]
moduleNames = Data.List.map fst . importedModules

-- | GHC imports the Haskell Prelude by default, but hint doesn't.
-- 
-- >>> let m name = (name, Nothing)
-- >>> :{
--   let testM exts modules = moduleNames m'
--         where
--           m0  = emptyModule
--           m1  = foldr addExtension m0 exts
--           m2  = foldr addImport m1 modules
--           m' = extendImports m2
-- :}
-- 
-- >>> testM [] []
-- ["Prelude"]
-- 
-- >>> testM [] [m "Data.Maybe"]
-- ["Prelude","Data.Maybe"]
-- 
-- >>> testM [] [m "Data.Maybe", m "Prelude", m "Data.Either"]
-- ["Data.Maybe","Prelude","Data.Either"]
-- 
-- >>> :{
-- testM [] [ ("Data.Maybe", Just "M")
--          , ("Prelude", Just "P")
--          , ("Data.Either", Just "E")
--          ]
-- :}
-- ["Data.Maybe","Prelude","Data.Either"]
-- 
-- >>> :{
-- testM ["OverloadedStrings","NoImplicitPrelude"]
--       [m "Data.Maybe"]
-- :}
-- ["Data.Maybe"]
extendImports :: HaskellModule -> HaskellModule
extendImports = until preludeOk
                    $ addImport unqualified_prelude
  where
    prelude = "Prelude"
    noPrelude = "NoImplicitPrelude"
    unqualified_prelude = (prelude, Nothing)
    
    preludeOk = liftA2 (||) hasPrelude noImplicitPrelude
    hasPrelude        m =   prelude `Data.List.elem` moduleNames m
    noImplicitPrelude m = noPrelude `Data.List.elem` languageExtensions m
-- | A computation which may raise warnings or fail in error.


type Warning = String
type Error = String

newtype UncertainT m a = UncertainT
  { unUncertainT :: ErrorT Error (WriterT [Warning] m) a }

type Uncertain a = UncertainT Identity a

instance Functor m => Functor (UncertainT m) where
  fmap f = UncertainT . fmap f . unUncertainT

instance (Functor m, Monad m) => Applicative (UncertainT m) where
  pure = UncertainT . pure
  UncertainT mf <*> UncertainT mx = UncertainT (mf <*> mx)

instance Monad m => Monad (UncertainT m) where
  return = UncertainT . return
  UncertainT mx >>= f = UncertainT (mx >>= f')
    where
      f' = unUncertainT . f
  fail s = UncertainT (fail s)

instance MonadTrans UncertainT where
  lift = UncertainT . lift . lift

instance MonadIO m => MonadIO (UncertainT m) where
  liftIO = lift . liftIO


warn :: Monad m => String -> UncertainT m ()
warn s = UncertainT $ lift $ tell [s]

fromRightM :: Monad m => Either String a -> UncertainT m a
fromRightM (Left e)  = fail e
fromRightM (Right x) = return x


multilineMsg :: String -> String
multilineMsg = Data.List.concat . Data.List.map (printf "\n  %s") . lines

-- | Indent a multiline warning message.
-- >>> :{
-- runUncertainIO $ do
--   multilineWarn "foo\nbar\n"
--   return 42
-- :}
-- warning: 
--   foo
--   bar
-- 42
multilineWarn :: Monad m => String -> UncertainT m ()
multilineWarn = warn . multilineMsg

-- | Indent a multiline error message.
-- >>> :{
-- runUncertainIO $ do
--   multilineFail "foo\nbar\n"
--   return 42
-- :}
-- error: 
--   foo
--   bar
-- *** Exception: ExitFailure 1
multilineFail :: Monad m => String -> UncertainT m a
multilineFail = fail . multilineMsg


mapUncertainT :: (forall a. m a -> m' a) -> UncertainT m b -> UncertainT m' b
mapUncertainT f = UncertainT . (mapErrorT . mapWriterT) f . unUncertainT

runUncertainT :: UncertainT m a -> m (Either Error a, [Warning])
runUncertainT = runWriterT . runErrorT . unUncertainT


-- | A version of `runWarnings` which allows you to interleave IO actions
--   with uncertain actions.
-- 
-- Note that the warnings are displayed after the IO's output.
-- 
-- >>> :{
-- runWarningsIO $ do
--   warn "before"
--   lift $ putStrLn "IO"
--   warn "after"
--   return 42
-- :}
-- IO
-- warning: before
-- warning: after
-- Right 42
-- 
-- >>> :{
-- runWarningsIO $ do
--   warn "before"
--   lift $ putStrLn "IO"
--   fail "fatal"
--   return 42
-- :}
-- IO
-- warning: before
-- Left "fatal"
runWarningsIO :: UncertainT IO a -> IO (Either String a)
runWarningsIO u = do
    (r, warnings) <- runUncertainT u
    mapM_ (System.IO.hPutStrLn stderr . printf "warning: %s") warnings
    return r

-- | A version of `runUncertain` which only prints the warnings, not the
--   errors. Unlike `runUncertain`, it doesn't terminate on error.
-- 
-- >>> :{
-- runWarnings $ do
--   warn "before"
--   warn "after"
--   return 42
-- :}
-- warning: before
-- warning: after
-- Right 42
-- 
-- >>> :{
-- runWarnings $ do
--   warn "before"
--   fail "fatal"
--   return 42
-- :}
-- warning: before
-- Left "fatal"
runWarnings :: Uncertain a -> IO (Either String a)
runWarnings = runWarningsIO . mapUncertainT (return . runIdentity)


-- | A version of `runUncertain` which allows you to interleave IO actions
--   with uncertain actions.
-- 
-- Note that the warnings are displayed after the IO's output.
-- 
-- >>> :{
-- runUncertainIO $ do
--   warn "before"
--   lift $ putStrLn "IO"
--   warn "after"
--   return 42
-- :}
-- IO
-- warning: before
-- warning: after
-- 42
-- 
-- >>> :{
-- runUncertainIO $ do
--   warn "before"
--   lift $ putStrLn "IO"
--   fail "fatal"
--   return 42
-- :}
-- IO
-- warning: before
-- error: fatal
-- *** Exception: ExitFailure 1
runUncertainIO :: UncertainT IO a -> IO a
runUncertainIO u = do
    r <- runWarningsIO u
    case r of
      Left e -> do
        System.IO.hPutStrLn stderr $ printf "error: %s" e
        exitFailure
      Right x -> return x

-- | Print warnings and errors, terminating on error.
-- 
-- Note that the warnings are displayed even if there is also an error.
-- 
-- >>> :{
-- runUncertainIO $ do
--   warn "first"
--   warn "second"
--   fail "fatal"
--   return 42
-- :}
-- warning: first
-- warning: second
-- error: fatal
-- *** Exception: ExitFailure 1
runUncertain :: Uncertain a -> IO a
runUncertain = runUncertainIO . mapUncertainT (return . runIdentity)


-- | Upgrade an `IO a -> IO a` wrapping function into a variant which uses
--   `UncertainT IO` instead of `IO`.
-- 
-- >>> :{
-- let wrap body = do { putStrLn "before"
--                    ; r <- body
--                    ; putStrLn "after"
--                    ; return r
--                    }
-- :}
-- 
-- >>> :{
-- wrap $ do { putStrLn "hello"
--           ; return 42
--           }
-- :}
-- before
-- hello
-- after
-- 42
-- 
-- >>> :{
-- runUncertainIO $ wrapUncertain wrap
--                $ do { lift $ putStrLn "hello"
--                     ; warn "be careful!"
--                     ; return 42
--                     }
-- :}
-- before
-- hello
-- after
-- warning: be careful!
-- 42
wrapUncertain :: (Monad m, Monad m')
              => (forall a. m a -> m' a)
              -> (UncertainT m b -> UncertainT m' b)
wrapUncertain wrap body = wrapUncertainArg wrap' body'
  where
    wrap' f = wrap $ f ()
    body' () = body

-- | A version of `wrapUncertain` for wrapping functions of type
--   `(Handle -> IO a) -> IO a`.
-- 
-- >>> :{
-- let wrap body = do { putStrLn "before"
--                    ; r <- body 42
--                    ; putStrLn "after"
--                    ; return r
--                    }
-- :}
-- 
-- >>> :{
-- wrap $ \x -> do { putStrLn "hello"
--                 ; return (x + 1)
--                 }
-- :}
-- before
-- hello
-- after
-- 43
-- 
-- >>> :{
-- runUncertainIO $ wrapUncertainArg wrap
--                $ \x -> do { lift $ putStrLn "hello"
--                           ; warn "be careful!"
--                           ; return (x + 1)
--                           }
-- :}
-- before
-- hello
-- after
-- warning: be careful!
-- 43
wrapUncertainArg :: (Monad m, Monad m')
                 => (forall a. (v -> m a) -> m' a)
                 -> ((v -> UncertainT m b) -> UncertainT m' b)
wrapUncertainArg wrap body = do
    (r, ws) <- lift $ wrap $ runUncertainT . body
    
    -- repackage the warnings and errors
    mapM_ warn ws
    fromRightM r