1 files changed, 20 insertions, 10 deletions
diff --git a/src/Language/Fiddle/Ast/Internal/Instances/Walk.hs b/src/Language/Fiddle/Ast/Internal/Instances/Walk.hs
index 221dd5b..fc77e1f 100644
--- a/src/Language/Fiddle/Ast/Internal/Instances/Walk.hs
+++ b/src/Language/Fiddle/Ast/Internal/Instances/Walk.hs
@@ -1,26 +1,30 @@
-module Language.Fiddle.Ast.Internal.Instances.Walk (Walk (..), GWalk (..), walk_) where
+module Language.Fiddle.Ast.Internal.Instances.Walk (Walk (..), GWalk (..), walk_, WalkContinuation (..)) where
 
 import Data.Typeable
 import GHC.Generics
 
--- | Like walk, but assumes no local state.
+-- | Like walk, but assumes no local state and always continue
 walk_ ::
   (Monad m, Traversable f, Typeable f, Typeable a, Walk t) =>
   (forall t'. (Walk t', Typeable t', Typeable f, Typeable a) => t' f a -> m ()) ->
   t f a ->
   m ()
-walk_ fn t = walk (\t _ -> fn t) t ()
+walk_ fn t = walk (\t _ -> fn t >> return (Continue ())) t ()
+
+data WalkContinuation s where
+  Continue :: s -> WalkContinuation s
+  Stop :: WalkContinuation s
 
 class (Typeable t) => Walk t where
   walk ::
     (Monad m, Traversable f, Typeable f, Typeable a) =>
-    (forall t'. (Walk t', Typeable t', Typeable f, Typeable a) => t' f a -> s -> m s) ->
+    (forall t'. (Walk t', Typeable t', Typeable f, Typeable a) => t' f a -> s -> m (WalkContinuation s)) ->
     t f a ->
     s ->
     m ()
   default walk ::
     (GWalk (Rep (t f a)) f a, Generic (t f a), Monad m, Traversable f, Typeable f, Typeable a) =>
-    (forall t'. (Walk t', Typeable t', Typeable f, Typeable a) => t' f a -> s -> m s) ->
+    (forall t'. (Walk t', Typeable t', Typeable f, Typeable a) => t' f a -> s -> m (WalkContinuation s)) ->
     t f a ->
     s ->
     m ()
@@ -29,7 +33,7 @@ class (Typeable t) => Walk t where
 class GWalk r f a where
   gwalk ::
     (Monad m, Typeable f, Typeable a, Traversable f) =>
-    (forall t'. (Walk t', Typeable t') => t' f a -> s -> m s) ->
+    (forall t'. (Walk t', Typeable t') => t' f a -> s -> m (WalkContinuation s)) ->
     r x ->
     s ->
     m ()
@@ -63,8 +67,11 @@ instance
   GWalk (Rec0 (t f a)) f a
   where
   gwalk fn (K1 k) s = do
-    s' <- fn k s
-    walk fn k s'
+    ( \case
+        Continue s' -> walk fn k s'
+        _ -> return ()
+      )
+      =<< fn k s
 
 instance
   ( Traversable f,
@@ -76,8 +83,11 @@ instance
   gwalk fn (K1 fk) s = do
     mapM_
       ( \tfa -> do
-          s' <- fn tfa s
-          walk fn tfa s'
+          ( \case
+              Continue s' -> walk fn tfa s'
+              _ -> return ()
+            )
+            =<< fn tfa s
       )
       fk