@@ -0,0 +1,290 @@
1+module Idris.TypeSearch (
2+ searchByType, searchPred, defaultScoreFunction
3+) where
4+
5+import Control.Applicative ((<$>), (<*>), (<|>))
6+import Control.Arrow (first, second, (&&&))
7+import Control.Monad (forM_, guard)
8+
9+import Data.Function (on)
10+import Data.List (find, sortBy, (\\))
11+import qualified Data.Map as M
12+import Data.Maybe (catMaybes, fromMaybe, isJust)
13+import Data.Monoid (Monoid (mempty, mappend))
14+import Data.Set (Set)
15+import qualified Data.Set as S
16+
17+import Idris.AbsSyntax (addUsingConstraints, addImpl, getContext, getIState, putIState, implicit)
18+import Idris.AbsSyntaxTree (class_instances, defaultSyntax, Idris,
19+ IState (idris_classes, tt_ctxt),
20+ implicitAllowed, prettyIst, PTerm, toplevel)
21+import Idris.Core.Evaluate (Context (definitions), Def (Function, TyDecl, CaseOp), normaliseC)
22+import Idris.Core.TT
23+import Idris.Core.Unify (match_unify)
24+import Idris.Delaborate (delab, delabTy)
25+import Idris.Docstrings (noDocs)
26+import Idris.ElabDecls (elabType')
27+import Idris.Output (ihRenderResult, ihPrintResult, ihPrintFunTypes)
28+
29+import System.IO (Handle)
30+
31+
32+searchByType :: (Ord a, Show a) => Handle -> (IState -> Type -> Type -> Maybe a) -> a -> PTerm -> Idris ()
33+searchByType h pred scoreLimit pterm = do
34+ pterm' <- addUsingConstraints syn emptyFC pterm
35+ pterm'' <- implicit toplevel syn n pterm'
36+ i <- getIState
37+ let pterm''' = addImpl i pterm''
38+ ty <- elabType' False toplevel syn (fst noDocs) (snd noDocs) emptyFC [] n pterm'
39+ putIState i -- don't actually make any changes
40+ ihRenderResult h (prettyIst i pterm)
41+ let names = searchUsing pred i ty
42+ let names' = take numLimit . takeWhile ((< scoreLimit) . snd . snd) $
43+ sortBy (compare `on` (snd . snd)) names
44+ forM_ names' $ \(name, (typ, val)) -> do
45+ ihPrintFunTypes h [] name [(name, delabTy i name)]
46+ ihPrintResult h ("\tScore: " ++ show val ++ "\n")
47+ where
48+ numLimit = 50
49+ syn = defaultSyntax { implicitAllowed = True } -- syntax
50+ n = sMN 0 "searchType" -- name
51+
52+
53+searchUsing :: (IState -> Type -> Type -> Maybe a) -> IState -> Type -> [(Name, (Type, a))]
54+searchUsing pred istate ty =
55+ concat . M.elems $ M.mapWithKey (\key -> M.toAscList . M.mapMaybe (f key)) (definitions ctxt)
56+ where
57+ ctxt = tt_ctxt istate
58+ f k x = do
59+ guard $ not (special k)
60+ y <- get (fst4 x)
61+ let ny = normaliseC ctxt [] y
62+ -- traceShow k False `seq` return ()
63+ val <- pred istate nty ny
64+ return (y, val)
65+ nty = normaliseC ctxt [] ty
66+ fst4 :: (a,b,c,d) -> a
67+ fst4 (w,x,y,z) = w
68+ get :: Def -> Maybe Type
69+ get (Function ty tm) = Just ty
70+ get (TyDecl _ ty) = Just ty
71+ -- get (Operator ty _ _) = Just ty
72+ get (CaseOp _ ty _ _ _ _) = Just ty
73+ get _ = Nothing
74+ special :: Name -> Bool
75+ special (SN _) = True
76+ special _ = False
77+
78+
79+tcToMaybe :: TC' e a -> Maybe a
80+tcToMaybe (OK x) = Just x
81+tcToMaybe (Error _) = Nothing
82+
83+
84+
85+searchPred :: (Score -> Int) -> IState -> Type -> Type -> Maybe Int
86+searchPred scoref istate ty1 = \ty2 -> case matcher ty2 of
87+ Nothing -> Nothing
88+ Just xs -> guard (not (null xs)) >> return (minimum (map scoref xs))
89+ where
90+ matcher = unifyWithHoles True istate ty1
91+
92+
93+
94+
95+reverseDag :: Ord k => [((k, a), Set k)] -> [((k, a), Set k)]
96+reverseDag xs = map f xs where
97+ f ((k, v), _) = ((k, v), S.fromList . map (fst . fst) $ filter (S.member k . snd) xs)
98+
99+-- run vToP first!
100+-- returns [(the name and type of the bound variable
101+-- the names in the type of the bound variable)]
102+computeDagP :: Ord n => TT n -> ([((n, TT n), Set n)], TT n)
103+computeDagP t = (reverse (map f args), retTy) where
104+
105+ f (n, t) = ((n, t), usedVars t)
106+
107+ (numArgs, args, retTy) = go 0 [] t
108+
109+ -- NOTE : args are in reverse order
110+ go k args (Bind n (Pi t) sc) = go (succ k) ( (n, t) : args ) sc
111+ go k args retTy = (k, args, retTy)
112+
113+ usedVars :: Ord n => TT n -> Set n
114+ usedVars (V j) = error "unexpected! run vToP first"
115+ usedVars (P Bound n t) = S.singleton n `S.union` usedVars t
116+ usedVars (Bind n binder t2) = (S.delete n (usedVars t2) `S.union`) $ case binder of
117+ Let t v -> usedVars t `S.union` usedVars v
118+ Guess t v -> usedVars t `S.union` usedVars v
119+ b -> usedVars (binderTy b)
120+ usedVars (App t1 t2) = usedVars t1 `S.union` usedVars t2
121+ usedVars (Proj t _) = usedVars t
122+ usedVars _ = S.empty
123+
124+
125+deleteFromDag :: Ord n => n -> [((n, TT n), (a, Set n))] -> [((n, TT n), (a, Set n))]
126+deleteFromDag name [] = []
127+deleteFromDag name (((name2, ty), (ix, set)) : xs) = (if name == name2
128+ then id
129+ else (((name2, ty) , (ix, S.delete name set)) :) ) (deleteFromDag name xs)
130+
131+
132+data Score = Score
133+ { transposition :: Int
134+ , leftApplied :: Int
135+ , rightApplied :: Int
136+ , leftTypeClass :: Int
137+ , rightTypeClass :: Int } deriving (Eq, Show)
138+
139+
140+scoreCriterion :: Score -> Bool
141+scoreCriterion (Score a b c d e) = True {- not
142+ ( (b > 0 && c > 0) || (b + c) > 2 ) -}
143+
144+defaultScoreFunction :: Score -> Int
145+defaultScoreFunction (Score a b c d e) = a + 9*b + 3*c + 12*d + 4*e + 100*(2*b + d)*(2*c + e)
146+ -- it's very bad to have *both* upcasting and downcasting
147+
148+instance Monoid Score where
149+ mempty = Score 0 0 0 0 0
150+ (Score a b c d e) `mappend` (Score a' b' c' d' e') = Score (a + a') (b + b') (c + c') (d + d') (e + e')
151+
152+
153+type ArgsDAG = [((Name, Type), (Int, Set Name))]
154+type ResType = ( [Name] , ArgsDAG , ArgsDAG )
155+
156+data State = State
157+ { holes :: ![Name]
158+ , args1 :: !ArgsDAG
159+ , args2 :: !ArgsDAG
160+ , score :: !Score
161+ }
162+
163+
164+--DONT run vToP first!
165+unifyWithHoles :: Bool -> IState -> Type -> Type -> Maybe [Score]
166+unifyWithHoles debugParam istate type1 = \type2 -> let
167+ (dag2, retTy2) = makeDag (uniqueBinders argNames1 type2)
168+ argNames2 = map (fst . fst) dag2
169+ startingHoles = argNames1 ++ argNames2
170+
171+ startingTypes = (retTy1, retTy2) : []
172+ in do
173+ state <- go (State startingHoles dag1 dag2 mempty) startingTypes
174+ return $ processDags state
175+ where
176+ ctxt = tt_ctxt istate
177+ classInfo = idris_classes istate
178+ (dag1, retTy1) = makeDag type1
179+ argNames1 = map (fst . fst) dag1
180+ makeDag = first (zipWith (\i (ty, deps) -> (ty, (i, deps))) [0..] . reverseDag) . computeDagP . vToP
181+
182+ matchf :: (Name, Term) -> Maybe (Name, Name)
183+ matchf (name, P Bound name2 _) = Just (name, name2)
184+ matchf _ = Nothing
185+
186+ -- update our state with the unification resolutions
187+ updateDags :: [(Name, Type)] -> ResType -> Maybe (ResType, [(Type, Type)], Score)
188+ updateDags [] res = Just (res, [], mempty)
189+ updateDags ((name, term@(P Bound name2 _)) : xs) (holes, args1, args2) | isJust findArgs = do
190+ ((ty1, ix1), (ty2, ix2)) <- findArgs
191+
192+ (res, queue, score) <- updateDags xs (holes', args1'', args2'')
193+ --traceShow (ty1, ty2) False `seq` return ()
194+ return $ (res, (ty1, ty2) : queue , score { transposition = transposition score + abs (ix2 - ix1) })
195+ where
196+ findArgs = ((,) <$> mgetType name args1 <*> mgetType name2 args2) <|>
197+ ((,) <$> mgetType name2 args1 <*> mgetType name args2)
198+ matchnames = [name, name2]
199+ holes' = holes \\ matchnames
200+ substf = deleteFromDag name . deleteFromDag name2
201+ args1' = substf args1
202+ args2' = substf args2
203+ args1'' = map (first . second $ subst name term) args1'
204+ args2'' = map (first . second $ subst name term) args2'
205+ mgetType name xs = fmap ((snd . fst) &&& (fst . snd)) . find ((name ==) . fst . fst) $ xs
206+
207+ updateDags ((name, term) : xs) (holes, args1, args2) = case (mgetType name args1, mgetType name args2) of
208+ (Just _, Nothing) -> thrd (\score -> score { leftApplied = succ (leftApplied score) }) <$>
209+ updateDags xs (holes', updatef args1, args2)
210+ (Nothing, Just _) -> thrd (\score -> score { rightApplied = succ (rightApplied score) }) <$>
211+ updateDags xs (holes', args1, updatef args2)
212+ _ -> error "Shouldn't happen. Watch the alpha conversion!"
213+ where
214+ thrd f (a,b,c) = (a,b,f c)
215+ holes' = holes \\ [name]
216+ updatef = map (first . second $ subst name term) . deleteFromDag name
217+ args1' = deleteFromDag name args1
218+ args1'' = map (first . second $ subst name term) args1'
219+ mgetType name xs = fmap (snd . fst) . find ((name ==) . fst . fst) $ xs
220+
221+
222+ go :: State -> [(Type, Type)] -> Maybe State
223+ --go (State holes args1 args2 score) queue | trace ("go\n\t" ++ show holes ++ "\n\t" ++ show args1 ++ "\n\t" ++ show args2 ++ "\n\t" ++ show queue) False = undefined
224+ go state [] = return state
225+ go (State holes args1 args2 score) ((ty1, ty2) : queue) = do
226+ res <- tcToMaybe $ match_unify ctxt [] ty1 ty2 [] holes []
227+ --trace ("UnifyResult: " ++ show (ty1, ty2, res, errors)) False `seq` return ()
228+ --guard (null errors)
229+ ((holes', args1', args2'), queueAdditions, scoreAdditions) <- updateDags res (holes, args1, args2)
230+ let newScore = score `mappend` scoreAdditions
231+ guard $ scoreCriterion newScore
232+ go (State holes' args1' args2' newScore) (queue ++ queueAdditions)
233+
234+
235+ processDags :: State -> [Score]
236+ processDags (State [] [] [] scoreAcc) = [scoreAcc]
237+ --processDags (State holes (_:_) [] scoreAcc) = []
238+ --processDags (State holes [] (_:_) scoreAcc) = []
239+ processDags (State holes dag1 dag2 scoreAcc) = concat [ processDags state | state <- allResults ] where
240+
241+ results = catMaybes [ go (State (holes \\ (map nameOf [ty1, ty2])) (deleteFromDag (nameOf ty1) dag1)
242+ (inArgTys (psubst (nameOf ty2) (P Bound (nameOf ty1) (typeOf ty1))) $ deleteFromDag (nameOf ty2) dag2) scoreAcc) [(typeOf ty1, typeOf ty2)]
243+ | ty1 <- canBeFirst dag1, ty2 <- canBeFirst dag2 {-, exactTypeEquality ctxt (typeOf ty1) (typeOf ty2) -} ]
244+
245+ results2 = [ State (holes \\ [nameOf ty])
246+ (deleteFromDag (nameOf ty) dag1) dag2
247+ (scoreAcc `mappend` (mempty { leftTypeClass = 1 }))
248+ | ty <- typeClassArgs1 ]
249+
250+ typeClassArgs1 = filter (isSaturatedClass . typeOf) dag1
251+ typeClassArgs2 = filter (isSaturatedClass . typeOf) dag2
252+
253+
254+ results3 = [ State (holes \\ [nameOf ty])
255+ dag1 (deleteFromDag (nameOf ty) dag2)
256+ (scoreAcc `mappend` (mempty { rightTypeClass = 1 }))
257+ | ty <- typeClassArgs2 ]
258+
259+ allResults :: [State]
260+ allResults = {- (if not (null typeClasses) then (traceShow typeClasses False `seq` id) else id ) -}
261+ (results ++ results2 ++ results3)
262+ where
263+ typeClasses = filter (isSaturatedClass . typeOf) (dag1 ++ dag2)
264+
265+
266+ -- check if the canBeFirst thing is losing any possibilities
267+
268+
269+ inArgTys = map . first . second
270+ typeOf ((name, ty), set) = ty
271+ nameOf ((name, ty), set) = name
272+
273+ -- XXX : debug stuff
274+ canBeFirst = if debugParam then filter (S.null . snd . snd) else id
275+ holes = map (fst . fst) dag1 ++ map (fst . fst) dag2
276+
277+ deleteIdx _ [] = []
278+ deleteIdx idx l@(x@(i,_,_) : xs) = if i == idx then xs else x : deleteIdx idx xs
279+
280+ isSaturatedClass :: Type -> Bool
281+ isSaturatedClass ty = fromMaybe False $ do
282+ className <- getClassName clss
283+ let possInstances = concatMap class_instances $ lookupCtxt className classInfo
284+ return $ (SN (sInstanceN className (map argToName args))) `elem` possInstances
285+ where
286+ (clss, args) = unApply ty
287+ getClassName (P (TCon _ _) className _) = Just className
288+ getClassName _ = Nothing
289+ argToName arg = show (delab istate arg)
290+