Fix various errors found by JET (#2908)

Author: odow

src/Bridges/Constraint/bridges/IntegerToZeroOneBridge.jl

@@ -152,7 +152,8 @@ function MOI.Bridges.final_touch(
     bridge::IntegerToZeroOneBridge{T},
     model::MOI.ModelLike,
 ) where {T}
-    ret = MOI.Utilities.get_bounds(model, T, bridge.x)
+    # Type annotation needed for JET
+    ret = MOI.Utilities.get_bounds(model, T, bridge.x)::Tuple{T,T}
     if ret === bridge.last_bounds
         return nothing  # final_touch already called
     elseif ret[1] == typemin(T) || ret[2] == typemax(T)

src/Bridges/Constraint/bridges/SplitIntervalBridge.jl

@@ -325,17 +325,15 @@ function MOI.get(
         upper_stat = MOI.get(model, attr, bridge.upper)
         if upper_stat == MOI.NONBASIC
             return MOI.NONBASIC_AT_UPPER
+        elseif bridge.lower === nothing
+            return upper_stat
         end
-    end
-    if bridge.lower === nothing
-        if bridge.upper === nothing
-            # The only case where the interval `[-∞, ∞]` is allowed is for
-            # `VariableIndex` constraints but `ConstraintBasisStatus` is not
-            # defined for `VariableIndex` constraints.
-            msg = "Cannot get `$attr` for a constraint in the interval `[-Inf, Inf]`."
-            throw(MOI.GetAttributeNotAllowed(attr, msg))
-        end
-        return upper_stat
+    elseif bridge.lower === nothing
+        # The only case where the interval `[-∞, ∞]` is allowed is for
+        # `VariableIndex` constraints but `ConstraintBasisStatus` is not
+        # defined for `VariableIndex` constraints.
+        msg = "Cannot get `$attr` for a constraint in the interval `[-Inf, Inf]`."
+        throw(MOI.GetAttributeNotAllowed(attr, msg))
     end
     lower_stat = MOI.get(model, attr, bridge.lower)
     if lower_stat == MOI.NONBASIC

src/Bridges/Constraint/single_bridge_optimizer.jl

@@ -48,15 +48,15 @@ mutable struct SingleBridgeOptimizer{BT<:AbstractBridge,OT<:MOI.ModelLike} <:
     map::Map # index of bridged constraint -> constraint bridge
     con_to_name::Dict{MOI.ConstraintIndex,String}
     name_to_con::Union{Dict{String,MOI.ConstraintIndex},Nothing}
-end
 
-function SingleBridgeOptimizer{BT}(model::OT) where {BT,OT<:MOI.ModelLike}
-    return SingleBridgeOptimizer{BT,OT}(
-        model,
-        Map(),
-        Dict{MOI.ConstraintIndex,String}(),
-        nothing,
-    )
+    function SingleBridgeOptimizer{BT}(model::OT) where {BT,OT<:MOI.ModelLike}
+        return new{BT,OT}(
+            model,
+            Map(),
+            Dict{MOI.ConstraintIndex,String}(),
+            nothing,
+        )
+    end
 end
 
 bridges(::MOI.Bridges.AbstractBridgeOptimizer) = EmptyMap()

src/Bridges/Objective/single_bridge_optimizer.jl

@@ -41,10 +41,10 @@ mutable struct SingleBridgeOptimizer{BT<:AbstractBridge,OT<:MOI.ModelLike} <:
                MOI.Bridges.AbstractBridgeOptimizer
     model::OT
     map::Map # `MOI.ObjectiveFunction` -> objective bridge
-end
 
-function SingleBridgeOptimizer{BT}(model::OT) where {BT,OT<:MOI.ModelLike}
-    return SingleBridgeOptimizer{BT,OT}(model, Map())
+    function SingleBridgeOptimizer{BT}(model::OT) where {BT,OT<:MOI.ModelLike}
+        return new{BT,OT}(model, Map())
+    end
 end
 
 bridges(::MOI.Bridges.AbstractBridgeOptimizer) = EmptyMap()

src/Bridges/Variable/single_bridge_optimizer.jl

@@ -50,17 +50,17 @@ mutable struct SingleBridgeOptimizer{BT<:AbstractBridge,OT<:MOI.ModelLike} <:
     name_to_var::Union{Dict{String,MOI.VariableIndex},Nothing}
     con_to_name::Dict{MOI.ConstraintIndex,String}
     name_to_con::Union{Dict{String,MOI.ConstraintIndex},Nothing}
-end
 
-function SingleBridgeOptimizer{BT}(model::OT) where {BT,OT<:MOI.ModelLike}
-    return SingleBridgeOptimizer{BT,OT}(
-        model,
-        Map(),
-        Dict{MOI.VariableIndex,String}(),
-        nothing,
-        Dict{MOI.ConstraintIndex,String}(),
-        nothing,
-    )
+    function SingleBridgeOptimizer{BT}(model::OT) where {BT,OT<:MOI.ModelLike}
+        return new{BT,OT}(
+            model,
+            Map(),
+            Dict{MOI.VariableIndex,String}(),
+            nothing,
+            Dict{MOI.ConstraintIndex,String}(),
+            nothing,
+        )
+    end
 end
 
 bridges(::MOI.Bridges.AbstractBridgeOptimizer) = EmptyMap()

src/Bridges/bridge_optimizer.jl

@@ -2389,7 +2389,10 @@ end
 Substitute any bridged [`MOI.VariableIndex`](@ref) in `value` by an equivalent
 expression in terms of variables of `b.model`.
 """
-function bridged_function(bridge::AbstractBridgeOptimizer, value)
+function bridged_function(
+    bridge::AbstractBridgeOptimizer,
+    value::V,
+)::V where {V}
     if !Variable.has_bridges(Variable.bridges(bridge))
         # Shortcut, this allows performance to be unaltered when no variable
         # bridges are used.
@@ -2398,10 +2401,9 @@ function bridged_function(bridge::AbstractBridgeOptimizer, value)
     # We assume that the type of `value` is not altered. This restricts
     # variable bridges to only return `ScalarAffineFunction` but otherwise,
     # the performance would be bad.
-    return MOI.Utilities.substitute_variables(
-        vi -> bridged_variable_function(bridge, vi),
-        value,
-    )::typeof(value)
+    return MOI.Utilities.substitute_variables(value) do vi
+        return bridged_variable_function(bridge, vi)
+    end
 end
 
 function bridged_function(b::AbstractBridgeOptimizer, func::MOI.VariableIndex)

src/FileFormats/LP/read.jl

@@ -511,6 +511,11 @@ end
 
 _is_number(c::Char) = isdigit(c) || c in ('.', 'e', 'E', '+', '-')
 
+function _is_negative_digit(state, c)
+    # The type annotation is needed for JET.
+    return c == '-' && isdigit(peek(state, Char)::Char)
+end
+
 function _peek_inner(state::_LexerState)
     while (c = peek(state, Char)) !== nothing
         pos = position(state.io)
@@ -523,7 +528,7 @@ function _peek_inner(state::_LexerState)
             while (c = peek(state, Char)) !== nothing && c != '\n'
                 _ = read(state, Char)
             end
-        elseif isdigit(c) || (c == '-' && isdigit(peek(state, Char))) # Number
+        elseif isdigit(c) || _is_negative_digit(state, c) # Number
             buf = IOBuffer()
             while (c = peek(state, Char)) !== nothing && _is_number(c)
                 write(buf, c)
@@ -826,7 +831,8 @@ function _parse_term(
         # <quadratic-expression>
         return _parse_quadratic_expression(state, cache, prefix)
     end
-    token = peek(state, _Token)
+    # The type annotation is needed for JET.
+    token = peek(state, _Token)::_Token
     return _throw_parse_error(
         state,
         token,
@@ -1050,7 +1056,8 @@ function _parse_constraint_sos(
     f, w = MOI.VectorOfVariables(MOI.VariableIndex[]), T[]
     while true
         if _next_token_is(state, _TOKEN_NEWLINE)
-            t = peek(state, _Token)
+            # The type annotation is needed for JET.
+            t = peek(state, _Token)::_Token
             _throw_parse_error(
                 state,
                 t,
@@ -1096,7 +1103,8 @@ function _parse_constraint_indicator(
     end
     _ = read(state, _Token, _TOKEN_IMPLIES)
     f = _parse_expression(state, cache)
-    set = _parse_set_suffix(state, cache)
+    # The type annotation is needed for JET.
+    set = _parse_set_suffix(state, cache)::MOI.AbstractScalarSet
     return MOI.add_constraint(
         cache.model,
         MOI.Utilities.operate(vcat, T, z, f),
@@ -1117,7 +1125,8 @@ function _parse_constraint(state::_LexerState, cache::_ReadCache)
         _parse_constraint_indicator(state, cache)
     else
         f = _parse_expression(state, cache)
-        set = _parse_set_suffix(state, cache)
+        # The type annotation is needed for JET.
+        set = _parse_set_suffix(state, cache)::MOI.AbstractScalarSet
         MOI.add_constraint(cache.model, f, set)
     end
     if name !== nothing

src/FileFormats/MPS/read.jl

@@ -432,7 +432,8 @@ function _add_indicator_constraint(
         scalar = MOI.ScalarAffineTerm(coef, variable_map[data.col_to_name[i]])
         push!(terms, MOI.VectorAffineTerm(2, scalar))
     end
-    f = MOI.VectorAffineFunction(terms, zeros(T, 2))
+    # The type annotation is needed for JET.
+    f = MOI.VectorAffineFunction(terms, zeros(T, 2)::Vector{T})
     c = MOI.add_constraint(model, f, MOI.Indicator{activate}(set))
     MOI.set(model, MOI.ConstraintName(), c, c_name)
     return

src/FileFormats/NL/read.jl

@@ -271,13 +271,15 @@ _try_scalar_affine_function(x::MOI.VariableIndex) = x
 function _try_scalar_affine_function(expr::Expr)
     if expr.args[1] == :+
         args = _try_scalar_affine_function.(expr.args[2:end])
-        if any(isnothing, args)
+        # The type annotation is needed for JET.
+        if any(isnothing, args)::Bool
             return nothing
         end
         return MOI.Utilities.operate(+, Float64, args...)
     elseif expr.args[1] == :*
         args = _try_scalar_affine_function.(expr.args[2:end])
-        if any(isnothing, args)
+        # The type annotation is needed for JET.
+        if any(isnothing, args)::Bool
             return nothing
         end
         n_affine_terms = 0

src/Nonlinear/SymbolicAD/SymbolicAD.jl

@@ -1333,7 +1333,12 @@ function Evaluator(
     constraint_index_by_hash = Dict{UInt64,Vector{Int}}()
     jac_offset, hess_offset = 0, 0
     if model.objective !== nothing
-        o_sym = _to_symbolic_form(model, model.objective, variable_to_column)
+        o_sym = _to_symbolic_form(
+            model,
+            # The type annotation is needed for JET.
+            model.objective::MOI.Nonlinear.Expression,
+            variable_to_column,
+        )
         x, ∇f, H, ∇²f = gradient_and_hessian(x -> x.value > 0, o_sym.f)
         dag = _DAG(model.operators, Any[o_sym.f; ∇f; ∇²f])
         hash_to_dag[o_sym.hash] = dag