version 0.2.0.3HEAD0.2.0.3master

16 files changed, 571 insertions, 104 deletions

diff --git a/src/HStyle.hs b/src/HStyle.hs
index 95762a9..abd70ef 100644
--- a/src/HStyle.hs
+++ b/src/HStyle.hs
@@ -1,12 +1,11 @@
-{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE OverloadedStrings, Rank2Types #-}
 module HStyle
-    ( checkStyle
-    , fixStyle
+    ( FileState (..)
+    , Options (..)
+    , checkStyle
     ) where
 
-import Control.Applicative ((<$>))
-import Control.Monad (forM, forM_)
-import Data.Char (isSpace)
+import Control.Monad (foldM, when)
 import Data.List (isPrefixOf)
 import Data.Maybe (fromMaybe)
 
@@ -14,62 +13,16 @@ import qualified Data.Text as T
 import qualified Data.Text.IO as T
 import qualified Language.Haskell.Exts.Annotated as H
 
-import HStyle.Alignment
 import HStyle.Block
-import HStyle.Selector
-import HStyle.Checker
-
--- | A selector and a check...
-type Rule = (Selector, Checker)
-
-runRule :: FilePath -> Block -> H.Module H.SrcSpanInfo -> Rule -> IO Bool
-runRule file block md (selector, check) = fmap and $
-    forM (selector md block) $ \block' -> do
-        let problems = check block'
-        forM_ problems $ \(i, problem) -> do
-            let line = absoluteLineNumber i block'
-            T.putStrLn $ T.pack file `T.append` ":" `T.append`
-                T.pack (show line) `T.append` ": " `T.append` problem
-            T.putStrLn ""
-            T.putStr   $ prettyBlock 4 block'
-            T.putStrLn ""
-        return $ null problems
-
-fromSrcSpanInfo :: H.SrcSpanInfo -> Block -> Block
-fromSrcSpanInfo ssi = subBlock start end
-  where
-    span' = H.srcInfoSpan ssi
-    start = H.srcSpanStartLine span'
-    end   = H.srcSpanEndLine span'
-
-typeSigSelector :: Selector
-typeSigSelector md block = map (flip fromSrcSpanInfo block) $ tss md
-  where
-    tss (H.Module _ _ _ _ decls) = [ssi | H.TypeSig ssi _ _ <- decls]
-    tss _                        = []
-
-typeSigCheck :: Checker
-typeSigCheck block = case checkAlignmentHead alignment of
-    Just t  -> [(1, t)]
-    Nothing -> []
-  where
-    alignment = alignmentOf ["::", "=>", "->"] $ toLines block
-
-tabsCheck :: Checker
-tabsCheck = checkLines $ \line -> case T.findIndex (== '\t') line of
-    Nothing -> Nothing
-    Just i  -> Just $ "\\t at column " `T.append` T.pack (show $ i + 1)
-
-lineLengthCheck :: Int -> Checker
-lineLengthCheck max' = checkLines $ \line -> if T.length line > max'
-    then Just $ "exceeds max line length of " `T.append` T.pack (show max')
-    else Nothing
-
-trailingWhiteSpace :: Checker
-trailingWhiteSpace = checkLines $ \line ->
-    if not (T.null line) && isSpace (T.last line)
-        then Just "trailing whitespace"
-        else Nothing
+import HStyle.Rule
+import HStyle.Rules.AppSpacing
+import HStyle.Rules.CaseAlignment
+import HStyle.Rules.DataAlignment
+import HStyle.Rules.EolComment
+import HStyle.Rules.LineLength
+import HStyle.Rules.Tabs
+import HStyle.Rules.TrailingWhiteSpace
+import HStyle.Rules.TypeSigAlignment
 
 -- | Filter out lines which use CPP macros
 unCPP :: String -> String
@@ -79,27 +32,33 @@ unCPP = unlines . map unCpp' . lines
         | "#" `isPrefixOf` x = ""
         | otherwise          = x
 
-checkStyle :: FilePath -> IO Bool
-checkStyle file = do
+checkStyle :: Options -> FilePath -> IO FileState
+checkStyle options file = do
     contents <- readFile file
-    let block      = fromText $ T.pack contents
+    let block     = fromText $ T.pack contents
         -- Determine the extensions used in the file, and update the parsing
         -- mode based upon those
         exts      = fromMaybe [] $ H.readExtensions contents
-        mode      = H.defaultParseMode {H.extensions = exts}
+        mode      = H.defaultParseMode
+            {H.extensions = exts, H.fixities = Nothing}
         -- Special handling for CPP, haskell-src-exts can't deal with it
         contents' = if H.CPP `elem` exts then unCPP contents else contents
-    case H.parseModuleWithMode mode contents' of
-        H.ParseOk md -> and <$> mapM (runRule file block md)
-            [ (typeSigSelector, typeSigCheck)
-            , (selectLines, tabsCheck)
-            , (selectLines, lineLengthCheck 78)
-            , (selectLines, trailingWhiteSpace)
-            ]
+        fs        = FileState block False True
+    case H.parseModuleWithComments mode contents' of
+        H.ParseOk x -> do
+            fs' <- foldM (runRule options file x) fs
+                [ typeSigAlignmentRule
+                , tabsRule 4
+                , lineLengthRule 78
+                , trailingWhiteSpaceRule
+                , eolCommentRule
+                , dataAlignmentRule
+                , appSpacingRule
+                , caseAlignmentRule
+                ]
+            when (fileUpdated fs') $ T.writeFile file $ toText $ fileBlock fs'
+            return fs'
         err         -> do
             putStrLn $ "HStyle.checkStyle: could not parse " ++
                 file ++ ": " ++ show err
-            return False
-
-fixStyle :: FilePath -> IO ()
-fixStyle = error "HStyle.fixStyle: Not implemented"
+            return fs
diff --git a/src/HStyle/Alignment.hs b/src/HStyle/Alignment.hs
index dd81382..7c8fdb1 100644
--- a/src/HStyle/Alignment.hs
+++ b/src/HStyle/Alignment.hs
@@ -12,15 +12,15 @@ type Alignment = [[(Int, Text)]]
 checkAlignmentHead :: Alignment
                   -> Maybe Text
 checkAlignmentHead alignment
-    | null alignment'       = Nothing
+    | null alignment'       = Nothing -- Isn't this comment to close?
     | equal (map fst heads) = Nothing
-    | otherwise             = Just $ "improper alignment of " `T.append`
+    | otherwise             = Just $ "Improper alignment of "`T.append`
         T.pack (show $ nub $ map snd heads)
   where
-    alignment' = filter (not . null) alignment
-    heads	   = map head alignment'
+    alignment' = filter(not . null) alignment
+    heads      = map head alignment'
 
-equal :: Eq a   
+equal :: Eq a
       => [a]
       -> Bool
 equal (x : y : r)
@@ -34,7 +34,7 @@ alignmentOf xs = map $ alignmentOf' 0
     alignmentOf' i t
         | T.null t  = []
         | otherwise = case find (`T.isPrefixOf` t) xs of
-            Nothing -> alignmentOf' (i + 1) (T.drop 1 t)
-            Just x  ->
+            Nothing -> alignmentOf' (i+1) (T.drop 1 t)
+            Just x ->
                 let len = T.length x
                 in (i, x) : alignmentOf' (i + len) (T.drop len t)
diff --git a/src/HStyle/Block.hs b/src/HStyle/Block.hs
index 2fe38a2..3b75ac6 100644
--- a/src/HStyle/Block.hs
+++ b/src/HStyle/Block.hs
@@ -3,11 +3,14 @@
 module HStyle.Block
     ( Block
     , fromText
+    , toText
     , prettyBlock
     , toLines
     , subBlock
+    , updateSubBlock
     , perLine
     , absoluteLineNumber
+    , mapLines
     ) where
 
 import Data.Text (Text)
@@ -18,7 +21,7 @@ import qualified Data.Text as T
 data Block = Block
     { blockOffset :: Int
     , blockLines  :: Vector Text
-    } deriving (Show)
+    } deriving (Eq, Show)
 
 fromText :: Text -> Block
 fromText text = Block
@@ -26,6 +29,9 @@ fromText text = Block
     , blockLines  = V.fromList $ T.lines text
     }
 
+toText :: Block -> Text
+toText = T.unlines . toLines
+
 prettyBlock :: Int -> Block -> Text
 prettyBlock indent block = T.unlines $
     map ((T.replicate indent " " `T.append`) . pretty) $
@@ -47,9 +53,34 @@ toLines = V.toList . blockLines
 -- | Subblock from start to end -- including both.
 subBlock :: Int -> Int -> Block -> Block
 subBlock start end block = Block
-    { blockOffset = blockOffset block + start - 1
-    , blockLines  = V.slice (start - 1) (end - start + 1) $ blockLines block
+    { blockOffset = blockOffset block + start'
+    , blockLines  = V.slice start' (end' - start') lines'
     }
+  where
+    -- Bounds checking
+    lines' = blockLines block
+    start' = start - 1
+    end'   = min (V.length lines') end
+
+-- | Update a subblock
+updateSubBlock :: Block  -- ^ Old
+               -> Block  -- ^ New
+               -> Block  -- ^ Block to update
+               -> Block  -- ^ Resulting block
+updateSubBlock old new block
+    | blockOffset old /= blockOffset new =
+        error "HStyle.Block.updateSubBlock: Internal error"
+    | otherwise                          = block
+        { blockLines = V.take subOffset lines' V.++ blockLines new V.++
+            V.drop (subOffset + V.length oldLines) lines'
+        }
+  where
+    subOffset
+        | blockOffset old == blockOffset new = blockOffset old
+        | otherwise                          = error
+            "HStyle.Block.updateSubBlock: Internal error"
+    oldLines = blockLines old
+    lines'   = blockLines block
 
 -- | Create a new block for every line.
 perLine :: Block -> [Block]
@@ -62,3 +93,7 @@ perLine (Block offset lines')  = map line $
 -- line number
 absoluteLineNumber :: Int -> Block -> Int
 absoluteLineNumber i = (+ i) . blockOffset
+
+-- | Map over the lines in a block
+mapLines :: (Text -> Text) -> Block -> Block
+mapLines f block = block {blockLines = V.map f (blockLines block)}
diff --git a/src/HStyle/Checker.hs b/src/HStyle/Checker.hs
index 684f15d..77368c1 100644
--- a/src/HStyle/Checker.hs
+++ b/src/HStyle/Checker.hs
@@ -9,12 +9,12 @@ import HStyle.Block
 
 -- | Takes a number of lines, and notifies of problems on each line. Indices
 -- in the result are 1-based.
-type Checker = Block -> [(Int, Text)]
+type Checker a = a -> Block -> [(Int, Text)]
 
 -- | Check every line of the block, possibly returning a problem description
-checkLines :: (Text -> Maybe Text) -> Checker
-checkLines checker block = do
+checkLines :: (a -> Text -> Maybe Text) -> Checker a
+checkLines checker x block = do
     (ln, text) <- zip [1 ..] (toLines block)
-    case checker text of
+    case checker x text of
         Nothing -> []
         Just p  -> [(ln , p)]
diff --git a/src/HStyle/Fixer.hs b/src/HStyle/Fixer.hs
new file mode 100644
index 0000000..176fb16
--- /dev/null
+++ b/src/HStyle/Fixer.hs
@@ -0,0 +1,18 @@
+module HStyle.Fixer
+    ( Fixer
+    , fixNothing
+    , fixLines
+    ) where
+
+import Data.Text (Text)
+
+import HStyle.Block
+
+-- | Takes a block and fixes it, if possible
+type Fixer a = a -> Block -> Maybe Block
+
+fixNothing :: Fixer a
+fixNothing = const $ const Nothing
+
+fixLines :: (a -> Text -> Text) -> Fixer a
+fixLines f x = Just . mapLines (f x)
diff --git a/src/HStyle/Rule.hs b/src/HStyle/Rule.hs
new file mode 100644
index 0000000..1407b21
--- /dev/null
+++ b/src/HStyle/Rule.hs
@@ -0,0 +1,90 @@
+{-# LANGUAGE ExistentialQuantification, OverloadedStrings #-}
+module HStyle.Rule
+    ( Rule (..)
+    , FileState (..)
+    , Options (..)
+    , runRule
+    ) where
+
+import Control.Monad (foldM, forM_, unless)
+
+import qualified Data.Text as T
+import qualified Data.Text.IO as T
+import qualified Language.Haskell.Exts.Annotated as H
+
+import HStyle.Block
+import HStyle.Selector
+import HStyle.Checker
+import HStyle.Fixer
+
+-- | Compose the elements of a rule. Use ExistentialQuantification so the
+-- internal state of a rule cannot be touched from the outside.
+data Rule = forall a. Rule (Selector a) (Checker a) (Fixer a)
+
+data FileState = FileState
+    { -- | A block holding the file contents
+      fileBlock   :: Block
+    , -- | Flag indicating whether or not the in-memory representation differs
+      -- from the file on disk
+      fileUpdated :: Bool
+    , -- | Flag indicating that all checks were OK
+      fileOk      :: Bool
+    } deriving (Show)
+
+-- | Options for checking files
+data Options = Options
+    { -- | Attempt to fix files
+      optionsFix   :: Bool
+    , -- | Be quiet
+      optionsQuiet :: Bool
+    } deriving (Show)
+
+-- | Represents fixing status
+data Fix
+    = DontFix      -- ^ User doesn't want to fix it
+    | CouldntFix   -- ^ Our library is unable to fix it
+    | Fixed        -- ^ Fixed, result
+    deriving (Eq, Show)
+
+runRule :: Options -> FilePath
+        -> (H.Module H.SrcSpanInfo, [H.Comment])
+        -> FileState -> Rule
+        -> IO FileState
+runRule options file mdc fileState (Rule selector checker fixer) =
+    foldM step fileState $ selector mdc $ fileBlock fileState
+  where
+    step fs (x, b) = checkBlock options file checker fixer fs x b
+
+checkBlock :: Options -> FilePath -> Checker a -> Fixer a -> FileState
+           -> a -> Block
+           -> IO FileState
+checkBlock options file checker fixer fs x block = do
+    -- Determine problems, and attempt to fix (lazily)
+    let problems      = checker x block
+        (fix, block') = case (optionsFix options, fixer x block) of
+            (False, _)      -> (DontFix, block)
+            (True, Nothing) -> (CouldntFix, block)
+            (True, Just b)  -> (Fixed, b)
+
+    -- Output our results for this check
+    forM_ problems $ \(i, problem) -> do
+        let line = absoluteLineNumber i block
+        T.putStrLn $ T.pack file `T.append` ":" `T.append`
+            T.pack (show line) `T.append` ": " `T.append` problem
+        unless (optionsQuiet options) $ do
+            T.putStrLn "    Found:"
+            T.putStr   $ prettyBlock 4 block
+            case fix of
+                DontFix    -> return ()
+                CouldntFix -> T.putStrLn "    (Couldn't automatically fix)"
+                Fixed      -> do
+                    T.putStrLn "    Fixed to:"
+                    T.putStr $ prettyBlock 4 block'
+            T.putStrLn ""
+
+    -- Return updated file state
+    return fs
+        { fileBlock   = updateSubBlock block block' (fileBlock fs)
+        , fileUpdated = fileUpdated fs || fix == Fixed
+        , fileOk      = fileOk fs      && null problems
+        }
diff --git a/src/HStyle/Rules/AppSpacing.hs b/src/HStyle/Rules/AppSpacing.hs
new file mode 100644
index 0000000..d9a2529
--- /dev/null
+++ b/src/HStyle/Rules/AppSpacing.hs
@@ -0,0 +1,58 @@
+-- | Checks spacing around function applications
+{-# LANGUAGE OverloadedStrings #-}
+module HStyle.Rules.AppSpacing
+    ( appSpacingRule
+    , appSelector
+    , appSpacingChecker
+    ) where
+
+import Data.Char (isSpace)
+
+import qualified Data.Text as T
+import qualified Language.Haskell.Exts.Annotated as H
+
+import HStyle.Block
+import HStyle.Checker
+import HStyle.Fixer
+import HStyle.Rule
+import HStyle.Selector
+
+exps :: H.Module H.SrcSpanInfo -> [H.Exp H.SrcSpanInfo]
+exps = everything
+
+data AppInfo
+    -- | Application of the form @x + y@
+    = Infix Snippet Snippet Snippet
+    -- | Application of the form @f x@
+    | Prefix Snippet Snippet
+    deriving (Show)
+
+appSpacingRule :: Rule
+appSpacingRule = Rule appSelector appSpacingChecker fixNothing
+
+appSelector :: Selector AppInfo
+appSelector (md, _) block =
+    [ (Infix (ts e1) (ts o) (ts e2), fromSrcSpanInfo l block)
+    | H.InfixApp l e1 o e2 <- exps md
+    ] ++
+    [ (Prefix (ts f) (ts x), fromSrcSpanInfo l block)
+    | H.App l f x <- exps md
+    ]
+  where
+    ts :: H.Annotated f => f H.SrcSpanInfo -> Snippet
+    ts = flip fromSrcSpanInfoSnippet block . H.ann
+
+appSpacingChecker :: Checker AppInfo
+appSpacingChecker app _ = case app of
+    (Infix _ so se2)
+        | spaceBefore so && spaceBefore se2 -> []
+        | otherwise                         ->
+            [(1, "Need spacing around " `T.append` snippetText so)]
+    (Prefix _ x)
+        | spaceBefore x -> []
+        | otherwise     ->
+            [(1, "Need spacing before " `T.append` snippetText x)]
+  where
+    spaceBefore (Snippet b _ c)
+        | c < 1     = True
+        | otherwise = isSpace $ T.index (toText b) (c - 2)
diff --git a/src/HStyle/Rules/CaseAlignment.hs b/src/HStyle/Rules/CaseAlignment.hs
new file mode 100644
index 0000000..a79b653
--- /dev/null
+++ b/src/HStyle/Rules/CaseAlignment.hs
@@ -0,0 +1,44 @@
+{-# LANGUAGE OverloadedStrings #-}
+module HStyle.Rules.CaseAlignment
+    ( caseAlignmentRule
+    , caseSelector
+    , caseAlignmentChecker
+    ) where
+
+import qualified Language.Haskell.Exts.Annotated as H
+
+import HStyle.Alignment
+import HStyle.Checker
+import HStyle.Fixer
+import HStyle.Rule
+import HStyle.Selector
+
+caseAlignmentRule :: Rule
+caseAlignmentRule = Rule caseSelector caseAlignmentChecker fixNothing
+
+caseSelector :: Selector [Snippet]
+caseSelector (md, _) block = do
+    -- Select a case statement
+    (l, alts) <- [(l, alts) | H.Case l _ alts <- exps]
+
+    -- Select all alternatives
+    let ls = [ gas
+             | H.Alt _ _ ga _ <- alts
+             , gas <- case ga of
+                        H.UnGuardedAlt l' _ -> return l'
+                        H.GuardedAlts _ ga' -> map H.ann ga'
+             ]
+
+    return (map snippet ls, fromSrcSpanInfo l block)
+  where
+    exps :: [H.Exp H.SrcSpanInfo]
+    exps = everything md
+    snippet :: H.SrcSpanInfo -> Snippet
+    snippet = flip fromSrcSpanInfoSnippet block
+
+caseAlignmentChecker :: Checker [Snippet]
+caseAlignmentChecker snippets _ = case checkAlignmentHead alignment of
+    Nothing -> []
+    Just t  -> [(1, t)]
+  where
+    alignment = [[(c, "->")] | Snippet _ _ c <- snippets]
diff --git a/src/HStyle/Rules/DataAlignment.hs b/src/HStyle/Rules/DataAlignment.hs
new file mode 100644
index 0000000..471317b
--- /dev/null
+++ b/src/HStyle/Rules/DataAlignment.hs
@@ -0,0 +1,34 @@
+{-# LANGUAGE OverloadedStrings #-}
+module HStyle.Rules.DataAlignment
+    ( dataAlignmentRule
+    , dataSelector
+    , dataAlignmentChecker
+    ) where
+
+import qualified Language.Haskell.Exts.Annotated as H
+
+import HStyle.Alignment
+import HStyle.Block
+import HStyle.Checker
+import HStyle.Fixer
+import HStyle.Rule
+import HStyle.Selector
+
+dataAlignmentRule :: Rule
+dataAlignmentRule = Rule dataSelector dataAlignmentChecker fixNothing
+
+dataSelector :: Selector ()
+dataSelector (md, _) block =
+    map (\ssi -> ((), fromSrcSpanInfo ssi block)) $ datas =<< everything md
+  where
+    datas :: H.Decl H.SrcSpanInfo -> [H.SrcSpanInfo]
+    datas decl = case decl of
+        d@(H.DataDecl _ _ _ _ _ _) -> [H.ann d]
+        _                          -> []
+
+dataAlignmentChecker :: Checker ()
+dataAlignmentChecker () block = case checkAlignmentHead alignment of
+    Just t  -> [(1, t)]
+    Nothing -> []
+  where
+    alignment = alignmentOf ["{", ",", "}"] $ toLines block
diff --git a/src/HStyle/Rules/EolComment.hs b/src/HStyle/Rules/EolComment.hs
new file mode 100644
index 0000000..3c009e4
--- /dev/null
+++ b/src/HStyle/Rules/EolComment.hs
@@ -0,0 +1,36 @@
+{-# LANGUAGE OverloadedStrings #-}
+module HStyle.Rules.EolComment
+    ( eolCommentRule
+    ) where
+
+import Control.Monad (guard)
+import Data.Char (isSpace)
+
+import qualified Data.Text as T
+import qualified Language.Haskell.Exts.Annotated as H
+
+import HStyle.Block
+import HStyle.Checker
+import HStyle.Fixer
+import HStyle.Rule
+import HStyle.Selector
+
+eolCommentRule :: Rule
+eolCommentRule = Rule eolCommentSelector eolCommentChecker fixNothing
+
+eolCommentSelector :: Selector Int
+eolCommentSelector (_, comments) block = do
+    H.Comment False ss _ <- comments
+    let start  = H.srcSpanStartLine ss
+        col    = H.srcSpanStartColumn ss
+        end    = H.srcSpanEndLine ss
+        block' = subBlock start end block
+    guard $ start == end && col > 2
+    -- Remember the start column of the comment
+    return (H.srcSpanStartColumn ss, block')
+
+eolCommentChecker :: Checker Int
+eolCommentChecker = checkLines $ \c t ->
+    if c > 2 && isSpace (T.index t (c - 3)) && isSpace (T.index t (c - 2))
+        then Nothing
+        else Just "Need two spaces between code and comment"
diff --git a/src/HStyle/Rules/LineLength.hs b/src/HStyle/Rules/LineLength.hs
new file mode 100644
index 0000000..5f793d4
--- /dev/null
+++ b/src/HStyle/Rules/LineLength.hs
@@ -0,0 +1,20 @@
+{-# LANGUAGE OverloadedStrings #-}
+module HStyle.Rules.LineLength
+    ( lineLengthRule
+    , lineLengthChecker
+    ) where
+
+import qualified Data.Text as T
+
+import HStyle.Checker
+import HStyle.Fixer
+import HStyle.Rule
+import HStyle.Selector
+
+lineLengthRule :: Int -> Rule
+lineLengthRule max' = Rule selectLines (lineLengthChecker max') fixNothing
+
+lineLengthChecker :: Int -> Checker ()
+lineLengthChecker max' = checkLines $ \() line -> if T.length line > max'
+    then Just $ "Exceeds max line length of " `T.append` T.pack (show max')
+    else Nothing
diff --git a/src/HStyle/Rules/Tabs.hs b/src/HStyle/Rules/Tabs.hs
new file mode 100644
index 0000000..7a24b90
--- /dev/null
+++ b/src/HStyle/Rules/Tabs.hs
@@ -0,0 +1,28 @@
+-- | Check for tabs in files
+{-# LANGUAGE OverloadedStrings #-}
+module HStyle.Rules.Tabs
+    ( tabsRule
+    , tabsChecker
+    , tabsFixer
+    ) where
+
+import qualified Data.Text as T
+
+import HStyle.Checker
+import HStyle.Fixer
+import HStyle.Rule
+import HStyle.Selector
+
+tabsRule :: Int -> Rule
+tabsRule i = Rule selectLines tabsChecker (tabsFixer i)
+
+tabsChecker :: Checker ()
+tabsChecker = checkLines $ \() line -> case T.findIndex (== '\t') line of
+    Nothing -> Nothing
+    Just i  -> Just $ "\\t at column " `T.append` T.pack (show $ i + 1)
+
+tabsFixer :: Int -> Fixer ()
+tabsFixer numSpaces = fixLines fixer
+  where
+    spaces   = T.replicate numSpaces " "
+    fixer () = T.intercalate spaces . T.split (== '\t')
diff --git a/src/HStyle/Rules/TrailingWhiteSpace.hs b/src/HStyle/Rules/TrailingWhiteSpace.hs
new file mode 100644
index 0000000..15babd8
--- /dev/null
+++ b/src/HStyle/Rules/TrailingWhiteSpace.hs
@@ -0,0 +1,28 @@
+{-# LANGUAGE OverloadedStrings #-}
+module HStyle.Rules.TrailingWhiteSpace
+    ( trailingWhiteSpaceRule
+    , trailingWhiteSpaceChecker
+    , trailingWhiteSpaceFixer
+    ) where
+
+import Data.Char (isSpace)
+
+import qualified Data.Text as T
+
+import HStyle.Checker
+import HStyle.Fixer
+import HStyle.Rule
+import HStyle.Selector
+
+trailingWhiteSpaceRule :: Rule
+trailingWhiteSpaceRule = Rule
+    selectLines trailingWhiteSpaceChecker trailingWhiteSpaceFixer
+
+trailingWhiteSpaceChecker :: Checker ()
+trailingWhiteSpaceChecker = checkLines $ \() line ->
+    if not (T.null line) && isSpace (T.last line)
+        then Just "Trailing whitespace"
+        else Nothing
+
+trailingWhiteSpaceFixer :: Fixer ()
+trailingWhiteSpaceFixer = fixLines $ const T.stripEnd
diff --git a/src/HStyle/Rules/TypeSigAlignment.hs b/src/HStyle/Rules/TypeSigAlignment.hs
new file mode 100644
index 0000000..1fc3b39
--- /dev/null
+++ b/src/HStyle/Rules/TypeSigAlignment.hs
@@ -0,0 +1,32 @@
+{-# LANGUAGE OverloadedStrings #-}
+module HStyle.Rules.TypeSigAlignment
+    ( typeSigAlignmentRule
+    , typeSigSelector
+    , typeSigAlignmentChecker
+    ) where
+
+import qualified Language.Haskell.Exts.Annotated as H
+
+import HStyle.Alignment
+import HStyle.Block
+import HStyle.Checker
+import HStyle.Fixer
+import HStyle.Rule
+import HStyle.Selector
+
+typeSigAlignmentRule :: Rule
+typeSigAlignmentRule = Rule typeSigSelector typeSigAlignmentChecker fixNothing
+
+typeSigSelector :: Selector ()
+typeSigSelector (md, _) block =
+    map (\ssi -> ((), fromSrcSpanInfo ssi block)) $ tss md
+  where
+    tss (H.Module _ _ _ _ decls) = [ssi | H.TypeSig ssi _ _ <- decls]
+    tss _                        = []
+
+typeSigAlignmentChecker :: Checker ()
+typeSigAlignmentChecker () block = case checkAlignmentHead alignment of
+    Just t  -> [(1, t)]
+    Nothing -> []
+  where
+    alignment = alignmentOf ["::", "=>", "->"] $ toLines block
diff --git a/src/HStyle/Selector.hs b/src/HStyle/Selector.hs
index 361ab4e..6a3dbf9 100644
--- a/src/HStyle/Selector.hs
+++ b/src/HStyle/Selector.hs
@@ -1,18 +1,60 @@
 module HStyle.Selector
     ( Selector
+    , Snippet (..)
     , selectAll
     , selectLines
+    , fromSrcSpanInfo
+    , fromSrcSpanInfoSnippet
+    , everything
     ) where
 
-import Language.Haskell.Exts.Annotated (Module, SrcSpanInfo)
+import Data.Maybe (maybeToList)
+
+import Data.Data (Data)
+import Data.Text (Text)
+import Data.Typeable (cast)
+import qualified Data.Generics as G
+import qualified Data.Text as T
+import qualified Language.Haskell.Exts.Annotated as H
 
 import HStyle.Block
 
 -- | Selects a portion from a haskell module
-type Selector = Module SrcSpanInfo -> Block -> [Block]
+type Selector a =
+    (H.Module H.SrcSpanInfo, [H.Comment]) -> Block -> [(a, Block)]
+
+data Snippet = Snippet
+    { snippetBlock :: Block
+    , snippetText  :: Text
+    , snippetCol   :: Int
+    } deriving (Show)
+
+selectAll :: Selector ()
+selectAll _ block = [((), block)]
+
+selectLines :: Selector ()
+selectLines _ block = [((), b) | b <- perLine block]
+
+fromSrcSpanInfo :: H.SrcSpanInfo -> Block -> Block
+fromSrcSpanInfo ssi = snippetBlock . fromSrcSpanInfoSnippet ssi
 
-selectAll :: Selector
-selectAll _ = return
+fromSrcSpanInfoSnippet :: H.SrcSpanInfo -> Block -> Snippet
+fromSrcSpanInfoSnippet ssi block = Snippet
+    { snippetBlock = block'
+    , snippetText  = text
+    , snippetCol   = startCol
+    }
+  where
+    (H.SrcSpan _ start startCol end endCol) = H.srcInfoSpan ssi
+    block' = subBlock start end block
+    lines' = toLines block'
+    text
+        | start == end = T.drop (startCol - 1) $
+            T.take (endCol - 1) $ head lines'
+        | otherwise    = T.unlines $
+            [T.drop (startCol - 1) (head lines')] ++
+            init (drop 1 lines') ++
+            [T.take endCol (last lines')]
 
-selectLines :: Selector
-selectLines _ = perLine
+everything :: Data d => H.Module H.SrcSpanInfo -> [d]
+everything = G.everything (++) (maybeToList . cast)
diff --git a/src/Main.hs b/src/Main.hs
index 4a8a4d3..64b5a77 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -1,24 +1,67 @@
 -- | This is a very simple project containing a few code style checks for use
 -- in git hooks for the Snap Framework. Hopefully we'll get some more
 -- sophisticated checks and automatic fixes implemented eventually.
+{-# LANGUAGE DeriveDataTypeable #-}
 module Main
     ( main
     ) where
 
 import Control.Applicative ((<$>))
-import System.Environment (getArgs)
+import System.Directory (doesDirectoryExist, getDirectoryContents)
 import System.Exit (ExitCode (..), exitWith)
+import System.FilePath (takeExtension, (</>))
+
+import System.Console.CmdArgs
 
 import HStyle
 
+-- | CmdArgs-enabled data-type
+data HStyle = HStyle
+    { fix   :: Bool
+    , quiet :: Bool
+    , files :: [FilePath]
+    } deriving (Show, Data, Typeable)
+
+-- | CmdArgs configuration
+hstyle :: HStyle
+hstyle = HStyle
+    { fix   = def &= help "Automatically fix (some) problems"
+    , quiet = def &= help "Print less output"
+    , files = def &= args 
+    }
+
+-- | Convert CmdArgs configuration to cleaner datatype
+toOptions :: HStyle -> Options
+toOptions hs = Options
+    { optionsFix   = fix hs
+    , optionsQuiet = quiet hs
+    }
+
+-- | Recursively list the contents of a directory. Only returns regular files.
+getFiles :: FilePath -> IO [FilePath]
+getFiles path = do
+    isDir <- doesDirectoryExist path
+    if not isDir
+        then return [path]
+        else do
+            contents <- filter proper <$> getDirectoryContents path
+            concat <$> mapM (getFiles . (path </>)) contents
+  where
+    proper = not . (`elem` [".", ".."])
+
+-- | Is a file haskell?
+isCheckable :: FilePath -> Bool
+isCheckable fp = case takeExtension fp of
+    ".hs"  -> True
+    ".lhs" -> True
+    _      -> False
+
 -- | Simple main that takes one command-line parameter of "check" or "fix" and
 -- a list of files to be checked.
 main :: IO ()
 main = do
-    args <- getArgs
-    case args of
-        "check" : files -> do
-            ok <- and <$> mapM checkStyle files
-            exitWith $ if ok then ExitSuccess else ExitFailure 1
-        "fix"   : files -> mapM_ fixStyle   files
-        _               -> error "Must specify 'check' or 'fix'"
+    config <- cmdArgs hstyle
+    -- Expand all files
+    files' <- filter isCheckable . concat <$> mapM getFiles (files config)
+    ok <- all fileOk <$> mapM (checkStyle $ toOptions config) files'
+    exitWith $ if ok then ExitSuccess else ExitFailure 1