pipe: block function

Author: MikeInnes

src/ir/ir.jl

@@ -845,6 +845,7 @@ mutable struct Pipe
   map::Dict{Any,Any}
   var::Int
   branch
+  block
 end
 
 var!(p::Pipe) = NewVariable(p.var += 1)
@@ -857,7 +858,7 @@ substitute(p::Pipe, x::Expr) = Expr(x.head, substitute.((p,), x.args)...)
 substitute(p::Pipe) = x -> substitute(p, x)
 
 function Pipe(ir)
-  p = Pipe(ir, IR(copy(ir.lines), meta = ir.meta), Dict(), 0, identity)
+  p = Pipe(ir, IR(copy(ir.lines), meta = ir.meta), Dict(), 0, identity, identity)
   for (x, T) in zip(p.from.blocks[1].args, p.from.blocks[1].argtypes)
     y = argument!(blocks(p.to)[end], nothing, T, insert = false)
     substitute!(p, x, y)
@@ -871,8 +872,10 @@ function pipestate(ir::IR)
 end
 
 branches(f, p::Pipe) = (p.branch = f)
+blocks(f, p::Pipe) = (p.block = f)
 
 function iterate(p::Pipe, (ks, b, i) = (pipestate(p.from), 1, 1))
+  i == 1 && b == 1 && p.block(b)
   if i == 1 && b != 1
     for (x, T) in zip(p.from.blocks[b].args, p.from.blocks[b].argtypes)
       y = argument!(blocks(p.to)[end], nothing, T, insert = false)
@@ -885,6 +888,7 @@ function iterate(p::Pipe, (ks, b, i) = (pipestate(p.from), 1, 1))
       br′ == nothing || push!(p.to.blocks[end].branches, map(substitute(p), br′))
     end
     b == length(ks) && return
+    p.block(b)
     block!(p.to)
     return iterate(p, (ks, b+1, 1))
   end