{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
module UseCounts.ProcessHie
  ( UsageCounter
  , UsageCount(..)
  , usageCounter
  ) where

import qualified Data.Map.Strict as M
import           Data.Map.Append.Strict (AppendMap(..))
import           Data.Maybe

import           GHC.Api
import           Utils

data UsageCount =
  UsageCount
    { UsageCount -> Int
usages :: !Int
    , UsageCount -> Bool
locallyDefined :: !Bool
    } deriving Int -> UsageCount -> ShowS
[UsageCount] -> ShowS
UsageCount -> String
(Int -> UsageCount -> ShowS)
-> (UsageCount -> String)
-> ([UsageCount] -> ShowS)
-> Show UsageCount
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [UsageCount] -> ShowS
$cshowList :: [UsageCount] -> ShowS
show :: UsageCount -> String
$cshow :: UsageCount -> String
showsPrec :: Int -> UsageCount -> ShowS
$cshowsPrec :: Int -> UsageCount -> ShowS
Show

instance Semigroup UsageCount where
  UsageCount Int
na Bool
da <> :: UsageCount -> UsageCount -> UsageCount
<> UsageCount Int
nb Bool
db
    = Int -> Bool -> UsageCount
UsageCount (Int
na Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nb) (Bool
da Bool -> Bool -> Bool
|| Bool
db)

instance Monoid UsageCount where
  mempty :: UsageCount
mempty = Int -> Bool -> UsageCount
UsageCount Int
0 Bool
False

type UsageCounter = AppendMap Name UsageCount

usageCounter :: HieAST a -> UsageCounter
usageCounter :: HieAST a -> UsageCounter
usageCounter HieAST a
node
  | String -> String -> HieAST a -> Bool
forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"FunBind" String
"HsBindLR" HieAST a
node
  = (HieAST a -> UsageCounter) -> [HieAST a] -> UsageCounter
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap HieAST a -> UsageCounter
forall a. HieAST a -> UsageCounter
findUsage (HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
node)
 UsageCounter -> UsageCounter -> UsageCounter
forall a. Semigroup a => a -> a -> a
<> (HieAST a -> UsageCounter) -> Maybe (HieAST a) -> UsageCounter
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap HieAST a -> UsageCounter
forall a. HieAST a -> UsageCounter
declaration ([HieAST a] -> Maybe (HieAST a)
forall a. [a] -> Maybe a
listToMaybe ([HieAST a] -> Maybe (HieAST a)) -> [HieAST a] -> Maybe (HieAST a)
forall a b. (a -> b) -> a -> b
$ HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
node)

  -- only get usages from instance declarations
  | ((FastString, FastString) -> Bool)
-> Set (FastString, FastString) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((FastString -> FastString -> Bool
forall a. Eq a => a -> a -> Bool
== FastString
"InstDecl") (FastString -> Bool)
-> ((FastString, FastString) -> FastString)
-> (FastString, FastString)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FastString, FastString) -> FastString
forall a b. (a, b) -> b
snd) (NodeInfo a -> Set (FastString, FastString)
forall a. NodeInfo a -> Set (FastString, FastString)
nodeAnnotations (NodeInfo a -> Set (FastString, FastString))
-> NodeInfo a -> Set (FastString, FastString)
forall a b. (a -> b) -> a -> b
$ HieAST a -> NodeInfo a
forall a. HieAST a -> NodeInfo a
getNodeInfo HieAST a
node)
  = (HieAST a -> UsageCounter) -> [HieAST a] -> UsageCounter
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap HieAST a -> UsageCounter
forall a. HieAST a -> UsageCounter
findUsage (HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
node)

  | Bool
otherwise
  = (HieAST a -> UsageCounter) -> [HieAST a] -> UsageCounter
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap HieAST a -> UsageCounter
forall a. HieAST a -> UsageCounter
declaration (HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
node)
 UsageCounter -> UsageCounter -> UsageCounter
forall a. Semigroup a => a -> a -> a
<> (HieAST a -> UsageCounter) -> [HieAST a] -> UsageCounter
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap HieAST a -> UsageCounter
forall a. HieAST a -> UsageCounter
findUsage (HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
node)

-- | Accrues all the top-level declarations if all different types
declaration :: HieAST a -> UsageCounter
declaration :: HieAST a -> UsageCounter
declaration HieAST a
node
  | ((FastString, FastString) -> Bool)
