[Type-constraints] Marginally better version

* Instead of listify, it walks the tree instead, so that constraints are only visible in relevant context.
* Uses Map instead of [(,)]
* Performs type substitutions in case of one child identifier found (relevant for constructor bindings)

Language/Haskell/GhcMod/Info.hs

@@ -5,13 +5,11 @@ module Language.Haskell.GhcMod.Info (
 
 import Data.Function (on)
 import Data.List (sortBy)
-import Data.Maybe (catMaybes, mapMaybe)
 import System.FilePath
 import Exception (ghandle, SomeException(..))
-import GHC (GhcMonad, LHsBind, LHsExpr, LPat, Id, TypecheckedModule(..), SrcSpan)
+import GHC (GhcMonad, SrcSpan)
 import Prelude
 import qualified GHC as G
-import qualified Var as G (varType)
 import qualified Language.Haskell.GhcMod.Gap as Gap
 
 import Language.Haskell.GhcMod.Convert
@@ -24,9 +22,6 @@ import Language.Haskell.GhcMod.SrcUtils
 import Language.Haskell.GhcMod.Types
 import Language.Haskell.GhcMod.Utils (mkRevRedirMapFunc)
 import Language.Haskell.GhcMod.FileMapping (fileModSummaryWithMapping)
-import Control.Applicative
-import Type
-import Control.Arrow
 
 ----------------------------------------------------------------
 
@@ -76,46 +71,8 @@ types file lineNo colNo =
      gmLog GmException "types" $ showDoc ex
      return []
 
-getSrcSpanType :: GhcMonad m => G.ModSummary -> Int -> Int -> m [(SrcSpan, Type)]
+getSrcSpanType :: (GhcMonad m) => G.ModSummary -> Int -> Int -> m [(SrcSpan, G.Type)]
-getSrcSpanType modSum lineNo colNo = do
+getSrcSpanType modSum lineNo colNo =
-  p <- G.parseModule modSum
+  G.parseModule modSum
-  tcm@TypecheckedModule{tm_typechecked_source = tcs} <- G.typecheckModule p
+  >>= G.typecheckModule
-  let bs = listifySpans tcs (lineNo, colNo) :: [LHsBind Id]
+  >>= flip collectSpansTypes (lineNo, colNo)
-      es = listifySpans tcs (lineNo, colNo) :: [LHsExpr Id]
-      ps = listifySpans tcs (lineNo, colNo) :: [LPat Id]
-      as = concatMap abtoex $ listifyAbsBinds tcs
-      abtoex (G.L _spn G.AbsBinds{abs_exports = es'})
-        = liftA2 (,) G.abe_mono (G.varType . G.abe_poly) `map` es'
-      abtoex _ = []
-      getType' b = getType tcm b >>= tryGetConstrainedType b
-      tryGetConstrainedType _ Nothing = return Nothing
-      tryGetConstrainedType b (Just gt) =
-        do
-          ids <- getId tcm b
-          return $ ct ids <|> Just gt
-        where
-          ct [pid] = (,) (fst gt) <$> lookup pid as
-          ct []    = Nothing
-          ct pids  =
-            let
-              ctys  = mapMaybe build pids
-              build x | Just cti <- x `lookup` as
-                      = let
-                          (preds', ctt) = getPreds cti
-                          vt = G.varType x
-                        in Just (preds', flip (,) vt <$> getTyVar_maybe ctt)
-                      | otherwise = Nothing
-              sty   = snd gt
-              preds = concatMap fst ctys
-              subs  = mkTopTvSubst $ mapMaybe snd ctys
-              ty    = substTy subs $ mkFunTys preds sty
-            in Just (fst gt, tidyTopType ty)
-          getPreds x | isForAllTy x = getPreds $ dropForAlls x
-                     | Just (c, t) <- splitFunTy_maybe x
-                     , isPredTy c = first (c:) $ getPreds t
-                     | otherwise = ([], x)
-
-  ets <- mapM getType' es
-  bts <- mapM getType' bs
-  pts <- mapM getType' ps
-  return $ catMaybes $ concat [ets, bts, pts]

Language/Haskell/GhcMod/SrcUtils.hs

@@ -1,4 +1,4 @@
-{-# LANGUAGE TupleSections, FlexibleInstances, Rank2Types #-}
+{-# LANGUAGE TupleSections, FlexibleInstances, Rank2Types, ImpredicativeTypes #-}
 {-# OPTIONS_GHC -fno-warn-orphans #-}
 
 module Language.Haskell.GhcMod.SrcUtils where
@@ -6,11 +6,14 @@ module Language.Haskell.GhcMod.SrcUtils where
 import Control.Applicative
 import CoreUtils (exprType)
 import Data.Generics
-import Data.Maybe (fromMaybe)
+import Data.Maybe
 import Data.Ord as O
 import GHC (LHsExpr, LPat, Id, DynFlags, SrcSpan, Type, Located, ParsedSource, RenamedSource, TypecheckedSource, GenLocated(L))
+import Var (Var)
 import qualified GHC as G
-import GHC.SYB.Utils (Stage(..), everythingStaged)
+import qualified Var as G
+import qualified Type as G
+import GHC.SYB.Utils
 import GhcMonad
 import qualified Language.Haskell.Exts.Annotated as HE
 import Language.Haskell.GhcMod.Doc
@@ -20,6 +23,9 @@ import OccName (OccName)
 import Outputable (PprStyle)
 import TcHsSyn (hsPatType)
 import Prelude
+import Control.Monad
+import Control.Arrow
+import qualified Data.Map as M
 
 ----------------------------------------------------------------
 
@@ -36,6 +42,48 @@ instance HasType (LPat Id) where
 
 ----------------------------------------------------------------
 
+type CstGenQS = M.Map Var Type
+type CstGenQT = forall m. GhcMonad m => CstGenQS -> (m [(SrcSpan, Type)], CstGenQS)
+
+collectSpansTypes :: (GhcMonad m) => G.TypecheckedModule -> (Int, Int) -> m [(SrcSpan, Type)]
+collectSpansTypes tcs lc =
+  everythingWithContext M.empty (liftM2 (++))
+    ((return [],) `mkQ` hsBind `extQ` hsExpr `extQ` hsPat)
+    (G.tm_typechecked_source tcs)
+  where
+    insExp x = M.insert (G.abe_mono x) (G.varType $ G.abe_poly x)
+    hsBind :: G.LHsBind Id -> CstGenQT
+    hsBind (L _ G.AbsBinds{abs_exports = es'}) s
+      = (return [], foldr insExp s es')
+    hsBind x@(L _ b) s = constrainedType' (G.collectHsBindBinders b) s x
+    hsExpr :: G.LHsExpr Id -> CstGenQT
+    hsExpr x s = (maybeToList <$> getType' x, s)
+    hsPat :: G.LPat Id -> CstGenQT
+    hsPat x@(L _ _) s = constrainedType' (G.collectPatBinders x) s x
+    getType' x@(L spn _)
+      | G.isGoodSrcSpan spn && spn `G.spans` lc
+      = getType tcs x
+      | otherwise = return Nothing
+    constrainedType' pids s x =
+        (maybe [] (uncurry $ constrainedType pids s) <$> getType' x, s)
+    constrainedType pids s spn genTyp =
+      let
+        ctys  = mapMaybe build pids
+        build x | Just cti <- x `M.lookup` s
+                = let
+                    (preds', ctt) = getPreds cti
+                    vt = G.varType x
+                  in (preds',) . (, vt) <$> G.getTyVar_maybe ctt
+                | otherwise = Nothing
+        preds = concatMap fst ctys
+        subs  = G.mkTopTvSubst $ map snd ctys
+        ty    = G.substTy subs $ G.mkFunTys preds genTyp
+      in [(spn, G.tidyTopType ty)]
+    getPreds x | G.isForAllTy x = getPreds $ G.dropForAlls x
+               | Just (c, t) <- G.splitFunTy_maybe x
+               , G.isPredTy c = first (c:) $ getPreds t
+               | otherwise = ([], x)
+
 listifySpans :: Typeable a => TypecheckedSource -> (Int, Int) -> [Located a]
 listifySpans tcs lc = listifyStaged TypeChecker p tcs
   where