stark 0.7.0 → 0.8.0

data/lib/stark/ruby.rb

@@ -12,12 +12,10 @@ module Stark
 
       @stream = stream
 
-
+      o "require 'set'"
       o "require 'stark/client'"
-      o "require 'stark/struct'"
-      o "require 'stark/field'"
-      o "require 'stark/converters'"
       o "require 'stark/processor'"
+      o "require 'stark/struct'"
       o "require 'stark/exception'"
     end
 
@@ -27,6 +25,7 @@ module Stark
     end
 
     def close
+      write_protocol
       if @namespace
         outdent
         o "end"
@@ -34,109 +33,74 @@ module Stark
     end
 
     def process_namespace(ns)
-      @namespace = ns.namespace if ns.lang == "rb"
-      o "module #{ns.namespace}"
+      return unless [nil, 'rb'].include?(ns.lang)
+      @namespace = ns.namespace.gsub('.', '::')
+      parts = @namespace.split('::')
+      if parts.length > 1
+        0.upto(parts.length - 2) do |i|
+          o "module #{parts[0..i].join('::')}; end unless defined?(#{parts[0..i].join('::')})"
+        end
+      end
+      o
+      o "module #{@namespace}"
       indent
     end
 
     def process_include(inc)
+      raise NotImplementedError, "include not implemented"
     end
 
     def process_enum(enum)
+      @enums[enum.name] = enum
       e = "Enum_#{enum.name}"
       o "#{e} = Hash.new { |h,k| p [:bad, k]; h[k] = -1 }"
-
       idx = 0
       enum.values.each do |f|
         o "#{e}[#{idx}] = :'#{f}'"
         o "#{e}[:'#{f}'] = #{idx}"
         idx += 1
       end
-
-      @enums[enum.name] = enum
     end
 
-    def converter(t)
-      if t.kind_of? Stark::Parser::AST::List
-        "Stark::Converters::List.new(#{converter(t.value)})"
-      elsif BUILTINS.include? t.downcase
-        "Stark::Converters::#{t.upcase}"
-      elsif desc = @structs[t]
-        "Stark::Converters::Struct.new(#{t})"
-      elsif desc = @enums[t]
-        "Stark::Converters::Enum.new(Enum_#{t})"
-      else
-        raise "Blah"
+    def write_field_declarations(fields)
+      max_field_len = fields.inject(0) {|max,f| f.name.length > max ? f.name.length : max }
+      max_index_len = fields.inject(0) {|max,f| f.index.to_s.length > max ? f.index.to_s.length : max }
+
+      current_index = 1
+      fields.each do |f|
+        if f.index != current_index
+          o "field_number #{f.index}"
+          current_index = f.index
+        end
+        current_index += 1
+        o("attr_accessor :%-*s  # %*s: %s" % [max_field_len, f.name, max_index_len, f.index, object_type(f.type)])
       end
     end
 
     def process_struct(str)
       @structs[str.name] = str
 
+      o
       o "class #{str.name} < Stark::Struct"
       indent
-      o "Fields = {"
-      indent
-
-      fields = str.fields.map do |f|
-        c = converter f.type
-        "#{f.index} => Stark::Field.new(#{f.index}, '#{f.name}', #{c})"
-      end
-
-      o "   #{fields.join(', ')}"
-
-      outdent
-      o "}"
-
-      str.fields.each do |f|
-        o "def #{f.name}; @fields['#{f.name}']; end"
-      end
-
+      write_field_declarations str.fields
       outdent
       o "end"
-
     end
 
-    BUILTINS = %w!bool byte i16 i32 i64 double string!
-
     def process_exception(str)
       @exceptions[str.name] = str
 
+      o
       o "class #{str.name} < Stark::Exception"
       indent
