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

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

  | Bool
otherwise
  = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall a. HieAST a -> UsageCounter
declaration (forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
node)
 forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall a. HieAST a -> UsageCounter
findUsage (forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
node)
  where
    annType :: NodeAnnotation -> FastString
annType =
#if MIN_VERSION_ghc(9,2,0)
      NodeAnnotation -> FastString
nodeAnnotType
#else
      snd
#endif

-- | Accrues all the top-level declarations if all different types
declaration :: HieAST a -> UsageCounter
declaration :: forall a. HieAST a -> UsageCounter
declaration HieAST a
node
  | forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((forall a. Eq a => a -> a -> Bool
== FastString
"ConDecl") forall b c a. (b -> c) -> (a -> b) -> a -> c
. NodeAnnotation -> FastString
annType) (forall a. NodeInfo a -> Set NodeAnnotation
nodeAnnotations forall a b. (a -> b) -> a -> b
$ forall a. HieAST a -> NodeInfo a
getNodeInfo HieAST a
node)
  = forall a. HieAST a -> UsageCounter
dataConDecl HieAST a
node
  where
    annType :: NodeAnnotation -> FastString
annType =
#if MIN_VERSION_ghc(9,2,0)
      NodeAnnotation -> FastString
nodeAnnotType
#else
      snd
#endif
declaration HieAST a
node = forall m k a. Monoid m => (k -> a -> m) -> Map k a -> m
M.foldMapWithKey forall {b} {a} {a}.
Ord b =>
Either a b -> IdentifierDetails a -> AppendMap b UsageCount
f 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
  where
    f :: Either a b -> IdentifierDetails a -> AppendMap b UsageCount
f (Right b
name) IdentifierDetails a
details = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ContextInfo -> AppendMap b UsageCount
g (forall a. IdentifierDetails a -> Set ContextInfo
identInfo IdentifierDetails a
details) where
      declare :: AppendMap b UsageCount
declare = 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 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
_                                   = forall a. Monoid a => a
mempty
    f Either a b
_ IdentifierDetails a
_ = 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 :: forall a. HieAST a -> UsageCounter
dataConDecl HieAST a
node = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall a. HieAST a -> UsageCounter
declaration [HieAST a]
dec
                forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall a. HieAST a -> UsageCounter
conField (forall a. HieAST a -> [HieAST a]
nodeChildren forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [HieAST a]
fields)
  where
    ([HieAST a]
dec, [HieAST a]
rest) = forall a. Int -> [a] -> ([a], [a])
splitAt Int
1 forall a b. (a -> b) -> a -> b
$ forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
node
    ([HieAST a]
fields, [HieAST a]
_) = forall a. Int -> [a] -> ([a], [a])
splitAt Int
1 [HieAST a]
rest
    conField :: HieAST a -> UsageCounter
conField HieAST a
n
      | forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"ConDeclField" String
"ConDeclField" HieAST a
n
      = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall a. HieAST a -> UsageCounter
declaration (forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
n)
      | Bool
otherwise = forall a. Monoid a => a
mempty

-- | Counts up the uses of all symbols in the AST.
findUsage :: HieAST a -> UsageCounter
findUsage :: forall a. HieAST a -> UsageCounter
findUsage HieAST a
node = (forall m k a. Monoid m => (k -> a -> m) -> Map k a -> m
M.foldMapWithKey forall {b} {a} {a}.
Ord b =>
Either a b -> IdentifierDetails a -> AppendMap b UsageCount
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. NodeInfo a -> NodeIdentifiers a
nodeIdentifiers forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. HieAST a -> NodeInfo a
getNodeInfo) HieAST a
node
              forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall a. HieAST a -> UsageCounter
findUsage (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 = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ContextInfo -> AppendMap b UsageCount
g (forall a. IdentifierDetails a -> Set ContextInfo
identInfo IdentifierDetails a
details) where
      use :: AppendMap b UsageCount
use = 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 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
_                                    = forall a. Monoid a => a
mempty
    f Either a b
_ IdentifierDetails a
_ = forall a. Monoid a => a
mempty