@@ -1,3 +1,5 @@
1+{-# LANGUAGE ScopedTypeVariables #-}
2+
13module Idris.TypeSearch (
24 searchByType, searchPred, defaultScoreFunction
35) where
@@ -6,12 +8,12 @@ import Control.Applicative ((<$>), (<*>), (<|>))
68import Control.Arrow (first, second, (&&&))
79import Control.Monad (forM_, guard)
810
9−import Data.Function (on)
10−import Data.List (find, minimumBy, partition, sortBy, (\\))
11+import Data.List (find, delete, deleteBy, minimumBy, partition, sortBy, (\\))
1112import Data.Map (Map)
1213import qualified Data.Map as M
13−import Data.Maybe (catMaybes, fromMaybe, isJust, maybeToList)
14+import Data.Maybe (catMaybes, fromMaybe, isJust, maybeToList, mapMaybe)
1415import Data.Monoid (Monoid (mempty, mappend))
16+import Data.Ord (comparing)
1517import qualified Data.PriorityQueue.FingerTree as Q
1618import Data.Set (Set)
1719import qualified Data.Set as S
@@ -24,14 +26,14 @@ import Idris.AbsSyntaxTree (class_instances, ClassInfo, defaultSyntax, Idris,
2426import Idris.Core.Evaluate (Context (definitions), Def (Function, TyDecl, CaseOp), normaliseC)
2527import Idris.Core.TT hiding (score)
2628import Idris.Core.Unify (match_unify)
27−import Idris.Delaborate (delab, delabTy)
29+import Idris.Delaborate (delabTy)
2830import Idris.Docstrings (noDocs, overview)
2931import Idris.ElabDecls (elabType')
30−import Idris.Output (ihRenderResult, ihPrintResult, ihPrintFunTypes)
32+import Idris.Output (ihRenderResult)
3133
3234import System.IO (Handle)
3335
34−import Util.Pretty (text, vsep, char, (<>), Doc)
36+import Util.Pretty (text, char, (<>), Doc)
3537
3638searchByType :: Handle -> PTerm -> Idris ()
3739searchByType h pterm = do
@@ -42,51 +44,41 @@ searchByType h pterm = do
4244 ty <- elabType' False toplevel syn (fst noDocs) (snd noDocs) emptyFC [] n pterm'
4345 putIState i -- don't actually make any changes
4446 let names = searchUsing searchPred i ty
45− let names' = take numLimit $
46− sortBy (compare `on` getScore) names
47+ let names' = take numLimit $ names
4748 let docs =
4849 [ let docInfo = (n, delabTy i n, fmap (overview . fst) (lookupCtxtExact n (idris_docstrings i))) in
4950 displayScore score <> char ' ' <> prettyDocumentedIst i docInfo
50− | (n, (_,score)) <- names']
51− ihRenderResult h $ vsep docs
51+ | (n, score) <- names']
52+ mapM_ (ihRenderResult h) docs
5253 where
53− getScore = defaultScoreFunction . snd . snd
5454 numLimit = 50
5555 syn = defaultSyntax { implicitAllowed = True } -- syntax
5656 n = sMN 0 "searchType" -- name
5757
5858-- | Conduct a type-directed search using a given match predicate
59−searchUsing :: (IState -> Type -> Type -> Maybe a) -> IState -> Type -> [(Name, (Type, a))]
60−searchUsing pred istate ty =
61− concat . M.elems $ M.mapWithKey (\key -> M.toAscList . M.mapMaybe (f key)) (definitions ctxt)
59+searchUsing :: (IState -> Type -> [(Name, Type)] -> [(Name, a)])
60+ -> IState -> Type -> [(Name, a)]
61+searchUsing pred istate ty = pred istate nty . concat . M.elems $
62+ M.mapWithKey (\key -> M.toAscList . M.mapMaybe (f key)) (definitions ctxt)
6263 where
64+ nty = normaliseC ctxt [] ty
6365 ctxt = tt_ctxt istate
6466 f k x = do
6567 guard $ not (special k)
66− y <- get (fst4 x)
67− let ny = normaliseC ctxt [] y
68− val <- pred istate nty ny
69− return (y, val)
70− nty = normaliseC ctxt [] ty
71− fst4 :: (a,b,c,d) -> a
72− fst4 (w,x,y,z) = w
73− get :: Def -> Maybe Type
74− get (Function ty tm) = Just ty
75− get (TyDecl _ ty) = Just ty
76− -- get (Operator ty _ _) = Just ty
77− get (CaseOp _ ty _ _ _ _) = Just ty
78− get _ = Nothing
68+ type2 <- typeFromDef x
69+ return $ normaliseC ctxt [] type2
7970 special :: Name -> Bool
8071 special (NS n ns) = special n
8172 special (SN _) = True
82− special (UN n) = T.pack "default#" `T.isPrefixOf` n
73+ special (UN n) = T.pack "default#" `T.isPrefixOf` n
74+ || n `elem` map T.pack ["believe_me", "really_believe_me"]
8375 special _ = False
8476
8577-- Our default search predicate.
86−searchPred :: IState -> Type -> Type -> Maybe Score
78+searchPred :: IState -> Type -> [(Name, Type)] -> [(Name, Score)]
8779searchPred istate ty1 = matcher where
8880 maxScore = 100
89− matcher = unifyWithHoles istate maxScore ty1
81+ matcher = matchTypesBulk istate maxScore ty1
9082
9183
9284typeFromDef :: (Def, b, c, d) -> Maybe Type
@@ -98,6 +90,8 @@ typeFromDef (def, _, _, _) = get def where
9890 get (CaseOp _ ty _ _ _ _) = Just ty
9991 get _ = Nothing
10092
93+
94+-- | reverse the edges for a directed acyclic graph
10195reverseDag :: Ord k => [((k, a), Set k)] -> [((k, a), Set k)]
10296reverseDag xs = map f xs where
10397 f ((k, v), _) = ((k, v), S.fromList . map (fst . fst) $ filter (S.member k . snd) xs)
@@ -112,7 +106,7 @@ computeDagP :: Ord n
112106 -> TT n
113107 -> ([((n, TT n), Set n)], [(n, TT n)], TT n)
114108computeDagP removePred t = (reverse (map f args), reverse removedArgs , retTy) where
115− f (n, t) = ((n, t), M.keysSet (usedVars t))
109+ f (n, t) = ((n, t), M.keysSet (usedVars False t))
116110
117111 (numArgs, args, removedArgs, retTy) = go 0 [] [] t
118112
@@ -124,17 +118,21 @@ computeDagP removePred t = (reverse (map f args), reverse removedArgs , retTy) w
124118 go k args removedArgs retTy = (k, args, removedArgs, retTy)
125119
126120-- | Collect the names and types of all the free variables
127−usedVars :: Ord n => TT n -> Map n (TT n)
128−usedVars (V j) = error "unexpected! run vToP first"
129−usedVars (P Bound n t) = M.singleton n t `M.union` usedVars t
130−usedVars (Bind n binder t2) = (M.delete n (usedVars t2) `M.union`) $ case binder of
131− Let t v -> usedVars t `M.union` usedVars v
132− Guess t v -> usedVars t `M.union` usedVars v
133− b -> usedVars (binderTy b)
134−usedVars (App t1 t2) = usedVars t1 `M.union` usedVars t2
135−usedVars (Proj t _) = usedVars t
136−usedVars _ = M.empty
137−
121+usedVars :: Ord n
122+ => Bool -- ^ only collect variables which must be determined due to injectivity
123+ -> TT n -> Map n (TT n)
124+usedVars inj = f where
125+ f (P Bound n t) = M.singleton n t `M.union` f t
126+ f (Bind n binder t2) = (M.delete n (f t2) `M.union`) $ case binder of
127+ Let t v -> f t `M.union` f v
128+ Guess t v -> f t `M.union` f v
129+ b -> f (binderTy b)
130+ f (App t1 t2) = f t1 `M.union` (if not inj || isInjective t1 then f t2 else M.empty)
131+ f (Proj t _) = f t
132+ f (V j) = error "unexpected! run vToP first"
133+ f _ = M.empty
134+
135+-- | Remove a node from a directed acyclic graph
138136deleteFromDag :: Ord n => n -> [((n, TT n), (a, Set n))] -> [((n, TT n), (a, Set n))]
139137deleteFromDag name [] = []
140138deleteFromDag name (((name2, ty), (ix, set)) : xs) = (if name == name2
@@ -145,33 +143,36 @@ deleteFromArgList :: Ord n => n -> [(n, TT n)] -> [(n, TT n)]
145143deleteFromArgList n = filter ((/= n) . fst)
146144
147145data Score = Score
148− { transposition :: Int
149− , leftApplied :: Int
150− , rightApplied :: Int
151− , leftTypeClass :: Int
152− , rightTypeClass :: Int } deriving (Eq, Show)
146+ { transposition :: !Int
147+ , leftApplied :: !Int
148+ , rightApplied :: !Int
149+ , leftTypeClassApp :: !Int
150+ , rightTypeClassApp :: !Int
151+ , leftTypeClassIntro :: !Int
152+ , rightTypeClassIntro :: !Int } deriving (Eq, Show)
153153
154154displayScore :: Score -> Doc a
155−displayScore (Score trans lapp rapp lclass rclass) = text $ case (lt, gt) of
155+displayScore (Score trans lapp rapp lclassapp rclassapp lclassintro rclassintro) = text $ case (lt, gt) of
156156 (True , True ) -> "=" -- types are isomorphic
157157 (True , False) -> "<" -- found type is more general than searched type
158158 (False, True ) -> ">" -- searched type is more general than found type
159− (False, False) -> " "
160− where lt = lapp + lclass == 0
161− gt = rapp + rclass == 0
159+ (False, False) -> "_"
160+ where lt = lapp + lclassapp + lclassintro == 0
161+ gt = rapp + rclassapp + rclassintro == 0
162162
163163
164164scoreCriterion :: Score -> Bool
165−scoreCriterion (Score a b c d e) = not
166− ( (b > 0 && c > 0) || (b + c) > 4 || d > 3 || e > 3 )
165+scoreCriterion (Score a b c d e f g) = not
166+ ( (b > 0 && c > 0) || (b + c) > 4 || any (> 3) [d,e,f,g])
167167
168168defaultScoreFunction :: Score -> Int
169−defaultScoreFunction (Score a b c d e) = a + 9*b + 3*c + 12*d + 4*e + 100*(2*b + d)*(2*c + e)
169+defaultScoreFunction (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)
170170 -- it's very bad to have *both* upcasting and downcasting
171171
172172instance Monoid Score where
173− mempty = Score 0 0 0 0 0
174− (Score a b c d e) `mappend` (Score a' b' c' d' e') = Score (a + a') (b + b') (c + c') (d + d') (e + e')
173+ mempty = Score 0 0 0 0 0 0 0
174+ (Score a b c d e f g) `mappend` (Score a' b' c' d' e' f' g') =
175+ Score (a + a') (b + b') (c + c') (d + d') (e + e') (f + f') (g + g')
175176
176177-- | A directed acyclic graph representing the arguments to a function
177178-- The 'Int' represents the position of the argument (1st argument, 2nd, etc.)
@@ -179,12 +180,12 @@ type ArgsDAG = [((Name, Type), (Int, Set Name))]
179180
180181-- | The state corresponding to an attempted match of two types.
181182data State = State
182− { holes :: ![Name] -- ^ names which have yet to be resolved
183− , args1 :: !ArgsDAG -- ^ arguments for the left type which have yet to be resolved
184− , args2 :: !ArgsDAG -- ^ arguments for the right type which have yet to be resolved
185− , classes1 :: ![(Name, Type)] -- ^ typeclass arguments for the left type which haven't been resolved
186− , classes2 :: ![(Name, Type)] -- ^ typeclass arguments for the right type which haven't been resolved
187− , score :: !Score -- ^ the score so far
183+ { holes :: ![(Name, Type)] -- ^ names which have yet to be resolved
184+ , args1 :: !ArgsDAG -- ^ arguments for the left type which have yet to be resolved
185+ , args2 :: !ArgsDAG -- ^ arguments for the right type which have yet to be resolved
186+ , classes1 :: ![(Name, Type)] -- ^ typeclass arguments for the left type which haven't been resolved
187+ , classes2 :: ![(Name, Type)] -- ^ typeclass arguments for the right type which haven't been resolved
188+ , score :: !Score -- ^ the score so far
188189 , usedNames :: ![Name] -- ^ all names that have been used
189190 } deriving Show
190191
@@ -237,30 +238,58 @@ isTypeClassArg classInfo ty = not (null (getClassName clss >>= flip lookupCtxt c
237238
238239
239240instance Ord Score where
240− compare = compare `on` defaultScoreFunction
241+ compare = comparing defaultScoreFunction
242+
243+-- | Compute the power set
244+subsets :: [a] -> [[a]]
245+subsets [] = [[]]
246+subsets (x : xs) = let ss = subsets xs in map (x :) ss ++ ss
247+
241248
242249--DONT run vToP first!
243250-- | Try to match two types together in a unification-like procedure.
244−-- Returns a list of possible scores representing ways in which the two
245−-- types can be matched.
246−unifyWithHoles :: IState -> Int -> Type -> Type -> Maybe Score
247−unifyWithHoles istate maxScore type1 = \type2 -> let
248− (dag2, typeClassArgs2, retTy2) = makeDag (uniqueBinders argNames1 type2)
249− argNames2 = map (fst . fst) dag2
250− startingHoles = argNames1 ++ argNames2
251−
252− startingTypes = (retTy1, retTy2) : []
253− in do
254− state <- unifyQueue (State startingHoles dag1 dag2
251+-- Returns a list of types and their minimum scores, sorted in order
252+-- of increasing score.
253+matchTypesBulk :: forall info. IState -> Int -> Type -> [(info, Type)] -> [(info, Score)]
254+matchTypesBulk istate maxScore type1 types = getAllResults startQueueOfQueues where
255+ getStartQueue :: (info, Type) -> Maybe (Score, (info, Q.PQueue Score State))
256+ getStartQueue nty@(info, type2) = do
257+ state <- unifyQueue (State startingHoles dag1 dag2
255258 typeClassArgs1 typeClassArgs2
256− mempty startingHoles) startingTypes
257− getResults (Q.singleton (score state) state)
258− where
259+ mempty usedNames) startingTypes
260+ let sc = score state
261+ return $ (sc, (info, Q.singleton sc state))
262+ where
263+ (dag2, typeClassArgs2, retTy2) = makeDag (uniqueBinders (map fst argNames1) type2)
264+ argNames2 = map fst dag2
265+ usedNames = map fst (argNames1 ++ argNames2)
266+ startingHoles = argNames1 ++ argNames2
267+
268+ startingTypes = (retTy1, retTy2) : []
269+
270+
271+ startQueueOfQueues :: Q.PQueue Score (info, Q.PQueue Score State)
272+ startQueueOfQueues = Q.fromList $ mapMaybe getStartQueue types
273+
274+ getAllResults :: Q.PQueue Score (info, Q.PQueue Score State) -> [(info, Score)]
275+ getAllResults q = case Q.minViewWithKey q of
276+ Nothing -> []
277+ Just ((nextScore, (info, stateQ)), q') ->
278+ if defaultScoreFunction nextScore <= maxScore
279+ then case nextStepsQueue stateQ of
280+ Nothing -> getAllResults q'
281+ Just (Left stateQ') -> case Q.minViewWithKey stateQ' of
282+ Nothing -> getAllResults q'
283+ Just ((newQscore,_), _) -> getAllResults (Q.add newQscore (info, stateQ') q')
284+ Just (Right score) -> (info, score) : getAllResults q'
285+ else []
286+
287+
259288 ctxt = tt_ctxt istate
260289 classInfo = idris_classes istate
261290
262291 (dag1, typeClassArgs1, retTy1) = makeDag type1
263− argNames1 = map (fst . fst) dag1
292+ argNames1 = map fst dag1
264293 makeDag :: Type -> (ArgsDAG, [(Name, Type)], Type)
265294 makeDag = first3 (zipWith (\i (ty, deps) -> (ty, (i, deps))) [0..] . reverseDag) .
266295 computeDagP (isTypeClassArg classInfo) . vToP
@@ -282,25 +311,27 @@ unifyWithHoles istate maxScore type1 = \type2 -> let
282311 findArgs = ((,) <$> findLeft name state <*> findRight name2 state) <|>
283312 ((,) <$> findLeft name2 state <*> findRight name state)
284313 matchnames = [name, name2]
285− holes' = holes \\ matchnames
314+ holes' = filter (not . (`elem` matchnames) . fst) holes
286315 deleteArgs = deleteLeft name . deleteLeft name2 . deleteRight name . deleteRight name2
287316 state' = modifyTypes (subst name term) $ deleteArgs
288317 (state { holes = holes'})
289318
290319 resolveUnis ((name, term) : xs)
291320 state@(State holes args1 args2 _ _ _ _) = case (findLeft name state, findRight name state) of
292− (Just (_,ix), Nothing) -> first (inScore (\score -> score { leftApplied = succ (leftApplied score) })) <$> nextStep
293− (Nothing, Just (_, ix)) -> first (inScore (\score -> score { rightApplied = succ (rightApplied score) })) <$> nextStep
294− (Nothing, Nothing) -> nextStep
295− _ -> error ("Shouldn't happen. Watch the alpha conversion!\n" ++ show args1 ++ "\n\n" ++ show args2)
321+ (Just (_,ix), Nothing) -> first (inScore (\score -> score { leftApplied = succ (leftApplied score) })) <$> nextStep
322+ (Nothing, Just (_, ix)) -> first (inScore (\score -> score { rightApplied = succ (rightApplied score) })) <$> nextStep
323+ (Nothing, Nothing) -> nextStep
324+ _ -> error ("Idris internal error: TypeSearch.resolveUnis")
296325 where
297− varsInTy = map fst $ M.toList (usedVars term) --[]
326+ -- find variables which are determined uniquely by the type
327+ -- due to injectivity
328+ varsInTy = M.keys $ usedVars True term
298329 toDelete = name : varsInTy
299330 deleteMany = foldr (.) id $ [ deleteLeft t . deleteRight t | t <- toDelete ]
300331
301332 inScore f state = state { score = f (score state) }
302333 state' = modifyTypes (subst name term) . deleteMany $
303− state { holes = holes \\ toDelete }
334+ state { holes = filter (not . (`elem` toDelete) . fst) holes }
304335 nextStep = resolveUnis xs state'
305336
306337
@@ -309,7 +340,7 @@ unifyWithHoles istate maxScore type1 = \type2 -> let
309340 unifyQueue state [] = return state
310341 unifyQueue state ((ty1, ty2) : queue) = do
311342 --trace ("go: \n" ++ show state) True `seq` return ()
312− res <- tcToMaybe $ match_unify ctxt [] ty1 ty2 [] (holes state) []
343+ res <- tcToMaybe $ match_unify ctxt [ (n, Pi ty) | (n, ty) <- holes state] ty1 ty2 [] (map fst $ holes state) []
313344 (state', queueAdditions) <- resolveUnis res state
314345 guard $ scoreCriterion (score state')
315346 unifyQueue state' (queue ++ queueAdditions)
@@ -328,18 +359,17 @@ unifyWithHoles istate maxScore type1 = \type2 -> let
328359 getClassName (P (TCon _ _) className _) = [className]
329360 getClassName _ = []
330361
331−
332− getResults :: Q.PQueue Score State -> Maybe Score
333− getResults queue = do
362+ -- Just if the computation hasn't totally failed yet, Nothing if it has
363+ -- Left if we haven't found a terminal state, Right if we have
364+ nextStepsQueue :: Q.PQueue Score State -> Maybe (Either (Q.PQueue Score State) Score)
365+ nextStepsQueue queue = do
334366 ((nextScore, next), rest) <- Q.minViewWithKey queue
335367 if isFinal next
336− then return nextScore
337− else do
338− guard (defaultScoreFunction nextScore <= maxScore)
339− let additions = if scoreCriterion nextScore
340− then Q.fromList [ (score state, state) | state <- nextSteps next ]
341− else Q.empty
342− getResults (Q.union rest additions)
368+ then Just $ Right nextScore
369+ else let additions = if scoreCriterion nextScore
370+ then Q.fromList [ (score state, state) | state <- nextSteps next ]
371+ else Q.empty in
372+ Just $ Left (Q.union rest additions)
343373 where
344374 isFinal (State [] [] [] [] [] score _) = True
345375 isFinal _ = False
@@ -349,36 +379,39 @@ unifyWithHoles istate maxScore type1 = \type2 -> let
349379 -- the search tree. Once we advance in a phase, there should be no going back.
350380 nextSteps :: State -> [State]
351381
352− -- Stage 2 - match typeclasses
382+ -- Stage 3 - match typeclasses
353383 nextSteps (State [] [] [] c1 c2 scoreAcc usedNames) =
354− if null results2 then results3 else results2
384+ if null results3 then results4 else results3
355385 where
356386 -- try to match a typeclass argument from the left with a typeclass argument from the right
357− results2 =
387+ results3 =
358388 catMaybes [ unifyQueue (State [] [] []
359389 (deleteFromArgList n1 c1) (map (second subst2for1) (deleteFromArgList n2 c2)) scoreAcc usedNames) [(ty1, ty2)]
360390 | (n1, ty1) <- c1, (n2, ty2) <- c2, let subst2for1 = psubst n2 (P Bound n1 ty1)]
361391
362392 -- try to hunt match a typeclass constraint by replacing it with an instance
363− results3 = results3A ++ results3B
393+ results4 = results4A ++ results4B
364394 typeClassArgs classes = [ ((n, ty), inst) | (n, ty) <- classes, inst <- possClassInstances usedNames ty ]
365− results3A = [ State [] [] []
395+ results4A = [ State [] [] []
366396 (deleteFromArgList n c1 ++ newClassArgs) c2
367− (scoreAcc `mappend` (mempty { leftTypeClass = 1 }))
397+ (scoreAcc `mappend` (mempty { leftTypeClassApp = 1 }))
368398 (usedNames ++ newHoles)
369399 | ((n, ty), inst) <- typeClassArgs c1
370400 , newClassArgs <- maybeToList $ typeclassUnify classInfo ctxt ty inst
371401 , let newHoles = map fst newClassArgs ]
372− results3B = [ State [] [] []
402+ results4B = [ State [] [] []
373403 c1 (deleteFromArgList n c2 ++ newClassArgs)
374− (scoreAcc `mappend` (mempty { rightTypeClass = 1 }))
404+ (scoreAcc `mappend` (mempty { rightTypeClassApp = 1 }))
375405 (usedNames ++ newHoles)
376406 | ((n, ty), inst) <- typeClassArgs c2
377407 , newClassArgs <- maybeToList $ typeclassUnify classInfo ctxt ty inst
378408 , let newHoles = map fst newClassArgs ]
379409
410+
380411 -- Stage 1 - match arguments
381− nextSteps (State holes dag1 dag2 c1 c2 scoreAcc usedNames) = results1 where
412+ nextSteps (State holes dag1 dag2 c1 c2 scoreAcc usedNames) = results where
413+
414+ results = concatMap takeSomeClasses results1
382415
383416 -- we only try to match arguments whose names don't appear in the types
384417 -- of any other arguments
@@ -386,6 +419,18 @@ unifyWithHoles istate maxScore type1 = \type2 -> let
386419 canBeFirst = map fst . filter (S.null . snd . snd)
387420
388421 -- try to match an argument from the left with an argument from the right
389− results1 = catMaybes [ unifyQueue (State (holes \\ [n1, n2]) (deleteFromDag n1 dag1)
422+ results1 = catMaybes [ unifyQueue (State (filter (not . (`elem` [n1,n2]) . fst) holes) (deleteFromDag n1 dag1)
390423 ((inArgTys subst2for1) $ deleteFromDag n2 dag2) c1 (map (second subst2for1) c2) scoreAcc usedNames) [(ty1, ty2)]
391424 | (n1, ty1) <- canBeFirst dag1, (n2, ty2) <- canBeFirst dag2, let subst2for1 = psubst n2 (P Bound n1 ty1)]
425+
426+
427+ -- Stage 2 - simply introduce a subset of the typeclasses
428+ -- we've finished, so take some classes
429+ takeSomeClasses (State [] [] [] c1 c2 scoreAcc usedNames) =
430+ let lc1 = length c1; lc2 = length c2 in
431+ [ State [] [] [] c1' c2' (scoreAcc `mappend`
432+ mempty { rightTypeClassIntro = lc1 - length c1',
433+ leftTypeClassIntro = lc2 - length c2' }) usedNames
434+ | c1' <- subsets c1, c2' <- subsets c2 ]
435+ -- still have arguments to match, so just be the identity
436+ takeSomeClasses s = [s]