@@ -316,6 +316,15 @@ data STrans : (m : Type -> Type) ->
316316 (val : ty') ->
317317 STrans m () ctxt (const (updateCtxt ctxt prf ty'))
318318
319+namespace Loop
320+ export
321+ data STransLoop : (m : Type -> Type) -> (ty : Type) ->
322+ Context -> (ty -> Context) -> Type where
323+ Bind : STrans m a st1 st2_fn ->
324+ ((result : a) -> Inf (STransLoop m b (st2_fn result) st3_fn)) ->
325+ STransLoop m b st1 st3_fn
326+ Pure : (result : ty) -> STransLoop m ty (out_fn result) out_fn
327+
319328export
320329dropEnv : Env ys -> SubCtxt xs ys -> Env xs
321330dropEnv [] SubNil = []
@@ -372,6 +381,20 @@ runST env (Call prog ctxt_prf) k
372381runST env (Read lbl prf) k = k (lookupEnv prf env) env
373382runST env (Write lbl prf val) k = k () (updateEnv prf env val)
374383
384+export
385+data Fuel = Empty | More (Lazy Fuel)
386+
387+export partial
388+forever : Fuel
389+forever = More forever
390+
391+runSTLoop : Fuel -> Env invars -> STransLoop m a invars outfn ->
392+ (k : (x : a) -> Env (outfn x) -> m b) ->
393+ (onDry : m b) -> m b
394+runSTLoop Empty _ _ _ onDry = onDry
395+runSTLoop (More x) env (Bind prog next) k onDry
396+ = runST env prog (\prog', env' => runSTLoop x env' (next prog') k onDry)
397+runSTLoop (More x) env (Pure result) k onDry = k result env
375398
376399export
377400pure : (result : ty) -> STrans m ty (out_fn result) out_fn
@@ -445,6 +468,17 @@ write : (lbl : Var) ->
445468 (val : ty') ->
446469 STrans m () ctxt (const (updateCtxt ctxt prf (State ty')))
447470write lbl {prf} val = Write lbl prf (Value val)
471+
472+namespace Loop
473+ export
474+ (>>=) : STrans m a st1 st2_fn ->
475+ ((result : a) -> Inf (STransLoop m b (st2_fn result) st3_fn)) ->
476+ STransLoop m b st1 st3_fn
477+ (>>=) = Bind
478+
479+ export
480+ pure : (result : ty) -> STransLoop m ty (out_fn result) out_fn
481+ pure = Pure
448482
449483public export %error_reduce
450484out_res : ty -> (as : List (Action ty)) -> Context
@@ -470,6 +504,13 @@ ST : (m : Type -> Type) ->
470504 List (Action ty) -> Type
471505ST m ty xs = STrans m ty (in_res xs) (\result : ty => out_res result xs)
472506
507+public export
508+%error_reduce -- always evaluate this before showing errors
509+STLoop : (m : Type -> Type) ->
510+ (ty : Type) ->
511+ List (Action ty) -> Type
512+STLoop m ty xs = STransLoop m ty (in_res xs) (\result : ty => out_res result xs)
513+
473514-- Console IO is useful sufficiently often that let's have it here
474515public export
475516interface ConsoleIO (m : Type -> Type) where
@@ -486,6 +527,13 @@ export
486527run : Applicative m => ST m a [] -> m a
487528run prog = runST [] prog (\res, env' => pure res)
488529
530+export
531+runLoop : Applicative m => Fuel -> STLoop m a [] ->
532+ (onEmpty : m a) ->
533+ m a
534+runLoop fuel prog onEmpty
535+ = runSTLoop fuel [] prog (\res, env' => pure res) onEmpty
536+
489537||| runWith allows running an STrans program with an initial environment,
490538||| which must be consumed.
491539||| It's only allowed in the IO monad, because it's inherently unsafe, so
@@ -499,6 +547,14 @@ runWith : {ctxtf : _} ->
499547 IO (res ** Env (ctxtf res))
500548runWith env prog = runST env prog (\res, env' => pure (res ** env'))
501549
550+export
551+runWithLoop : {ctxtf : _} ->
552+ Env ctxt -> Fuel -> STransLoop IO a ctxt (\res => ctxtf res) ->
553+ IO (Maybe (res ** Env (ctxtf res)))
554+runWithLoop env fuel prog
555+ = runSTLoop fuel env prog (\res, env' => pure (Just (res ** env')))
556+ (pure Nothing)
557+
502558export
503559runPure : ST Basics.id a [] -> a
504560runPure prog = runST [] prog (\res, env' => res)