Preserve method-related meta forms through lowering (#90)

Co-authored-by: Shuhei Kadowaki <[email protected]>

Author: mlechu

src/ast.jl

@@ -519,6 +519,14 @@ end
 # the middle of a pass.
 const CompileHints = Base.ImmutableDict{Symbol,Any}
 
+function CompileHints(d::Dict{Symbol, Any})
+    id = CompileHints()
+    for (k, v) in d
+        id = CompileHints(id, k, v)
+    end
+    id
+end
+
 function setmeta!(ex::SyntaxTree; kws...)
     @assert length(kws) == 1 # todo relax later ?
     key = first(keys(kws))

src/compat.jl

@@ -459,27 +459,39 @@ function _insert_convert_expr(@nospecialize(e), graph::SyntaxGraph, src::SourceA
         st_k = K"hygienic_scope"
     elseif e.head === :meta
         # Messy and undocumented.  Only sometimes we want a K"meta".
-        @assert e.args[1] isa Symbol
-        if e.args[1] === :nospecialize
-            if nargs > 2
-                st_k = K"block"
-                # Kick the can down the road (should only be simple atoms?)
-                child_exprs = map(c->Expr(:meta, :nospecialize, c), child_exprs[2:end])
+        if e.args[1] isa Expr && e.args[1].head === :purity
+            st_k = K"meta"
+            child_exprs = [Expr(:quoted_symbol, :purity), Base.EffectsOverride(e.args[1].args...)]
+        else
+            @assert e.args[1] isa Symbol
+            if e.args[1] === :nospecialize
+                if nargs > 2
+                    st_k = K"block"
+                    # Kick the can down the road (should only be simple atoms?)
+                    child_exprs = map(c->Expr(:meta, :nospecialize, c), child_exprs[2:end])
+                else
+                    st_id, src = _insert_convert_expr(e.args[2], graph, src)
+                    setmeta!(SyntaxTree(graph, st_id); nospecialize=true)
+                    return st_id, src
+                end
+            elseif e.args[1] in (:inline, :noinline, :generated, :generated_only,
+                                 :max_methods, :optlevel, :toplevel, :push_loc, :pop_loc,
+                                 :no_constprop, :aggressive_constprop, :specialize, :compile, :infer,
+                                 :nospecializeinfer, :force_compile, :propagate_inbounds, :doc)
+                # TODO: Some need to be handled in lowering
+                for (i, ma) in enumerate(e.args)
+                    if ma isa Symbol
+                        # @propagate_inbounds becomes (meta inline
+                        # propagate_inbounds), but usually(?) only args[1] is
+                        # converted here
+                        child_exprs[i] = Expr(:quoted_symbol, e.args[i])
+                    end
+                end
             else
-                st_id, src = _insert_convert_expr(e.args[2], graph, src)
-                setmeta!(SyntaxTree(graph, st_id); nospecialize=true)
-                return st_id, src
+                # Can't throw a hard error; it is explicitly tested that meta can take arbitrary keys.
+                @error("Unknown meta form at $src: `$e`\n$(sprint(dump, e))")
+                child_exprs[1] = Expr(:quoted_symbol, e.args[1])
             end
-        elseif e.args[1] in (:inline, :noinline, :generated, :generated_only,
-                             :max_methods, :optlevel, :toplevel, :push_loc, :pop_loc,
-                             :no_constprop, :aggressive_constprop, :specialize, :compile, :infer,
-                             :nospecializeinfer, :force_compile, :doc)
-            # TODO: Some need to be handled in lowering
-            child_exprs[1] = Expr(:quoted_symbol, e.args[1])
-        else
-            # Can't throw a hard error; it is explicitly tested that meta can take arbitrary keys.
-            @error("Unknown meta form at $src: `$e`\n$(sprint(dump, e))")
-            child_exprs[1] = Expr(:quoted_symbol, e.args[1])
         end
     elseif e.head === :scope_layer
         @assert nargs === 2

src/eval.jl

@@ -218,7 +218,7 @@ end
 
 # Convert SyntaxTree to the CodeInfo+Expr data stuctures understood by the
 # Julia runtime
-function to_code_info(ex::SyntaxTree, slots::Vector{Slot})
+function to_code_info(ex::SyntaxTree, slots::Vector{Slot}, meta::CompileHints)
     stmts = Any[]
 
     current_codelocs_stack = ir_debug_info_state(ex)
@@ -267,18 +267,22 @@ function to_code_info(ex::SyntaxTree, slots::Vector{Slot})
     # - call site @assume_effects
     ssaflags = zeros(UInt32, length(stmts))
 
-    # TODO: Set true for @propagate_inbounds
-    propagate_inbounds  = false
+    propagate_inbounds =
+        get(meta, :propagate_inbounds, false)
     # TODO: Set true if there's a foreigncall
-    has_fcall           = false
-    # TODO: Set for @nospecializeinfer
-    nospecializeinfer   = false
-    # TODO: Set based on @inline -> 0x01 or @noinline -> 0x02
-    inlining            = 0x00
-    # TODO: Set based on @constprop :aggressive -> 0x01 or @constprop :none -> 0x02
-    constprop           = 0x00
-    # TODO: Set based on Base.@assume_effects
-    purity              = 0x0000
+    has_fcall = false
+    nospecializeinfer =
+        get(meta, :nospecializeinfer, false)
+    inlining =
+        get(meta, :inline, false) ? 0x01 :
+        get(meta, :noinline, false) ? 0x02 : 0x00
+    constprop =
+        get(meta, :aggressive_constprop, false) ? 0x01 :
+        get(meta, :no_constprop, false) ? 0x02 : 0x00
+    purity =
+        let eo = get(meta, :purity, nothing)
+            isnothing(eo) ? 0x0000 : Base.encode_effects_override(eo)
+        end
 
     # The following CodeInfo fields always get their default values for
     # uninferred code.
@@ -367,7 +371,7 @@ function _to_lowered_expr(ex::SyntaxTree, stmt_offset::Int)
         e1 = ex[1]
         getmeta(ex, :as_Expr, false) ? QuoteNode(Expr(e1)) : e1
     elseif k == K"code_info"
-        ir = to_code_info(ex[1], ex.slots)
+        ir = to_code_info(ex[1], ex.slots, ex.meta)
         if ex.is_toplevel_thunk
             Expr(:thunk, ir) # TODO: Maybe nice to just return a CodeInfo here?
         else

src/linear_ir.jl

@@ -78,6 +78,7 @@ struct LinearIRContext{GraphType} <: AbstractLoweringContext
     finally_handlers::Vector{FinallyHandler{GraphType}}
     symbolic_jump_targets::Dict{String,JumpTarget{GraphType}}
     symbolic_jump_origins::Vector{JumpOrigin{GraphType}}
+    meta::Dict{Symbol, Any}
     mod::Module
 end
 
@@ -89,7 +90,7 @@ function LinearIRContext(ctx, is_toplevel_thunk, lambda_bindings, return_type)
                     is_toplevel_thunk, lambda_bindings, rett,
                     Dict{String,JumpTarget{GraphType}}(), SyntaxList(ctx), SyntaxList(ctx),
                     Vector{FinallyHandler{GraphType}}(), Dict{String,JumpTarget{GraphType}}(),
-                    Vector{JumpOrigin{GraphType}}(), ctx.mod)
+                    Vector{JumpOrigin{GraphType}}(), Dict{Symbol, Any}(), ctx.mod)
 end
 
 function current_lambda_bindings(ctx::LinearIRContext)
