{-# 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
(DefType -> DefType -> Bool)
-> (DefType -> DefType -> Bool) -> Eq DefType
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
Eq DefType
-> (DefType -> DefType -> Ordering)
-> (DefType -> DefType -> Bool)
-> (DefType -> DefType -> Bool)
-> (DefType -> DefType -> Bool)
-> (DefType -> DefType -> Bool)
-> (DefType -> DefType -> DefType)
-> (DefType -> DefType -> DefType)
-> Ord 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
$cp1Ord :: Eq DefType
Ord, Int -> DefType -> ShowS
[DefType] -> ShowS
DefType -> String
(Int -> DefType -> ShowS)
-> (DefType -> String) -> ([DefType] -> ShowS) -> Show DefType
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 :: SourceCode -> HieAST a -> DefCounter
declLines SourceCode
src HieAST a
node
  | String -> String -> HieAST a -> Bool
forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"ClsInstD" String
"InstDecl" HieAST a
node
  Bool -> Bool -> Bool
|| String -> String -> HieAST a -> Bool
forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"DerivDecl" String
"DerivDecl" HieAST a
node
  = Map DefType (Sum Int, Sum Int) -> DefCounter
forall k v. Map k v -> AppendMap k v
AppendMap (Map DefType (Sum Int, Sum Int) -> DefCounter)
-> Map DefType (Sum Int, Sum Int) -> DefCounter
forall a b. (a -> b) -> a -> b
$ DefType -> (Sum Int, Sum Int) -> Map DefType (Sum Int, Sum Int)
forall k a. k -> a -> Map k a
M.singleton DefType
ClassInst (Span -> Sum Int
numLines (Span -> Sum Int) -> Span -> Sum Int
forall a b. (a -> b) -> a -> b
$ HieAST a -> Span
forall a. HieAST a -> Span
nodeSpan HieAST a
node, Sum Int
1)

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

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

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

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

  | Bool
otherwise = (HieAST a -> DefCounter) -> [HieAST a] -> DefCounter
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ( (IdentifierDetails a -> DefCounter)
-> Map Identifier (IdentifierDetails a) -> DefCounter
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((ContextInfo -> DefCounter) -> Set ContextInfo -> DefCounter
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (SourceCode -> ContextInfo -> DefCounter
tyDeclLines SourceCode
src) (Set ContextInfo -> DefCounter)
-> (IdentifierDetails a -> Set ContextInfo)
-> IdentifierDetails a
-> DefCounter
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IdentifierDetails a -> Set ContextInfo
forall a. IdentifierDetails a -> Set ContextInfo
identInfo)
                        (Map Identifier (IdentifierDetails a) -> DefCounter)
-> (HieAST a -> Map Identifier (IdentifierDetails a))
-> HieAST a
-> DefCounter
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NodeInfo a -> Map Identifier (IdentifierDetails a)
forall a. NodeInfo a -> NodeIdentifiers a
nodeIdentifiers
                        (NodeInfo a -> Map Identifier (IdentifierDetails a))
-> (HieAST a -> NodeInfo a)
-> HieAST a
-> Map Identifier (IdentifierDetails a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HieAST a -> NodeInfo a
forall a. HieAST a -> NodeInfo a
getNodeInfo )
              ([HieAST a] -> DefCounter) -> [HieAST a] -> DefCounter
forall a b. (a -> b) -> a -> b
$ HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
node

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

tyDeclLines :: SourceCode -> ContextInfo -> DefCounter
tyDeclLines :: SourceCode -> ContextInfo -> DefCounter
tyDeclLines SourceCode
src = \case
  Decl DeclType
declTy (Just Span
srcSpan)
    | Just DefType
defTy <- Span -> DeclType -> Maybe DefType
toDefType Span
srcSpan DeclType
declTy
    -> Map DefType (Sum Int, Sum Int) -> DefCounter
forall k v. Map k v -> AppendMap k v
AppendMap (Map DefType (Sum Int, Sum Int) -> DefCounter)
-> Map DefType (Sum Int, Sum Int) -> DefCounter
forall a b. (a -> b) -> a -> b
$ DefType -> (Sum Int, Sum Int) -> Map DefType (Sum Int, Sum Int)
forall k a. k -> a -> Map k a
M.singleton DefType
defTy (Span -> Sum Int
numLines Span
srcSpan, Sum Int
1)
  ContextInfo
_ -> DefCounter
forall a. Monoid a => a
mempty
  where
    toDefType :: Span -> DeclType -> Maybe DefType
toDefType Span
srcSpan = \case
      DeclType
FamDec           -> DefType -> Maybe DefType
forall a. a -> Maybe a
Just DefType
Fam
      DeclType
SynDec           -> DefType -> Maybe DefType
forall a. a -> Maybe a
Just DefType
Syn
      DeclType
DataDec
        | Bool
isNewtypeDec -> DefType -> Maybe DefType
forall a. a -> Maybe a
Just DefType
Newtype
        | Bool
otherwise    -> DefType -> Maybe DefType
forall a. a -> Maybe a
Just DefType
Data
      DeclType
PatSynDec        -> DefType -> Maybe DefType
forall a. a -> Maybe a
Just DefType
PatSyn
      DeclType
ClassDec         -> DefType -> Maybe DefType
forall a. a -> Maybe a
Just DefType
Class
      DeclType
InstDec          -> DefType -> Maybe DefType
forall a. a -> Maybe a
Just DefType
TyFamInst
      DeclType
_                -> Maybe DefType
forall a. Maybe a
Nothing

      where
        isNewtypeDec :: Bool
isNewtypeDec =
          let ln :: Int
ln = Span -> Int
srcSpanStartLine Span
srcSpan Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
              col :: Int
col = Span -> Int
srcSpanStartCol Span
srcSpan Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
              (Int
lBnd, Int
uBnd) = SourceCode -> (Int, Int)
forall i e. Array i e -> (i, i)
A.bounds SourceCode
src
           in Int
ln Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
lBnd Bool -> Bool -> Bool
&& Int
ln Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
uBnd
           Bool -> Bool -> Bool
&& ByteString
"newtype" ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== (Int -> ByteString -> ByteString
BS.take Int
7 (ByteString -> ByteString)
-> (ByteString -> ByteString) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ByteString -> ByteString
BS.drop Int
col (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ SourceCode
src SourceCode -> Int -> ByteString
forall i e. Ix i => Array i e -> i -> e
A.! Int
ln)