-> Set (FastString, FastString) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((FastString -> FastString -> Bool
forall a. Eq a => a -> a -> Bool
== FastString
"ConDecl") (FastString -> Bool)
-> ((FastString, FastString) -> FastString)
-> (FastString, FastString)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FastString, FastString) -> FastString
forall a b. (a, b) -> b
snd) (NodeInfo a -> Set (FastString, FastString)
forall a. NodeInfo a -> Set (FastString, FastString)
nodeAnnotations (NodeInfo a -> Set (FastString, FastString))
-> NodeInfo a -> Set (FastString, FastString)
forall a b. (a -> b) -> a -> b
$ HieAST a -> NodeInfo a
forall a. HieAST a -> NodeInfo a
getNodeInfo HieAST a
node)
  = HieAST a -> UsageCounter
forall a. HieAST a -> UsageCounter
dataConDecl HieAST a
node
declaration HieAST a
node = (Either ModuleName Name -> IdentifierDetails a -> UsageCounter)
-> Map (Either ModuleName Name) (IdentifierDetails a)
-> UsageCounter
forall m k a. Monoid m => (k -> a -> m) -> Map k a -> m
M.foldMapWithKey Either ModuleName Name -> IdentifierDetails a -> UsageCounter
forall b a a.
Ord b =>
Either a b -> IdentifierDetails a -> AppendMap b UsageCount
f (Map (Either ModuleName Name) (IdentifierDetails a)
 -> UsageCounter)
-> (NodeInfo a
    -> Map (Either ModuleName Name) (IdentifierDetails a))
-> NodeInfo a
-> UsageCounter
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NodeInfo a -> Map (Either ModuleName Name) (IdentifierDetails a)
forall a. NodeInfo a -> NodeIdentifiers a
nodeIdentifiers (NodeInfo a -> UsageCounter) -> NodeInfo a -> UsageCounter
forall a b. (a -> b) -> a -> b
$ HieAST a -> NodeInfo a
forall a. HieAST a -> NodeInfo a
getNodeInfo HieAST a
node
  where
    f :: Either a b -> IdentifierDetails a -> AppendMap b UsageCount
f (Right b
name) IdentifierDetails a
details = (ContextInfo -> AppendMap b UsageCount)
-> Set ContextInfo -> AppendMap b UsageCount
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ContextInfo -> AppendMap b UsageCount
g (IdentifierDetails a -> Set ContextInfo
forall a. IdentifierDetails a -> Set ContextInfo
identInfo IdentifierDetails a
details) where
      declare :: AppendMap b UsageCount
declare = Map b UsageCount -> AppendMap b UsageCount
forall k v. Map k v -> AppendMap k v
AppendMap (Map b UsageCount -> AppendMap b UsageCount)
-> Map b UsageCount -> AppendMap b UsageCount
forall a b. (a -> b) -> a -> b
$ b -> UsageCount -> Map b UsageCount
forall k a. k -> a -> Map k a
M.singleton b
name (Int -> Bool -> UsageCount
UsageCount Int
0 Bool
True)
      g :: ContextInfo -> AppendMap b UsageCount
g (ValBind BindType
RegularBind Scope
ModuleScope Maybe Span
_) = AppendMap b UsageCount
declare
      g (PatternBind Scope
ModuleScope Scope
_ Maybe Span
_)       = AppendMap b UsageCount
declare
      g (Decl DeclType
t Maybe Span
_) | DeclType -> Bool
checkDeclType DeclType
t        = AppendMap b UsageCount
declare
      g ContextInfo
TyDecl                              = AppendMap b UsageCount
declare
      g ClassTyDecl{}                       = AppendMap b UsageCount
declare
      g ContextInfo
_                                   = AppendMap b UsageCount
forall a. Monoid a => a
mempty
    f Either a b
_ IdentifierDetails a
_ = AppendMap b UsageCount
forall a. Monoid a => a
mempty

    checkDeclType :: DeclType -> Bool
checkDeclType = \case
      DeclType
InstDec -> Bool
False -- type fam instance is not a declaration
      DeclType
_       -> Bool
True

-- | Handles data constructor declarations
dataConDecl :: HieAST a -> UsageCounter
dataConDecl :: HieAST a -> UsageCounter
dataConDecl HieAST a
node = (HieAST a -> UsageCounter) -> [HieAST a] -> UsageCounter
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap HieAST a -> UsageCounter
forall a. HieAST a -> UsageCounter
declaration [HieAST a]
dec
                UsageCounter -> UsageCounter -> UsageCounter