-
-      o "class Struct < Stark::Struct"
-      indent
-      o "Fields = {"
-      indent
-
-      fields = str.fields.map do |f|
-        c = converter f.type
-        "#{f.index} => Stark::Field.new(#{f.index}, '#{f.name}', #{c})"
-      end
-
-      o "   #{fields.join(', ')}"
-
-      outdent
-      o "}"
-
-      str.fields.each do |f|
-        o "def #{f.name}; @fields['#{f.name}']; end"
-      end
-
-      outdent
-      o "end"
-
-      str.fields.each do |f|
-        o "def #{f.name}; @struct.#{f.name}; end"
-      end
-
+      write_field_declarations str.fields
       outdent
       o "end"
     end
 
-    def o(str)
-      @stream.print(" " * @indent)
+    def o(str = nil)
+      @stream.print(" " * @indent) if str
       @stream.puts str
     end
 
@@ -156,6 +120,7 @@ module Stark
       'i32' => "::Thrift::Types::I32",
       'i64' => "::Thrift::Types::I64",
       'string' => '::Thrift::Types::STRING',
+      'binary' => '::Thrift::Types::STRING',
       'struct' => '::Thrift::Types::STRUCT',
       'map' => '::Thrift::Types::MAP',
       'set' => '::Thrift::Types::SET',
@@ -170,6 +135,7 @@ module Stark
       'i32' => "read_i32",
       'i64' => "read_i64",
       'string' => 'read_string',
+      'binary' => 'read_string',
     }
 
     WriteFunc = {
@@ -180,6 +146,7 @@ module Stark
       'i32' => "write_i32",
       'i64' => "write_i64",
       'string' => 'write_string',
+      'binary' => 'write_string',
     }
 
     def type(t)
@@ -187,7 +154,7 @@ module Stark
     end
 
     def wire_type(t)
-      return "::Thrift::Types::STRUCT" if @structs[t]
+      return "::Thrift::Types::STRUCT" if @structs[t] || @exceptions[t]
       return "::Thrift::Types::I32" if @enums[t]
 
       case t
@@ -202,9 +169,23 @@ module Stark
       end
     end
 
+    def wire_type_and_ruby_type(t)
+      return "::Thrift::Types::STRUCT, #{t}" if @structs[t] || @exceptions[t]
+      return "::Thrift::Types::I32, Enum_#{t}" if @enums[t]
+      wire_type(t)
+    end
+
     def object_type(t)
-      return t if @structs[t]
-      type t
+      case t
+      when Stark::Parser::AST::Map
+        "map<#{object_type(t.key)},#{object_type(t.value)}>"
+      when Stark::Parser::AST::List
+        "list<#{object_type(t.value)}>"
+      when Stark::Parser::AST::Set
+        "set<#{object_type(t.value)}>"
+      else
+        t
+      end
     end
 
     def read_func(t)
@@ -215,146 +196,160 @@ module Stark
       WriteFunc[t] || raise("unknown type - #{t}")
     end
 
