WIP: allow @generated begin ... end inside a function to provide an optional optimizer

JeffBezanson · JeffBezanson · commit bbb9284106f6 · 2017-08-09T16:32:10.000-04:00
diff --git a/base/expr.jl b/base/expr.jl
@@ -334,8 +334,17 @@ end
 
 macro generated(f)
     if isa(f, Expr) && (f.head === :function || is_short_function_def(f))
-        pushmeta!(f, :generated)
-        return Expr(:escape, f)
+        body = f.args[2]
+        lno = body.args[1]
+        return Expr(:escape,
+                    Expr(f.head, f.args[1],
+                         Expr(:block,
+                              lno,
+                              Expr(:meta, :generator, body),
+                              Expr(:meta, :generated_only),
+                              Expr(:return, nothing))))
+    elseif isa(f, Expr) && f.head === :block
+        return Expr(:escape, Expr(:meta, :generator, f))
     else
         error("invalid syntax; @generated must be used with a function definition")
     end
diff --git a/base/methodshow.jl b/base/methodshow.jl
@@ -42,6 +42,16 @@ 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)
+        gm = first(methods(m.generator))
+        return method_argnames(gm)[length(m.sparam_syms)+2 : end]
+    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 +62,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)
diff --git a/base/reflection.jl b/base/reflection.jl
@@ -682,7 +682,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) :
+                                                     uncompressed_ast(first(methods(m.generator)))
 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
 
diff --git a/src/ast.c b/src/ast.c
@@ -55,7 +55,8 @@ jl_sym_t *meta_sym; jl_sym_t *compiler_temp_sym;
 jl_sym_t *inert_sym; jl_sym_t *vararg_sym;
 jl_sym_t *unused_sym; jl_sym_t *static_parameter_sym;
 jl_sym_t *polly_sym; jl_sym_t *inline_sym;
-jl_sym_t *propagate_inbounds_sym; jl_sym_t *generated_sym;
+jl_sym_t *propagate_inbounds_sym; jl_sym_t *generator_sym;
+jl_sym_t *generated_only_sym;
 jl_sym_t *isdefined_sym; jl_sym_t *nospecialize_sym;
 
 static uint8_t flisp_system_image[] = {
@@ -437,7 +438,8 @@ void jl_init_frontend(void)
     propagate_inbounds_sym = jl_symbol("propagate_inbounds");
     isdefined_sym = jl_symbol("isdefined");
     nospecialize_sym = jl_symbol("nospecialize");
-    generated_sym = jl_symbol("generated");
+    generator_sym = jl_symbol("generator");
+    generated_only_sym = jl_symbol("generated_only");
 }
 
 JL_DLLEXPORT void jl_lisp_prompt(void)
diff --git a/src/ast.scm b/src/ast.scm
@@ -345,8 +345,11 @@
   (and (if one (length= e 3) (length> e 2))
        (eq? (car e) 'meta) (eq? (cadr e) 'nospecialize)))
 
-(define (generated-meta? e)
-  (and (pair? e) (eq? (car e) 'meta) (any (lambda (x) (eq? x 'generated)) (cdr e))))
+(define (generator-meta? e)
+  (and (length= e 3) (eq? (car e) 'meta) (eq? (cadr e) 'generator)))
+
+(define (generated_only-meta? e)
+  (and (length= e 2) (eq? (car e) 'meta) (eq? (cadr e) 'generated_only)))
 
 ;; flatten nested expressions with the given head
 ;; (op (op a b) c) => (op a b c)
diff --git a/src/codegen.cpp b/src/codegen.cpp
@@ -1197,8 +1197,6 @@ jl_llvm_functions_t jl_compile_linfo(jl_method_instance_t **pli, jl_code_info_t
                 li->inferred &&
                 // and there is something to delete (test this before calling jl_ast_flag_inlineable)
                 li->inferred != jl_nothing &&
-                // don't delete the code for the generator
-                li != li->def.method->generator &&
                 // don't delete inlineable code, unless it is constant
                 (li->jlcall_api == 2 || !jl_ast_flag_inlineable((jl_array_t*)li->inferred)) &&
                 // don't delete code when generating a precompile file
diff --git a/src/dump.c b/src/dump.c
@@ -1417,7 +1417,7 @@ static jl_value_t *jl_deserialize_value_method(jl_serializer_state *s, jl_value_
     m->unspecialized = (jl_method_instance_t*)jl_deserialize_value(s, (jl_value_t**)&m->unspecialized);
     if (m->unspecialized)
         jl_gc_wb(m, m->unspecialized);
-    m->generator = (jl_method_instance_t*)jl_deserialize_value(s, (jl_value_t**)&m->generator);
+    m->generator = jl_deserialize_value(s, (jl_value_t**)&m->generator);
     if (m->generator)
         jl_gc_wb(m, m->generator);
     m->invokes.unknown = jl_deserialize_value(s, (jl_value_t**)&m->invokes);
diff --git a/src/jltypes.c b/src/jltypes.c
@@ -2045,7 +2045,7 @@ void jl_init_types(void)
                             jl_simplevector_type,
                             jl_any_type,
                             jl_any_type, // jl_method_instance_type
-                            jl_any_type, // jl_method_instance_type
+                            jl_any_type,
                             jl_array_any_type,
                             jl_any_type,
                             jl_int32_type,
@@ -2158,7 +2158,6 @@ void jl_init_types(void)
 #endif
     jl_svecset(jl_methtable_type->types, 8, jl_int32_type); // uint32_t
     jl_svecset(jl_method_type->types, 10, jl_method_instance_type);
-    jl_svecset(jl_method_type->types, 11, jl_method_instance_type);
     jl_svecset(jl_method_instance_type->types, 12, jl_voidpointer_type);
     jl_svecset(jl_method_instance_type->types, 13, jl_voidpointer_type);
     jl_svecset(jl_method_instance_type->types, 14, jl_voidpointer_type);
diff --git a/src/julia-syntax.scm b/src/julia-syntax.scm
@@ -294,6 +294,17 @@
                (map (lambda (x) (replace-outer-vars x renames))
                     (cdr e))))))
 
