Clean up test case

Author: Keno

test/compiler/contextual.jl

@@ -1,85 +1,110 @@
-using Core.Compiler: method_instances, retrieve_code_info, CodeInfo,
-    MethodInstance, SSAValue, GotoNode, Slot, SlotNumber
-using Base.Meta: isexpr
-using Test
+module MiniCassette
+    # A minimal demonstration of the cassette mechanism. Doesn't support all the
+    # fancy features, but sufficient to exercise this code path in the compiler.
 
-struct Ctx; end
+    using Core.Compiler: method_instances, retrieve_code_info, CodeInfo,
+        MethodInstance, SSAValue, GotoNode, Slot, SlotNumber, quoted,
+        signature_type
+    using Base: _methods_by_ftype
+    using Base.Meta: isexpr
+    using Test
 
-# A no-op cassette-like transform
-function transform_expr(expr, map_slot_number, map_ssa_value)
-    transform(expr) = transform_expr(expr, map_slot_number, map_ssa_value)
-    if isexpr(expr, :call)
-        return Expr(:call, overdub, SlotNumber(1), map(transform, expr.args)...)
-    elseif isexpr(expr, :gotoifnot)
-        return Expr(:gotoifnot, transform(expr), map_ssa_value(SSAValue(expr.args[2])).id)
-    elseif isa(expr, GotoNode)
-        return GotoNode(map_ssa_value(SSAValue(expr.label)).id)
-    elseif isa(expr, Slot)
-        return map_slot_number(expr.id)
-    elseif isa(expr, SSAValue)
-        return map_ssa_value(expr)
-    else
-        return expr
-    end
-end
+    export Ctx, overdub
 
-function transform!(ci, nargs)
-    code = ci.code
-    ci.slotnames = Symbol[Symbol("#self#"), :ctx, :f, :args, ci.slotnames[nargs+1:end]...]
-    ci.slotflags = UInt8[(0x00 for i = 1:4)..., ci.slotflags[nargs+1:end]...]
-    # Insert one SSAValue for every argument statement
-    for i = 1:nargs
-        pushfirst!(code, Expr(:getfield, SlotNumber(3), i))
-    end
-    function map_slot_number(slot)
-        if slot == 1
-            # self in the original function is now `f`
-            return SlotNumber(3)
-        elseif 2 <= slot <= nargs + 1
-            # Arguments get inserted as ssa values at the top of the function
-            return SSAValue(slot - 2)
+    struct Ctx; end
+
+    # A no-op cassette-like transform
+    function transform_expr(expr, map_slot_number, map_ssa_value, sparams)
+        transform(expr) = transform_expr(expr, map_slot_number, map_ssa_value, sparams)
+        if isexpr(expr, :call)
+            return Expr(:call, overdub, SlotNumber(2), map(transform, expr.args)...)
+        elseif isexpr(expr, :gotoifnot)
+            return Expr(:gotoifnot, transform(expr.args[1]), map_ssa_value(SSAValue(expr.args[2])).id)
+        elseif isexpr(expr, :static_parameter)
+            return quoted(sparams[expr.args[1]])
+        elseif isa(expr, Expr)
+            return Expr(expr.head, map(transform, expr.args)...)
+        elseif isa(expr, GotoNode)
+            return GotoNode(map_ssa_value(SSAValue(expr.label)).id)
+        elseif isa(expr, Slot)
+            return map_slot_number(expr.id)
+        elseif isa(expr, SSAValue)
+            return map_ssa_value(expr)
         else
-            # The first non-argument slot will be 5
-            return SlotNumber(slot - (nargs + 1) + 4)
+            return expr
         end
     end