@@ -807,7 +808,17 @@ function compile(ctx::LinearIRContext, ex, needs_value, in_tail_pos)
             nothing
         end
     elseif k == K"meta"
-        emit(ctx, ex)
+        @chk numchildren(ex) >= 1
+        if ex[1].name_val in ("inline", "noinline", "propagate_inbounds",
+                              "nospecializeinfer", "aggressive_constprop", "no_constprop")
+            for c in children(ex)
+                ctx.meta[Symbol(c.name_val)] = true
+            end
+        elseif ex[1].name_val === "purity"
+            ctx.meta[Symbol(ex[1].name_val)] = ex[2].value::Base.EffectsOverride
+        else
+            emit(ctx, ex)
+        end
         if needs_value
             val = @ast ctx ex "nothing"::K"core"
             if in_tail_pos
@@ -1099,10 +1110,9 @@ function compile_lambda(outer_ctx, ex)
         @assert info.kind == :static_parameter
         slot_rewrites[id] = i
     end
-    # @info "" @ast ctx ex [K"block" ctx.code...]
     code = renumber_body(ctx, ctx.code, slot_rewrites)
     @ast ctx ex [K"code_info"(is_toplevel_thunk=ex.is_toplevel_thunk,
-                              slots=slots)
+                              slots=slots, meta=CompileHints(ctx.meta))
         [K"block"(ex[3])
             code...
         ]
