Add better support for auto, with several options searched

Language/Haskell/GhcMod/FillSig.hs

@@ -9,20 +9,26 @@ module Language.Haskell.GhcMod.FillSig (
 import Data.Char (isSymbol)
 import Data.Function (on)
 import Data.List (find, nub, sortBy)
-import Data.Maybe (isJust)
+import qualified Data.Map.Strict as M
+import Data.Maybe (isJust, catMaybes)
 import Exception (ghandle, SomeException(..))
 import GHC (GhcMonad, Id, ParsedModule(..), TypecheckedModule(..), DynFlags, SrcSpan, Type, GenLocated(L))
 import qualified GHC as G
+import qualified ConLike as G
+import qualified Name as G
 import qualified Language.Haskell.GhcMod.Gap as Gap
 import Language.Haskell.GhcMod.Convert
 import Language.Haskell.GhcMod.Monad
 import Language.Haskell.GhcMod.SrcUtils
 import Language.Haskell.GhcMod.Types
 import CoreMonad (liftIO)
-import Outputable (PprStyle)
+import Outputable (PprStyle, ppr)
 import qualified Type as Ty
 import qualified HsBinds as Ty
 import qualified Class as Ty
+import qualified Var as Ty
+import qualified DataCon as Ty
+import qualified HsPat as Ty
 import qualified Language.Haskell.Exts.Annotated as HE
 import Djinn.GHC
 
@@ -330,9 +336,87 @@ auto file lineNo colNo = ghandle handler body
         opt <- options
         modSum <- Gap.fileModSummary file
         p <- G.parseModule modSum
-        tcm@TypecheckedModule{tm_typechecked_source = tcs} <- G.typecheckModule p
+        tcm@TypecheckedModule{tm_typechecked_source = tcs
+                             ,tm_checked_module_info = minfo} <- G.typecheckModule p
         whenFound' opt (findVar dflag style tcm tcs lineNo colNo) $ \(loc, _name, rty, paren) -> do
-          djinns <- djinn True rty
+          topLevel <- getEverythingInTopLevel minfo
-          return (fourInts loc, map (doParen paren) (nub djinns))
+          let (f,pats) = getPatsForVariable tcs (lineNo,colNo)
+              -- Remove self function to prevent recursion, and id to trim cases
+              filterFn = (\(n,_) -> let funName = G.getOccString n
+                                        recName = G.getOccString (G.getName f)
+                                     in not $ funName `elem` recName:notWantedFuns)
+              -- Find without using other functions in top-level
+              localBnds = M.unions $ map (\(L _ pat) -> getBindingsForPat pat) pats
+              lbn = filter filterFn (M.toList localBnds)
+          djinnsEmpty <- djinn True (Just minfo) lbn rty (Max 10) 100000
+          let -- Find with the entire top-level
+              almostEnv = M.toList $ M.union localBnds topLevel
+              env = filter filterFn almostEnv
+          djinns <- djinn True (Just minfo) env rty (Max 10) 100000
+          return (fourInts loc, map (doParen paren) $ nub (djinnsEmpty ++ djinns))
             
     handler (SomeException _) = emptyResult =<< options
+
+-- Functions we do not want in completions
+notWantedFuns :: [String]
+notWantedFuns = ["id", "asTypeOf", "const"]
+
+-- Get all things defined in top-level
+getEverythingInTopLevel :: GhcMonad m => G.ModuleInfo -> m (M.Map G.Name Type)
+getEverythingInTopLevel m = do
+  let modInfo  = tyThingsToInfo (G.modInfoTyThings m)
+      topNames = G.modInfoTopLevelScope m
+  case topNames of
+    Just topNames' -> do topThings <- mapM G.lookupGlobalName topNames'
+                         let topThings' = catMaybes topThings
+                             topInfo    = tyThingsToInfo topThings'
+                         return $ M.union modInfo topInfo
+    Nothing -> return modInfo
+
+tyThingsToInfo :: [Ty.TyThing] -> M.Map G.Name Type
+tyThingsToInfo [] = M.empty
+tyThingsToInfo (G.AnId i : xs) = M.insert (G.getName i) (Ty.varType i) (tyThingsToInfo xs)
+-- Getting information about constructors is not needed
+-- because they will be added by djinn-ghc when traversing types
+-- #if __GLASGOW_HASKELL__ >= 708
+-- tyThingToInfo (G.AConLike (G.RealDataCon con)) = return [(Ty.dataConName con, Ty.dataConUserType con)]
+-- #else
+-- tyThingToInfo (G.AConLike con) = return [(Ty.dataConName con, Ty.dataConUserType con)]
+-- #endif
+tyThingsToInfo (_:xs) = tyThingsToInfo xs
+
+-- Find the Id of the function and the pattern where the hole is located
+getPatsForVariable :: G.TypecheckedSource -> (Int,Int) -> (Id, [Ty.LPat Id])
+getPatsForVariable tcs (lineNo, colNo) =
+  let (L _ bnd:_) = sortBy (cmp `on` G.getLoc) $ listifySpans tcs (lineNo, colNo) :: [G.LHsBind Id]
+   in case bnd of
+        G.PatBind { Ty.pat_lhs = L ploc pat }  -> case pat of
+          Ty.ConPatIn (L _ i) _ -> (i, [L ploc pat])
+        G.FunBind { Ty.fun_id = L _ funId } ->
+          let m = sortBy (cmp `on` G.getLoc) $ listifySpans tcs (lineNo, colNo) :: [G.LMatch Id (G.LHsExpr Id)]
+              (L _ (G.Match pats _ _):_) = m
+           in (funId, pats)
+        _ -> (error "This should never happen", [])
+
+getBindingsForPat :: Ty.Pat Id -> M.Map G.Name Type
+getBindingsForPat (Ty.VarPat i) = M.singleton (G.getName i) (Ty.varType i)
+getBindingsForPat (Ty.LazyPat (L _ l)) = getBindingsForPat l
+getBindingsForPat (Ty.BangPat (L _ b)) = getBindingsForPat b
+getBindingsForPat (Ty.AsPat (L _ a) (L _ i)) = M.insert (G.getName a) (Ty.varType a) (getBindingsForPat i)
+getBindingsForPat (Ty.ListPat  l _ _) = M.unions $ map (\(L _ i) -> getBindingsForPat i) l
+getBindingsForPat (Ty.TuplePat l _ _) = M.unions $ map (\(L _ i) -> getBindingsForPat i) l
+getBindingsForPat (Ty.PArrPat  l _)   = M.unions $ map (\(L _ i) -> getBindingsForPat i) l
+getBindingsForPat (Ty.ViewPat _ (L _ i) _) = getBindingsForPat i
+getBindingsForPat (Ty.SigPatIn  (L _ i) _) = getBindingsForPat i
+getBindingsForPat (Ty.SigPatOut (L _ i) _) = getBindingsForPat i
+getBindingsForPat (Ty.ConPatIn (L _ i) d) = M.insert (G.getName i) (Ty.varType i) (getBindingsForRecPat d)
+getBindingsForPat (Ty.ConPatOut { Ty.pat_args = d }) = getBindingsForRecPat d
+getBindingsForPat _ = M.empty
+
+getBindingsForRecPat :: Ty.HsConPatDetails Id -> M.Map G.Name Type
+getBindingsForRecPat (Ty.PrefixCon args) = M.unions $ map (\(L _ i) -> getBindingsForPat i) args
+getBindingsForRecPat (Ty.InfixCon (L _ a1) (L _ a2)) = M.union (getBindingsForPat a1) (getBindingsForPat a2)
+getBindingsForRecPat (Ty.RecCon (Ty.HsRecFields { Ty.rec_flds = fields })) = getBindingsForRecFields fields
+  where getBindingsForRecFields [] = M.empty
+        getBindingsForRecFields (Ty.HsRecField { Ty.hsRecFieldArg = (L _ a) } : fs) =
+          M.union (getBindingsForPat a) (getBindingsForRecFields fs)

ghc-mod.cabal

@@ -107,7 +107,7 @@ Library
                       , split
                       , haskell-src-exts
                       , text
-                      , djinn-ghc
+                      , djinn-ghc >= 0.0.2
   if impl(ghc >= 7.8)
     Build-Depends:      Cabal >= 1.18
   else
@@ -192,7 +192,7 @@ Test-Suite spec
                       , split
                       , haskell-src-exts
                       , text
-                      , djinn-ghc
+                      , djinn-ghc >= 0.0.2
   if impl(ghc >= 7.8)
     Build-Depends:      Cabal >= 1.18
   else