git.iko.soy/Idris-dev
on 2aed52a help sign in
Idris-dev/src/Idris/TypeSearch.hs
444 lines 19 kB history raw
1{-# LANGUAGE ScopedTypeVariables #-}
3module Idris.TypeSearch (
4 searchByType, searchPred, defaultScoreFunction
5) where
7import Control.Applicative ((<$>), (<*>), (<|>))
8import Control.Arrow (first, second, (&&&))
9import Control.Monad (forM_, guard)
11import Data.List (find, delete, deleteBy, minimumBy, partition, sortBy, (\\))
12import Data.Map (Map)
13import qualified Data.Map as M
14import Data.Maybe (catMaybes, fromMaybe, isJust, maybeToList, mapMaybe)
15import Data.Monoid (Monoid (mempty, mappend))
16import Data.Ord (comparing)
17import qualified Data.PriorityQueue.FingerTree as Q
18import Data.Set (Set)
19import qualified Data.Set as S
20import qualified Data.Text as T (pack, isPrefixOf)
22import Idris.AbsSyntax (addUsingConstraints, addImpl, getContext, getIState, putIState, implicit)
23import Idris.AbsSyntaxTree (class_instances, ClassInfo, defaultSyntax, Idris,
24 IState (idris_classes, idris_docstrings, tt_ctxt, idris_outputmode),
25 implicitAllowed, OutputMode(..), prettyDocumentedIst, prettyIst, PTerm, toplevel)
26import Idris.Core.Evaluate (Context (definitions), Def (Function, TyDecl, CaseOp), normaliseC)
27import Idris.Core.TT hiding (score)
28import Idris.Core.Unify (match_unify)
29import Idris.Delaborate (delabTy)
30import Idris.Docstrings (noDocs, overview)
31import Idris.ElabDecls (elabType')
32import Idris.Output (ihRenderOutput, ihPrintResult, ihRenderResult)
34import System.IO (Handle)
36import Util.Pretty (text, char, vsep, (<>), Doc)
38searchByType :: Handle -> PTerm -> Idris ()
39searchByType h pterm = do
40 pterm' <- addUsingConstraints syn emptyFC pterm
41 pterm'' <- implicit toplevel syn n pterm'
42 i <- getIState
43 let pterm''' = addImpl i pterm''
44 ty <- elabType' False toplevel syn (fst noDocs) (snd noDocs) emptyFC [] n pterm'
45 putIState i -- don't actually make any changes
46 let names = searchUsing searchPred i ty
47 let names' = take numLimit $ names
48 let docs =
49 [ let docInfo = (n, delabTy i n, fmap (overview . fst) (lookupCtxtExact n (idris_docstrings i))) in
50 displayScore score <> char ' ' <> prettyDocumentedIst i docInfo
51 | (n, score) <- names']
52 case idris_outputmode i of
53 RawOutput -> do mapM_ (ihRenderOutput h) docs
54 ihPrintResult h ""
55 IdeSlave n -> ihRenderResult h (vsep docs)
56 where
57 numLimit = 50
58 syn = defaultSyntax { implicitAllowed = True } -- syntax
59 n = sMN 0 "searchType" -- name
61-- | Conduct a type-directed search using a given match predicate
62searchUsing :: (IState -> Type -> [(Name, Type)] -> [(Name, a)])
63 -> IState -> Type -> [(Name, a)]
64searchUsing pred istate ty = pred istate nty . concat . M.elems $
65 M.mapWithKey (\key -> M.toAscList . M.mapMaybe (f key)) (definitions ctxt)
66 where
67 nty = normaliseC ctxt [] ty
68 ctxt = tt_ctxt istate
69 f k x = do
70 guard $ not (special k)
71 type2 <- typeFromDef x
72 return $ normaliseC ctxt [] type2
73 special :: Name -> Bool
74 special (NS n ns) = special n
75 special (SN _) = True
76 special (UN n) = T.pack "default#" `T.isPrefixOf` n
77 || n `elem` map T.pack ["believe_me", "really_believe_me"]
78 special _ = False
80-- Our default search predicate.
81searchPred :: IState -> Type -> [(Name, Type)] -> [(Name, Score)]
82searchPred istate ty1 = matcher where
83 maxScore = 100
84 matcher = matchTypesBulk istate maxScore ty1
87typeFromDef :: (Def, b, c, d) -> Maybe Type
88typeFromDef (def, _, _, _) = get def where
89 get :: Def -> Maybe Type
90 get (Function ty tm) = Just ty
91 get (TyDecl _ ty) = Just ty
92 -- get (Operator ty _ _) = Just ty
93 get (CaseOp _ ty _ _ _ _) = Just ty
94 get _ = Nothing
97-- | reverse the edges for a directed acyclic graph
98reverseDag :: Ord k => [((k, a), Set k)] -> [((k, a), Set k)]
99reverseDag xs = map f xs where
100 f ((k, v), _) = ((k, v), S.fromList . map (fst . fst) $ filter (S.member k . snd) xs)
102-- run vToP first!
103-- | Compute a directed acyclic graph corresponding to the
104-- arguments of a function.
105-- returns [(the name and type of the bound variable
106-- the names in the type of the bound variable)]
107computeDagP :: Ord n
108 => (TT n -> Bool) -- ^ filter to remove some arguments
109 -> TT n
110 -> ([((n, TT n), Set n)], [(n, TT n)], TT n)
111computeDagP removePred t = (reverse (map f args), reverse removedArgs , retTy) where
112 f (n, t) = ((n, t), M.keysSet (usedVars t))
114 (numArgs, args, removedArgs, retTy) = go 0 [] [] t
116 -- NOTE : args are in reverse order
117 go k args removedArgs (Bind n (Pi t) sc) = let arg = (n, t) in
118 if removePred t
119 then go k args (arg : removedArgs) sc
120 else go (succ k) (arg : args) removedArgs sc
121 go k args removedArgs retTy = (k, args, removedArgs, retTy)
123-- | Collect the names and types of all the free variables
124-- The Boolean indicates those variables which are determined due to injectivity
125-- I have not given nearly enough thought to know whether this is correct
126usedVars :: Ord n => TT n -> Map n (TT n, Bool)
127usedVars = f True where
128 f b (P Bound n t) = M.singleton n (t, b) `M.union` f b t
129 f b (Bind n binder t2) = (M.delete n (f b t2) `M.union`) $ case binder of
130 Let t v -> f b t `M.union` f b v
131 Guess t v -> f b t `M.union` f b v
132 bind -> f b (binderTy bind)
133 f b (App t1 t2) = f b t1 `M.union` f (b && isInjective t1) t2
134 f b (Proj t _) = f b t
135 f _ (V j) = error "unexpected! run vToP first"
136 f _ _ = M.empty
138-- | Remove a node from a directed acyclic graph
139deleteFromDag :: Ord n => n -> [((n, TT n), (a, Set n))] -> [((n, TT n), (a, Set n))]
140deleteFromDag name [] = []
141deleteFromDag name (((name2, ty), (ix, set)) : xs) = (if name == name2
142 then id
143 else (((name2, ty) , (ix, S.delete name set)) :) ) (deleteFromDag name xs)
145deleteFromArgList :: Ord n => n -> [(n, TT n)] -> [(n, TT n)]
146deleteFromArgList n = filter ((/= n) . fst)
148data Score = Score
149 { transposition :: !Int
150 , leftApplied :: !Int
151 , rightApplied :: !Int
152 , leftTypeClassApp :: !Int
153 , rightTypeClassApp :: !Int
154 , leftTypeClassIntro :: !Int
155 , rightTypeClassIntro :: !Int } deriving (Eq, Show)
157displayScore :: Score -> Doc a
158displayScore (Score trans lapp rapp lclassapp rclassapp lclassintro rclassintro) = text $ case (lt, gt) of
159 (True , True ) -> "=" -- types are isomorphic
160 (True , False) -> "<" -- found type is more general than searched type
161 (False, True ) -> ">" -- searched type is more general than found type
162 (False, False) -> "_"
163 where lt = lapp + lclassapp + lclassintro == 0
164 gt = rapp + rclassapp + rclassintro == 0
167scoreCriterion :: Score -> Bool
168scoreCriterion (Score a b c d e f g) = not
169 ( (b > 0 && c > 0) || (b + c) > 4 || any (> 3) [d,e,f,g])
171defaultScoreFunction :: Score -> Int
172defaultScoreFunction (Score a b c d e f g) = a + 9*b + 3*c + 12*d + 4*e + 6*f + 2*g + 100*(2*b + d + f)*(2*c + e + g)
173 -- it's very bad to have *both* upcasting and downcasting
175instance Monoid Score where
176 mempty = Score 0 0 0 0 0 0 0
177 (Score a b c d e f g) `mappend` (Score a' b' c' d' e' f' g') =
178 Score (a + a') (b + b') (c + c') (d + d') (e + e') (f + f') (g + g')
180-- | A directed acyclic graph representing the arguments to a function
181-- The 'Int' represents the position of the argument (1st argument, 2nd, etc.)
182type ArgsDAG = [((Name, Type), (Int, Set Name))]
184-- | The state corresponding to an attempted match of two types.
185data State = State
186 { holes :: ![(Name, Type)] -- ^ names which have yet to be resolved
187 , args1 :: !ArgsDAG -- ^ arguments for the left type which have yet to be resolved
188 , args2 :: !ArgsDAG -- ^ arguments for the right type which have yet to be resolved
189 , classes1 :: ![(Name, Type)] -- ^ typeclass arguments for the left type which haven't been resolved
190 , classes2 :: ![(Name, Type)] -- ^ typeclass arguments for the right type which haven't been resolved
191 , score :: !Score -- ^ the score so far
192 , usedNames :: ![Name] -- ^ all names that have been used
193 } deriving Show
195modifyTypes :: (Type -> Type) -> State -> State
196modifyTypes f (State h a1 a2 c1 c2 s un) =
197 State h (modifyDag a1) (modifyDag a2)
198 (modifyList c1) (modifyList c2)
199 s un
200 where
201 modifyDag = map (first (second f))
202 modifyList = map (second f)
204findNameInArgsDAG :: Name -> ArgsDAG -> Maybe (Type, Maybe Int)
205findNameInArgsDAG name xs = fmap ((snd . fst) &&& (Just . fst . snd)) . find ((name ==) . fst . fst) $ xs
207findLeft, findRight :: Name -> State -> Maybe (Type, Maybe Int)
208findLeft n (State _ a1 a2 c1 c2 _ _) = findNameInArgsDAG n a1 <|> ((,) <$> lookup n c1 <*> Nothing)
209findRight n (State _ a1 a2 c1 c2 _ _) = findNameInArgsDAG n a2 <|> ((,) <$> lookup n c2 <*> Nothing)
211deleteLeft, deleteRight :: Name -> State -> State
212deleteLeft n state = state { args1 = deleteFromDag n (args1 state) , classes1 = filter ((/= n) . fst) (classes1 state) }
213deleteRight n state = state { args2 = deleteFromDag n (args2 state) , classes2 = filter ((/= n) . fst) (classes2 state) }
216tcToMaybe :: TC' e a -> Maybe a
217tcToMaybe (OK x) = Just x
218tcToMaybe (Error _) = Nothing
220inArgTys :: (Type -> Type) -> ArgsDAG -> ArgsDAG
221inArgTys = map . first . second
224typeclassUnify :: Ctxt ClassInfo -> Context -> Type -> Type -> Maybe [(Name, Type)]
225typeclassUnify classInfo ctxt ty tyTry = do
226 res <- tcToMaybe $ match_unify ctxt [] ty retTy [] holes []
227 guard $ null (holes \\ map fst res)
228 let argTys' = map (second $ foldr (.) id [ subst n t | (n, t) <- res ]) tcArgs
229 return argTys'
230 where
231 tyTry' = vToP tyTry
232 holes = map fst nonTcArgs
233 retTy = getRetTy tyTry'
234 (tcArgs, nonTcArgs) = partition (isTypeClassArg classInfo . snd) $ getArgTys tyTry'
236isTypeClassArg :: Ctxt ClassInfo -> Type -> Bool
237isTypeClassArg classInfo ty = not (null (getClassName clss >>= flip lookupCtxt classInfo)) where
238 (clss, args) = unApply ty
239 getClassName (P (TCon _ _) className _) = [className]
240 getClassName _ = []
243instance Ord Score where
244 compare = comparing defaultScoreFunction
246-- | Compute the power set
247subsets :: [a] -> [[a]]
248subsets [] = [[]]
249subsets (x : xs) = let ss = subsets xs in map (x :) ss ++ ss
252--DONT run vToP first!
253-- | Try to match two types together in a unification-like procedure.
254-- Returns a list of types and their minimum scores, sorted in order
255-- of increasing score.
256matchTypesBulk :: forall info. IState -> Int -> Type -> [(info, Type)] -> [(info, Score)]
257matchTypesBulk istate maxScore type1 types = getAllResults startQueueOfQueues where
258 getStartQueue :: (info, Type) -> Maybe (Score, (info, Q.PQueue Score State))
259 getStartQueue nty@(info, type2) = do
260 state <- unifyQueue (State startingHoles dag1 dag2
261 typeClassArgs1 typeClassArgs2
262 mempty usedNames) startingTypes
263 let sc = score state
264 return $ (sc, (info, Q.singleton sc state))
265 where
266 (dag2, typeClassArgs2, retTy2) = makeDag (uniqueBinders (map fst argNames1) type2)
267 argNames2 = map fst dag2
268 usedNames = map fst (argNames1 ++ argNames2)
269 startingHoles = argNames1 ++ argNames2
271 startingTypes = (retTy1, retTy2) : []
274 startQueueOfQueues :: Q.PQueue Score (info, Q.PQueue Score State)
275 startQueueOfQueues = Q.fromList $ mapMaybe getStartQueue types
277 getAllResults :: Q.PQueue Score (info, Q.PQueue Score State) -> [(info, Score)]
278 getAllResults q = case Q.minViewWithKey q of
279 Nothing -> []
280 Just ((nextScore, (info, stateQ)), q') ->
281 if defaultScoreFunction nextScore <= maxScore
282 then case nextStepsQueue stateQ of
283 Nothing -> getAllResults q'
284 Just (Left stateQ') -> case Q.minViewWithKey stateQ' of
285 Nothing -> getAllResults q'
286 Just ((newQscore,_), _) -> getAllResults (Q.add newQscore (info, stateQ') q')
287 Just (Right score) -> (info, score) : getAllResults q'
288 else []
291 ctxt = tt_ctxt istate
292 classInfo = idris_classes istate
294 (dag1, typeClassArgs1, retTy1) = makeDag type1
295 argNames1 = map fst dag1
296 makeDag :: Type -> (ArgsDAG, [(Name, Type)], Type)
297 makeDag = first3 (zipWith (\i (ty, deps) -> (ty, (i, deps))) [0..] . reverseDag) .
298 computeDagP (isTypeClassArg classInfo) . vToP
299 first3 f (a,b,c) = (f a, b, c)
301 -- update our state with the unification resolutions
302 resolveUnis :: [(Name, Type)] -> State -> Maybe (State, [(Type, Type)])
303 resolveUnis [] state = Just (state, [])
304 resolveUnis ((name, term@(P Bound name2 _)) : xs)
305 state@(State holes args1 args2 _ _ _ _) | isJust findArgs = do
306 ((ty1, ix1), (ty2, ix2)) <- findArgs
308 (state'', queue) <- resolveUnis xs state'
309 let transScore = fromMaybe 0 (abs <$> ((-) <$> ix1 <*> ix2))
310 return $ (inScore (\s -> s { transposition = transposition s + transScore }) state'', (ty1, ty2) : queue)
311 where
312 mgetType name xs = fmap ((snd . fst) &&& (fst . snd)) . find ((name ==) . fst . fst) $ xs
313 inScore f state = state { score = f (score state) }
314 findArgs = ((,) <$> findLeft name state <*> findRight name2 state) <|>
315 ((,) <$> findLeft name2 state <*> findRight name state)
316 matchnames = [name, name2]
317 holes' = filter (not . (`elem` matchnames) . fst) holes
318 deleteArgs = deleteLeft name . deleteLeft name2 . deleteRight name . deleteRight name2
319 state' = modifyTypes (subst name term) $ deleteArgs
320 (state { holes = holes'})
322 resolveUnis ((name, term) : xs)
323 state@(State holes args1 args2 _ _ _ _) = case (findLeft name state, findRight name state) of
324 (Just (_,ix), Nothing) -> first (addScore (mempty { leftApplied = 1, rightApplied = otherApplied})) <$> nextStep
325 (Nothing, Just (_,ix)) -> first (addScore (mempty { rightApplied = 1, leftApplied = otherApplied})) <$> nextStep
326 (Nothing, Nothing) -> nextStep
327 _ -> error ("Idris internal error: TypeSearch.resolveUnis")
328 where
329 -- find variables which are determined uniquely by the type
330 -- due to injectivity
331 matchedVarMap = usedVars term
332 both f (x, y) = (f x, f y)
333 (injUsedVars, notInjUsedVars) = both M.keys . M.partition id . M.filterWithKey (\k _-> k `elem` map fst holes) $ M.map snd matchedVarMap
334 varsInTy = injUsedVars ++ notInjUsedVars
335 toDelete = name : varsInTy
336 deleteMany = foldr (.) id $ [ deleteLeft t . deleteRight t | t <- toDelete ]
338 otherApplied = length notInjUsedVars
340 addScore additions state = state { score = score state `mappend` additions }
341 state' = modifyTypes (subst name term) . deleteMany $
342 state { holes = filter (not . (`elem` toDelete) . fst) holes }
343 nextStep = resolveUnis xs state'
346 -- | resolve a queue of unification constraints
347 unifyQueue :: State -> [(Type, Type)] -> Maybe State
348 unifyQueue state [] = return state
349 unifyQueue state ((ty1, ty2) : queue) = do
350 --trace ("go: \n" ++ show state) True `seq` return ()
351 res <- tcToMaybe $ match_unify ctxt [ (n, Pi ty) | (n, ty) <- holes state] ty1 ty2 [] (map fst $ holes state) []
352 (state', queueAdditions) <- resolveUnis res state
353 guard $ scoreCriterion (score state')
354 unifyQueue state' (queue ++ queueAdditions)
356 possClassInstances :: [Name] -> Type -> [Type]
357 possClassInstances usedNames ty = do
358 className <- getClassName clss
359 classDef <- lookupCtxt className classInfo
360 n <- class_instances classDef
361 def <- lookupCtxt n (definitions ctxt)
362 ty <- normaliseC ctxt [] <$> (case typeFromDef def of Just x -> [x]; Nothing -> [])
363 let ty' = vToP (uniqueBinders usedNames ty)
364 return ty'
365 where
366 (clss, _) = unApply ty
367 getClassName (P (TCon _ _) className _) = [className]
368 getClassName _ = []
370 -- Just if the computation hasn't totally failed yet, Nothing if it has
371 -- Left if we haven't found a terminal state, Right if we have
372 nextStepsQueue :: Q.PQueue Score State -> Maybe (Either (Q.PQueue Score State) Score)
373 nextStepsQueue queue = do
374 ((nextScore, next), rest) <- Q.minViewWithKey queue
375 if isFinal next
376 then Just $ Right nextScore
377 else let additions = if scoreCriterion nextScore
378 then Q.fromList [ (score state, state) | state <- nextSteps next ]
379 else Q.empty in
380 Just $ Left (Q.union rest additions)
381 where
382 isFinal (State [] [] [] [] [] score _) = True
383 isFinal _ = False
385 -- | Find all possible matches starting from a given state.
386 -- We go in stages rather than concatenating all three results in hopes of narrowing
387 -- the search tree. Once we advance in a phase, there should be no going back.
388 nextSteps :: State -> [State]
390 -- Stage 3 - match typeclasses
391 nextSteps (State [] [] [] c1 c2 scoreAcc usedNames) =
392 if null results3 then results4 else results3
393 where
394 -- try to match a typeclass argument from the left with a typeclass argument from the right
395 results3 =
396 catMaybes [ unifyQueue (State [] [] []
397 (deleteFromArgList n1 c1) (map (second subst2for1) (deleteFromArgList n2 c2)) scoreAcc usedNames) [(ty1, ty2)]
398 | (n1, ty1) <- c1, (n2, ty2) <- c2, let subst2for1 = psubst n2 (P Bound n1 ty1)]
400 -- try to hunt match a typeclass constraint by replacing it with an instance
401 results4 = results4A ++ results4B
402 typeClassArgs classes = [ ((n, ty), inst) | (n, ty) <- classes, inst <- possClassInstances usedNames ty ]
403 results4A = [ State [] [] []
404 (deleteFromArgList n c1 ++ newClassArgs) c2
405 (scoreAcc `mappend` (mempty { leftTypeClassApp = 1 }))
406 (usedNames ++ newHoles)
407 | ((n, ty), inst) <- typeClassArgs c1
408 , newClassArgs <- maybeToList $ typeclassUnify classInfo ctxt ty inst
409 , let newHoles = map fst newClassArgs ]
410 results4B = [ State [] [] []
411 c1 (deleteFromArgList n c2 ++ newClassArgs)
412 (scoreAcc `mappend` (mempty { rightTypeClassApp = 1 }))
413 (usedNames ++ newHoles)
414 | ((n, ty), inst) <- typeClassArgs c2
415 , newClassArgs <- maybeToList $ typeclassUnify classInfo ctxt ty inst
416 , let newHoles = map fst newClassArgs ]
419 -- Stage 1 - match arguments
420 nextSteps (State holes dag1 dag2 c1 c2 scoreAcc usedNames) = results where
422 results = concatMap takeSomeClasses results1
424 -- we only try to match arguments whose names don't appear in the types
425 -- of any other arguments
426 canBeFirst :: ArgsDAG -> [(Name, Type)]
427 canBeFirst = map fst . filter (S.null . snd . snd)
429 -- try to match an argument from the left with an argument from the right
430 results1 = catMaybes [ unifyQueue (State (filter (not . (`elem` [n1,n2]) . fst) holes) (deleteFromDag n1 dag1)
431 ((inArgTys subst2for1) $ deleteFromDag n2 dag2) c1 (map (second subst2for1) c2) scoreAcc usedNames) [(ty1, ty2)]
432 | (n1, ty1) <- canBeFirst dag1, (n2, ty2) <- canBeFirst dag2, let subst2for1 = psubst n2 (P Bound n1 ty1)]
435 -- Stage 2 - simply introduce a subset of the typeclasses
436 -- we've finished, so take some classes
437 takeSomeClasses (State [] [] [] c1 c2 scoreAcc usedNames) =
438 let lc1 = length c1; lc2 = length c2 in
439 [ State [] [] [] c1' c2' (scoreAcc `mappend`
440 mempty { rightTypeClassIntro = lc1 - length c1',
441 leftTypeClassIntro = lc2 - length c2' }) usedNames
442 | c1' <- subsets c1, c2' <- subsets c2 ]
443 -- still have arguments to match, so just be the identity
444 takeSomeClasses s = [s]