@@ -1131,7 +1141,8 @@ loops, etc) to gotos and exception handling to enter/leave. We also convert
                            SyntaxList(graph), SyntaxList(graph),
                            Vector{FinallyHandler{GraphType}}(),
                            Dict{String, JumpTarget{GraphType}}(),
-                           Vector{JumpOrigin{GraphType}}(), ctx.mod)
+                           Vector{JumpOrigin{GraphType}}(),
+                           Dict{Symbol, Any}(), ctx.mod)
     res = compile_lambda(_ctx, reparent(_ctx, ex))
     _ctx, res
 end

test/macros.jl

@@ -362,4 +362,48 @@ end
     @test test_mod.X1 isa Enum
 end
 
+@testset "macros producing meta forms" begin
+    function find_method_ci(thunk)
+        ci = thunk.args[1]::Core.CodeInfo
+        m = findfirst(x->(x isa Expr && x.head === :method && length(x.args) === 3), ci.code)
+        ci.code[m].args[3]
+    end
+    jlower_e(s) = JuliaLowering.to_lowered_expr(
+        JuliaLowering.lower(
+            test_mod, JuliaLowering.parsestmt(
+                JuliaLowering.SyntaxTree, s);
+            expr_compat_mode=true))
+
+    prog = "Base.@assume_effects :foldable function foo(); end"
+    ref = Meta.lower(test_mod, Meta.parse(prog))
+    our = jlower_e(prog)
+    @test find_method_ci(ref).purity === find_method_ci(our).purity
+
+    prog = "Base.@inline function foo(); end"
+    ref = Meta.lower(test_mod, Meta.parse(prog))
+    our = jlower_e(prog)
+    @test find_method_ci(ref).inlining === find_method_ci(our).inlining
+
+    prog = "Base.@noinline function foo(); end"
+    ref = Meta.lower(test_mod, Meta.parse(prog))
+    our = jlower_e(prog)
+    @test find_method_ci(ref).inlining === find_method_ci(our).inlining
+
+    prog = "Base.@constprop :none function foo(); end"
+    ref = Meta.lower(test_mod, Meta.parse(prog))
+    our = jlower_e(prog)
+    @test find_method_ci(ref).constprop === find_method_ci(our).constprop
+
+    prog = "Base.@nospecializeinfer function foo(); end"
+    ref = Meta.lower(test_mod, Meta.parse(prog))
+    our = jlower_e(prog)
+    @test find_method_ci(ref).nospecializeinfer === find_method_ci(our).nospecializeinfer
+
+    prog = "Base.@propagate_inbounds function foo(); end"
+    ref = Meta.lower(test_mod, Meta.parse(prog))
+    our = jlower_e(prog)
+    @test find_method_ci(ref).propagate_inbounds === find_method_ci(our).propagate_inbounds
+
+end
+
 end