3 files changed, 37 insertions, 13 deletions

diff --git a/src/Main.hs b/src/Main.hs
index 06cc6ed..f4a863f 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -2,7 +2,7 @@
 
 import Control.Monad
 
-import Control.Lens
+import Control.Lens hiding (Context(..))
 
 import System.Console.Shell
 import System.Console.Shell.ShellMonad
@@ -146,7 +146,7 @@ rollTest = withArg $ enactTest' >=> maybe (return ()) (shellPutStrLn . ppResult)
 
 enactTest' :: (FormulaM Stats Test) -> Sh GameState (Maybe TestResult)
 enactTest' test = withFocus' $ \focus -> do
-  (newStats, result) <- evalFormula (view eStats focus) (enactTest =<< test)
-  gFocus'.eStats .= newStats
+  (newFocus, result) <- evalFormula focus (enactTest =<< test)
+  gFocus' .= newFocus
   return result
       
diff --git a/src/Sequence/Formula.hs b/src/Sequence/Formula.hs
index 0d99fc0..9000c58 100644
--- a/src/Sequence/Formula.hs
+++ b/src/Sequence/Formula.hs
@@ -1,13 +1,14 @@
-{-# LANGUAGE RecordWildCards, RankNTypes, TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses, ViewPatterns, TypeFamilies #-}
+{-# LANGUAGE RecordWildCards, RankNTypes, TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses, ViewPatterns, TypeFamilies, GADTs, TypeOperators, ExistentialQuantification, FlexibleContexts #-}
 
 module Sequence.Formula
   ( FormulaM, Formula, quot'
+  , (:<:)(..), Context(..)
   , evalFormula
   , val
   , d, z
   ) where
 
-import Control.Lens
+import Control.Lens hiding (Context(..))
 import Data.Data.Lens
 
 import Control.Monad
@@ -30,8 +31,27 @@ import Data.Either
 import Data.Set (Set)
 import qualified Data.Set as Set
 
+class (:<:) small large where
+  ctx' :: Traversal' large small
 
-type FormulaM input a = StateT (Set String) (ReaderT input (ExceptT (Question input) EventM)) a
+instance a :<: a where
+  ctx' = simple
+
+instance a :<: (a, a) where
+  ctx' = both
+
+instance a :<: (a, b) where
+  ctx' = _1
+
+instance a :<: (b, a) where
+  ctx' = _2
+
+data Context small = forall large. (small :<: large) => Context large
+
+ctx :: Traversal' (Context input) input
+ctx modifySmall (Context large) = Context <$> ctx' modifySmall large
+
+type FormulaM input a = StateT (Set String) (ReaderT (Context input) (ExceptT (Question input) EventM)) a
 
 type Formula input = FormulaM input Int
 
@@ -66,22 +86,22 @@ instance Integral a => Num (FormulaM input a) where
 quot' :: Integral a => FormulaM input a -> FormulaM input a -> FormulaM input a
 quot' = liftM2 quot
 
-askQuestion :: MonadIO m => input -> (Question input) -> m input
-askQuestion input q@(Question{..}) = flip (set answer) input . maybe (throwError q) return <$> askQ prompt (join . fmap readMaybe)
+askQuestion :: (MonadIO m, sInput :<: lInput) => lInput -> Question sInput -> m lInput
+askQuestion input q@(Question{..}) = flip (set $ ctx' . answer) input . maybe (throwError q) return <$> askQ prompt (join . fmap readMaybe)
 
-evalFormula :: (MonadIO m, Ord a, Show a) => input -> FormulaM input a -> m (input, a)
+evalFormula :: (MonadIO m, sInput :<: lInput) => lInput -> FormulaM sInput a -> m (lInput, a)
 evalFormula = evalFormula' []
   where
     evalFormula' finalChanges input formula = do
-      result <- liftIO . enact . runExceptT . (runReaderT ?? input) . (evalStateT ?? Set.empty) $ formula
+      result <- liftIO . enact . runExceptT . (runReaderT ?? (Context input)) . (evalStateT ?? Set.empty) $ formula
       case result of
-        Left q@(Question{..}) -> askQuestion input q >>= flip (evalFormula' $ bool (pure . set answer $ throwError q) mempty keepResult ++ finalChanges) formula
+        Left q@(Question{..}) -> askQuestion input q >>= flip (evalFormula' $ bool (pure . set (ctx' . answer) $ throwError q) mempty keepResult ++ finalChanges) formula
         Right result -> return (foldr ($) input finalChanges, result)
 
 val :: Integral a => 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{..})
-  preview answer >>= maybe (throwError Question{..}) id
+  preview (ctx . answer) >>= maybe (throwError Question{..}) id
 
 d, z :: Integral a => Int -> FormulaM input a
 d n = liftBase . fmap fromIntegral $ D.d n
diff --git a/src/Sequence/Types.hs b/src/Sequence/Types.hs
index 541505c..480dfee 100644
--- a/src/Sequence/Types.hs
+++ b/src/Sequence/Types.hs
@@ -1,4 +1,4 @@
-{-# LANGUAGE FlexibleInstances, FlexibleContexts, MultiParamTypeClasses, TypeSynonymInstances, ViewPatterns, OverloadedStrings, TemplateHaskell, GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE FlexibleInstances, FlexibleContexts, MultiParamTypeClasses, TypeSynonymInstances, ViewPatterns, OverloadedStrings, TemplateHaskell, GeneralizedNewtypeDeriving, TypeOperators #-}
 
 module Sequence.Types
   ( GameState, gEntities, gEntityNames, gFocus, gNextId'
@@ -33,6 +33,7 @@ import Data.Tuple
 import Data.Ord
 
 import Sequence.Contact.Types
+import Sequence.Formula ((:<:)(..))
 
 import Text.Read (readMaybe)
 
@@ -77,6 +78,9 @@ instance Default Entity where
         , _eStats = def
         }
 
+instance Stats :<: Entity where
+  ctx' = eStats
+
 newtype EntityName = EntityName { _entityName :: CI String }
   deriving (Show, Eq, Ord)