Add Problem as type parameter to SweepIterator

RegionPlan is ommited as this is just vector of kwargs whos type is
unimportant

Author: Jack Dunham

src/solvers/iterators.jl

@@ -1,44 +1,3 @@
-#
-# SweepIterator
-#
-
-mutable struct SweepIterator
-  sweep_kws
-  region_iter
-  which_sweep::Int
-end
-
-problem(S::SweepIterator) = problem(S.region_iter)
-
-Base.length(S::SweepIterator) = length(S.sweep_kws)
-
-function Base.iterate(S::SweepIterator, which=nothing)
-  if isnothing(which)
-    sweep_kws_state = iterate(S.sweep_kws)
-  else
-    sweep_kws_state = iterate(S.sweep_kws, which)
-  end
-  isnothing(sweep_kws_state) && return nothing
-  current_sweep_kws, next = sweep_kws_state
-
-  if !isnothing(which)
-    S.region_iter = region_iterator(
-      problem(S.region_iter); sweep=S.which_sweep, current_sweep_kws...
-    )
-  end
-  S.which_sweep += 1
-  return S.region_iter, next
-end
-
-function sweep_iterator(problem, sweep_kws)
-  region_iter = region_iterator(problem; sweep=1, first(sweep_kws)...)
-  return SweepIterator(sweep_kws, region_iter, 1)
-end
-
-function sweep_iterator(problem, nsweeps::Integer; sweep_kws...)
-  return sweep_iterator(problem, Iterators.repeated(sweep_kws, nsweeps))
-end
-
 #
 # RegionIterator
 #
@@ -54,10 +13,14 @@ current_region_plan(R::RegionIterator) = R.region_plan[R.which_region]
 current_region(R::RegionIterator) = current_region_plan(R)[1]
 region_kwargs(R::RegionIterator) = current_region_plan(R)[2]
 function previous_region(R::RegionIterator)
-  R.which_region==1 ? nothing : R.region_plan[R.which_region - 1][1]
+  return R.which_region == 1 ? nothing : R.region_plan[R.which_region - 1][1]
 end
 function next_region(R::RegionIterator)
-  R.which_region==length(R.region_plan) ? nothing : R.region_plan[R.which_region + 1][1]
+  return if R.which_region == length(R.region_plan)
+    nothing
+  else
+    R.region_plan[R.which_region + 1][1]
+  end
 end
 is_last_region(R::RegionIterator) = isnothing(next_region(R))
 
@@ -98,3 +61,44 @@ end
 function region_plan(problem; kws...)
   return euler_sweep(state(problem); kws...)
 end
+
+#
+# SweepIterator
+#
+
+mutable struct SweepIterator{Problem}
+  sweep_kws
+  region_iter::RegionIterator{Problem}
+  which_sweep::Int
+end
+
+problem(S::SweepIterator) = problem(S.region_iter)
+
+Base.length(S::SweepIterator) = length(S.sweep_kws)
+
+function Base.iterate(S::SweepIterator, which=nothing)
+  if isnothing(which)
+    sweep_kws_state = iterate(S.sweep_kws)
+  else
+    sweep_kws_state = iterate(S.sweep_kws, which)
+  end
+  isnothing(sweep_kws_state) && return nothing
+  current_sweep_kws, next = sweep_kws_state
+
+  if !isnothing(which)
+    S.region_iter = region_iterator(
+      problem(S.region_iter); sweep=S.which_sweep, current_sweep_kws...
+    )
+  end
+  S.which_sweep += 1
+  return S.region_iter, next
+end
+
+function sweep_iterator(problem, sweep_kws)
+  region_iter = region_iterator(problem; sweep=1, first(sweep_kws)...)
+  return SweepIterator(sweep_kws, region_iter, 1)
+end
+
+function sweep_iterator(problem, nsweeps::Integer; sweep_kws...)
+  return sweep_iterator(problem, Iterators.repeated(sweep_kws, nsweeps))
+end