-    def write_field(ft, name, idx)
-      if desc = @structs[ft]
-        o "op.write_field_begin '#{name}', ::Thrift::Types::STRUCT, #{idx}"
-        output_struct desc, name
-        o "op.write_field_end"
+    def read_type(t, lhs, found_type = 'rtype')
+      o "#{lhs} = expect ip, #{wire_type(t)}, #{found_type} do"
+      indent
+      if desc = @structs[t]
+        o "read_#{desc.name}(ip)"
+      elsif desc = @enums[t]
+        o "Enum_#{desc.name}[ip.read_i32]"
+      elsif t.kind_of? Stark::Parser::AST::Map
+        o "expect_map ip, #{wire_type(t.key)}, #{wire_type(t.value)} do |kt,vt,size|"
+        indent
+        o   "{}.tap do |_hash|"
+        indent
+        o     "size.times do"
+        indent
+        read_type(t.key, "k", "kt")
+        read_type(t.value, "v", "vt")
+        o       "_hash[k] = v"
+        outdent
+        o     "end"
+        outdent
+        o   "end"
+        outdent
+        o "end"
+      elsif t.kind_of? Stark::Parser::AST::List
+        o "expect_list ip, #{wire_type(t.value)} do |vt,size|"
+        indent
+        o   "Array.new(size) do"
+        indent
+        read_type t.value, "_elem", "vt"
+        outdent
+        o   "end"
+        outdent
+        o "end"
+      elsif t.kind_of? Stark::Parser::AST::Set
+        o "expect_set ip, #{wire_type(t.value)} do |vt,size|"
+        indent
+        o   "_arr = Array.new(size) do"
+        indent
+        read_type t.value, "element", "vt"
+        o     "element"
+        outdent
+        o   "end"
+        o   "::Set.new(_arr)"
+        outdent
+        o "end"
+      else
+        o "ip.#{read_func(t)}"
+      end
+      outdent
+      o "end"
+    end
+
+    def write_type(ft, name)
+      if desc = (@structs[ft] || @exceptions[ft])
+        o "write_#{desc.name} op, #{name}"
       elsif desc = @enums[ft]
-        o "op.write_field_begin '#{name}', ::Thrift::Types::I32, #{idx}"
         o "op.write_i32 Enum_#{desc.name}[#{name}.to_sym]"
-
-        o "op.write_field_end"
       elsif ft.kind_of? Stark::Parser::AST::Map
         o "#{name} = hash_cast #{name}"
-        o "op.write_field_begin '#{name}', ::Thrift::Types::MAP, #{idx}"
         o "op.write_map_begin(#{wire_type(ft.key)}, #{wire_type(ft.value)}, #{name}.size)"
-
         o "#{name}.each do |k,v|"
         indent
-        o "op.#{write_func(ft.key)} k"
-        o "op.#{write_func(ft.value)} v"
+        write_type ft.key, "k"
+        write_type ft.value, "v"
         outdent
         o "end"
-
         o "op.write_map_end"
-        o "op.write_field_end"
       elsif ft.kind_of? Stark::Parser::AST::List
         o "#{name} = Array(#{name})"
-        o "op.write_field_begin '#{name}', ::Thrift::Types::LIST, #{idx}"
         o "op.write_list_begin(#{wire_type(ft.value)}, #{name}.size)"
-        o "#{name}.each { |v| op.#{write_func(ft.value)}(v) }"
+        o "#{name}.each do |v|"
+        indent
+        write_type ft.value, "v"
+        outdent
+        o "end"
         o "op.write_list_end"
-
-        o "op.write_field_end"
+      elsif ft.kind_of? Stark::Parser::AST::Set
+        o "#{name} = Set.new(#{name})"
+        o "op.write_list_begin(#{wire_type(ft.value)}, #{name}.size)"
+        o "#{name}.each do |v|"
+        indent
+        write_type ft.value, "v"
+        outdent
+        o "end"
+        o "op.write_set_end"
       elsif ft != "void"
-        o "op.write_field_begin '#{name}', #{type(ft)}, #{idx}"
         o "op.#{write_func(ft)} #{name}"
-        o "op.write_field_end"
       end
-
     end
 
-    def output_struct(desc, obj)
-      o "op.write_struct_begin '#{desc.name}'"
-
-      desc.fields.each do |f|
-        o "#{f.name} = #{obj}.#{f.name}"
-        write_field f.type, f.name, f.index
-      end
-
-      o "op.write_field_stop"
-      o "op.write_struct_end"
+    def write_field(ft, name, idx)
+      o "op.write_field_begin '#{name}', #{wire_type(ft)}, #{idx}"
+      write_type ft, name
+      o "op.write_field_end"
     end
 
     def write_processor(serv)
       o "class Processor < Stark::Processor"
       indent
+      o "include Protocol"
 
       serv.functions.each do |func|
+        o
         o "def process_#{func.name}(seqid, ip, op)"
         indent
 
         o "ip.read_struct_begin"
         args = Array(func.arguments)