+(define (make-generator-function name sp-names arg-names body)
+  (let ((arg-names (append sp-names
+                           (map (lambda (n)
+                                  (if (eq? n '|#self#|) (gensy) n))
+                                arg-names))))
+    (let ((body (insert-after-meta body  ;; don't specialize on generator arguments
+                                   `((meta nospecialize ,@arg-names)))))
+      `(block
+        (global ,name)
+        (function (call ,name ,@arg-names) ,body)))))
+
 ;; construct the (method ...) expression for one primitive method definition,
 ;; assuming optional and keyword args are already handled
 (define (method-def-expr- name sparams argl body (rett '(core Any)))
@@ -328,7 +339,14 @@
            (error "function argument and static parameter names must be distinct")))
      (if (or (and name (not (sym-ref? name))) (eq? name 'true) (eq? name 'false))
          (error (string "invalid function name \"" (deparse name) "\"")))
-     (let* ((types (llist-types argl))
+     (let* ((generator (let ((found (find generator-meta? body)))
+                         (if found
+                             (let* ((gname (symbol (string (gensy) "#" (current-julia-module-counter))))
+                                    (gf (make-generator-function gname names (llist-vars argl) (caddr (car found)))))
+                               (set-car! (cddar found) gname)
+                               (list gf))
+                             '())))
+            (types (llist-types argl))
             (body  (method-lambda-expr argl body rett))
             ;; HACK: the typevars need to be bound to ssavalues, since this code
             ;; might be moved to a different scope by closure-convert.
@@ -360,8 +378,10 @@
                                  (call (core svec) ,@temps)))
                           ,body))))
        (if (symbol? name)
-           `(block (method ,name) ,mdef (unnecessary ,name))  ;; return the function
-           mdef)))))
+           `(block ,@generator (method ,name) ,mdef (unnecessary ,name))  ;; return the function
+           (if (not (null? generator))
+               `(block ,@generator ,mdef)
+               mdef))))))
 
 ;; wrap expr in nested scopes assigning names to vals
 (define (scopenest names vals expr)
@@ -411,11 +431,8 @@
                                 keynames))
          ;; list of function's initial line number and meta nodes (empty if none)
          (prologue (extract-method-prologue body))
-         (annotations (append (if (any generated-meta? prologue)
-                                  '((meta generated))
-                                  '())
-                              (map (lambda (a) `(meta nospecialize ,(arg-name (cadr (caddr a)))))
-                                   (filter nospecialize-meta? kargl))))
+         (annotations (map (lambda (a) `(meta nospecialize ,(arg-name (cadr (caddr a)))))
+                           (filter nospecialize-meta? kargl)))
          ;; body statements
          (stmts (cdr body))
          (positional-sparams
@@ -565,10 +582,9 @@
       '()))
 
 (define (without-generated stmts)
-  (map (lambda (x) (if (generated-meta? x)
-                       (filter (lambda (e) (not (eq? e 'generated))) x)
-                       x))
-       stmts))
+  (filter (lambda (x) (not (or (generator-meta? x)
+                               (generated_only-meta? x))))
+          stmts))
 
 ;; keep only sparams used by `expr` or other sparams
 (define (filter-sparams expr sparams)
diff --git a/src/julia.h b/src/julia.h
@@ -247,7 +247,7 @@ typedef struct _jl_method_t {
     jl_svec_t *sparam_syms;  // symbols giving static parameter names
     jl_value_t *source;  // original code template (jl_code_info_t, but may be compressed), null for builtins
     struct _jl_method_instance_t *unspecialized;  // unspecialized executable method instance, or null
-    struct _jl_method_instance_t *generator;  // executable code-generating function if available
+    jl_value_t *generator;  // executable code-generating function if available
     jl_array_t *roots;  // pointers in generated code (shared to reduce memory), or null
 
     // cache of specializations of this method for invoke(), i.e.
diff --git a/src/julia_internal.h b/src/julia_internal.h
@@ -994,7 +994,8 @@ extern jl_sym_t *pure_sym; extern jl_sym_t *simdloop_sym;
 extern jl_sym_t *meta_sym; extern jl_sym_t *list_sym;
 extern jl_sym_t *inert_sym; extern jl_sym_t *static_parameter_sym;
 extern jl_sym_t *polly_sym; extern jl_sym_t *inline_sym;
-extern jl_sym_t *propagate_inbounds_sym; extern jl_sym_t *generated_sym;
+extern jl_sym_t *propagate_inbounds_sym; extern jl_sym_t *generator_sym;
+extern jl_sym_t *generated_only_sym;
 extern jl_sym_t *isdefined_sym; extern jl_sym_t *nospecialize_sym;
 
 void jl_register_fptrs(uint64_t sysimage_base, const char *base, const int32_t *offsets,
diff --git a/src/method.c b/src/method.c
@@ -245,24 +245,23 @@ jl_code_info_t *jl_new_code_info_from_ast(jl_expr_t *ast)
 }
 
 // invoke (compiling if necessary) the jlcall function pointer for a method template
-STATIC_INLINE jl_value_t *jl_call_staged(jl_svec_t *sparam_vals, jl_method_instance_t *generator,
+STATIC_INLINE jl_value_t *jl_call_staged(jl_method_t *def, jl_value_t *generator, jl_svec_t *sparam_vals,
                                          jl_value_t **args, uint32_t nargs)
 {
-    jl_generic_fptr_t fptr;
-    fptr.fptr = generator->fptr;
-    fptr.jlcall_api = generator->jlcall_api;
-    if (__unlikely(fptr.fptr == NULL || fptr.jlcall_api == 0)) {
-        size_t world = generator->def.method->min_world;
-        const char *F = jl_compile_linfo(&generator, (jl_code_info_t*)generator->inferred, world, &jl_default_cgparams).functionObject;
-        fptr = jl_generate_fptr(generator, F, world);
+    size_t spl = jl_svec_len(sparam_vals);
+    jl_value_t **gargs;
+    size_t totargs = 1 + spl + nargs + def->isva;
+    JL_GC_PUSHARGS(gargs, totargs);
+    gargs[0] = generator;
+    memcpy(&gargs[1], jl_svec_data(sparam_vals), spl * sizeof(void*));
+    memcpy(&gargs[1+spl], args, nargs * sizeof(void*));
+    if (def->isva) {
+        gargs[totargs-1] = jl_f_tuple(NULL, &gargs[1+spl+def->nargs-1], nargs - (def->nargs-1));
+        gargs[1+spl+def->nargs-1] = gargs[totargs-1];
     }
-    assert(jl_svec_len(generator->def.method->sparam_syms) == jl_svec_len(sparam_vals));
-    if (fptr.jlcall_api == 1)
-        return fptr.fptr1(args[0], &args[1], nargs-1);
-    else if (fptr.jlcall_api == 3)
-        return fptr.fptr3(sparam_vals, args[0], &args[1], nargs-1);
-    else
-        abort(); // shouldn't have inferred any other calling convention
+    jl_value_t *code = jl_apply(gargs, 1+spl+def->nargs);
+    JL_GC_POP();
+    return code;
 }
 
 // return a newly allocated CodeInfo for the function signature
@@ -275,9 +274,11 @@ JL_DLLEXPORT jl_code_info_t *jl_code_for_staged(jl_method_instance_t *linfo)
     jl_expr_t *ex = NULL;
     jl_value_t *linenum = NULL;
     jl_svec_t *sparam_vals = env;
-    jl_method_instance_t *generator = linfo->def.method->generator;
+    jl_value_t *generator = linfo->def.method->generator;
+    jl_method_t *gen_meth = jl_gf_mtable(generator)->defs.leaf->func.method;
     assert(generator != NULL);
     assert(linfo != generator);
+    assert(jl_is_method(gen_meth));
     jl_code_info_t *func = NULL;
     JL_GC_PUSH4(&ex, &linenum, &sparam_vals, &func);
     jl_ptls_t ptls = jl_get_ptls_states();
@@ -292,13 +293,14 @@ JL_DLLEXPORT jl_code_info_t *jl_code_for_staged(jl_method_instance_t *linfo)
         // need to eval macros in the right module
         ptls->current_task->current_module = ptls->current_module = linfo->def.method->module;
         // and the right world
-        ptls->world_age = generator->def.method->min_world;
+        ptls->world_age = gen_meth->min_world;
 
         ex = jl_exprn(lambda_sym, 2);
 
-        jl_array_t *argnames = jl_alloc_vec_any(linfo->def.method->nargs);
+        jl_array_t *argnames = jl_alloc_vec_any(linfo->def.method->nargs + jl_svec_len(sparam_vals) + 1);
         jl_array_ptr_set(ex->args, 0, argnames);
-        jl_fill_argnames((jl_array_t*)generator->inferred, argnames);
+        jl_fill_argnames((jl_array_t*)gen_meth->source, argnames);
+        jl_array_del_beg(argnames, jl_svec_len(sparam_vals) + 1);
 
         // build the rest of the body to pass to expand
         jl_expr_t *scopeblock = jl_exprn(jl_symbol("scope-block"), 1);
@@ -319,7 +321,7 @@ JL_DLLEXPORT jl_code_info_t *jl_code_for_staged(jl_method_instance_t *linfo)
         // invoke code generator
         assert(jl_nparams(tt) == jl_array_len(argnames) ||
                (linfo->def.method->isva && (jl_nparams(tt) >= jl_array_len(argnames) - 1)));
-        jl_value_t *generated_body = jl_call_staged(sparam_vals, generator, jl_svec_data(tt->parameters), jl_nparams(tt));
+        jl_value_t *generated_body = jl_call_staged(linfo->def.method, generator, sparam_vals, jl_svec_data(tt->parameters), jl_nparams(tt));
         jl_array_ptr_set(body->args, 2, generated_body);
 
         if (jl_is_code_info(generated_body)) {
@@ -398,7 +400,7 @@ jl_method_instance_t *jl_get_specialized(jl_method_t *m, jl_value_t *types, jl_s
     return new_linfo;
 }
 
-static void jl_method_set_source(jl_method_t *m, jl_code_info_t *src, int *isstaged)
+static void jl_method_set_source(jl_method_t *m, jl_code_info_t *src, jl_value_t **generator, int *gen_only)
 {
     uint8_t j;
     uint8_t called = 0;
@@ -470,12 +472,17 @@ static void jl_method_set_source(jl_method_t *m, jl_code_info_t *src, int *issta
                 }
                 st = jl_nothing;
             }
-            else {
-                for (size_t j=0; j < jl_expr_nargs(st); j++) {
-                    if (jl_exprarg(st, j) == (jl_value_t*)generated_sym) {
-                        *isstaged = 1; break;
-                    }
+            else if (jl_expr_nargs(st) == 2 && jl_exprarg(st, 0) == (jl_value_t*)generator_sym) {
+                jl_value_t *gname = jl_exprarg(st, 1);
+                *generator = jl_get_global(m->module, (jl_sym_t*)gname);
+                if (*generator == NULL) {
+                    jl_error("invalid @generated function; try placing it in global scope");
                 }
+                st = jl_nothing;
+            }
+            else if (jl_expr_nargs(st) == 1 && jl_exprarg(st, 0) == (jl_value_t*)generated_only_sym) {
+                *gen_only = 1;
+                st = jl_nothing;
             }
         }
         else {
@@ -530,7 +537,6 @@ static jl_method_t *jl_new_method(
         jl_svec_t *tvars)
 {
     size_t i, l = jl_svec_len(tvars);
-    int isstaged = 0;
     jl_svec_t *sparam_syms = jl_alloc_svec_uninit(l);
     for (i = 0; i < l; i++) {
         jl_svecset(sparam_syms, i, ((jl_tvar_t*)jl_svecref(tvars, i))->name);
@@ -547,12 +553,14 @@ static jl_method_t *jl_new_method(
     m->sig = (jl_value_t*)sig;
     m->isva = isva;
     m->nargs = nargs;
-    jl_method_set_source(m, definition, &isstaged);
-    if (isstaged) {
-        // create and store generator for generated functions
-        m->generator = jl_get_specialized(m, (jl_value_t*)jl_anytuple_type, jl_emptysvec);
+    jl_value_t *gen = NULL; int gen_only = 0;
+    jl_method_set_source(m, definition, &gen, &gen_only);
+    if (gen) {
+        m->generator = gen;
         jl_gc_wb(m, m->generator);
-        m->generator->inferred = (jl_value_t*)m->source;
+    }
+    if (gen_only) {
+        assert(gen);
         m->source = NULL;
     }
 
diff --git a/src/utils.scm b/src/utils.scm
@@ -84,3 +84,8 @@
                                           (without (cdr alst) remove)))))
 
 (define (caddddr x) (car (cdr (cdr (cdr (cdr x))))))
+
+(define (find p lst)
+  (cond ((atom? lst) #f)
+	((p (car lst)) lst)
+	(else (find p (cdr lst)))))
