Generating Bytecode from an AST

generating_bytecode.rb

def cancel
  authorize OrderIssue

  order_issue = OrderIssue.find(params[:id])
  order_issue.cancel!
  render_results(:results    => order_issue,
                 :status     => :ok,
                 :serializer => OrderIssueSerializer)
end
code = File.read(__FILE__).lines[0...__LINE__-1].join

# Just fkn around after showing someone at work how to see Ruby's bytecode

require 'parser/ruby23'

def self.get_locals(ast, locals)
  return locals unless ast
  get_locals_for_all = lambda do |asts|
    asts.each { |child| get_locals child, locals }
    locals
  end

  case ast.type
  when :begin
    get_locals_for_all.(ast.children)
  when :sym, :lvar, :int
    locals
  when :lvasgn
    name, value = ast.children
    locals << name
    get_locals value, locals
  when :send
    receiver, _name, *args = ast.children
    get_locals receiver, locals
    get_locals_for_all.(args)
  when :hash
    get_locals_for_all.call ast.children.flat_map(&:children)
  when :const
    namespace, _name = ast.children
    get_locals namespace, locals
  else
    raise "No local getting code for #{ast.inspect}"
  end
end

def self.lineno_generator
  instruction_no = 1
  lambda do |offset|
    begin
      sprintf "%04d", instruction_no
    ensure
      instruction_no += offset
    end
  end
end

class Locals
  def initialize(offset)
    @offset = offset
    @names_to_offsets = {}
  end
  def <<(name)
    @names_to_offsets[name] = @offset
  ensure
    @offset += 1
    return self
  end
  def [](name)
    @names_to_offsets.fetch name
  end
  def length
    @names_to_offsets.length
  end
  def to_declarations
    @names_to_offsets.keys.reverse.each.with_index(1).map { |name, index| "[ #{index}] #{name}" }.reverse
  end
end

def self.generate_bytecode(ast, locals=Locals.new(3), leave_on_stack=false)
  case ast.type
  when :def
    name, args_ast, body_ast = ast.children
    args = args_ast.children
    raise "Expected no args" unless args.empty?
    get_locals body_ast, locals
    "== disassemble the instruction sequence for #{name}\n" +
      "local table (size: #{locals.length}, argc: #{args.length} [opts: 0, rest: -1, post: 0, block: -1, kw: -1@-1, kwrest: -1])\n" +
      locals.to_declarations.join('        ') + "\n" +
      generate_bytecode(body_ast, locals, true) +
      "leave\n"
  when :begin
    *children, last_child = ast.children
    children.map { |child| generate_bytecode child, locals }.join + 
      generate_bytecode(last_child, locals, leave_on_stack)
  when :send
    receiver_ast, name, *arg_asts = ast.children
    if receiver_ast
      receiver_bytecode = generate_bytecode(receiver_ast, locals, true)
    else
      receiver_bytecode = "putself\n"
    end
    if arg_asts.last&.type == :hash && arg_asts.last.children.map(&:children).transpose.first.all? { |key| key.type == :sym }
      key_asts, values = arg_asts.pop.children.map(&:children).transpose
      keyword_names    = key_asts.map { |key| key.children.first }
      keyword_bytecode = values.map { |v| generate_bytecode v, locals, true }.join
    else
      keyword_names = []
      keyword_bytecode = ''
    end
    args_bytecode = arg_asts.map { |arg| generate_bytecode arg, locals, true }.join + keyword_bytecode
    kw_callinfo   = " kw:[#{keyword_names.join ?,}]" if keyword_names.any?
    callinfo      = "<call!#{name}, argc:#{arg_asts.length + keyword_names.length}#{kw_callinfo}>"
    if name == :[]
      call_bytecode = "opt_aref #{callinfo}\n"
    else
      call_bytecode = "opt_send_without_block #{callinfo}\n"
    end
    if leave_on_stack
      cleanup = ''
    else
      cleanup = "pop\n"
    end
    receiver_bytecode + args_bytecode + call_bytecode + cleanup
  when :lvar
    name = ast.children[0]
    "getlocal_OP__WC__0 #{locals[name]}\n"
  when :lvasgn
    name, value = ast.children
    generate_bytecode(value, locals, true) + "setlocal_OP__WC__0 #{locals[name]}\n"
  when :const
    namespace, name = ast.children
    raise "YO GET THE NAMESPACE!" if namespace
    "getconstant #{name.inspect}\n"
  when :hash
    "FIXME :hash\n"
  when :sym
    "putobject #{ast.children.first.inspect}\n"
  else
    raise "Unknkown ast: #{ast.inspect}"
  end
