Skip to content
Permalink
Browse files

Type stability fixes in NestedIterator

  • Loading branch information...
blegat committed Oct 2, 2019
1 parent 37f9a0c commit 1192fd86cefd31c11b43c1a8210c85c5f284f3af
Showing with 23 additions and 25 deletions.
  1. +21 −23 src/Containers/nested_iterator.jl
  2. +2 −2 test/perf/axis_constraints.jl
@@ -24,49 +24,47 @@ for i1 in iterators[1]()
end
```
"""
struct NestedIterator{T}
struct NestedIterator{T, C}
iterators::T # Tuple of functions
condition::Function
condition::C
end
function nested(iterators...; condition = (args...) -> true)
return NestedIterator(iterators, condition)
end
Base.IteratorSize(::Type{<:NestedIterator}) = Base.SizeUnknown()
Base.IteratorEltype(::Type{<:NestedIterator}) = Base.EltypeUnknown()
function next_iterate(it::NestedIterator, i, elems, states, iterator, elem_state)
function next_iterate(iterators, condition, elems, states, iterator, elem_state)
if elem_state === nothing
return nothing
end
elem, state = elem_state
elems_states = first_iterate(
it, i + 1, (elems..., elem),
Base.tail(iterators), condition, (elems..., elem),
(states..., (iterator, state, elem)))
if elems_states !== nothing
return elems_states
end
return next_iterate(it, i, elems, states, iterator, iterate(iterator, state))
return next_iterate(iterators, condition, elems, states, iterator, iterate(iterator, state))
end
function first_iterate(it::NestedIterator, i, elems, states)
if i > length(it.iterators)
if it.condition(elems...)
return elems, states
else
return nothing
end
end
iterator = it.iterators[i](elems...)
return next_iterate(it, i, elems, states, iterator, iterate(iterator))
end
function tail_iterate(it::NestedIterator, i, elems, states)
if i > length(it.iterators)
function first_iterate(::Tuple{}, condition, elems, states)
if condition(elems...)
return elems, states
else
return nothing
end
next = tail_iterate(it, i + 1, (elems..., states[i][3]), states)
end
function first_iterate(iterators, condition, elems, states)
iterator = iterators[1](elems...)
return next_iterate(iterators, condition, elems, states, iterator, iterate(iterator))
end
tail_iterate(::Tuple{}, condition, elems, states, prev_states) = nothing
function tail_iterate(iterators, condition, elems, states, prev_states)
next = tail_iterate(Base.tail(iterators), condition, (elems..., states[1][3]), Base.tail(states), (prev_states..., states[1]))
if next !== nothing
return next
end
iterator = states[i][1]
next_iterate(it, i, elems, states[1:(i - 1)], iterator, iterate(iterator, states[i][2]))
iterator = states[1][1]
next_iterate(iterators, condition, elems, prev_states, iterator, iterate(iterator, states[1][2]))
end
Base.iterate(it::NestedIterator) = first_iterate(it, 1, tuple(), tuple())
Base.iterate(it::NestedIterator, states) = tail_iterate(it, 1, tuple(), states)
Base.iterate(it::NestedIterator) = first_iterate(it.iterators, it.condition, tuple(), tuple())
Base.iterate(it::NestedIterator, states) = tail_iterate(it.iterators, it.condition, tuple(), states, tuple())
@@ -28,7 +28,7 @@ function dense_axis_constraints(n)

@variable(model, x[1:n])
set = MOI.EqualTo(0.0)
@time @constraint(model, con_refs[i = 2:n], x[i] in set)
con_refs = @time @constraint(model, [i = 2:n], x[i] in set)
optimize!(model)
@assert sum_iterate(con_refs) == n - 1
@btime sum_iterate($con_refs)
@@ -47,7 +47,7 @@ function sparse_axis_constraints(n)

@variable(model, x[1:n])
set = MOI.EqualTo(0.0)
@time @constraint(model, con_refs[i = 1:n; iseven(i)], x[i] in set)
con_refs = @time @constraint(model, [i = 1:n; iseven(i)], x[i] in set)
optimize!(model)
@assert sum_iterate(con_refs) == div(n, 2)
@btime sum_iterate($con_refs)

0 comments on commit 1192fd8

Please sign in to comment.
You can’t perform that action at this time.