-    map_ssa_value(ssa::SSAValue) = SSAValue(ssa.id + nargs)
-    for i = (nargs+1:length(code))
-        code[i] = transform_expr(code[i], map_slot_number, map_ssa_value)
+
+    function transform!(ci, nargs, sparams)
+        code = ci.code
+        ci.slotnames = Symbol[Symbol("#self#"), :ctx, :f, :args, ci.slotnames[nargs+1:end]...]
+        ci.slotflags = UInt8[(0x00 for i = 1:4)..., ci.slotflags[nargs+1:end]...]
+        # Insert one SSAValue for every argument statement
+        prepend!(code, [Expr(:call, getfield, SlotNumber(4), i) for i = 1:nargs])
+        prepend!(ci.codelocs, [0 for i = 1:nargs])
+        ci.ssavaluetypes += nargs
+        function map_slot_number(slot)
+            if slot == 1
+                # self in the original function is now `f`
+                return SlotNumber(3)
+            elseif 2 <= slot <= nargs + 1
+                # Arguments get inserted as ssa values at the top of the function
+                return SSAValue(slot - 1)
+            else
+                # The first non-argument slot will be 5
+                return SlotNumber(slot - (nargs + 1) + 4)
+            end
+        end
+        map_ssa_value(ssa::SSAValue) = SSAValue(ssa.id + nargs)
+        for i = (nargs+1:length(code))
+            code[i] = transform_expr(code[i], map_slot_number, map_ssa_value, sparams)
+        end
     end
-end
 
-function overdub_generator(self, c, f, args)
-    mis = method_instances(f.instance, args)
-    @assert length(mis) == 1
-    mi = mis[1]
-    # Unsupported in this mini-cassette
-    @assert !mi.def.isva
-    code_info = retrieve_code_info(mi)
-    @assert isa(code_info, CodeInfo)
-    code_info = copy(code_info)
-    if isdefined(code_info, :edges)
-        code_info.edges = MethodInstance[mi]
+    function overdub_generator(self, c, f, args)
+        if f <: Core.Builtin || !isdefined(f, :instance)
+            return :(return f(args...))
+        end
+
+        tt = Tuple{f, args...}
+        mthds = _methods_by_ftype(tt, -1, typemax(UInt))
+        @assert length(mthds) == 1
+        mtypes, msp, m = mthds[1]
+        mi = ccall(:jl_specializations_get_linfo, Ref{MethodInstance}, (Any, Any, Any), m, mtypes, msp)
+        # Unsupported in this mini-cassette
+        @assert !mi.def.isva
+        code_info = retrieve_code_info(mi)
+        @assert isa(code_info, CodeInfo)
+        code_info = copy(code_info)
+        if isdefined(code_info, :edges)
+            code_info.edges = MethodInstance[mi]
+        end
+        transform!(code_info, length(args), msp)
+        code_info
     end
-    transform!(code_info, length(args))
-    code_info
-end
 
-@eval function overdub(c::Ctx, f, args...)
-    $(Expr(:meta, :generated_only))
-    $(Expr(:meta,
-            :generated,
-            Expr(:new,
-                Core.GeneratedFunctionStub,
-                :overdub_generator,
-                Any[:overdub, :ctx, :f, :args],
-                Any[],
-                @__LINE__,
-                QuoteNode(Symbol(@__FILE__)),
-                true)))
+    @eval function overdub(c::Ctx, f, args...)
+        $(Expr(:meta, :generated_only))
+        $(Expr(:meta,
+                :generated,
+                Expr(:new,
+                    Core.GeneratedFunctionStub,
+                    :overdub_generator,
+                    Any[:overdub, :ctx, :f, :args],
+                    Any[],
+                    @__LINE__,
+                    QuoteNode(Symbol(@__FILE__)),
+                    true)))
+    end
 end
 
+using .MiniCassette
+
+# Test #265 for Cassette
 f() = 1
 @test overdub(Ctx(), f) === 1
 f() = 2
 @test overdub(Ctx(), f) === 2
+
+# Test that MiniCassette is at least somewhat capable by overdubbing gcd
+@test overdub(Ctx(), gcd, 10, 20) === gcd(10, 20)