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