{-# 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 (DefType -> DefType -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: DefType -> DefType -> Bool
$c/= :: DefType -> DefType -> Bool
== :: DefType -> DefType -> Bool
$c== :: DefType -> DefType -> Bool
Eq, Eq DefType
DefType -> DefType -> Bool
DefType -> DefType -> Ordering
DefType -> DefType -> DefType
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: DefType -> DefType -> DefType
$cmin :: DefType -> DefType -> DefType
max :: DefType -> DefType -> DefType
$cmax :: DefType -> DefType -> DefType
>= :: DefType -> DefType -> Bool
$c>= :: DefType -> DefType -> Bool
> :: DefType -> DefType -> Bool
$c> :: DefType -> DefType -> Bool
<= :: DefType -> DefType -> Bool
$c<= :: DefType -> DefType -> Bool
< :: DefType -> DefType -> Bool
$c< :: DefType -> DefType -> Bool
compare :: DefType -> DefType -> Ordering
$ccompare :: DefType -> DefType -> Ordering
Ord, Int -> DefType -> ShowS
[DefType] -> ShowS
DefType -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DefType] -> ShowS
$cshowList :: [DefType] -> ShowS
show :: DefType -> String
$cshow :: DefType -> String
showsPrec :: Int -> DefType -> ShowS
$cshowsPrec :: Int -> DefType -> ShowS
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 :: forall a. SourceCode -> HieAST a -> DefCounter
declLines SourceCode
src HieAST a
node
  | forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"ClsInstD" String
"InstDecl" HieAST a
node
  Bool -> Bool -> Bool
|| forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"DerivDecl" String
"DerivDecl" HieAST a
node
  = forall k v. Map k v -> AppendMap k v
AppendMap forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton DefType
ClassInst (Span -> Sum Int
numLines forall a b. (a -> b) -> a -> b
$ forall a. HieAST a -> Span
nodeSpan HieAST a
node, Sum Int
1)

  | forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"TypeSig" String
"Sig" HieAST a
node
  = forall k v. Map k v -> AppendMap k v
AppendMap forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton DefType
Func (Span -> Sum Int
numLines forall a b. (a -> b) -> a -> b
$ forall a. HieAST a -> Span
nodeSpan HieAST a
node, Sum Int
0)

  | forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"FunBind" String
"HsBindLR" HieAST a
node
  = forall k v. Map k v -> AppendMap k v
AppendMap forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton DefType
Func (Span -> Sum Int
numLines forall a b. (a -> b) -> a -> b
$ forall a. HieAST a -> Span
nodeSpan HieAST a
node, Sum Int
1)

  | forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"ImportDecl" String
"ImportDecl" HieAST a
node
  = forall k v. Map k v -> AppendMap k v
AppendMap forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton DefType
ModImport (Span -> Sum Int
numLines forall a b. (a -> b) -> a -> b
$ forall a. HieAST a -> Span
nodeSpan HieAST a
node, Sum Int
1)

  | forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"IEName" String
"IEWrappedName" HieAST a
node
  = forall k v. Map k v -> AppendMap k v
AppendMap forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton DefType
ExportThing (Span -> Sum Int
numLines forall a b. (a -> b) -> a -> b
$ forall a. HieAST a -> Span
nodeSpan HieAST a
node, Sum Int
1)

  | forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"IEThingAll" String
"IE" HieAST a
node
  = forall k v. Map k v -> AppendMap k v
AppendMap forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton DefType
ExportThing (Span -> Sum Int
numLines forall a b. (a -> b) -> a -> b
$ forall a. HieAST a -> Span
nodeSpan HieAST a
node, Sum Int
1)

  | forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"IEPattern" String
"IEWrappedName" HieAST a
node
  = forall k v. Map k v -> AppendMap k v
AppendMap forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton DefType
ExportThing (Span -> Sum Int
numLines forall a b. (a -> b) -> a -> b
$ forall a. HieAST a -> Span
nodeSpan HieAST a
node, Sum Int
1)

  | forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"IEVar" String
"IE" HieAST a
node
  = forall k v. Map k v -> AppendMap k v
