Inspect entities

1 files changed, 21 insertions, 5 deletions

diff --git a/src/Sequence/Formula.hs b/src/Sequence/Formula.hs
index a3675c6..2ac1210 100644
--- a/src/Sequence/Formula.hs
+++ b/src/Sequence/Formula.hs
@@ -4,6 +4,7 @@ module Sequence.Formula
   ( FormulaM, Formula, quot'
   , (:<:)(..), Context(..), ctx
   , evalFormula, evalFormula'
+  , findDistribution, findDistribution'
   , val
   , d, z
   , Table, table
@@ -109,18 +110,33 @@ askQuestion promptPref input q@(Question{..}) = flip (if keepResult then set $ _
     sep = "»"
 
 evalFormula :: (MonadIO m, sInput :<: lInput) => [String] -> lInput -> FormulaM sInput a -> m (lInput, a)
-evalFormula promptPref input = evalFormula' [] promptPref (input, Nothing)
+evalFormula = primEvalFormula $ liftIO . enact
+
+findDistribution :: (MonadIO m, sInput :<: lInput, Ord a) => [String] -> lInput -> FormulaM sInput a -> m (lInput, (Map a Rational))
+findDistribution = primEvalFormula $ return . fmap Map.fromList . seqEither . outcomes
   where
-    evalFormula' :: (MonadIO m, sInput :<: lInput) => [lInput -> lInput] -> [String] -> (lInput, Maybe (Formula sInput)) -> FormulaM sInput a -> m (lInput, a)
-    evalFormula' finalChanges promptPref (input, fSt) formula = do
-      result <- liftIO . enact . runExceptT . (runReaderT ?? (Context input fSt)) . (evalStateT ?? Set.empty) $ formula
+    seqEither :: [(Either q a, Rational)] -> Either q [(a, Rational)]
+    seqEither = mapM seqEither'
+
+    seqEither' (Left  a, _) = Left a
+    seqEither' (Right b, c) = Right (b, c)
+
+primEvalFormula :: (MonadIO m, sInput :<: lInput) => (EventM (Either (Question sInput) a) -> m (Either (Question sInput) r)) -> [String] -> lInput -> FormulaM sInput a -> m (lInput, r)
+primEvalFormula fromOutcomes promptPref input = evalFormula' fromOutcomes [] promptPref (input, Nothing)
+  where
+    evalFormula' :: (MonadIO m, sInput :<: lInput) => (EventM (Either (Question sInput) a) -> m (Either (Question sInput) r)) -> [lInput -> lInput] -> [String] -> (lInput, Maybe (Formula sInput)) -> FormulaM sInput a -> m (lInput, r)
+    evalFormula' fromOutcomes finalChanges promptPref (input, fSt) formula = do
+      result <- fromOutcomes . runExceptT . (runReaderT ?? (Context input fSt)) . (evalStateT ?? Set.empty) $ formula
       case result of
-        Left q@(Question{..}) -> askQuestion promptPref (input, fSt) q >>= flip (flip evalFormula' promptPref $ set (ctx' . answer) (throwError q) : finalChanges) formula
+        Left q@(Question{..}) -> askQuestion promptPref (input, fSt) q >>= flip (flip (evalFormula' fromOutcomes) promptPref $ set (ctx' . answer) (throwError q) : finalChanges) formula
         Right result -> return (foldr ($) input finalChanges, result)
 
 evalFormula' :: (MonadIO m, sInput :<: lInput, MonadState lInput m) => [String] -> FormulaM sInput a -> m a
 evalFormula' promptPref formula = uncurry (<$) . over _2 put . swap =<< flip (evalFormula promptPref) formula =<< get
 
+findDistribution' :: (MonadIO m, sInput :<: lInput, MonadState lInput m, Ord a) => [String] -> FormulaM sInput a -> m (Map a Rational)
+findDistribution' promptPref formula = uncurry (<$) . over _2 put . swap =<< flip (findDistribution promptPref) formula =<< get
+
 val :: Traversal' input (Formula input) -> [String] -> Bool -> Formula input
 val answer prompt keepResult = do
   gets (Set.member prompt) >>= bool (modify $ Set.insert prompt) (modify (Set.delete prompt) >> throwError Question{..})