restructure nondeterministic choice in inference monad

examples/qselect.imp

@@ -9,9 +9,10 @@ def qselect():
     hi = N - 1
 
 
-    while [lo >= 0 ] lo < hi:
+    while [ lo >= 0 ] lo < hi:
         tick(hi - lo)
-        p = rand(N)
+        p = rand(hi-lo)
+        p = lo + p
 
         if p = pos:
             hi = lo

src/Data/PWhile/CostExpression.hs

@@ -233,8 +233,9 @@ norms c = S.fromList [n | GNorm _ _ n <- gNorms c]
 -- pretty printers
 
 instance PP.Pretty Norm where
-  -- pretty (Norm _ e) = PP.text "<" PP.<> PP.hang 0 (PP.pretty e) PP.<> PP.text ">"
-  pretty (Norm g e) = PP.brackets (PP.hang 0 (PP.group (PP.pretty e) PP.</> PP.text "|" PP.<+> PP.group (PP.pretty g)))
+  pretty (Norm _ e) = PP.text "<" PP.<> PP.hang 0 (PP.pretty e) PP.<> PP.text ">"
+
+--   pretty (Norm g e) = PP.brackets (PP.hang 0 (PP.group (PP.pretty e) PP.</> PP.text "|" PP.<+> PP.group (PP.pretty g)))
 
 instance (PP.Pretty c, Eq c, Num c) => PP.Pretty (CExp c) where
   pretty = pp id where

src/PWhile/InferEt.hs

 {-# OPTIONS_GHC -Wno-unused-local-binds -Wno-orphans -Wno-unused-top-binds -Wno-type-defaults #-}
-{-# LANGUAGE DeriveFunctor, FlexibleContexts, PatternSynonyms, ViewPatterns
+{-# LANGUAGE DeriveFunctor, FlexibleContexts, PatternSynonyms, ViewPatterns, TupleSections
            , NoOverloadedStrings #-}
 module PWhile.InferEt
   (
@@ -16,7 +16,7 @@ import           Data.List                    (nub)
 import           Control.Applicative
 import           Control.Monad.Except
 import qualified Control.Monad.State          as St
-import           Data.Maybe                   (listToMaybe)
+import           Data.Maybe                   (maybeToList, catMaybes, listToMaybe)
 import qualified Data.Map.Strict              as M
 import           Data.Tree                    (Forest)
 import qualified Data.Set                     as S
@@ -36,6 +36,7 @@ import qualified Data.PWhile.Expression       as E
 import qualified Data.PWhile.CostExpression   as C
 import           Data.PWhile.Program
 import qualified PWhile.Util                  as U
+import Data.Either (isRight)
 
 
 --- * Solver Monad ---------------------------------------------------------------------------------------------------
@@ -66,7 +67,7 @@ runSolveM' :: SolveM a -> L.ListT (MemoT IO) (Either String a, Forest String)
 runSolveM' = runTraceT . runExceptT . runSolveM_
 
 runSolveM :: SolveM a -> IO [(Either String a, Forest String)]
-runSolveM = runMemoT . L.toList . runTraceT . runExceptT . runSolveM_
+runSolveM = runMemoT . L.toList . runSolveM'
 
 
 
@@ -97,6 +98,19 @@ first p = go where
       Nothing    -> return []
       Just (a,l) -> if p a then return [a] else go l
 
+choose :: SolveM a -> SolveM a
+choose ma = do
+  res <- liftMemoT (go (runSolveM' ma) [])
+  case res of
+    Left ls -> U.logBlk "unsucessful trial(s)" $ mapM putTrace ls >> throwError "All branches failed"
+    Right (r, l) -> putTrace l >> return r
+  where
+    go s ls = do
+      a <- L.uncons s
+      case a of
+        Nothing -> return (Left ls)
+        Just ((Left _,l),s') -> go s' (l:ls)
+        Just ((Right r,l),_) -> return (Right (r,l))
 
 -- ** Logging
 
@@ -148,6 +162,7 @@ run t = fmap (map toResult) . runMemoT . select t . runSolveM' where
     success _                = False
   select (Take i) = L.toList . L.take i
   select (Pick i) = L.toList . L.drop (pred i) . L.take i
+
   select All      = L.toList
   toResult (Left err, l) = Failure err l
   toResult (Right a, l) = Success a l
@@ -261,6 +276,7 @@ boundedSum (i,f) (l,o) = bs f
             , cont = C.zero
             , limt = ns }
       (ui,s) <- alternative [(uiDown, o), (uiUp, l)]
+      logData "Invariant" ui
       E.substitute i s <$> (debugM "result" =<< solveInvariant ui)
 
 expectation :: E.Dist E.Exp -> (E.Exp -> CExp) -> SolveM CExp
@@ -402,12 +418,12 @@ ect c f =
   logData "f" f >> logData "Program" c *> etM Ect c f
 
 evt c f =
-  logBlk "Expected Cost" $
+  logBlk "Expected Value" $
   logData "f" f >> logData "Program" c *> etM Evt c f
 
-etM :: Knd -> C -> CExp -> SolveM CExp
-etM = memoized et
 
+etM :: Knd -> C -> CExp -> SolveM CExp
+etM k c f = choose (memoized et k c f)
 
 et :: Knd -> C -> CExp -> SolveM CExp
 et Ect c f
@@ -434,8 +450,10 @@ et t d@(While _  i b c) f  = logBlk "While.step" $ do
   logData "Problem" d
   logData "f" f
   g <- case t of {Evt -> return C.zero; Ect -> logBlk "Expected Cost Body" $ et Ect c C.zero}
-  ns <- logResult "Norms" $ extractRanking c i b g f
-  hs <- traverse (et Evt c . C.fromNorm) ns
+  candidates <- logResult "Norms" $ extractRanking c i b g f
+  let evnorm n = (Just <$> (n,) <$> evt c (C.fromNorm n)) `catchError` (\ _ -> return Nothing)
+  (ns,hs) <- unzip <$> catMaybes <$> traverse evnorm candidates
+--  hs <- traverse (et Evt c . C.fromNorm) ns
   let ui = UpperInvariant { inv  = i
                           , cnd = b
                           , cost = g