[Containers] document nested and vectorized_product (#2751)

Author: odow

Diff for: docs/src/reference/containers.md

@@ -8,7 +8,9 @@ Containers.DenseAxisArray
 Containers.SparseAxisArray
 Containers.container
 Containers.default_container
+Containers.@container
 Containers.VectorizedProductIterator
+Containers.vectorized_product
 Containers.NestedIterator
-Containers.@container
+Containers.nested
 ```

Diff for: src/Containers/nested_iterator.jl

@@ -10,9 +10,20 @@
     end
 
 Iterators over the tuples that are produced by a nested for loop.
-For instance, if `length(iterators) == 3` , this corresponds to the tuples
-`(i1, i2, i3)` produced by:
-```
+
+Construct a `NestedIterator` using [`nested`](@ref).
+
+## Example
+
+If `length(iterators) == 3` ,
+```julia
+x = NestedIterator(iterators, condition)
+for (i1, i2, i3) in x
+    # produces (i1, i2, i3)
+end
+````
+is the same as
+```julia
 for i1 in iterators[1]()
     for i2 in iterator[2](i1)
         for i3 in iterator[3](i1, i2)
@@ -28,40 +39,69 @@ struct NestedIterator{T,C}
     iterators::T # Tuple of functions
     condition::C
 end
+
+"""
+    nested(iterators...; condition = (args...) -> true)
+
+Create a [`NestedIterator`](@ref).
+
+## Example
+
+```julia
+nested(1:2, ["A", "B"]; condition = (i, j) -> isodd(i) || j == "B")
+```
+"""
 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(iterators, condition, elems, states, iterator, elem_state)
+
+function _next_iterate(
+    iterators,
+    condition,
+    elems,
+    states,
+    iterator,
+    elem_state,
+)
     while true
-        elem_state === nothing && return nothing
+        if elem_state === nothing
+            return
+        end
         elem, state = elem_state
-        elems_states = first_iterate(
+        elems_states = _first_iterate(
             Base.tail(iterators),
             condition,
             (elems..., elem),
             (states..., (iterator, state, elem)),
         )
-        elems_states !== nothing && return elems_states
-        # This could be written as a recursive function where we call `next_iterate`
-        # here with this new value of `next_iterate` instead of the `while` loop`.
-        # However, if there are too many consecutive elements for which `condition`
-        # is `false` for the last iterator, this will result in a `StackOverflow`.
-        # See https://github.com/jump-dev/JuMP.jl/issues/2335
+        if elems_states !== nothing
+            return elems_states
+        end
+        # This could be written as a recursive function where we call
+        # `_next_iterate` here with this new value of `_next_iterate` instead of
+        # the `while` loop`.
+        #
+        # However, if there are too many consecutive elements for which
+        # `condition`is `false` for the last iterator, this will result in a
+        # `StackOverflow`. See https://github.com/jump-dev/JuMP.jl/issues/2335
         elem_state = iterate(iterator, state)
     end
 end
-function first_iterate(::Tuple{}, condition, elems, states)
+
+function _first_iterate(::Tuple{}, condition, elems, states)
     if condition(elems...)
         return elems, states
-    else
-        return nothing
     end
+    return
 end
-function first_iterate(iterators, condition, elems, states)
+
+function _first_iterate(iterators, condition, elems, states)
     iterator = iterators[1](elems...)
-    return next_iterate(
+    return _next_iterate(
         iterators,
         condition,
         elems,
@@ -70,9 +110,11 @@ function first_iterate(iterators, condition, elems, states)
         iterate(iterator),
     )
 end
-tail_iterate(::Tuple{}, condition, elems, states, prev_states) = nothing
-function tail_iterate(iterators, condition, elems, states, prev_states)
-    next = tail_iterate(
+
+_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]),
@@ -83,7 +125,7 @@ function tail_iterate(iterators, condition, elems, states, prev_states)
         return next
     end
     iterator = states[1][1]
-    return next_iterate(
+    return _next_iterate(
         iterators,
         condition,
         elems,
@@ -92,12 +134,15 @@ function tail_iterate(iterators, condition, elems, states, prev_states)
         iterate(iterator, states[1][2]),
     )
 end
+
 function Base.iterate(it::NestedIterator)
-    return first_iterate(it.iterators, it.condition, tuple(), tuple())
+    return _first_iterate(it.iterators, it.condition, tuple(), tuple())
 end
+
 function Base.iterate(it::NestedIterator, states)
-    return tail_iterate(it.iterators, it.condition, tuple(), states, tuple())
+    return _tail_iterate(it.iterators, it.condition, tuple(), states, tuple())
 end
+
 function _eltype_or_any(::NestedIterator{<:Tuple{Vararg{Any,N}}}) where {N}
     return NTuple{N,Any}
 end

Diff for: src/Containers/vectorized_product_iterator.jl

@@ -27,50 +27,73 @@
 # Long story short, we want to tried everything as an iterator without shape
 # while `Iterators.ProductIterator` does care about preserving the shape
 # when doing the Cartesian product.
+
 """
     struct VectorizedProductIterator{T}
         prod::Iterators.ProductIterator{T}
     end
 
-Cartesian product of the iterators `prod.iterators`. It is the same iterator as
-`Base.Iterators.ProductIterator` except that it is independent of the
-`IteratorSize` of the elements of `prod.iterators`.
-For instance:
-* `size(Iterators.product(1, 2))` is `tuple()` while
-  `size(VectorizedProductIterator(1, 2))` is `(1, 1)`.
-* `size(Iterators.product(ones(2, 3)))` is `(2, 3)` while
-  `size(VectorizedProductIterator(ones(2, 3)))` is `(1, 1)`.
+A wrapper type for `Iterators.ProuctIterator` that discards shape information
+and returns a `Vector`.
+
+Construct a `VectorizedProductIterator` using [`vectorized_product`](@ref).
 """
 struct VectorizedProductIterator{T}
     prod::Iterators.ProductIterator{T}
 end
 
 # Collect iterators with unknown size so they can be used as axes.
 _collect(::Base.SizeUnknown, x) = collect(x)
+
 _collect(::Any, x) = x
+
 function _collect(::Base.IsInfinite, x)
     return error("Unable to form a container. Axis $(x) has infinite size!")
 end
+
 _collect(x) = _collect(Base.IteratorSize(x), x)
 
+"""
+    vectorized_product(iterators...)
+
+Created a [`VectorizedProductIterator`](@ref).
+
+## Examples
+
+```julia
+vectorized_product(1:2, ["A", "B"])
+```
+"""
 function vectorized_product(iterators...)
     return VectorizedProductIterator(Iterators.product(_collect.(iterators)...))
 end
+
 function Base.IteratorSize(
     ::Type{<:VectorizedProductIterator{<:Tuple{Vararg{Any,N}}}},
 ) where {N}
     return Base.HasShape{N}()
 end
+
 Base.IteratorEltype(::Type{<:VectorizedProductIterator}) = Base.EltypeUnknown()
+
 Base.size(it::VectorizedProductIterator) = _prod_size(it.prod.iterators)
+
 _prod_size(::Tuple{}) = ()
+
 _prod_size(t::Tuple) = (length(t[1]), _prod_size(Base.tail(t))...)
+
 Base.axes(it::VectorizedProductIterator) = _prod_indices(it.prod.iterators)
+
 _prod_indices(::Tuple{}) = ()
+
 function _prod_indices(t::Tuple)
     return (Base.OneTo(length(t[1])), _prod_indices(Base.tail(t))...)
 end
+
 Base.ndims(it::VectorizedProductIterator) = length(axes(it))
+
 Base.length(it::VectorizedProductIterator) = prod(size(it))
+
 Base.iterate(it::VectorizedProductIterator, args...) = iterate(it.prod, args...)
+
 _eltype_or_any(it::VectorizedProductIterator) = eltype(it.prod)