Merge branch 'master' into ox/colprac

Author: DhairyaLGandhi

Project.toml

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

README.md

@@ -1,7 +1,7 @@
 # Mjolnir
 
+[![Build Status](https://travis-ci.org/FluxML/Mjolnir.jl.svg?branch=master)](https://travis-ci.org/FluxML/Mjolnir.jl)
 [![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor's%20Guide-blueviolet)](https://github.com/SciML/ColPrac)
-[![Build Status](https://travis-ci.org/MikeInnes/Mjolnir.jl.svg?branch=master)](https://travis-ci.org/MikeInnes/Mjolnir.jl)
 
 Mjolnir is a hybrid approach to partial evaluation / abstract interpretation,
 with an implementation in Julia. It can be thought of as a blend of
@@ -26,7 +26,7 @@ pow (generic function with 1 method)
 
 julia> using Mjolnir
 
-julia> @trace pow(Int, 3)
+julia> @trace pow(::Int, 3)
 1: (%1 :: const(pow), %2 :: Int64, %3 :: const(3))
   %4 = (*)(1, %2) :: Int64
   %5 = (*)(%4, %2) :: Int64
@@ -44,7 +44,7 @@ can generate diagnostics when there are issues. Mjolnir can thus compile a much
 wider range of Julia programs than OO approaches.
 
 ```julia
-julia> @trace pow(Int, Int)
+julia> @trace pow(::Int, ::Int)
 1: (%1 :: const(pow), %2 :: Int64, %3 :: Int64)
   %4 = (>)(%3, 0) :: Bool
   br 3 (1) unless %4

docs/types.md

@@ -47,7 +47,7 @@ functions are called _primitives_.
 julia> f(x) = 3x^2 + 2x + 1
 f (generic function with 1 method)
 
-julia> @trace f(Int)
+julia> @trace f(::Int)
 1: (%1 :: const(f), %2 :: Int64)
   %3 = (*)(%2, %2) :: Int64
   %4 = (*)(3, %3) :: Int64
@@ -144,7 +144,7 @@ julia> @abstract MyPrimitives (a::AType{T} * b::AType{T}) where T<:Union{Float64
 
 julia> MyDefaults() = Multi(MyPrimitives(), Mjolnir.Basic())
 
-julia> @trace MyDefaults() Int * Float64
+julia> @trace MyDefaults() ::Int * ::Float64
 1: (%1 :: const(*), %2 :: Int64, %3 :: Float64)
   %4 = ($(QuoteNode(Float64)))(%2) :: Float64
   %5 = (*)(%4, %3) :: Float64

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
@@ -273,12 +269,17 @@ function trace(P, Ts...)
   end
 end
 
+to_type_level(x) = :(Mjolnir.Const($x))
+to_type_level(ex::Expr) = isexpr(ex, :(::)) && length(ex.args) == 1 ? (ex.args[1]) : :(Mjolnir.Const($ex)) 
+
 atype(T::AType) = T
 atype(x) = Const(x)
+#atype.(($(esc(f)), $(esc.(args)...)))...)
 
 function tracem(P, ex)
   @capture(ex, f_(args__)) || error("@trace f(args...)")
-  :(trace($(esc(P)), atype.(($(esc(f)), $(esc.(args)...)))...))
+  sig = ((esc ∘ to_type_level).((f, args...)))  
+  :(trace($(esc(P)), $(sig...)))
 end
 
 """
@@ -288,7 +289,7 @@ end
 Get a typed trace for `f`, analagous to `@code_typed`. Note that unlike
 `@code_typed`, you probably want to pass types rather than values, e.g.
 
-    julia> @trace Int+Int
+    julia> @trace ::Int + ::Int
     1: (%1 :: const(+), %2 :: Int64, %3 :: Int64)
       %4 = (+)(%2, %3) :: Int64
       return %4

test/trace.jl

@@ -64,24 +64,24 @@ ir = @code_ir fact(1)
 tr = @trace f()
 @test returntype(tr) == String
 
-tr = @trace pow(Int, 3)
+tr = @trace pow(::Int, 3)
 @test length(tr.blocks) == 1
 @test returntype(tr) == Int
 
 tr = @trace pow(2, 3)
 @test length(tr.blocks) == 1
 @test returntype(tr) == Const(8)
 
-tr = @trace pow(2, Int)
+tr = @trace pow(2, ::Int)
 @test returntype(tr) == Int
 
-tr = @trace pow(2.0, Int)
+tr = @trace pow(2.0, ::Int)
 @test returntype(tr) == Union{Float64,Int}
 
-tr = @trace pow(1, Int)
+tr = @trace pow(1, ::Int)
 @test returntype(tr) == Const(1)
 
-tr = @trace bar(Int, Int)
+tr = @trace bar(::Int, ::Int)
 @test returntype(tr) == Int
 
 function foo(x)
@@ -93,7 +93,7 @@ end
 tr = @trace foo(1)
 @test returntype(tr) == Const(1)
 
-tr = @trace foo(Int)
+tr = @trace foo(::Int)
 @test returntype(tr) == Int
 
 function foo(x)
@@ -142,7 +142,7 @@ end
 tr = @trace pow(2, 3)
 @test returntype(tr) == Const(8)
 
-tr = @trace pow(Int, Int)
+tr = @trace pow(::Int, ::Int)
 @test returntype(tr) == Int
 
 function sumabs2(xs)
@@ -153,14 +153,26 @@ function sumabs2(xs)
   return s
 end
 
-tr = @trace sumabs2(arrayshape(Float64, 3))
+tr = @trace sumabs2(::arrayshape(Float64, 3))
 @test length(blocks(tr)) == 1
 @test returntype(tr) == Float64
 
 f(xs) = sum(xs)
-tr = @trace f(Matrix{Int32})
+tr = @trace f(::Matrix{Int32})
 @test returntype(tr) == Int32
 
 f(xs) = sum(xs, dims = 1)
-tr = @trace f(Matrix{Int32})
+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