Merge pull request #21 from MasonProtter/fix-19

Fix #19

Author: MasonProtter

Project.toml

@@ -1,7 +1,7 @@
 name = "Mjolnir"
 uuid = "1154507a-7ac2-44fe-9f15-34617bca9db5"
 authors = ["Mike J Innes <[email protected]>"]
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 IRTools = "7869d1d1-7146-5819-86e3-90919afe41df"

src/infer.jl

@@ -86,7 +86,6 @@ mutable struct Frame
   edges::Vector{Any}
   stmts::Vector{Vector{Variable}}
   rettype::AType
-  outer::Bool
 end
 
 getblock(fr::Frame, b, follow) =
@@ -96,7 +95,7 @@ getblock(fr::Frame, b, follow) =
 # TODO clear up inlined edges
 uninline!(fr::Frame, b, follow) = deleteat!(fr.inlined[b], follow+1:length(fr.inlined[b]))
 
-Frame(ir::IR; outer = false) = Frame(ir, [IR[] for _ = 1:length(blocks(ir))], [], keys.(blocks(ir)), Union{}, outer)
+Frame(ir::IR) = Frame(ir, [IR[] for _ = 1:length(blocks(ir))], [], keys.(blocks(ir)), Union{})
 
 function frame(ir::IR, args...)
   prepare_ir!(ir)
@@ -201,7 +200,7 @@ function step!(inf::Inference)
       error("Unrecognised expression $(st.expr)")
     end
   elseif (brs = openbranches(block); length(brs) == 1 && !isreturn(brs[1])
-          && !(frame.outer && brs[1].block == length(frame.ir.blocks)))
+          && !(brs[1].block == length(frame.ir.blocks)))
     inferbranch!(inf, frame, b, f, brs[1])
   else
     for br in brs
@@ -250,19 +249,3 @@ function return_type(ir::IR, args...)
   infer!(inf)
   return fr.rettype
 end
-
-function inlineable(ir::IR, inf::Inference)
-  for (v, st) in ir
-    Ts = exprtype.((ir,), st.expr.args)
-    haskey(inf.frames, Ts) && return (v, Ts)
-  end
-end
-
-function inlineall(ir::IR, inf::Inference)
-  while (next = inlineable(ir, inf)) != nothing
-    v, Ts = next
-    subir = inf.frames[Ts].ir
-    ir = inline(ir, v, subir)
-  end
-  return ir
-end

src/trace.jl

@@ -133,11 +133,7 @@ function abstract!(tr, bl, env)
   for i = 1:length(args)
     argtypes(ir)[i] = exprtype(tr.ir, rename(env, args[i]))
   end
-  inf = Inference(Frame(ir, outer = true), tr.primitives)
-  infer!(inf)
-  # TODO weird that we need to expand here; perhaps inlining is impicitly pruning
-  # somehow.
-  ir = inlineall(ir, inf) |> expand!
+  infer!(Inference(Frame(ir), tr.primitives))
   inline!(tr.ir, ir, rename.((env,), args))
   for (k, v) in zip(arguments(block(bl.ir, after)), arguments(blocks(tr.ir)[end]))
     env[k] = v

test/trace.jl

@@ -164,3 +164,15 @@ tr = @trace f(::Matrix{Int32})
 f(xs) = sum(xs, dims = 1)
 tr = @trace f(::Matrix{Int32})
 @test returntype(tr) == Matrix{Int32}
+
+
+function negsquare(x)
+    if x > 0
+        return x^2
+    else
+        return -x^2
+    end
+end
+
+tr = @trace negsquare(::Float64)
+@test returntype(tr) == Float64