Support GotoIfNot & ReturnNode (#38)

xref JuliaLang/julia#36318

Author: timholy

Project.toml

@@ -1,13 +1,13 @@
 name = "LoweredCodeUtils"
 uuid = "6f1432cf-f94c-5a45-995e-cdbf5db27b0b"
 authors = ["Tim Holy <[email protected]>"]
-version = "1.1.2"
+version = "1.1.3"
 
 [deps]
 JuliaInterpreter = "aa1ae85d-cabe-5617-a682-6adf51b2e16a"
 
 [compat]
-JuliaInterpreter = "0.7.16"
+JuliaInterpreter = "0.7.23"
 julia = "1"
 
 [extras]

src/LoweredCodeUtils.jl

@@ -12,7 +12,8 @@ module LoweredCodeUtils
 using JuliaInterpreter
 using JuliaInterpreter: SSAValue, SlotNumber, Frame
 using JuliaInterpreter: @lookup, moduleof, pc_expr, step_expr!, is_global_ref, is_quotenode, whichtt,
-                        next_until!, finish_and_return!, get_return, nstatements, codelocation, linetable
+                        next_until!, finish_and_return!, get_return, nstatements, codelocation, linetable,
+                        is_return, lookup_return, is_GotoIfNot, is_ReturnNode
 
 include("packagedef.jl")
 

src/codeedges.jl

@@ -304,6 +304,10 @@ function add_links!(target::Pair{Union{SSAValue,SlotNumber,NamedVar},Links}, @no
         end
     elseif isa(stmt, QuoteNode)
         add_links!(target, stmt.value, cl)
+    elseif is_GotoIfNot(stmt)
+        add_links!(target, (stmt::Core.GotoIfNot).cond, cl)
+    elseif is_ReturnNode(stmt)
+        add_links!(target, (stmt::Core.ReturnNode).val, cl)
     end
     return nothing
 end
@@ -787,7 +791,7 @@ function selective_eval!(@nospecialize(recurse), frame::Frame, isrequired::Abstr
     pcexec = pcexec === nothing ? pclast : pcexec
     frame.pc = pcexec
     node = pc_expr(frame)
-    isexpr(node, :return) && return @lookup(frame, (node::Expr).args[1])
+    is_return(node) && return lookup_return(frame, node)
     isassigned(frame.framedata.ssavalues, pcexec) && return frame.framedata.ssavalues[pcexec]
     return nothing
 end

src/signatures.jl

@@ -142,13 +142,11 @@ function identify_framemethod_calls(frame)
             tsrc = stmt.args[1]::CodeInfo
             if length(tsrc.code) == 1
                 tstmt = tsrc.code[1]
-                if isa(tstmt, Expr)
-                    if tstmt.head === :return && length(tstmt.args) == 1
-                        tex = tstmt.args[1]
-                        if isa(tex, Expr)
-                            if tex.head === :method && (methname = tex.args[1]; isa(methname, Symbol))
-                                push!(refs, methname=>i)
-                            end
+                if is_return(tstmt)
+                    tex = isa(tstmt, Expr) ? tstmt.args[1] : tstmt.val
+                    if isa(tex, Expr)
+                        if tex.head === :method && (methname = tex.args[1]; isa(methname, Symbol))
+                            push!(refs, methname=>i)
                         end
                     end
                 end

src/utils.jl

@@ -122,6 +122,7 @@ function typedef_range(src::CodeInfo, idx)
         if isa(stmt, Expr)
             (stmt.head === :global || stmt.head === :return) && break
         end
+        is_return(stmt) && break
         iend += 1
     end
     iend <= n || (@show src; error("no final :global found"))