IRTools: how can I insert a shortcycling `@ifsomething` into IR?

General Usage
,
1

Hi there,
I am experimenting with IRTools and IR.
So far I understood that branches, including return statements, always have to be last in a block.

Hence if I want to create a shortcycling statement like realised by the @ifsomething macro

macro ifsomething(ex)
    quote
        result = $(esc(ex))
        result === nothing && return nothing
        result
    end
end

I somehow need to split the original IR block into two

  1. the result === nothing && return nothing needs to be added to the original block.
  2. everything else in the original block needs to move to a new block, passing all necessary parameters.

Looking into documentation, examples, tests and source code of IRTools I still have no clue how to do this as of now.
Any help or hint where else I might get help is highly appreciated!

2

after getting into the internals of the IR representation, I was able to create a solution.

function shortcycle_if_nothing!(ir, var)
  oldblock, i_var_oldblock = IRTools.Inner.blockidx(ir, var)
  newblock = IRTools.block!(ir, oldblock.id+1)  # insert right after given block

  # move next statements to new block
  oldblock_n_statements = length(IRTools.BasicBlock(oldblock).stmts)
  for j in (i_var_oldblock+1):oldblock_n_statements
      push!(newblock, oldblock[j])
      oldblock[j] = nothing
  end
  # move variable ids to new block
  for (i, (i_block, i_var_block)) in enumerate(ir.defs)
      if i_block == oldblock.id && i_var_block > i_var_oldblock
          ir.defs[i] = (newblock.id, i_var_block - i_var_oldblock)
      # we need to delete the original references which have been created for newblock
      elseif i_block == newblock.id
          ir.defs[i] = (-1, -1)
      end
  end

  # move branches to new block
  append!(IRTools.branches(newblock), IRTools.branches(oldblock))
  empty!(IRTools.branches(oldblock))

  # add shortcycling in the original branch if we found `nothing`
  var_isnothing = IRTools.insertafter!(ir, var, IRTools.xcall(===, var, nothing))
  IRTools.branch!(oldblock, newblock, unless = var_isnothing)
  IRTools.return!(oldblock, GlobalRef(Main, :nothing))
  return ir
end
1 Like