Fixes #258: Support colon indexing in function calls in generators

src/nlp.jl

@@ -919,7 +919,7 @@ end
 # `Var` wrapping an arithmetic node that would be re-evaluated per row.
 @inline _indexed_var(i::AbstractArgNode, o) = Var(ArgLeaf(i + o))
 @inline _indexed_var(i, o) = Var(i + o)
-@inline function Base.getindex(v::V, is...) where {V<:AbstractVariable}
+@inline function _getindex(v::V, is, ::Val{false}) where {V<:AbstractVariable}
     @assert(length(is) == length(v.size), "Variable index dimension error")
     _bound_check(v.size, is)
     Var(v.offset + idxx(is .- (_start.(v.size) .- 1), _length.(v.size)))
@@ -944,7 +944,7 @@ end
     idx = idxx(is .- (_start.(s.size) .- 1), _length.(s.size))
     return _reindex(s.f, s.iter[idx])
 end
-@inline function Base.getindex(s::Expression, is...)
+@inline function _getindex(s::Expression, is, ::Val{false})
     # Symbolic indices case - the symbolic indices ARE the iterator elements
     # No adjustment needed; the indices are used directly in expression building
     @assert(length(is) == length(s.size), "Expression index dimension error")
@@ -955,12 +955,31 @@ end
     _bound_check(p.size, i)
     ParameterNode(i + (p.offset - _start(p.size[1]) + 1))
 end
-@inline function Base.getindex(p::P, is...) where {P<:Parameter}
+@inline function _getindex(p::P, is, ::Val{false}) where {P<:Parameter}
     @assert(length(is) == length(p.size), "Parameter index dimension error")
     _bound_check(p.size, is)
     ParameterNode(p.offset + idxx(is .- (_start.(p.size) .- 1), _length.(p.size)))
 end
 
+# `x[:, i]` expands each colon axis to its entries, column-major as in Base.
+const _Indexable = Union{AbstractVariable,Expression,Parameter}
+@inline Base.getindex(x::_Indexable, is...) = _getindex(x, is, _hascolon(is))
+@inline Base.getindex(v::V, ::Colon) where {V<:AbstractVariable} = _allentries(v)
+@inline Base.getindex(s::Expression, ::Colon) = _allentries(s)
+@inline Base.getindex(p::P, ::Colon) where {P<:Parameter} = _allentries(p)
+@inline _allentries(x) = _getindex(x, map(_ -> :, x.size), Val(true))
+
+@inline _hascolon(::Tuple{}) = Val(false)
+@inline _hascolon(::Tuple{Colon,Vararg{Any}}) = Val(true)
+@inline _hascolon(is::Tuple) = _hascolon(Base.tail(is))
+
+@inline _axis(n, ::Colon) = _start(n):_start(n)+_length(n)-1
+@inline _axis(n, i) = (i,)
+@inline function _getindex(x, is, ::Val{true})
+    @assert(length(is) == length(x.size), "Colon index dimension error")
+    Tuple(x[I...] for I in Iterators.product(map(_axis, x.size, is)...))
+end
+
 
 function _bound_check(sizes, i::I) where {I<:Integer}
     __bound_check(sizes[1], i)

test/NLPTest/feature_test.jl

@@ -198,6 +198,49 @@ function test_nonunit_expr(backend)
     end
 end
 
+function test_colon_index(backend)
+    @testset "Colon expands to the tuple of entries along an axis" begin
+        N, K = 4, 3
+        itr = [(i, k, exp10(i / N)) for i in 1:N, k in 1:K]
+        f(x...) = -x[1] * 2 + x[2] * x[3] - x[4] + x[5]
+        g(x, θ) = x[1] * θ[1] - x[5] * θ[2]
+        function build(con)
+            c = ExaCore(; backend, concrete = Val(true))
+            c, z = add_var(c, 1:5, 1:N, 0:K; start = 1.0)
+            c, θ = add_par(c, [2.0, 3.0])
+            c, _ = con(c, z, θ)
+            ExaModel(c)
+        end
+        m = build(
+            (c, z, θ) ->
+                add_con(c, -hi * (f(z[:, i, k]...) + g(z[:, i, k], θ[:])) for (i, k, hi) in itr),
+        )
+        mref = build(
+            (c, z, θ) -> add_con(
+                c,
+                -hi * (
+                    f(z[1, i, k], z[2, i, k], z[3, i, k], z[4, i, k], z[5, i, k]) +
+                    g((z[1, i, k], z[2, i, k], z[3, i, k], z[4, i, k], z[5, i, k]), (θ[1], θ[2]))
+                ) for (i, k, hi) in itr
+            ),
+        )
+        x0 = ExaModels.convert_array([sin(i) for i in 1:m.meta.nvar], backend)
+        @test Array(NLPModels.cons(m, x0)) ≈ Array(NLPModels.cons(mref, x0))
+
+        # column-major order, non-unit starts, `x[:]` as in Base, Expression, wrong rank
+        c = ExaCore(; backend, concrete = Val(true))
+        c, x = add_var(c, 2, 3)
+        c, y = add_var(c, 2:4, 0:1)
+        c, s = add_expr(c, x[i, j]^2 for (i, j) in Iterators.product(1:2, 1:3))
+        @test [v.i for v in x[:, 2]] == [3, 4]
+        @test [v.i for v in x[2, :]] == [2, 4, 6]
+        @test [v.i for v in y[:, 1]] == [10, 11, 12]
+        @test [v.i for v in x[:]] == [v.i for v in x[:, :]] == 1:6
+        @test length(s[:, 1]) == 2
+        @test_throws Exception x[:, 1, 1]
+    end
+end
+
 function test_generator_free_constr(backend)
     @testset "add_con(core, expr, itr)" begin
         c = ExaCore(; backend, concrete = Val(true))
@@ -225,6 +268,9 @@ function test_features(backend)
     @testset "Non-unit expression indexing" begin
         test_nonunit_expr(backend)
     end
+    @testset "Colon indexing" begin
+        test_colon_index(backend)
+    end
     @testset "Generator-free constraint" begin
         test_generator_free_constr(backend)
     end