AppendMap forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton DefType
ExportThing (Span -> Sum Int
numLines forall a b. (a -> b) -> a -> b
$ forall a. HieAST a -> Span
nodeSpan HieAST a
node, Sum Int
1)

  | Bool
otherwise = forall m a. Monoid m => (HieAST a -> m) -> HieAST a -> m
foldNodeChildren (forall a. SourceCode -> HieAST a -> DefCounter
tyDeclLines SourceCode
src) HieAST a
node

numLines :: Span -> Sum Int
numLines :: Span -> Sum Int
numLines Span
s = forall a. a -> Sum a
Sum forall a b. (a -> b) -> a -> b
$ Span -> Int
srcSpanEndLine Span
s forall a. Num a => a -> a -> a
- Span -> Int
srcSpanStartLine Span
s forall a. Num a => a -> a -> a
+ Int
1

tyDeclLines :: SourceCode -> HieAST a -> DefCounter
tyDeclLines :: forall a. SourceCode -> HieAST a -> DefCounter
tyDeclLines SourceCode
src HieAST a
node = DefCounter
fromCurrentNode forall a. Semigroup a => a -> a -> a
<> DefCounter
fromChildren
  where
    fromCurrentNode :: DefCounter
fromCurrentNode =
      forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ContextInfo -> DefCounter
go forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IdentifierDetails a -> Set ContextInfo
identInfo) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. NodeInfo a -> NodeIdentifiers a
nodeIdentifiers forall a b. (a -> b) -> a -> b
$ forall a. HieAST a -> NodeInfo a
getNodeInfo HieAST a
node
    fromChildren :: DefCounter
fromChildren = forall m a. Monoid m => (HieAST a -> m) -> HieAST a -> m
foldNodeChildren (forall a. SourceCode -> HieAST a -> DefCounter
tyDeclLines SourceCode
src) HieAST a
node

    go :: ContextInfo -> DefCounter
go = \case
      Decl DeclType
declTy (Just Span
srcSpan)
        | Just DefType
defTy <- Span -> DeclType -> Maybe DefType
toDefType Span
srcSpan DeclType
declTy
        -> forall k v. Map k v -> AppendMap k v
AppendMap forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton DefType
defTy (Span -> Sum Int
numLines Span
srcSpan, Sum Int
1)
      ContextInfo
_ -> forall a. Monoid a => a
mempty
      where
        toDefType :: Span -> DeclType -> Maybe DefType
toDefType Span
srcSpan = \case
          DeclType
FamDec           -> forall a. a -> Maybe a
Just DefType
Fam
          DeclType
SynDec           -> forall a. a -> Maybe a
Just DefType
Syn
          DeclType
DataDec
            | Bool
isNewtypeDec -> forall a. a -> Maybe a
Just DefType
Newtype
            | Bool
otherwise    -> forall a. a -> Maybe a
Just DefType
Data
          DeclType
PatSynDec        -> forall a. a -> Maybe a
Just DefType
PatSyn
          DeclType
ClassDec         -> forall a. a -> Maybe a
Just DefType
Class
          DeclType
InstDec          -> forall a. a -> Maybe a
Just DefType
TyFamInst
          DeclType
_                -> forall a. Maybe a
Nothing
          where
            isNewtypeDec :: Bool
isNewtypeDec =
              let ln :: Int
ln = Span -> Int
srcSpanStartLine Span
srcSpan forall a. Num a => a -> a -> a
- Int
1
                  col :: Int
col = Span -> Int
srcSpanStartCol Span
srcSpan forall a. Num a => a -> a -> a
- Int
1
                  (Int
lBnd, Int
uBnd) = forall i e. Array i e -> (i, i)
A.bounds SourceCode
src
               in Int
ln forall a. Ord a => a -> a -> Bool
>= Int
lBnd Bool -> Bool -> Bool
&& Int
ln forall a. Ord a => a -> a -> Bool
<= Int
uBnd
               Bool -> Bool -> Bool
&& ByteString
"newtype" forall a. Eq a => a -> a -> Bool
== (Int -> ByteString -> ByteString
BS.take Int
7 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ByteString -> ByteString
BS.drop Int
col forall a b. (a -> b) -> a -> b
$ SourceCode
src forall i e. Ix i => Array i e -> i -> e
A.! Int
ln)