-        o "args = Array.new(#{args.size})"
-
-        args.each do |arg|
-          if desc = @structs[arg.type]
-            o "_, rtype, rid = ip.read_field_begin"
-            o "if rtype != #{wire_type(arg.type)}"
-            o "  handle_unexpected rtype"
-            o "else"
-            o "  args[#{arg.index - 1}] = read_struct ip, rtype, rid, #{desc.name}"
-            o "end"
-            o "ip.read_field_end"
-          elsif desc = @enums[arg.type]
-            o "_, rtype, _ = ip.read_field_begin"
-            o "if rtype != #{wire_type(arg.type)}"
-            o "  handle_unexpected rtype"
-            o "else"
-            o "  args[#{arg.index - 1}] = Enum_#{desc.name}[ip.read_i32]"
-            o "end"
-            o "ip.read_field_end"
-
-          elsif arg.type.kind_of? Stark::Parser::AST::Map
-            ft = arg.type
-
-            o "_, rtype, _ = ip.read_field_begin"
-            o "if rtype != #{wire_type(arg.type)}"
-            o "  handle_unexpected rtype"
-            o "else"
-            o "  kt, vt, size = ip.read_map_begin"
-            o "  if kt == #{wire_type(ft.key)} && vt == #{wire_type(ft.value)}"
-            o "    result = {}"
-            o "    size.times do"
-            o "      result[ip.#{read_func(ft.key)}] = ip.#{read_func(ft.value)}"
-            o "    end"
-            o "    args[#{arg.index - 1}] = result"
-            o "  else"
-            o "    handle_bad_map size"
-            o "  end"
-            o "  ip.read_map_end"
-            o "end"
-            o "ip.read_field_end"
-          elsif arg.type.kind_of? Stark::Parser::AST::List
-            ft = arg.type
-            o "_, rtype, _ = ip.read_field_begin"
-            o "if rtype == ::Thrift::Types::LIST"
-            o "  vt, size = ip.read_list_begin"
-            o "  if vt == #{wire_type(ft.value)}"
-            o "    args[#{arg.index - 1}] = Array.new(size) { |n| ip.#{read_func(ft.value)} }"
-            o "  else"
-            o "    handle_bad_list size"
-            o "  end"
-            o "  ip.read_list_end"
-            o "else"
-            o "  handle_unexpected rtype"
-            o "end"
-            o "ip.read_field_end"
-          else
-            o "_, rtype, _ = ip.read_field_begin"
-            o "if rtype == #{type(arg.type)}"
-            o "  args[#{arg.index - 1}]= ip.#{read_func(arg.type)}"
-            o "else"
-            o "  handle_unexpected rtype"
-            o "end"
-            o "ip.read_field_end"
+        o "fields = {}"
+
+        o "while true"
+        o "  _, ftype, fid = ip.read_field_begin"
+        o "  break if ftype == ::Thrift::Types::STOP"
+
+        if args.empty?
+          o "  ip.skip ftype"
+        else
+          o "  case fid"
+          args.each do |arg|
+            o "  when #{arg.index}"
+            indent; indent
+            read_type arg.type, "fields[#{arg.index}]", "ftype"
+            outdent; outdent
           end
+          o "  else"
+          o "    ip.skip ftype"
+          o "  end"
         end
-
-        o "_, rtype, _ = ip.read_field_begin"
-        o "fail unless rtype == ::Thrift::Types::STOP"
+        o "  ip.read_field_end"
+        o "end"
         o "ip.read_struct_end"
         o "ip.read_message_end"
 
-        if t = func.throws
-          o "result = check_raise_specific('#{func.name}', seqid, op, #{t.first.type}) do"
-        o "  @handler.#{func.name}(*args)"
-        o "end"
-
-        o "return unless result"
+        o "args = #{args.map(&:index).inspect}.inject([]) {|arr,i| arr << fields[i] }"
 
