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