generated functions, next generation

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

@@ -430,4 +430,32 @@ show(@nospecialize a) = show(STDOUT, a)
 print(@nospecialize a...) = print(STDOUT, a...)
 println(@nospecialize a...) = println(STDOUT, a...)
 
+struct GeneratedFunctionStub
+    gen
+    argnames::Array{Any,1}
+    spnames::Union{Void, Array{Any,1}}
+    line::Int
+    file::Symbol
+end
+
+# invoke and wrap the results of @generated
+function (g::GeneratedFunctionStub)(@nospecialize 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

@@ -577,12 +577,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

@@ -597,12 +597,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.
@@ -737,7 +738,9 @@ 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) :
+                              isdefined(m,:generator) ? error("Method is @generated; try `code_lowered` instead.") :
+                              error("Code for this Method is not available.")
 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)
@@ -851,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
@@ -861,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