forall a. Semigroup a => a -> a -> a
<> (HieAST a -> UsageCounter) -> [HieAST a] -> UsageCounter
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap HieAST a -> UsageCounter
forall a. HieAST a -> UsageCounter
conField (HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
nodeChildren (HieAST a -> [HieAST a]) -> [HieAST a] -> [HieAST a]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [HieAST a]
fields)
  where
    ([HieAST a]
dec, [HieAST a]
rest) = Int -> [HieAST a] -> ([HieAST a], [HieAST a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
1 ([HieAST a] -> ([HieAST a], [HieAST a]))
-> [HieAST a] -> ([HieAST a], [HieAST a])
forall a b. (a -> b) -> a -> b
$ HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
node
    ([HieAST a]
fields, [HieAST a]
_) = Int -> [HieAST a] -> ([HieAST a], [HieAST a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
1 [HieAST a]
rest
    conField :: HieAST a -> UsageCounter
conField HieAST a
n
      | String -> String -> HieAST a -> Bool
forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"ConDeclField" String
"ConDeclField" HieAST a
n
      = (HieAST a -> UsageCounter) -> [HieAST a] -> UsageCounter
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap HieAST a -> UsageCounter
forall a. HieAST a -> UsageCounter
declaration (HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
n)
      | Bool
otherwise = UsageCounter
forall a. Monoid a => a
mempty

-- | Counts up the uses of all symbols in the AST.
findUsage :: HieAST a -> UsageCounter
findUsage :: HieAST a -> UsageCounter
findUsage HieAST a
node = ((Either ModuleName Name -> IdentifierDetails a -> UsageCounter)
-> Map (Either ModuleName Name) (IdentifierDetails a)
-> UsageCounter
forall m k a. Monoid m => (k -> a -> m) -> Map k a -> m
M.foldMapWithKey Either ModuleName Name -> IdentifierDetails a -> UsageCounter
forall b a a.
Ord b =>
Either a b -> IdentifierDetails a -> AppendMap b UsageCount
f (Map (Either ModuleName Name) (IdentifierDetails a)
 -> UsageCounter)
-> (HieAST a -> Map (Either ModuleName Name) (IdentifierDetails a))
-> HieAST a
-> UsageCounter
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NodeInfo a -> Map (Either ModuleName Name) (IdentifierDetails a)
forall a. NodeInfo a -> NodeIdentifiers a
nodeIdentifiers (NodeInfo a -> Map (Either ModuleName Name) (IdentifierDetails a))
-> (HieAST a -> NodeInfo a)
-> HieAST a
-> Map (Either ModuleName Name) (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
node
              UsageCounter -> UsageCounter -> UsageCounter
forall a. Semigroup a => a -> a -> a
<> (HieAST a -> UsageCounter) -> [HieAST a] -> UsageCounter
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap HieAST a -> UsageCounter
forall a. HieAST a -> UsageCounter
findUsage (HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
node)
  where
    f :: Either a b -> IdentifierDetails a -> AppendMap b UsageCount
f (Right b
name) IdentifierDetails a
details = (ContextInfo -> AppendMap b UsageCount)
-> Set ContextInfo -> AppendMap b UsageCount
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ContextInfo -> AppendMap b UsageCount
g (IdentifierDetails a -> Set ContextInfo
forall a. IdentifierDetails a -> Set ContextInfo
identInfo IdentifierDetails a
details) where
      use :: AppendMap b UsageCount
use = Map b UsageCount -> AppendMap b UsageCount
forall k v. Map k v -> AppendMap k v
AppendMap (Map b UsageCount -> AppendMap b UsageCount)
-> Map b UsageCount -> AppendMap b UsageCount
forall a b. (a -> b) -> a -> b
$ b -> UsageCount -> Map b UsageCount
forall k a. k -> a -> Map k a
M.singleton b
name (Int -> Bool -> UsageCount
UsageCount Int
1 Bool
False)
      g :: ContextInfo -> AppendMap b UsageCount
g ContextInfo
Use                                  = AppendMap b UsageCount
use
      g (ValBind BindType
InstanceBind Scope
ModuleScope Maybe Span
_) = AppendMap b UsageCount
use
      g (Decl DeclType
InstDec Maybe Span
_)                     = AppendMap b UsageCount
use
      g (RecField RecFieldContext
RecFieldAssign Maybe Span
_)          = AppendMap b UsageCount
use
      g (RecField RecFieldContext
RecFieldMatch Maybe Span
_)           = AppendMap b UsageCount
use
      g ContextInfo
_                                    = AppendMap b UsageCount
forall a. Monoid a => a
mempty
    f Either a b
_ IdentifierDetails a
_ = AppendMap b UsageCount
forall a. Monoid a => a
mempty