feat: add functional version of nif for easier @generated

Author: MilesCranmer

src/Utils.jl

@@ -153,4 +153,55 @@ struct ResultOk2{A<:AbstractArray,B<:AbstractArray}
     ok::Bool
 end
 
+#! format: on
+"""
+    nif(condition, expression, [else_expression,] ::Val{N}) where {N}
+
+Generate a sequence of `if ... elseif ... else ... end` statements.
+This is copied from https://github.com/JuliaLang/julia/pull/55093.
+
+# Arguments
+- `condition`: A function that takes an integer between `1` and `N-1` and
+    returns a boolean condition.
+- `expression`: A function that takes an integer between `1` and `N` (or,
+    only up to `N-1`, if `else_expression` is provided) and is called if
+    the condition is true.
+- `else_expression`: (optional) A function that takes `N` as input
+    returns an expression to be evaluated if all conditions are false.
+- `N`: The number of conditions to check, passed as a `Val{N}` instance.
+This function is similar to the `@nif` macro but can be used in cases
+where `N` is not known at parse time.
+# Examples
+For example, here we find the first index of a positive element in a
+fixed-size tuple using `nif`:
+```jldoctest
+julia> x = (0, -1, 1, 0)
+(0, -1, 1, 0)
+julia> Base.Cartesian.nif(d -> x[d] > 0, d -> d, Val(4))
+3
+```
+"""
+@inline function nif(condition::F, expression::G, ::Val{N}) where {F,G,N}
+    return nif(condition, expression, expression, Val(N))
+end
+@inline function nif(
+    condition::F, expression::G, else_expression::H, ::Val{N}
+) where {F,G,H,N}
+    n = N::Int  # Can improve inference; see #54544
+    (n >= 0) || throw(ArgumentError("if statement length should be ≥ 0, got $n"))
+    if @generated
+        :(Base.Cartesian.@nif $N d -> condition(d) d -> expression(d) d ->
+            else_expression(d))
+    else
+        for d in 1:(n - 1)
+            if condition(d)
+                return expression(d)
+            end
+        end
+        return else_expression(n)
+    end
+end
+typeof(function nif end).name.max_methods = UInt8(2)
+#! format: off
+
 end

test/test_nif.jl

@@ -0,0 +1,49 @@
+# Copied from https://github.com/JuliaLang/julia/pull/55093
+@testitem "nif" begin
+    using DynamicExpressions.UtilsModule: nif
+
+    x = (0, -1, 1, 0)
+    @test nif(d -> x[d] > 0, d -> d, Val(4)) == 3
+
+    @test nif(d -> d > 1, d -> "A", d -> "B", Val(1)) == "B"
+    @test nif(d -> d > 3, d -> "A", d -> "B", Val(3)) == "B"
+
+    # Test with N = 0
+    @test nif(d -> d > 0, d -> "", d -> "A", Val(0)) == "A"
+
+    # Specific branch true
+    @test nif(d -> d == 2, d -> d, d -> "else", Val(3)) == 2
+
+    # Test with condition only true for last branch
+    @test nif(d -> d == 5, d -> "A", d -> "B", Val(5)) == "B"
+
+    # Test with bad input:
+    @test_throws ArgumentError("if statement length should be ≥ 0, got -1") nif(
+        identity, identity, Val(-1)
+    )
+
+    # Non-Int64 also throws
+    @test_throws TypeError nif(identity, identity, Val(1.5))
+
+    # Make sure all conditions are actually evaluated
+    result = let c = Ref(0)
+        nif(d -> (c[] += 1; false), d -> 1, Val(4))
+        c[]
+    end
+    @test result == 3
+
+    # Test inference is good
+    t = ("i am not an int", ntuple(d -> d, Val(10))...)
+    function extract_from_tuple(t::Tuple, i)
+        return nif(
+            d -> d == i,
+            d -> t[d + 1],  # We skip the non-integer element
+            Val(length(t) - 1),
+        )
+    end
+    # Normally, had we used getindex here, inference would have
+    # not been able to infer that the return type never includes
+    # the first element. But since we used an `nif`, the compiler
+    # knows all possible branches and can infer the correct type.
+    @test @inferred(extract_from_tuple(t, 3)) == 3
+end

test/unittest.jl

@@ -129,3 +129,4 @@ include("test_operator_construction_edgecases.jl")
 include("test_node_interface.jl")
 include("test_expression_math.jl")
 include("test_structured_expression.jl")
+include("test_nif.jl")