Improve the type stability of constant powers (#4016)

Author: odow

src/operators.jl

@@ -226,6 +226,17 @@ function Base.:/(lhs::GenericAffExpr, rhs::_Constant)
     return map_coefficients(c -> c / rhs, lhs)
 end
 
+# We need an implementation for both Base.literal_pow and Base.:^ because of the
+# difference between `x^2` and `a = 2; x^a`.
+
+function Base.literal_pow(::typeof(^), x::AbstractVariableRef, ::Val{0})
+    return one(value_type(typeof(x)))
+end
+
+Base.literal_pow(::typeof(^), x::AbstractVariableRef, ::Val{1}) = x
+
+Base.literal_pow(::typeof(^), x::AbstractVariableRef, ::Val{2}) = x * x
+
 function Base.:^(lhs::V, rhs::Integer) where {V<:AbstractVariableRef}
     if rhs == 0
         return one(value_type(V))
@@ -238,6 +249,12 @@ function Base.:^(lhs::V, rhs::Integer) where {V<:AbstractVariableRef}
     end
 end
 
+Base.literal_pow(::typeof(^), x::GenericAffExpr{T}, ::Val{0}) where {T} = one(T)
+
+Base.literal_pow(::typeof(^), x::GenericAffExpr, ::Val{1}) = x
+
+Base.literal_pow(::typeof(^), x::GenericAffExpr, ::Val{2}) = x * x
+
 function Base.:^(lhs::GenericAffExpr{T,V}, rhs::Integer) where {T,V}
     if rhs == 0
         return one(T)

test/test_operator.jl

@@ -758,4 +758,41 @@ function test_is_hermitian()
     return
 end
 
+function test_literal_pow_AbstractVariableRef()
+    model = Model()
+    @variable(model, x)
+    # We need a function barrier to allow Julia to optimize the constant power
+    p0(x) = x^0
+    p1(x) = x^1
+    p2(x) = x^2
+    @test (@inferred p0(x)) === 1.0
+    @test isequal_canonical(@inferred(p1(x)), x)
+    @test isequal_canonical(@inferred(p2(x)), 1.0 * x * x)
+    # Test the Base.:^ fallback
+    pa(x, a) = x^a
+    @test pa(x, 0) === 1.0
+    @test isequal_canonical(pa(x, 1), x)
+    @test isequal_canonical(pa(x, 2), 1.0 * x * x)
+    return
+end
+
+function test_literal_pow_GenericAffExpr()
+    model = Model()
+    @variable(model, y)
+    x = 2.0 * y + 1.0
+    # We need a function barrier to allow Julia to optimize the constant power
+    p0(x) = x^0
+    p1(x) = x^1
+    p2(x) = x^2
+    @test (@inferred p0(x)) === 1.0
+    @test isequal_canonical(@inferred(p1(x)), x)
+    @test isequal_canonical(@inferred(p2(x)), 1.0 * x * x)
+    # Test the Base.:^ fallback
+    pa(x, a) = x^a
+    @test pa(x, 0) === 1.0
+    @test isequal_canonical(pa(x, 1), x)
+    @test isequal_canonical(pa(x, 2), 1.0 * x * x)
+    return
+end
+
 end  # module