ProcessHie.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
module DefCounts.ProcessHie
  ( DefCounter
  , DefType(..)
  , declLines
  ) where

import qualified Data.Array as A
import qualified Data.ByteString as BS
import           Data.Map.Append.Strict (AppendMap(..))
import qualified Data.Map.Strict as M
import           Data.Monoid

import           GHC.Api
import           Utils

-- TODO standalone kind sigs
data DefType
  = Func
  | Data
  | Newtype
  | Class
  | ClassInst
  | Fam
  | TyFamInst
  | Syn
  | PatSyn
  | ModImport
  | ExportThing
  deriving (Eq, Ord, Show)

type DefCounter =
  AppendMap DefType
            ( Sum Int -- num lines
            , Sum Int -- num occurrences
            )

-- | Supports indexing into the source code by line number
type SourceCode = A.Array Int BS.ByteString

-- | Counts up the different types of definitions in the given 'HieAST'.
declLines :: SourceCode -> HieAST a -> DefCounter
declLines src node
  | nodeHasAnnotation "ClsInstD" "InstDecl" node
  || nodeHasAnnotation "DerivDecl" "DerivDecl" node
  = AppendMap $ M.singleton ClassInst (numLines $ nodeSpan node, 1)

  | nodeHasAnnotation "TypeSig" "Sig" node
  = AppendMap $ M.singleton Func (numLines $ nodeSpan node, 0)

  | nodeHasAnnotation "FunBind" "HsBindLR" node
  = AppendMap $ M.singleton Func (numLines $ nodeSpan node, 1)

  | nodeHasAnnotation "ImportDecl" "ImportDecl" node
  = AppendMap $ M.singleton ModImport (numLines $ nodeSpan node, 1)

  | nodeHasAnnotation "IEName" "IEWrappedName" node
  = AppendMap $ M.singleton ExportThing (numLines $ nodeSpan node, 1)

  | otherwise = foldMap ( foldMap (foldMap (tyDeclLines src) . identInfo)
                        . nodeIdentifiers
                        . getNodeInfo )
              $ nodeChildren node

numLines :: Span -> Sum Int
numLines s = Sum $ srcSpanEndLine s - srcSpanStartLine s + 1

tyDeclLines :: SourceCode -> ContextInfo -> DefCounter
tyDeclLines src = \case
  Decl declTy (Just srcSpan)
    | Just defTy <- toDefType srcSpan declTy
    -> AppendMap $ M.singleton defTy (numLines srcSpan, 1)
  _ -> mempty
  where
    toDefType srcSpan = \case
      FamDec           -> Just Fam
      SynDec           -> Just Syn
      DataDec
        | isNewtypeDec -> Just Newtype
        | otherwise    -> Just Data
      PatSynDec        -> Just PatSyn
      ClassDec         -> Just Class
      InstDec          -> Just TyFamInst
      _                -> Nothing

      where
        isNewtypeDec =
          let ln = srcSpanStartLine srcSpan - 1
              col = srcSpanStartCol srcSpan - 1
              (lBnd, uBnd) = A.bounds src
           in ln >= lBnd && ln <= uBnd
           && "newtype" == (BS.take 7 . BS.drop col $ src A.! ln)