+        if t = func.throws
+          o "begin"
+          indent
+          o "result = @handler.#{func.name}(*args)"
+          outdent
+          t.each do |ex|
+            o "rescue #{ex.type} => ex#{ex.index}"
+            indent
+            o   "op.write_message_begin '#{func.name}', ::Thrift::MessageTypes::REPLY, seqid"
+            o   "op.write_struct_begin '#{func.name}_result'"
+            write_field ex.type, "ex#{ex.index}", ex.index
+            o   "op.write_field_stop"
+            o   "op.write_struct_end"
+            o   "op.write_message_end"
+            o   "op.trans.flush"
+            o   "return"
+            outdent
+          end
+          o "end"
         else
           o "result = @handler.#{func.name}(*args)"
         end
@@ -366,12 +361,19 @@ module Stark
           next
         end
 
+        ft = func.return_type
+        o "result = value_for_write(result, #{wire_type_and_ruby_type(ft)})" if ft != "void"
+
         o "op.write_message_begin '#{func.name}', ::Thrift::MessageTypes::REPLY, seqid"
         o "op.write_struct_begin '#{func.name}_result'"
 
-        ft = func.return_type
-
-        write_field ft, 'result', 0
+        if ft != "void"
+          o "if result"
+          indent
+          write_field ft, 'result', 0
+          outdent
+          o "end"
+        end
 
         o "op.write_field_stop"
         o "op.write_struct_end"
@@ -387,31 +389,69 @@ module Stark
       o "end"
     end
 
-    def process_service(serv)
-      o "module #{serv.name}"
-      indent
-      o "class Client < Stark::Client"
+    def write_protocol
+      o
+      o "module Protocol"
       indent
+      @structs.merge(@exceptions).each do |name, struct|
+        o
+        o "def read_#{name}(ip)"
+        indent
+        o "obj = #{name}.new"
+        o "ip.read_struct_begin"
 
-      o "Functions = {}"
-      serv.functions.each do |func|
-        o "Functions[\"#{func.name}\"] = {"
-
-        o "    :args => {"
+        o "while true"
+        o "  _, ftype, fid = ip.read_field_begin"
+        o "  break if ftype == ::Thrift::Types::STOP"
 
-        mapped_args = Array(func.arguments).map do |a|
-          "#{a.index} => #{wire_type(a.type)}"
+        o "  case fid"
+        struct.fields.each do |f|
+          o "  when #{f.index}"
+          indent; indent
+          read_type f.type, "obj.#{f.name}", 'ftype'
+          outdent; outdent
         end
+        o "  else"
+        o "    ip.skip ftype"
+        o "  end"
+        o "  ip.read_field_end"
+        o "end"
 
-        o "      #{mapped_args.join(', ')}"
+        o "ip.read_struct_end"
+        o "obj"
+        outdent
+        o "end"
+        o
+        o "def write_#{name}(op, str)"
+        indent
+        o "op.write_struct_begin '#{name}'"
 
-        o "    }"
-        o "  }"
+        struct.fields.each do |f|
+          o "if #{f.name} = value_for_write(str.#{f.name}, #{wire_type_and_ruby_type(f.type)})"
+          indent
+          write_field f.type, f.name, f.index
+          outdent
+          o "end"
+        end
+
+        o "op.write_field_stop"
+        o "op.write_struct_end"
+        outdent
+        o "end"
       end
+      outdent
+      o "end"
+    end
+
+    def write_client(serv)
+      o "class Client < Stark::Client"
+      indent
+      o "include Protocol"
 
       serv.functions.each do |func|
         names = Array(func.arguments).map { |f| f.name }.join(", ")
 
+        o
         o "def #{func.name}(#{names})"
         indent
         o "op = @oprot"
@@ -419,6 +459,7 @@ module Stark
         o "op.write_struct_begin \"#{func.name}_args\""
 
         Array(func.arguments).each do |arg|
