support for proof search with high-order functions

3 files changed, 84 insertions, 72 deletions

diff --git a/src/runtime/haskell/PGF/Forest.hs b/src/runtime/haskell/PGF/Forest.hs
index 0b481d0a2..5d06f4e97 100644
--- a/src/runtime/haskell/PGF/Forest.hs
+++ b/src/runtime/haskell/PGF/Forest.hs
@@ -203,13 +203,9 @@ foldForest f g b fcat forest =
 

 instance Selector FId where

   splitSelector s = (s,s)

-  select cat dp = TcM (\abstr s ms -> case Map.lookup cat (cats abstr) of

-                                        Just (_,fns) -> iter abstr s ms fns

-                                        Nothing      -> Fail s (UnknownCat cat))

+  select cat scope dp = do

+    gens <- typeGenerators scope cat

+    TcM (\abstr s ms -> iter s ms gens)

     where

-      iter abstr s ms []           = Zero

-      iter abstr s ms ((_,fn):fns) = Plus (select_helper fn abstr s ms) (iter abstr s ms fns)

-

-select_helper fn = unTcM $ do

-  ty <- lookupFunType fn

-  return (EFun fn,ty)

+      iter s ms []              = Zero

+      iter s ms ((_,e,tty):fns) = Plus (Ok s ms (e,tty)) (iter s ms fns)

diff --git a/src/runtime/haskell/PGF/Generate.hs b/src/runtime/haskell/PGF/Generate.hs
index 5ba822b3b..e79b6071f 100644
--- a/src/runtime/haskell/PGF/Generate.hs
+++ b/src/runtime/haskell/PGF/Generate.hs
@@ -82,18 +82,19 @@ generate sel pgf ty dp =
                                 sel emptyMetaStore]
 
 prove :: Selector sel => Maybe Int -> Scope -> TType -> TcM sel Expr
-prove dp scope (TTyp env1 (DTyp [] cat es1)) = do
-  (fe,DTyp hypos _ es2) <- select cat dp
+prove dp scope (TTyp env1 (DTyp hypos1 cat es1)) = do
+  vs1 <- mapM (PGF.TypeCheck.eval env1) es1
+  let scope' = exScope scope env1 hypos1
+  (fe,TTyp env2 (DTyp hypos2 _ es2)) <- select cat scope' dp
   if fe == EFun (mkCId "plus") then mzero else return ()
   case dp of
-    Just 0 | not (null hypos) -> mzero
-    _                         -> return ()
-  (env2,args) <- mkEnv [] hypos
-  vs1 <- mapM (PGF.TypeCheck.eval env1) es1
+    Just 0 | not (null hypos2) -> mzero
+    _                          -> return ()
+  (env2,args) <- mkEnv scope' env2 hypos2
   vs2 <- mapM (PGF.TypeCheck.eval env2) es2
