2 files changed, 127 insertions, 61 deletions

diff --git a/Spear/Step.hs b/Spear/Step.hs
index 5df873d..f1aef59 100644
--- a/Spear/Step.hs
+++ b/Spear/Step.hs
@@ -2,14 +2,20 @@
 module Spear.Step
 (
     -- * Definitions
-    Step(..)
+    Step
 ,   Elapsed
 ,   Dt
+    -- * Running
+,   runStep
     -- * Constructors
+,   step
 ,   sid
 ,   spure
 ,   sfst
 ,   ssnd
+,   switch
+,   multiSwitch
+,   sfold
     -- * Combinators
 ,   (.>)
 ,   (<.)
@@ -17,59 +23,133 @@ module Spear.Step
 )
 where
 
+import Data.List (foldl')
+import qualified Data.Map as Map
+import Data.Map (Map)
 import Data.Monoid
 
 type Elapsed = Double
 type Dt = Float
 
 -- | A step function.
-data Step a b = Step { step :: Elapsed -> Dt -> a -> (b, Step a b) }
+data Step s e a b =
+     Step { runStep :: Elapsed -> Dt -> s -> e -> a -> (b, Step s e a b) }
+
+-- | Construct a step from a function.
+step :: (Elapsed -> Dt -> s -> e -> a -> (b, Step s e a b)) -> Step s e a b
+step = Step
 
 -- | Step identity.
-sid :: Step a a
-sid = Step $ \_ _ a -> (a, sid)
+sid :: Step s e a a
+sid = Step $ \_ _ _ _ a -> (a, sid)
+
+-- | Construct a step from a pure function.
+spure :: (a -> b) -> Step s e a b
+spure f = Step $ \_ _ _ _ x -> (f x, spure f)
 
 -- | The step that returns the first component in the tuple.
-sfst :: Step (a,b) a
+sfst :: Step s e (a,b) a
 sfst = spure $ \(a,_) -> a
 
 -- | The step that returns the second component in the tuple.
-ssnd :: Step (a,b) b
+ssnd :: Step s e (a,b) b
 ssnd = spure $ \(_,b) -> b
 
--- | Construct a step from a pure function.
-spure :: (a -> b) -> Step a b
-spure f = Step $ \_ _ x -> (f x, spure f)
+-- | Construct a step that switches between two steps based on input.
+--
+-- The initial step is the first one.
+switch :: Eq e
+       => e -> (Step s (Maybe e) a a)
+       -> e -> (Step s (Maybe e) a a)
+       -> Step s (Maybe e) a a
+switch flag1 s1 flag2 s2 = switch' s1 flag1 s1 flag2 s2
+
+switch' :: Eq e
+        => (Step s (Maybe e) a a)
+        -> e -> (Step s (Maybe e) a a)
+        -> e -> (Step s (Maybe e) a a)
+        -> Step s (Maybe e) a a
+switch' cur flag1 s1 flag2 s2 = Step $ \elapsed dt g e a ->
+        case e of
+             Nothing ->
+                     let (a',s') = runStep cur elapsed dt g Nothing a
+                     in (a', switch' s' flag1 s1 flag2 s2)
+             Just e' ->
+                     let next = if e' == flag1 then s1
+                                else if e' == flag2 then s2
+                                else cur
+                         (a',s') = runStep next elapsed dt g e a
+                         in (a', switch' s' flag1 s1 flag2 s2)
+
+-- | Construct a step that switches among multiple steps based on input.
+multiSwitch :: (Eq e, Ord e) => [(e, Step s (Maybe e) a a)] -> Step s (Maybe e) a a
+multiSwitch xs = multiSwitch' Nothing sid (Map.fromList xs)
+
+multiSwitch' :: (Eq e, Ord e)
+             => Maybe e -> Step s (Maybe e) a a -> Map e (Step s (Maybe e) a a)
+             -> Step s (Maybe e) a a
+multiSwitch' curKey cur m = Step $ \elapsed dt g e a ->
+             let singleStep = let (a',s') = runStep cur elapsed dt g e a
+                              in (a', multiSwitch' curKey s' m)
+             in case e of
+                     Nothing -> singleStep
+                     Just e' -> case Map.lookup e' m of
+                          Nothing -> singleStep
+                          Just s ->
+                               let (a',s') = runStep s elapsed dt g e a
+                                   m' = case curKey of
+                                             Nothing  -> m
+                                             Just key -> Map.insert key cur m
+                               in (a', multiSwitch' e s' m')
+
+-- | Construct a step that folds a given list of inputs.
+--
+-- The step is run N+1 times, where N is the size of the input list.
+sfold :: Step s (Maybe e) a a -> Step s [e] a a
+sfold s = Step $ \elapsed dt g es a ->
+      case es of
+           [] ->
+              let (b',s') = runStep s elapsed dt g Nothing a
+              in (b', sfold s')
+           es ->
+              let (b',s') = sfold' elapsed dt g s a es
+              in (b', sfold s')
+
+sfold' :: Elapsed -> Dt -> s -> Step s (Maybe e) a a -> a -> [e]
+       -> (a, Step s (Maybe e) a a)
+sfold' elapsed dt g s a es = foldl' f (a',s') es
+       where f (a,s) e = runStep s elapsed dt g (Just e) a
+             (a',s') = runStep s elapsed dt g Nothing a
 
-instance Functor (Step a) where
-         fmap f (Step s1) = Step $ \elapsed dt x ->
-              let (a, s') = s1 elapsed dt x
+instance Functor (Step s e a) where
+         fmap f (Step s1) = Step $ \elapsed dt g e x ->
+              let (a, s') = s1 elapsed dt g e x
               in (f a, fmap f s')
 
-instance Monoid (Step a a) where
+instance Monoid (Step s e a a) where
          mempty = sid
 
-         mappend (Step s1) (Step s2) = Step $ \elapsed dt a ->
-                 let (b, s1') = s1 elapsed dt a
-                     (c, s2') = s2 elapsed dt b
+         mappend (Step s1) (Step s2) = Step $ \elapsed dt g e a ->
+                 let (b, s1') = s1 elapsed dt g e a
+                     (c, s2') = s2 elapsed dt g e b
                  in (c, mappend s1' s2')
 
 -- Combinators
 
--- | Chain two steps.
-(.>) :: Step a b -> Step b c -> Step a c
-(Step s1) .> (Step s2) = Step $ \elapsed dt a ->
-      let (b, s1') = s1 elapsed dt a
-          (c, s2') = s2 elapsed dt b
+-- | Compose two steps.
+(.>) :: Step s e a b -> Step s e b c -> Step s e a c
+(Step s1) .> (Step s2) = Step $ \elapsed dt g e a ->
+      let (b, s1') = s1 elapsed dt g e a
+          (c, s2') = s2 elapsed dt g e b
       in (c, s1' .> s2')
 
--- | Chain two steps.
-(<.) :: Step a b -> Step c a -> Step c b
+-- | Compose two steps.
+(<.) :: Step s e a b -> Step s e c a -> Step s e c b
 (<.) = flip (.>)
 
 -- | Evaluate two steps and zip their results.
-szip :: (a -> b -> c) -> Step d a -> Step d b -> Step d c
-szip f (Step s1) (Step s2) = Step $ \elapsed dt d ->
-     let (a, s1') = s1 elapsed dt d
-         (b, s2') = s2 elapsed dt d
+szip :: (a -> b -> c) -> Step s e d a -> Step s e d b -> Step s e d c
+szip f (Step s1) (Step s2) = Step $ \elapsed dt g e d ->
+     let (a, s1') = s1 elapsed dt g e d
+         (b, s2') = s2 elapsed dt g e d
      in (f a b, szip f s1' s2')
\ No newline at end of file
diff --git a/demos/pong/Pong.hs b/demos/pong/Pong.hs
index b323aa2..6b2f888 100644
--- a/demos/pong/Pong.hs
+++ b/demos/pong/Pong.hs
@@ -24,14 +24,14 @@ data GameEvent
      | MoveRight
      | StopLeft
      | StopRight
-     deriving Eq
+     deriving (Eq, Ord)
 
 -- Game objects
 
 data GameObject = GameObject
      { aabb   :: AABB2
      , obj    :: Obj2
-     , gostep :: Step ([GameEvent], [GameObject], GameObject) GameObject
+     , gostep :: Step [GameObject] [GameEvent] GameObject GameObject
      }
 
 instance Spatial2 GameObject where
@@ -43,7 +43,7 @@ stepWorld elapsed dt evts gos = map (update elapsed dt evts gos) gos
 
 update :: Elapsed -> Dt -> [GameEvent] -> [GameObject] -> GameObject -> GameObject
 update elapsed dt evts gos go =
-       let (go', s') = step (gostep go) elapsed dt (evts, gos, go)
+       let (go', s') = runStep (gostep go) elapsed dt gos evts go
        in go' { gostep = s' }
 
 ballBox :: AABB2
@@ -63,23 +63,12 @@ newWorld =
          , GameObject padBox  (obj2 0.5 0.1) stepPlayer
          ]
 
--- Generic steppers
-
-ignore :: Step ([GameEvent], [GameObject], GameObject) GameObject
-ignore = spure $ \(_,_,go) -> go
-
-ignoreEvts :: Step ([GameEvent], [GameObject], GameObject) ([GameObject], GameObject)
-ignoreEvts = spure $ \(_, world, go) -> (world, go)
-
-ignoreGOs :: Step ([GameEvent], [GameObject], GameObject) ([GameEvent], GameObject)
-ignoreGOs = spure $ \(evts, _, go) -> (evts, go)
-
 -- Ball steppers
 
-stepBall vel = ignoreEvts .> collideBall vel .> moveBall
+stepBall vel = collideBall vel .> moveBall
 
-collideBall :: Vector2 -> Step ([GameObject], GameObject) (Vector2, GameObject)
-collideBall vel = Step $ \_ _ (gos, ball) ->
+collideBall :: Vector2 -> Step [GameObject] e GameObject (Vector2, GameObject)
+collideBall vel = step $ \_ _ gos _ ball ->
             let (AABB2 pmin pmax) = aabb ball `aabbAdd` pos ball
                 collideCol = x pmin < 0 || x pmax > 1
                 collideRow = y pmin < 0 || y pmax > 1
@@ -99,15 +88,15 @@ collide go1 go2 =
 
 aabbAdd (AABB2 pmin pmax) p = AABB2 (p+pmin) (p+pmax)
 
-moveBall :: Step (Vector2, GameObject) GameObject
-moveBall = Step $ \_ dt (vel,ball) -> (move (scale dt vel) ball, moveBall)
+moveBall :: Step s e (Vector2, GameObject) GameObject
+moveBall = step $ \_ dt _ _ (vel,ball) -> (move (scale dt vel) ball, moveBall)
 
 -- Enemy stepper
 
-stepEnemy = ignore .> movePad
+stepEnemy = movePad
 
-movePad :: Step GameObject GameObject
-movePad = Step $ \elapsed _ pad ->
+movePad :: Step s e GameObject GameObject
+movePad = step $ \elapsed _ _ _ pad ->
         let p  = vec2 px 0.9
             px = double2Float (sin elapsed * 0.5 + 0.5)
                * (1 - 2 * x padSize)
@@ -116,20 +105,17 @@ movePad = Step $ \elapsed _ pad ->
 
 -- Player stepper
 
-stepPlayer = ignoreGOs
-           .> moveGO False MoveLeft StopLeft
-           .> moveGO False MoveRight StopRight
-           .> ssnd
-           .> clamp
+stepPlayer = sfold moveGO .> clamp
+
+moveGO = mconcat
+       [ switch StopLeft  sid MoveLeft  (moveGO' $ vec2 (-1) 0)
+       , switch StopRight sid MoveRight (moveGO' $ vec2 1 0)
+       ]
 
-moveGO :: Bool -> GameEvent -> GameEvent
-       -> Step ([GameEvent], GameObject) ([GameEvent], GameObject)
-moveGO moving start stop = Step $ \_ dt (evts, go) ->
-       let moving' = (moving || any (==start) evts) && not (any (==stop) evts)
-           dir = scale dt $ toDir moving' start
-       in ((evts, move dir go), moveGO moving' start stop)
+moveGO' :: Vector2 -> Step s e GameObject GameObject
+moveGO' dir = step $ \_ dt _ _ go -> (move (scale dt dir) go, moveGO' dir)
 
-clamp :: Step GameObject GameObject
+clamp :: Step s e GameObject GameObject
 clamp = spure $ \go ->
       let p' = vec2 (clamp' x s (1 - s)) y
           (Vector2 x y) = pos go