+          o "#{arg.name} = value_for_write #{arg.name}, #{wire_type_and_ruby_type(arg.type)}"
           write_field arg.type, arg.name, arg.index
         end
 
@@ -444,68 +485,32 @@ module Stark
         o "_, rtype, rid = ip.read_field_begin"
 
         if t = func.throws
-          o "if rid == 1"
-          o "  handle_throw #{t.first.type}"
+          o "case rid"
+          t.each do |ex|
+            o "when #{ex.index}"
+            o "  _ex = read_#{ex.type}(ip)"
+            o "  ip.read_field_end"
+            o "  _, rtype, _ = ip.read_field_begin"
+            o "  fail if rtype != ::Thrift::Types::STOP"
+            o "  ip.read_struct_end"
+            o "  ip.read_message_end"
+            o "  raise _ex"
+          end
           o "end"
         end
 
-        o "fail unless rid == 0"
-
         o "result = nil"
 
-        if func.return_type != "void"
-          if desc = @structs[func.return_type]
-            o "if rtype != #{wire_type(func.return_type)}"
-            o "  handle_unexpected rtype"
-            o "else"
-            o "  result = read_generic rtype, rid, #{desc.name}"
-            o "end"
-          elsif desc = @enums[func.return_type]
-            o "if rtype != #{wire_type(func.return_type)}"
-            o "  handle_unexpected rtype"
-            o "else"
-            o "  result = Enum_#{desc.name}[ip.read_i32]"
-            o "end"
-
-          elsif func.return_type.kind_of? Stark::Parser::AST::Map
-            ft = func.return_type
-
-            o "if rtype != #{wire_type(func.return_type)}"
-            o "  handle_unexpected rtype"
-            o "else"
-            o "  kt, vt, size = ip.read_map_begin"
-            o "  if kt == #{wire_type(ft.key)} && vt == #{wire_type(ft.value)}"
-            o "    result = {}"
-            o "    size.times do"
-            o "      result[ip.#{read_func(ft.key)}] = ip.#{read_func(ft.value)}"
-            o "    end"
-            o "  else"
-            o "    handle_bad_map size"
-            o "  end"
-            o "  ip.read_map_end"
-            o "end"
-          elsif func.return_type.kind_of? Stark::Parser::AST::List
-            ft = func.return_type
-            o "if rtype == ::Thrift::Types::LIST"
-            o "  vt, size = ip.read_list_begin"
-            o "  if vt == #{wire_type(ft.value)}"
-            o "    result = Array.new(size) { |n| ip.#{read_func(ft.value)} }"
-            o "  else"
-            o "    handle_bad_list size"
-            o "  end"
-            o "  ip.read_list_end"
-            o "else"
-            o "  handle_unexpected rtype"
-            o "end"
-          else
-            o "if rtype == #{type(func.return_type)}"
-            o "  result = ip.#{read_func(func.return_type)}"
-            o "else"
-            o "  handle_unexpected rtype"
-            o "end"
-          end
+        o "fail unless rid == 0"
 
-          o "_, rtype, rid = ip.read_field_begin unless rtype == ::Thrift::Types::STOP"
+        if func.return_type != "void"
+          o "if rtype != ::Thrift::Types::STOP"
+          indent
+          read_type func.return_type, "result"
+          o "ip.read_field_end"
+          o "_, rtype, rid = ip.read_field_begin"
+          outdent
+          o "end"
         end
 
         o "fail if rtype != ::Thrift::Types::STOP"
@@ -521,7 +526,21 @@ module Stark
 
       outdent
       o "end"
+    end
+
+    def process_service(serv)
+      unless @protocol_declared
+        o
+        o "module Protocol; end"
+        @protocol_declared = true
+      end
+
+      o
+      o "module #{serv.name}"
+      indent
 
+      write_client serv
+      o
       write_processor serv
 
       outdent