-  sequence_ [eqValue mzero suspend (scopeSize scope) v1 v2 | (v1,v2) <- zip vs1 vs2]
-  es <- mapM descend args
-  return (foldl EApp fe es)
+  sequence_ [eqValue mzero suspend (scopeSize scope') v1 v2 | (v1,v2) <- zip vs1 vs2]
+  es <- mapM (descend scope') args
+  return (abs hypos1 (foldl EApp fe es))
   where
     suspend i c = do
       mv <- getMeta i
@@ -103,23 +104,33 @@ prove dp scope (TTyp env1 (DTyp [] cat es1)) = do
                                       setMeta i (MBound e)
                                       sequence_ [c e | c <- (c:cs)]
 
-    mkEnv env []                = return (env,[])
-    mkEnv env ((bt,x,ty):hypos) = do
+    abs []                e = e
+    abs ((bt,x,ty):hypos) e = EAbs bt x (abs hypos e)
+
+    exScope scope env []                = scope
+    exScope scope env ((bt,x,ty):hypos) = 
+       let env' | x /= wildCId = VGen (scopeSize scope) [] : env
+                | otherwise    = env
+       in exScope (addScopedVar x (TTyp env ty) scope) env' hypos
+
+    mkEnv scope env []                = return (env,[])
+    mkEnv scope env ((bt,x,ty):hypos) = do
       (env,arg) <- if x /= wildCId
                     then do i <- newMeta scope (TTyp env ty)
                             return (VMeta i (scopeEnv scope) [] : env,Right (EMeta i))
                     else return (env,Left (TTyp env ty))
-      (env,args) <- mkEnv env hypos
+      (env,args) <- mkEnv scope env hypos
       return (env,(bt,arg):args)
 
-    descend (bt,arg) = do let dp' = fmap (flip (-) 1) dp
-                          e <- case arg of
-                                 Right e  -> return e
-                                 Left tty -> prove dp' scope tty
-                          e <- case bt of
-                                 Implicit -> return (EImplArg e)
-                                 Explicit -> return e
-                          return e
+    descend scope (bt,arg) = do
+      let dp' = fmap (flip (-) 1) dp
+      e <- case arg of
+             Right e  -> return e
+             Left tty -> prove dp' scope tty
+      e <- case bt of
+             Implicit -> return (EImplArg e)
+             Explicit -> return e
+      return e
 
 
 -- Helper function for random generation. After every
@@ -137,50 +148,30 @@ restart g f =
 
 instance Selector () where
   splitSelector s = (s,s)
-  select cat dp
-    | cat == cidInt    = return (ELit (LInt 999),  DTyp [] cat [])
-    | cat == cidFloat  = return (ELit (LFlt 3.14), DTyp [] cat [])
-    | cat == cidString = return (ELit (LStr "Foo"),DTyp [] cat [])
-    | otherwise        = TcM (\abstr s ms -> case Map.lookup cat (cats abstr) of
-                                               Just (_,fns) -> iter abstr ms fns
-                                               Nothing      -> Fail s (UnknownCat cat))
+  select cat scope dp = do
+    gens <- typeGenerators scope cat
+    TcM (\abstr s ms -> iter ms gens)
     where
-      iter abstr ms []           = Zero
-      iter abstr ms ((_,fn):fns) = Plus (select_helper fn abstr () ms) (iter abstr ms fns)
+      iter ms []              = Zero
+      iter ms ((_,e,tty):fns) = Plus (Ok () ms (e,tty)) (iter ms fns)
+
 
 instance RandomGen g => Selector (Identity g) where
   splitSelector (Identity g) = let (g1,g2) = split g
                                in (Identity g1, Identity g2)
 
-  select cat dp
-    | cat == cidInt    = TcM (\abstr (Identity g) ms ->
-                                  let (n,g') = maybe random (\d -> randomR ((-10)*d,10*d)) dp g
-                                  in Ok (Identity g) ms (ELit (LInt n),DTyp [] cat []))
-    | cat == cidFloat  = TcM (\abstr (Identity g) ms ->
-                                  let (d,g') = maybe random (\d' -> let d = fromIntegral d'
-                                                                    in randomR ((-pi)*d,pi*d)) dp g
-                                  in Ok (Identity g) ms (ELit (LFlt d),DTyp [] cat []))
-    | cat == cidString = TcM (\abstr (Identity g) ms ->
-                                  let (g1,g2) = split g
-                                      s = take (fromMaybe 10 dp) (randomRs ('A','Z') g1)
-                                  in Ok (Identity g2) ms (ELit (LStr s),DTyp [] cat []))
-    | otherwise        = TcM (\abstr (Identity g) ms ->
-                                  case Map.lookup cat (cats abstr) of
-                                    Just (_,fns) -> do_rand abstr g ms 1.0 fns
-                                    Nothing      -> Fail (Identity g) (UnknownCat cat))
+  select cat scope dp = do
+    gens <- typeGenerators scope cat
+    TcM (\abstr (Identity g) ms -> do_rand abstr g ms 1.0 gens)
     where
-      do_rand abstr g ms p []  = Zero
-      do_rand abstr g ms p fns = let (d,g')    = randomR (0.0,p) g
-                                     (g1,g2)   = split g'
-                                     (p',fn,fns') = hit d fns
-                                 in Plus (select_helper fn abstr (Identity g1) ms) (do_rand abstr g2 ms (p-p') fns')
-
-      hit :: Double -> [(Double,a)] -> (Double,a,[(Double,a)])
-      hit d (px@(p,x):xs)
-        | d < p     = (p,x,xs)
-        | otherwise = let (p',x',xs') = hit (d-p) xs
-                      in (p,x',px:xs')
-
-select_helper fn = unTcM $ do
-  ty <- lookupFunType fn
-  return (EFun fn,ty)
+      do_rand abstr g ms p []   = Zero
+      do_rand abstr g ms p gens = let (d,g')    = randomR (0.0,p) g
+                                      (g1,g2)   = split g'
+                                      (p',e_ty,gens') = hit d gens
+                                  in Plus (Ok (Identity g1) ms e_ty) (do_rand abstr g2 ms (p-p') gens')
+
+      hit :: Double -> [(Double,Expr,TType)] -> (Double,(Expr,TType),[(Double,Expr,TType)])
+      hit d (gen@(p,e,ty):gens)
+        | d < p     = (p,(e,ty),gens)
+        | otherwise = let (p',e_ty',gens') = hit (d-p) gens
+                      in (p',e_ty',gen:gens')
diff --git a/src/runtime/haskell/PGF/TypeCheck.hs b/src/runtime/haskell/PGF/TypeCheck.hs
index fcfb58a9f..28ec016bc 100644
--- a/src/runtime/haskell/PGF/TypeCheck.hs
+++ b/src/runtime/haskell/PGF/TypeCheck.hs
@@ -23,14 +23,14 @@ module PGF.TypeCheck ( checkType, checkExpr, inferExpr
                      , Scope, emptyScope, scopeSize, scopeEnv, addScopedVar
                      , TcM(..), TcResult(..), runTcM, TType(..), Selector(..)
                      , tcExpr, infExpr, eqType, eqValue
-                     , lookupFunType, eval
+                     , lookupFunType, typeGenerators, eval
                      , generateForMetas, generateForForest, checkResolvedMetaStore
                      ) where
 
 import PGF.Data
 import PGF.Expr hiding (eval, apply, value2expr)
 import qualified PGF.Expr as Expr
-import PGF.Macros (typeOfHypo)
+import PGF.Macros (typeOfHypo, cidInt, cidFloat, cidString)
 import PGF.CId
 
 import Data.Map as Map
@@ -93,7 +93,7 @@ data TcResult s a
 
 class Selector s where
   splitSelector :: s -> (s,s)
-  select        :: CId -> Maybe Int -> TcM s (Expr,Type)
+  select        :: CId -> Scope -> Maybe Int -> TcM s (Expr,TType)
 
 instance Monad (TcM s) where
   return x = TcM (\abstr s ms -> Ok s ms x)
@@ -148,6 +148,31 @@ lookupFunType fun = TcM (\abstr s ms -> case Map.lookup fun (funs abstr) of
                                           Just (ty,_,_,_) -> Ok s ms ty
                                           Nothing         -> Fail s (UnknownFun fun))
 
+typeGenerators :: Scope -> CId -> TcM s [(Double,Expr,TType)]
+typeGenerators scope cat = fmap normalize (liftM2 (++) x y)
+  where
+    x = return
+           [(0.25,EVar i,tty) | (i,(_,tty@(TTyp _ (DTyp _ cat' _)))) <- zip [0..] gamma
+                              , cat == cat']
+         where
+           Scope gamma = scope
+
+    y | cat == cidInt    = return [(1.0,ELit (LInt 999),  TTyp [] (DTyp [] cat []))]
+      | cat == cidFloat  = return [(1.0,ELit (LFlt 3.14), TTyp [] (DTyp [] cat []))]
+      | cat == cidString = return [(1.0,ELit (LStr "Foo"),TTyp [] (DTyp [] cat []))]
+      | otherwise        = TcM (\abstr s ms ->
+                                    case Map.lookup cat (cats abstr) of
+                                      Just (_,fns) -> unTcM (mapM helper fns) abstr s ms
+                                      Nothing      -> Fail s (UnknownCat cat))
+
+    helper (p,fn) = do
+      ty <- lookupFunType fn
+      return (p,EFun fn,TTyp [] ty)
+      
+    normalize gens = [(p/s,e,tty) | (p,e,tty) <- gens]
+      where
+        s = sum [p | (p,_,_) <- gens]
+
 emptyMetaStore :: MetaStore s
 emptyMetaStore = IntMap.empty