end

ast = Parser::Ruby23.parse code
generate_bytecode ast
# => "== disassemble the instruction sequence for cancel\n" +
#    "local table (size: 1, argc: 0 [opts: 0, rest: -1, post: 0, block: -1, kw: -1@-1, kwrest: -1])\n" +
#    "[ 1] order_issue\n" +
#    "putself\n" +
#    "getconstant :OrderIssue\n" +
#    "opt_send_without_block <call!authorize, argc:1>\n" +
#    "pop\n" +
#    "getconstant :OrderIssue\n" +
#    "putself\n" +
#    "opt_send_without_block <call!params, argc:0>\n" +
#    "putobject :id\n" +
#    "opt_aref <call![], argc:1>\n" +
#    "opt_send_without_block <call!find, argc:1>\n" +
#    "setlocal_OP__WC__0 3\n" +
#    "getlocal_OP__WC__0 3\n" +
#    "opt_send_without_block <call!cancel!, argc:0>\n" +
#    "pop\n" +
#    "putself\n" +
#    "getlocal_OP__WC__0 3\n" +
#    "putobject :ok\n" +
#    "getconstant :OrderIssueSerializer\n" +
#    "opt_send_without_block <call!render_results, argc:3 kw:[results,status,serializer]>\n" +
#    "leave\n"


RubyVM::InstructionSequence.disasm(method :cancel).lines.grep_v(/^\d+\s+(trace|[sg]etinlinecache)\b/).join
# => "== disasm: #<ISeq:cancel@/var/folders/yj/23xw_th91pl6xkdnyc733r01t48gqj/T/seeing_is_believing_temp_dir20170721-81675-krkwmd/program.rb>\n" +
#    "local table (size: 1, argc: 0 [opts: 0, rest: -1, post: 0, block: -1, kw: -1@-1, kwrest: -1])\n" +
#    "[ 1] order_issue\n" +
#    "0004 putself          \n" +
#    "0008 getconstant      :OrderIssue\n" +
#    "0012 opt_send_without_block <callinfo!mid:authorize, argc:1, FCALL|ARGS_SIMPLE>, <callcache>\n" +
#    "0015 pop              \n" +
#    "0021 getconstant      :OrderIssue\n" +
#    "0025 putself          \n" +
#    "0026 opt_send_without_block <callinfo!mid:params, argc:0, FCALL|VCALL|ARGS_SIMPLE>, <callcache>\n" +
#    "0029 putobject        :id\n" +
#    "0031 opt_aref         <callinfo!mid:[], argc:1, ARGS_SIMPLE>, <callcache>\n" +
#    "0034 opt_send_without_block <callinfo!mid:find, argc:1, ARGS_SIMPLE>, <callcache>\n" +
#    "0037 setlocal_OP__WC__0 3\n" +
#    "0041 getlocal_OP__WC__0 3\n" +
#    "0043 opt_send_without_block <callinfo!mid:cancel!, argc:0, ARGS_SIMPLE>, <callcache>\n" +
#    "0046 pop              \n" +
#    "0049 putself          \n" +
#    "0050 getlocal_OP__WC__0 3\n" +
#    "0052 putobject        :ok                                             (   7)\n" +
#    "0057 getconstant      :OrderIssueSerializer\n" +
#    "0061 opt_send_without_block <callinfo!mid:render_results, argc:3, kw:[results,status,serializer], FCALL|KWARG>, <callcache>(   6)\n" +
#    "0066 leave                                                            (   6)\n"