add if @generated ... else ... end inside functions to provide optional optimizers

use meta nodes instead of `stagedfunction` expression head

Author: JeffBezanson

NEWS.md

@@ -19,6 +19,10 @@ New language features
  * The macro call syntax `@macroname[args]` is now available and is parsed
    as `@macroname([args])` ([#23519]).
 
+  * The construct `if @generated ...; else ...; end` can be used to provide both
+    `@generated` and normal implementations of part of a function. Surrounding code
+    will be common to both versions ([#23168]).
+
 Language changes
 ----------------
 

base/boot.jl

@@ -432,4 +432,32 @@ show(@nospecialize a) = show(STDOUT, a)
 print(@nospecialize a...) = print(STDOUT, a...)
 println(@nospecialize a...) = println(STDOUT, a...)
 
+struct GeneratedFunctionStub
+    gen
+    argnames
+    spnames
+    line
+    file
+end
+
+# invoke and wrap the results of @generated
+function (g::GeneratedFunctionStub)(args...)
+    body = g.gen(args...)
+    if body isa CodeInfo
+        return body
+    end
+    lam = Expr(:lambda, g.argnames,
+               Expr(Symbol("scope-block"),
+                    Expr(:block,
+                         LineNumberNode(g.line, g.file),
+                         Expr(:meta, :push_loc, g.file, Symbol("@generated body")),
+                         Expr(:return, body),
+                         Expr(:meta, :pop_loc))))
+    if g.spnames === nothing
+        return lam
+    else
+        return Expr(Symbol("with-static-parameters"), lam, g.spnames...)
+    end
+end
+
 ccall(:jl_set_istopmod, Void, (Any, Bool), Core, true)

base/docs/Docs.jl

@@ -642,7 +642,7 @@ finddoc(λ, def) = false
 
 # Predicates and helpers for `docm` expression selection:
 
-const FUNC_HEADS    = [:function, :stagedfunction, :macro, :(=)]
+const FUNC_HEADS    = [:function, :macro, :(=)]
 const BINDING_HEADS = [:typealias, :const, :global, :(=)]  # deprecation: remove `typealias` post-0.6
 # For the special `:@mac` / `:(Base.@mac)` syntax for documenting a macro after definition.
 isquotedmacrocall(x) =

base/expr.jl

@@ -332,10 +332,23 @@ function remove_linenums!(ex::Expr)
     return ex
 end
 
+macro generated()
+    return Expr(:generated)
+end
+
 macro generated(f)
-     if isa(f, Expr) && (f.head === :function || is_short_function_def(f))
-        f.head = :stagedfunction
-        return Expr(:escape, f)
+    if isa(f, Expr) && (f.head === :function || is_short_function_def(f))
+        body = f.args[2]
+        lno = body.args[1]
+        return Expr(:escape,
+                    Expr(f.head, f.args[1],
+                         Expr(:block,
+                              lno,
+                              Expr(:if, Expr(:generated),
+                                   body,
+                                   Expr(:block,
+                                        Expr(:meta, :generated_only),
+                                        Expr(:return, nothing))))))
     else
         error("invalid syntax; @generated must be used with a function definition")
     end

base/linalg/bidiag.jl

@@ -573,12 +573,18 @@ _valuefields(::Type{<:AbstractTriangular}) = [:data]
 
 const SpecialArrays = Union{Diagonal,Bidiagonal,Tridiagonal,SymTridiagonal,AbstractTriangular}
 
-@generated function fillslots!(A::SpecialArrays, x)
-    ex = :(xT = convert(eltype(A), x))
-    for field in _valuefields(A)
-        ex = :($ex; fill!(A.$field, xT))
+function fillslots!(A::SpecialArrays, x)
+    xT = convert(eltype(A), x)
+    if @generated
+        quote
+            $([ :(fill!(A.$field, xT)) for field in _valuefields(A) ]...)
+        end
+    else
+        for field in _valuefields(A)
+            fill!(getfield(A, field), xT)
+        end
     end
-    :($ex;return A)
+    return A
 end
 
 # for historical reasons:

base/methodshow.jl

@@ -42,6 +42,15 @@ function argtype_decl(env, n, sig::DataType, i::Int, nargs, isva::Bool) # -> (ar
     return s, string_with_env(env, t)
 end
 
+function method_argnames(m::Method)
+    if !isdefined(m, :source) && isdefined(m, :generator)
+        return m.generator.argnames
+    end
+    argnames = Vector{Any}(m.nargs)
+    ccall(:jl_fill_argnames, Void, (Any, Any), m.source, argnames)
+    return argnames
+end
+
 function arg_decl_parts(m::Method)
     tv = Any[]
     sig = m.sig
@@ -52,8 +61,7 @@ function arg_decl_parts(m::Method)
     file = m.file
     line = m.line
     if isdefined(m, :source) || isdefined(m, :generator)
-        argnames = Vector{Any}(m.nargs)
-        ccall(:jl_fill_argnames, Void, (Any, Any), isdefined(m, :source) ? m.source : m.generator.inferred, argnames)
+        argnames = method_argnames(m)
         show_env = ImmutableDict{Symbol, Any}()
         for t in tv
             show_env = ImmutableDict(show_env, :unionall_env => t)

base/multidimensional.jl

@@ -549,14 +549,11 @@ end
 @noinline throw_checksize_error(A, sz) = throw(DimensionMismatch("output array is the wrong size; expected $sz, got $(size(A))"))
 
 ## setindex! ##
-@generated function _setindex!(l::IndexStyle, A::AbstractArray, x, I::Union{Real, AbstractArray}...)
-    N = length(I)
-    quote
-        @_inline_meta
-        @boundscheck checkbounds(A, I...)
-        _unsafe_setindex!(l, _maybe_reshape(l, A, I...), x, I...)
-        A
-    end
+function _setindex!(l::IndexStyle, A::AbstractArray, x, I::Union{Real, AbstractArray}...)
+    @_inline_meta
+    @boundscheck checkbounds(A, I...)
+    _unsafe_setindex!(l, _maybe_reshape(l, A, I...), x, I...)
+    A
 end
 
 _iterable(v::AbstractArray) = v
@@ -916,28 +913,29 @@ function copy!(dest::AbstractArray{T,N}, src::AbstractArray{T,N}) where {T,N}
     dest
 end
 
-@generated function copy!(dest::AbstractArray{T1,N},
-                          Rdest::CartesianRange{N},
-                          src::AbstractArray{T2,N},
-                          Rsrc::CartesianRange{N}) where {T1,T2,N}
-    quote
-        isempty(Rdest) && return dest
-        if size(Rdest) != size(Rsrc)
-            throw(ArgumentError("source and destination must have same size (got $(size(Rsrc)) and $(size(Rdest)))"))
+function copy!(dest::AbstractArray{T1,N}, Rdest::CartesianRange{N},
+               src::AbstractArray{T2,N}, Rsrc::CartesianRange{N}) where {T1,T2,N}
+    isempty(Rdest) && return dest
+    if size(Rdest) != size(Rsrc)
+        throw(ArgumentError("source and destination must have same size (got $(size(Rsrc)) and $(size(Rdest)))"))
+    end
+    @boundscheck checkbounds(dest, first(Rdest))
+    @boundscheck checkbounds(dest, last(Rdest))
+    @boundscheck checkbounds(src, first(Rsrc))
+    @boundscheck checkbounds(src, last(Rsrc))
+    ΔI = first(Rdest) - first(Rsrc)
+    if @generated
+        quote
+            @nloops $N i (n->Rsrc.indices[n]) begin
+                @inbounds @nref($N,dest,n->i_n+ΔI[n]) = @nref($N,src,i)
+            end
         end
-        @boundscheck checkbounds(dest, first(Rdest))
-        @boundscheck checkbounds(dest, last(Rdest))
-        @boundscheck checkbounds(src, first(Rsrc))
-        @boundscheck checkbounds(src, last(Rsrc))
-        ΔI = first(Rdest) - first(Rsrc)
-        # TODO: restore when #9080 is fixed
-        # for I in Rsrc
-        #     @inbounds dest[I+ΔI] = src[I]
-        @nloops $N i (n->Rsrc.indices[n]) begin
-            @inbounds @nref($N,dest,n->i_n+ΔI[n]) = @nref($N,src,i)
+    else
+        for I in Rsrc
+            @inbounds dest[I + ΔI] = src[I]
         end
-        dest
     end
+    dest
 end
 
 """

base/reflection.jl

@@ -598,12 +598,13 @@ end
 """
     code_lowered(f, types, expand_generated = true)
 
-Return an array of lowered ASTs for the methods matching the given generic function and type signature.
+Return an array of the lowered forms (IR) for the methods matching the given generic function
+and type signature.
 
-If `expand_generated` is `false`, then the `CodeInfo` instances returned for `@generated`
-methods will correspond to the generators' lowered ASTs. If `expand_generated` is `true`,
-these `CodeInfo` instances will correspond to the lowered ASTs of the method bodies yielded
-by expanding the generators.
+If `expand_generated` is `false`, the returned `CodeInfo` instances will correspond to fallback
+implementations. An error is thrown if no fallback implementation exists.
+If `expand_generated` is `true`, these `CodeInfo` instances will correspond to the method bodies
+yielded by expanding the generators.
 
 Note that an error will be thrown if `types` are not leaf types when `expand_generated` is
 `true` and the corresponding method is a `@generated` method.
@@ -738,7 +739,8 @@ function length(mt::MethodTable)
 end
 isempty(mt::MethodTable) = (mt.defs === nothing)
 
-uncompressed_ast(m::Method) = uncompressed_ast(m, isdefined(m, :source) ? m.source : m.generator.inferred)
+uncompressed_ast(m::Method) = isdefined(m,:source) ? uncompressed_ast(m, m.source) :
+    error("Method is @generated; try `code_lowered` instead.")
 uncompressed_ast(m::Method, s::CodeInfo) = s
 uncompressed_ast(m::Method, s::Array{UInt8,1}) = ccall(:jl_uncompress_ast, Any, (Any, Any), m, s)::CodeInfo
 uncompressed_ast(m::Core.MethodInstance) = uncompressed_ast(m.def)
@@ -852,7 +854,7 @@ code_native(::IO, ::Any, ::Symbol) = error("illegal code_native call") # resolve
 
 # give a decent error message if we try to instantiate a staged function on non-leaf types
 function func_for_method_checked(m::Method, @nospecialize types)
-    if isdefined(m,:generator) && !isdefined(m,:source) && !_isleaftype(types)
+    if isdefined(m,:generator) && !_isleaftype(types)
         error("cannot call @generated function `", m, "` ",
               "with abstract argument types: ", types)
     end
@@ -862,7 +864,7 @@ end
 """
     code_typed(f, types; optimize=true)
 
-Returns an array of lowered and type-inferred ASTs for the methods matching the given
+Returns an array of type-inferred lowered form (IR) for the methods matching the given
 generic function and type signature. The keyword argument `optimize` controls whether
 additional optimizations, such as inlining, are also applied.
 """

base/sysimg.jl

@@ -236,21 +236,27 @@ include("broadcast.jl")
 using .Broadcast
 
 # define the real ntuple functions
-@generated function ntuple(f::F, ::Val{N}) where {F,N}
-    Core.typeassert(N, Int)
-    (N >= 0) || return :(throw($(ArgumentError(string("tuple length should be ≥0, got ", N)))))
-    return quote
-        $(Expr(:meta, :inline))
-        @nexprs $N i -> t_i = f(i)
-        @ncall $N tuple t
+@inline function ntuple(f::F, ::Val{N}) where {F,N}
+    N::Int
+    (N >= 0) || throw(ArgumentError(string("tuple length should be ≥0, got ", N)))
+    if @generated
+        quote
+            @nexprs $N i -> t_i = f(i)
+            @ncall $N tuple t
+        end
+    else
+        Tuple(f(i) for i = 1:N)
     end
 end
-@generated function fill_to_length(t::Tuple, val, ::Val{N}) where {N}
-    M = length(t.parameters)
-    M > N  && return :(throw($(ArgumentError("input tuple of length $M, requested $N"))))
-    return quote
-        $(Expr(:meta, :inline))
-        (t..., $(Any[ :val for i = (M + 1):N ]...))
+@inline function fill_to_length(t::Tuple, val, ::Val{N}) where {N}
+    M = length(t)
+    M > N && throw(ArgumentError("input tuple of length $M, requested $N"))
+    if @generated
+        quote
+            (t..., $(fill(:val, N-length(t.parameters))...))
+        end
+    else
+        (t..., fill(val, N-M)...)
     end
 end
 

doc/src/manual/metaprogramming.md

@@ -1012,17 +1012,16 @@ syntax tree.
 A very special macro is `@generated`, which allows you to define so-called *generated functions*.
 These have the capability to generate specialized code depending on the types of their arguments
 with more flexibility and/or less code than what can be achieved with multiple dispatch. While
-macros work with expressions at parsing-time and cannot access the types of their inputs, a generated
+macros work with expressions at parse time and cannot access the types of their inputs, a generated
 function gets expanded at a time when the types of the arguments are known, but the function is
 not yet compiled.
 
 Instead of performing some calculation or action, a generated function declaration returns a quoted
 expression which then forms the body for the method corresponding to the types of the arguments.
-When called, the body expression is first evaluated and compiled, then the returned expression
-is compiled and run. To make this efficient, the result is often cached. And to make this inferable,
-only a limited subset of the language is usable. Thus, generated functions provide a flexible
-framework to move work from run-time to compile-time, at the expense of greater restrictions on
-the allowable constructs.
+When a generated function is called, the expression it returns is compiled and then run.
+To make this efficient, the result is usually cached. And to make this inferable, only a limited
+subset of the language is usable. Thus, generated functions provide a flexible way to move work from
+run time to compile time, at the expense of greater restrictions on allowed constructs.
 
 When defining generated functions, there are four main differences to ordinary functions:
 
@@ -1038,7 +1037,7 @@ When defining generated functions, there are four main differences to ordinary f
    This means they can only read global constants, and cannot have any side effects.
    In other words, they must be completely pure.
    Due to an implementation limitation, this also means that they currently cannot define a closure
-   or untyped generator.
+   or generator.
 
 It's easiest to illustrate this with an example. We can declare a generated function `foo` as
 
@@ -1053,9 +1052,8 @@ foo (generic function with 1 method)
 Note that the body returns a quoted expression, namely `:(x * x)`, rather than just the value
 of `x * x`.
 
-From the caller's perspective, they are very similar to regular functions; in fact, you don't
-have to know if you're calling a regular or generated function - the syntax and result of the
-call is just the same. Let's see how `foo` behaves:
+From the caller's perspective, this is identical to a regular function; in fact, you don't
+have to know whether you're calling a regular or generated function. Let's see how `foo` behaves:
 
 ```jldoctest generated
 julia> x = foo(2); # note: output is from println() statement in the body
@@ -1199,7 +1197,7 @@ end and at the call site; however, *don't copy them*, for the following reasons:
     when, how often or how many times these side-effects will occur
   * the `bar` function solves a problem that is better solved with multiple dispatch - defining `bar(x) = x`
     and `bar(x::Integer) = x ^ 2` will do the same thing, but it is both simpler and faster.
-  * the `baz` function is pathologically insane
+  * the `baz` function is pathological
 
 Note that the set of operations that should not be attempted in a generated function is unbounded,
 and the runtime system can currently only detect a subset of the invalid operations. There are
@@ -1317,3 +1315,59 @@ the two tuples, multiplication and addition/subtraction. All the looping is perf
 and we avoid looping during execution entirely. Thus, we only loop *once per type*, in this case
 once per `N` (except in edge cases where the function is generated more than once - see disclaimer
 above).
+
+### Optionally-generated functions
+
+Generated functions can achieve high efficiency at run time, but come with a compile time cost:
+a new function body must be generated for every combination of concrete argument types.
+Typically, Julia is able to compile "generic" versions of functions that will work for any
+arguments, but with generated functions this is impossible.
+This means that programs making heavy use of generated functions might be impossible to
+statically compile.
+
+To solve this problem, the language provides syntax for writing normal, non-generated
+alternative implementations of generated functions.
+Applied to the `sub2ind` example above, it would look like this:
+
+```julia
+function sub2ind_gen(dims::NTuple{N}, I::Integer...) where N
+    if N != length(I)
+        throw(ArgumentError("Number of dimensions must match number of indices."))
+    end
+    if @generated
+        ex = :(I[$N] - 1)
+        for i = (N - 1):-1:1
+            ex = :(I[$i] - 1 + dims[$i] * $ex)
+        end
+        return :($ex + 1)
+    else
+        ind = I[N] - 1
+        for i = (N - 1):-1:1
+            ind = I[i] - 1 + dims[i]*ind
+        end
+        return ind + 1
+    end
+end
+```
+
+Internally, this code creates two implementations of the function: a generated one where
+the first block in `if @generated` is used, and a normal one where the `else` block is used.
+Inside the `then` part of the `if @generated` block, code has the same semantics as other
+generated functions: argument names refer to types, and the code should return an expression.
+Multiple `if @generated` blocks may occur, in which case the generated implementation uses
+all of the `then` blocks and the alternate implementation uses all of the `else` blocks.
+
+Notice that we added an error check to the top of the function.
+This code will be common to both versions, and is run-time code in both versions
+(it will be quoted and returned as an expression from the generated version).
+That means that the values and types of local variables are not available at code generation
+time --- the code-generation code can only see the types of arguments.
+
+In this style of definition, the code generation feature is essentially an optional
+optimization.
+The compiler will use it if convenient, but otherwise may choose to use the normal
+implementation instead.
+This style is preferred, since it allows the compiler to make more decisions and compile
+programs in more ways, and since normal code is more readable than code-generating code.
+However, which implementation is used depends on compiler implementation details, so it
+is essential for the two implementations to behave identically.