timlawrenz/gnn-ruby-code-study

1#!/usr/bin/env ruby2# frozen_string_literal: true3 4require 'json'5require 'parser/current'6 7# Pretty-printing script to convert AST JSON back to formatted Ruby code8#9# This script takes the JSON representation of an AST (as used in the project's dataset)10# and converts it back into formatted Ruby code that is syntactically equivalent11# to the original source.12 13class ASTPrettyPrinter14  # Main method to convert AST JSON to pretty-printed Ruby code15  #16  # @param ast_json [String, Hash] JSON string or parsed hash representing the AST17  # @return [String] Pretty-printed Ruby code18  def self.ast_to_ruby(ast_json)19    ast_hash = ast_json.is_a?(String) ? JSON.parse(ast_json) : ast_json20    new.convert_node(ast_hash)21  end22 23  # Convert a single AST node to Ruby code24  #25  # @param node [Hash, String, nil] AST node to convert26  # @param indent_level [Integer] Current indentation level27  # @return [String] Ruby code representation28  def convert_node(node, indent_level = 0)29    return '' if node.nil?30    return node.to_s if node.is_a?(String) || node.is_a?(Numeric)31 32    return '' unless node.is_a?(Hash) && node['type']33 34    case node['type']35    when 'def'36      convert_def(node, indent_level)37    when 'args'38      convert_args(node)39    when 'begin'40      convert_begin(node, indent_level)41    when 'send'42      convert_send(node)43    when 'block'44      convert_block(node, indent_level)45    when 'self'46      'self'47    when 'nil'48      'nil'49    when 'true'50      'true'51    when 'false'52      'false'53    when 'str'54      convert_string(node)55    when 'int'56      convert_int(node)57    when 'float'58      convert_float(node)59    when 'sym'60      convert_symbol(node)61    when 'lvar'62      convert_lvar(node)63    when 'ivar'64      convert_ivar(node)65    when 'cvar'66      convert_cvar(node)67    when 'gvar'68      convert_gvar(node)69    when 'const'70      convert_const(node)71    when 'if'72      convert_if(node, indent_level)73    when 'unless'74      convert_unless(node, indent_level)75    when 'while'76      convert_while(node, indent_level)77    when 'until'78      convert_until(node, indent_level)79    when 'for'80      convert_for(node, indent_level)81    when 'case'82      convert_case(node, indent_level)83    when 'when'84      convert_when(node, indent_level)85    when 'return'86      convert_return(node)87    when 'break'88      convert_break(node)89    when 'next'90      convert_next(node)91    when 'yield'92      convert_yield(node)93    when 'and'94      convert_and(node)95    when 'or'96      convert_or(node)97    when 'not'98      convert_not(node)99    when 'array'100      convert_array(node)101    when 'hash'102      convert_hash(node)103    when 'pair'104      convert_pair(node)105    when 'splat'106      convert_splat(node)107    when 'kwsplat'108      convert_kwsplat(node)109    when 'class'110      convert_class(node, indent_level)111    when 'module'112      convert_module(node, indent_level)113    when 'sclass'114      convert_sclass(node, indent_level)115    when 'lvasgn'116      convert_lvasgn(node)117    when 'ivasgn'118      convert_ivasgn(node)119    when 'cvasgn'120      convert_cvasgn(node)121    when 'gvasgn'122      convert_gvasgn(node)123    when 'casgn'124      convert_casgn(node)125    when 'dstr'126      convert_dstr(node)127    when 'dnum'128      convert_dnum(node)129    when 'regexp'130      convert_regexp(node)131    when 'resbody'132      convert_resbody(node, indent_level)133    when 'ensure'134      convert_ensure(node, indent_level)135    when 'retry'136      'retry'137    when 'kwarg'138      convert_kwarg(node)139    when 'kwoptarg'140      convert_kwoptarg(node)141    when 'kwrestarg'142      convert_kwrestarg(node)143    when 'blockarg'144      convert_blockarg(node)145    when 'optarg'146      convert_optarg(node)147    when 'restarg'148      convert_restarg(node)149    when 'arg'150      convert_arg(node)151    else152      # Fallback for unknown node types153      convert_unknown(node)154    end155  end156 157  private158 159  def indent(level)160    '  ' * level161  end162 163  def convert_def(node, indent_level)164    children = node['children'] || []165    method_name = children[0]166    args_node = children[1]167    body_node = children[2]168 169    result = "#{indent(indent_level)}def #{method_name}"170    171    if args_node && args_node['children'] && !args_node['children'].empty?172      result += "(#{convert_node(args_node)})"173    end174    175    result += "\n"176    177    if body_node178      body_code = convert_node(body_node, indent_level + 1)179      result += body_code unless body_code.strip.empty?180    end181    182    result += "\n#{indent(indent_level)}end"183    result184  end185 186  def convert_args(node)187    children = node['children'] || []188    children.map { |child| convert_node(child) }.join(', ')189  end190 191  def convert_begin(node, indent_level)192    children = node['children'] || []193    children.map { |child| 194      code = convert_node(child, indent_level)195      # Add indentation if the code doesn't already have it196      if code.start_with?(indent(indent_level))197        code198      else199        "#{indent(indent_level)}#{code}"200      end201    }.join("\n")202  end203 204  def convert_send(node)205    children = node['children'] || []206    receiver = children[0]207    method_name = children[1]208    args = children[2..-1]209 210    result = ''211    212    if receiver213      result += convert_node(receiver)214      215      # Special case for array access216      if method_name.to_s == '[]'217        result += "[#{args.map { |arg| convert_node(arg) }.join(', ')}]"218        return result219      end220      221      result += '.'222    end223    224    if method_name.to_s.end_with?('=') && args && !args.empty?225      # For assignment methods, handle specially226      base_name = method_name.to_s.chomp('=')227      result += "#{base_name} = #{args.map { |arg| convert_node(arg) }.join(', ')}"228    else229      # Regular method calls230      result += method_name.to_s231      232      if args && !args.empty?233        # For common methods like raise, don't use parentheses234        if should_omit_parentheses?(method_name.to_s)235          result += " #{args.map { |arg| convert_node(arg) }.join(', ')}"236        else237          result += "(#{args.map { |arg| convert_node(arg) }.join(', ')})"238        end239      end240    end241    242    result243  end244 245  def should_omit_parentheses?(method_name)246    # Common methods that look better without parentheses247    common_methods = %w[raise puts print p require require_relative include extend]248    # Operators that should not have parentheses249    operators = %w[+ - * / % ** == != < > <= >= <=> === =~ !~ << >> & | ^ ~ && ||]250    251    common_methods.include?(method_name) || operators.include?(method_name)252  end253 254  def convert_block(node, indent_level)255    children = node['children'] || []256    call_node = children[0]257    block_args = children[1]258    block_body = children[2]259 260    result = convert_node(call_node)261    262    # Determine if we should use do...end (for multi-line) or { } (for single-line)263    use_do_end = should_use_do_end_block?(block_body)264    265    if use_do_end266      # Multi-line do...end block267      if block_args && block_args['children'] && !block_args['children'].empty?268        result += " do |#{convert_node(block_args)}|\n"269      else270        result += " do\n"271      end272 273      if block_body274        body_code = convert_node(block_body, indent_level + 1)275        result += body_code276        result += "\n"277      end278 279      result += "#{indent(indent_level)}end"280    else281      # Single-line { } block282      if block_args && block_args['children'] && !block_args['children'].empty?283        result += " { |#{convert_node(block_args)}|"284      else285        result += " {"286      end287 288      if block_body289        body_code = convert_node(block_body)290        result += " #{body_code} "291      end292 293      result += "}"294    end295    296    result297  end298 299  def should_use_do_end_block?(block_body)300    return false unless block_body301    302    # Use do...end for multi-statement blocks303    if block_body['type'] == 'begin'304      children = block_body['children'] || []305      return children.length > 1306    end307    308    # Use do...end for complex single statements (control flow, etc.)309    case block_body['type']310    when 'if', 'unless', 'case', 'while', 'until', 'for', 'def', 'class', 'module'311      true312    else313      false314    end315  end316 317  def convert_string(node)318    children = node['children'] || []319    "\"#{children[0]}\""320  end321 322  def convert_int(node)323    children = node['children'] || []324    children[0].to_s325  end326 327  def convert_float(node)328    children = node['children'] || []329    children[0].to_s330  end331 332  def convert_symbol(node)333    children = node['children'] || []334    ":#{children[0]}"335  end336 337  def convert_lvar(node)338    children = node['children'] || []339    children[0].to_s340  end341 342  def convert_ivar(node)343    children = node['children'] || []344    "@#{children[0]}"345  end346 347  def convert_cvar(node)348    children = node['children'] || []349    "@@#{children[0]}"350  end351 352  def convert_gvar(node)353    children = node['children'] || []354    "$#{children[0]}"355  end356 357  def convert_const(node)358    children = node['children'] || []359    scope = children[0]360    const_name = children[1]361    362    result = ""363    if scope364      result += "#{convert_node(scope)}::"365    end366    result += const_name.to_s367    result368  end369 370  def convert_lvasgn(node)371    children = node['children'] || []372    var_name = children[0]373    value = children[1]374    375    if value376      "#{var_name} = #{convert_node(value)}"377    else378      var_name.to_s379    end380  end381 382  def convert_ivasgn(node)383    children = node['children'] || []384    var_name = children[0]385    value = children[1]386    387    if value388      "@#{var_name} = #{convert_node(value)}"389    else390      "@#{var_name}"391    end392  end393 394  def convert_cvasgn(node)395    children = node['children'] || []396    var_name = children[0]397    value = children[1]398    399    if value400      "@@#{var_name} = #{convert_node(value)}"401    else402      "@@#{var_name}"403    end404  end405 406  def convert_gvasgn(node)407    children = node['children'] || []408    var_name = children[0]409    value = children[1]410    411    if value412      "$#{var_name} = #{convert_node(value)}"413    else414      "$#{var_name}"415    end416  end417 418  def convert_casgn(node)419    children = node['children'] || []420    scope = children[0]421    const_name = children[1]422    value = children[2]423    424    result = ""425    if scope426      result += "#{convert_node(scope)}::"427    end428    result += const_name.to_s429    430    if value431      result += " = #{convert_node(value)}"432    end433    434    result435  end436 437  def convert_dstr(node)438    children = node['children'] || []439    result = "\""440    441    children.each do |child|442      if child.is_a?(String)443        result += child444      else445        case child['type']446        when 'begin'447          # String interpolation448          result += "\#{#{convert_node(child)}}"449        else450          result += convert_node(child)451        end452      end453    end454    455    result += "\""456    result457  end458 459  def convert_dnum(node)460    children = node['children'] || []461    children[0].to_s462  end463 464  def convert_regexp(node)465    children = node['children'] || []466    pattern = children[0]467    flags = children[1]468    469    result = "/#{convert_node(pattern)}/"470    if flags && flags != 0471      result += convert_regexp_flags(flags)472    end473    result474  end475 476  def convert_regexp_flags(flags)477    flag_str = ""478    flag_str += "i" if flags & 1 != 0   # IGNORECASE479    flag_str += "x" if flags & 2 != 0   # EXTENDED480    flag_str += "m" if flags & 4 != 0   # MULTILINE481    flag_str482  end483 484  def convert_if(node, indent_level)485    children = node['children'] || []486    condition = children[0]487    then_body = children[1]488    else_body = children[2]489 490    result = "#{indent(indent_level)}if #{convert_node(condition)}\n"491    492    if then_body493      result += convert_node(then_body, indent_level + 1)494      result += "\n"495    end496    497    if else_body498      result += "#{indent(indent_level)}else\n"499      result += convert_node(else_body, indent_level + 1)500      result += "\n"501    end502    503    result += "#{indent(indent_level)}end"504    result505  end506 507  def convert_unless(node, indent_level)508    children = node['children'] || []509    condition = children[0]510    then_body = children[1]511    else_body = children[2]512 513    result = "#{indent(indent_level)}unless #{convert_node(condition)}\n"514    515    if then_body516      result += convert_node(then_body, indent_level + 1)517      result += "\n"518    end519    520    if else_body521      result += "#{indent(indent_level)}else\n"522      result += convert_node(else_body, indent_level + 1)523      result += "\n"524    end525    526    result += "#{indent(indent_level)}end"527    result528  end529 530  def convert_return(node)531    children = node['children'] || []532    if children.empty?533      'return'534    else535      "return #{children.map { |child| convert_node(child) }.join(', ')}"536    end537  end538 539  def convert_break(node)540    children = node['children'] || []541    if children.empty?542      'break'543    else544      "break #{children.map { |child| convert_node(child) }.join(', ')}"545    end546  end547 548  def convert_next(node)549    children = node['children'] || []550    if children.empty?551      'next'552    else553      "next #{children.map { |child| convert_node(child) }.join(', ')}"554    end555  end556 557  def convert_yield(node)558    children = node['children'] || []559    if children.empty?560      'yield'561    else562      "yield(#{children.map { |child| convert_node(child) }.join(', ')})"563    end564  end565 566  def convert_and(node)567    children = node['children'] || []568    left = convert_node_with_precedence(children[0], 'and')569    right = convert_node_with_precedence(children[1], 'and')570    "#{left} && #{right}"571  end572 573  def convert_or(node)574    children = node['children'] || []575    left = convert_node_with_precedence(children[0], 'or')576    right = convert_node_with_precedence(children[1], 'or')577    "#{left} || #{right}"578  end579 580  def convert_node_with_precedence(node, parent_op)581    return convert_node(node) unless node.is_a?(Hash) && node['type']582    583    # Define operator precedence (higher number = higher precedence)584    precedence = {585      'or' => 1,   # ||586      'and' => 2   # &&587    }588    589    node_op = node['type']590    591    # Add parentheses if the child operator has lower precedence than parent592    if precedence[node_op] && precedence[parent_op] && 593       precedence[node_op] < precedence[parent_op]594      "(#{convert_node(node)})"595    else596      convert_node(node)597    end598  end599 600  def convert_not(node)601    children = node['children'] || []602    "!#{convert_node(children[0])}"603  end604 605  def convert_array(node)606    children = node['children'] || []607    "[#{children.map { |child| convert_node(child) }.join(', ')}]"608  end609 610  def convert_hash(node)611    children = node['children'] || []612    "{ #{children.map { |child| convert_node(child) }.join(', ')} }"613  end614 615  def convert_pair(node)616    children = node['children'] || []617    key = children[0]618    value = children[1]619    620    # Use modern syntax for symbol keys (key: value instead of :key => value)621    if key && key['type'] == 'sym'622      key_name = key['children'][0]623      "#{key_name}: #{convert_node(value)}"624    else625      "#{convert_node(key)} => #{convert_node(value)}"626    end627  end628 629  def convert_splat(node)630    children = node['children'] || []631    "*#{convert_node(children[0])}"632  end633 634  def convert_kwsplat(node)635    children = node['children'] || []636    "**#{convert_node(children[0])}"637  end638 639  def convert_arg(node)640    children = node['children'] || []641    children[0].to_s642  end643 644  def convert_optarg(node)645    children = node['children'] || []646    name = children[0]647    default = children[1]648    "#{name} = #{convert_node(default)}"649  end650 651  def convert_restarg(node)652    children = node['children'] || []653    if children.empty?654      '*'655    else656      "*#{children[0]}"657    end658  end659 660  def convert_kwarg(node)661    children = node['children'] || []662    "#{children[0]}:"663  end664 665  def convert_kwoptarg(node)666    children = node['children'] || []667    name = children[0]668    default = children[1]669    "#{name}: #{convert_node(default)}"670  end671 672  def convert_kwrestarg(node)673    children = node['children'] || []674    if children.empty?675      '**'676    else677      "**#{children[0]}"678    end679  end680 681  def convert_blockarg(node)682    children = node['children'] || []683    "&#{children[0]}"684  end685 686  def convert_class(node, indent_level)687    children = node['children'] || []688    name = children[0]689    superclass = children[1]690    body = children[2]691 692    result = "#{indent(indent_level)}class #{convert_node(name)}"693    694    if superclass695      result += " < #{convert_node(superclass)}"696    end697    698    result += "\n"699    700    if body701      result += convert_node(body, indent_level + 1)702      result += "\n"703    end704    705    result += "#{indent(indent_level)}end"706    result707  end708 709  def convert_module(node, indent_level)710    children = node['children'] || []711    name = children[0]712    body = children[1]713 714    result = "#{indent(indent_level)}module #{convert_node(name)}\n"715    716    if body717      result += convert_node(body, indent_level + 1)718      result += "\n"719    end720    721    result += "#{indent(indent_level)}end"722    result723  end724 725  def convert_while(node, indent_level)726    children = node['children'] || []727    condition = children[0]728    body = children[1]729 730    result = "#{indent(indent_level)}while #{convert_node(condition)}\n"731    732    if body733      result += convert_node(body, indent_level + 1)734      result += "\n"735    end736    737    result += "#{indent(indent_level)}end"738    result739  end740 741  def convert_until(node, indent_level)742    children = node['children'] || []743    condition = children[0]744    body = children[1]745 746    result = "#{indent(indent_level)}until #{convert_node(condition)}\n"747    748    if body749      result += convert_node(body, indent_level + 1)750      result += "\n"751    end752    753    result += "#{indent(indent_level)}end"754    result755  end756 757  def convert_case(node, indent_level)758    children = node['children'] || []759    expr = children[0]760    branches = children[1..-1]761 762    result = "#{indent(indent_level)}case #{convert_node(expr)}\n"763    764    # Separate when branches from else clause765    when_branches = []766    else_clause = nil767    768    branches.each do |branch|769      if branch && branch['type'] == 'when'770        when_branches << branch771      elsif branch && branch['type'] != 'when'772        # This is likely the else clause (any non-when node at the end)773        else_clause = branch774      end775    end776    777    # Process when branches778    when_branches.each do |when_node|779      result += convert_node(when_node, indent_level)780      result += "\n"781    end782    783    # Process else clause if present784    if else_clause785      result += "#{indent(indent_level)}else\n"786      result += convert_node(else_clause, indent_level + 1)787      result += "\n"788    end789    790    result += "#{indent(indent_level)}end"791    result792  end793 794  def convert_when(node, indent_level)795    children = node['children'] || []796    conditions = children[0..-2]797    body = children[-1]798 799    result = "#{indent(indent_level)}when #{conditions.map { |cond| convert_node(cond) }.join(', ')}\n"800    801    if body802      result += convert_node(body, indent_level + 1)803    end804    805    result806  end807 808  def convert_for(node, indent_level)809    children = node['children'] || []810    var = children[0]811    collection = children[1]812    body = children[2]813 814    result = "#{indent(indent_level)}for #{convert_node(var)} in #{convert_node(collection)}\n"815    816    if body817      result += convert_node(body, indent_level + 1)818      result += "\n"819    end820    821    result += "#{indent(indent_level)}end"822    result823  end824 825  def convert_sclass(node, indent_level)826    children = node['children'] || []827    expr = children[0]828    body = children[1]829 830    result = "#{indent(indent_level)}class << #{convert_node(expr)}\n"831    832    if body833      result += convert_node(body, indent_level + 1)834      result += "\n"835    end836    837    result += "#{indent(indent_level)}end"838    result839  end840 841  def convert_rescue(node, indent_level)842    children = node['children'] || []843    body = children[0]844    rescue_clauses = children[1..-2]845    else_clause = children[-1]846 847    result = ""848    849    if body850      result += convert_node(body, indent_level)851      result += "\n"852    end853    854    rescue_clauses.each do |rescue_clause|855      if rescue_clause && rescue_clause['type'] == 'resbody'856        result += convert_resbody(rescue_clause, indent_level)857        result += "\n"858      end859    end860    861    if else_clause862      result += "#{indent(indent_level)}else\n"863      result += convert_node(else_clause, indent_level + 1)864      result += "\n"865    end866    867    result868  end869 870  def convert_resbody(node, indent_level)871    children = node['children'] || []872    exception_types = children[0]873    exception_var = children[1]874    body = children[2]875 876    result = "#{indent(indent_level)}rescue"877    878    if exception_types && exception_types['children'] && !exception_types['children'].empty?879      result += " #{exception_types['children'].map { |type| convert_node(type) }.join(', ')}"880    end881    882    if exception_var883      result += " => #{convert_node(exception_var)}"884    end885    886    result += "\n"887    888    if body889      result += convert_node(body, indent_level + 1)890    end891    892    result893  end894 895  def convert_ensure(node, indent_level)896    children = node['children'] || []897    body = children[0]898    ensure_body = children[1]899 900    result = ""901    902    if body903      result += convert_node(body, indent_level)904      result += "\n"905    end906    907    result += "#{indent(indent_level)}ensure\n"908    909    if ensure_body910      result += convert_node(ensure_body, indent_level + 1)911      result += "\n"912    end913    914    result915  end916 917  def convert_unknown(node)918    # Fallback: try to handle unknown nodes gracefully919    children = node['children'] || []920    if children.empty?921      node['type']922    else923      "#{node['type']}(#{children.map { |child| convert_node(child) }.join(', ')})"924    end925  end926end927 928# Command-line interface929if __FILE__ == $0930  begin931    if ARGV.empty? && !$stdin.tty?932      # Read from stdin when no arguments and input is piped933      input = $stdin.read934    elsif ARGV.length == 1935      if ARGV[0] == '-'936        # Explicit stdin937        input = $stdin.read938      else939        # Read from file940        input = File.read(ARGV[0])941      end942    else943      puts "Usage: #{$0} <ast_json_file>"944      puts "   or: echo '<ast_json>' | #{$0}"945      puts "   or: #{$0} -"946      exit 1947    end948 949    result = ASTPrettyPrinter.ast_to_ruby(input)950    puts result951  rescue => e952    $stderr.puts "Error: #{e.message}"953    exit 1954  end955end