{-# LANGUAGE OverloadedStrings #-}
module MatchSigs.ProcessHie
  ( SigMap
  , MatchedSigs(..)
  , mkSigMap
  ) where

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

import           GHC.Api
import           MatchSigs.Matching (MatchedSigs(..))
import           MatchSigs.Sig (Sig, sigFingerprint, sigsFromHie)
import           Utils

type SigMap = AppendMap [Sig ()] MatchedSigs

-- | Collect all the function definitions in the 'HieAST' that have isomorphic
-- type signatures.
mkSigMap :: DynFlags -> HieAST HieTypeFix -> SigMap
mkSigMap :: DynFlags -> HieAST HieTypeFix -> SigMap
mkSigMap DynFlags
dynFlags HieAST HieTypeFix
node =
  let renderedSigs :: Map Name String
renderedSigs = forall m a. Monoid m => (HieAST a -> m) -> HieAST a -> m
foldNodeChildren (DynFlags -> HieAST HieTypeFix -> Map Name String
nameSigRendered DynFlags
dynFlags) HieAST HieTypeFix
node
      sigReps :: Map Name [Sig FreeVarIdx]
sigReps = forall m a. Monoid m => (HieAST a -> m) -> HieAST a -> m
foldNodeChildren forall a. HieAST a -> Map Name [Sig FreeVarIdx]
sigsFromHie HieAST HieTypeFix
node
      mkMatch :: Name -> String -> [Sig FreeVarIdx] -> ([Sig ()], MatchedSigs)
mkMatch Name
n String
s [Sig FreeVarIdx]
r = (forall a. [Sig a] -> [Sig ()]
sigFingerprint [Sig FreeVarIdx]
r, [SigMatches] -> MatchedSigs
MatchedSigs [([Sig FreeVarIdx]
r, String
s, [Name
n])])
      sigMatches :: [([Sig ()], MatchedSigs)]
sigMatches = forall k a. Map k a -> [a]
M.elems forall a b. (a -> b) -> a -> b
$ forall k a b c.
Ord k =>
(k -> a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWithKey Name -> String -> [Sig FreeVarIdx] -> ([Sig ()], MatchedSigs)
mkMatch Map Name String
renderedSigs Map Name [Sig FreeVarIdx]
sigReps
   in forall k v. Map k v -> AppendMap k v
AppendMap forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith forall a. Semigroup a => a -> a -> a
(<>) [([Sig ()], MatchedSigs)]
sigMatches

-- | Produce a 'Map' from function 'Name's to their rendered type signatures
nameSigRendered :: DynFlags -> HieAST HieTypeFix -> M.Map Name String
nameSigRendered :: DynFlags -> HieAST HieTypeFix -> Map Name String
nameSigRendered DynFlags
dynFlags HieAST HieTypeFix
node
  | forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"FunBind" String
"HsBindLR" HieAST HieTypeFix
node
  , Just HieAST HieTypeFix
ident <- Maybe (HieAST HieTypeFix)
mIdent
  , Right Name
name : [Identifier]
_ <- forall k a. Map k a -> [k]
M.keys 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 HieTypeFix
ident
  , let renderedTy :: String
renderedTy = [String] -> String
unwords
                   forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (DynFlags -> HieTypeFix -> String
renderHieType DynFlags
dynFlags)
                   forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. NodeInfo a -> [a]
nodeType
                   forall a b. (a -> b) -> a -> b
$ forall a. HieAST a -> NodeInfo a
getNodeInfo HieAST HieTypeFix
node
  = forall k a. k -> a -> Map k a
M.singleton Name
name String
renderedTy

  | Bool
otherwise = forall a. Monoid a => a
mempty
  where
    mIdent :: Maybe (HieAST HieTypeFix)
mIdent
      | HieAST HieTypeFix
c : [HieAST HieTypeFix]
_ <- forall a. HieAST a -> [HieAST a]
nodeChildren HieAST HieTypeFix
node
      -- multiple decls result in Match nodes
      , forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
"Match" String
"Match" HieAST HieTypeFix
c
      , HieAST HieTypeFix
i : [HieAST HieTypeFix]
_ <- forall a. HieAST a -> [HieAST a]
nodeChildren HieAST HieTypeFix
c
      = forall a. a -> Maybe a
Just HieAST HieTypeFix
i

      | HieAST HieTypeFix
i : [HieAST HieTypeFix]
_ <- forall a. HieAST a -> [HieAST a]
nodeChildren HieAST HieTypeFix
node
      = forall a. a -> Maybe a
Just HieAST HieTypeFix
i

      | Bool
otherwise = forall a. Maybe a
Nothing