{-# 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 = (HieAST HieTypeFix -> Map Name String)
-> HieAST HieTypeFix -> Map Name String
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 = (HieAST HieTypeFix -> Map Name [Sig FreeVarIdx])
-> HieAST HieTypeFix -> Map Name [Sig FreeVarIdx]
forall m a. Monoid m => (HieAST a -> m) -> HieAST a -> m
foldNodeChildren HieAST HieTypeFix -> Map Name [Sig FreeVarIdx]
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 = ([Sig FreeVarIdx] -> [Sig ()]
forall a. [Sig a] -> [Sig ()]
sigFingerprint [Sig FreeVarIdx]
r, [SigMatches] -> MatchedSigs
MatchedSigs [([Sig FreeVarIdx]
r, String
s, [Name
n])])
      sigMatches :: [([Sig ()], MatchedSigs)]
sigMatches = Map Name ([Sig ()], MatchedSigs) -> [([Sig ()], MatchedSigs)]
forall k a. Map k a -> [a]
M.elems (Map Name ([Sig ()], MatchedSigs) -> [([Sig ()], MatchedSigs)])
-> Map Name ([Sig ()], MatchedSigs) -> [([Sig ()], MatchedSigs)]
forall a b. (a -> b) -> a -> b
$ (Name -> String -> [Sig FreeVarIdx] -> ([Sig ()], MatchedSigs))
-> Map Name String
-> Map Name [Sig FreeVarIdx]
-> Map Name ([Sig ()], MatchedSigs)
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 Map [Sig ()] MatchedSigs -> SigMap
forall k v. Map k v -> AppendMap k v
AppendMap (Map [Sig ()] MatchedSigs -> SigMap)
-> Map [Sig ()] MatchedSigs -> SigMap
forall a b. (a -> b) -> a -> b
$ (MatchedSigs -> MatchedSigs -> MatchedSigs)
-> [([Sig ()], MatchedSigs)] -> Map [Sig ()] MatchedSigs
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith MatchedSigs -> MatchedSigs -> MatchedSigs
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
  | String -> String -> HieAST HieTypeFix -> Bool
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]
_ <- Map Identifier (IdentifierDetails HieTypeFix) -> [Identifier]
forall k a. Map k a -> [k]
M.keys (Map Identifier (IdentifierDetails HieTypeFix) -> [Identifier])
-> (NodeInfo HieTypeFix
    -> Map Identifier (IdentifierDetails HieTypeFix))
-> NodeInfo HieTypeFix
-> [Identifier]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NodeInfo HieTypeFix
-> Map Identifier (IdentifierDetails HieTypeFix)
forall a. NodeInfo a -> NodeIdentifiers a
nodeIdentifiers (NodeInfo HieTypeFix -> [Identifier])
-> NodeInfo HieTypeFix -> [Identifier]
forall a b. (a -> b) -> a -> b
$ HieAST HieTypeFix -> NodeInfo HieTypeFix
forall a. HieAST a -> NodeInfo a
getNodeInfo HieAST HieTypeFix
ident
  , let renderedTy :: String
renderedTy = [String] -> String
unwords
                   ([String] -> String)
-> (NodeInfo HieTypeFix -> [String])
-> NodeInfo HieTypeFix
-> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (HieTypeFix -> String) -> [HieTypeFix] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (DynFlags -> HieTypeFix -> String
renderHieType DynFlags
dynFlags)
                   ([HieTypeFix] -> [String])
-> (NodeInfo HieTypeFix -> [HieTypeFix])
-> NodeInfo HieTypeFix
-> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NodeInfo HieTypeFix -> [HieTypeFix]
forall a. NodeInfo a -> [a]
nodeType
                   (NodeInfo HieTypeFix -> String) -> NodeInfo HieTypeFix -> String
forall a b. (a -> b) -> a -> b
$ HieAST HieTypeFix -> NodeInfo HieTypeFix
forall a. HieAST a -> NodeInfo a
getNodeInfo HieAST HieTypeFix
node
  = Name -> String -> Map Name String
forall k a. k -> a -> Map k a
M.singleton Name
name String
renderedTy

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

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

      | Bool
otherwise = Maybe (HieAST HieTypeFix)
forall a. Maybe a
Nothing