database.rb in sypex_geo-0.2.0

- old
+ new
@@ -3,161 +3,128 @@
 module SypexGeo
   class DatabaseError < StandardError
   end
 
   class Database
-    attr_reader :version
+    TYPE_COUNTRY = 1
+    TYPE_CITY    = 3
 
+    attr_reader :version, :time
+
     def initialize(path)
       @file = File.open(path, 'rb')
 
-      setup!
+      parse_header
+      setup
     end
 
-    def lookup(ip, full = false)
-      if seek = search(ip)
-        read_location(seek, full)
+    def query(ip)
+      if position = search(ip)
+        Result.new(position, self)
       end
     end
 
+    def read_country(position)
+      @country_parser ||= Pack.new(@country_pack)
+      @country_parser.parse(@cities_db[position, @country_size])
+    end
+
+    def read_region(position)
+      @region_parser ||= Pack.new(@region_pack)
+      @region_parser.parse(@regions_db[position, @region_size])
+    end
+
+    def read_city(position)
+      @city_parser ||= Pack.new(@city_pack)
+      @city_parser.parse(@cities_db[position, @city_size])
+    end
+
+    def country?
+      @type == TYPE_COUNTRY
+    end
+
+    def city?
+      @type == TYPE_CITY
+    end
+
     def inspect
       "#<#{self.class}:0x#{object_id} @version=#{@version}>"
     end
 
     protected
 
-    def setup!
+    def parse_header
       if header = @file.read(40)
         id, @version, @time, @type, @charset,
-        @b_idx_len, @m_idx_len, @range, @db_items, @id_len,
-        @max_region, @max_city, @region_size, @city_size,
-        @max_country, @country_size,
+        @block_idx_size, @main_idx_size, @range, @db_records_count, @id_size,
+        @region_size, @city_size, @regions_db_size, @cities_db_size,
+        @country_size, @countries_db_size,
         @pack_size = header.unpack('a3CNCCCnnNCnnNNnNn')
       end
 
       raise DatabaseError.new, 'Wrong file format' unless id == 'SxG'
+    end
 
+    def setup
       @pack = @file.read(@pack_size).split("\0")
-      @b_idx_arr = @file.read(@b_idx_len * 4).unpack('N*')
-      @m_idx_arr = @file.read(@m_idx_len * 4).scan(/.{1,4}/m)
+      @country_pack, @region_pack, @city_pack = @pack
 
-      @block_len = 3 + @id_len
-      @db_begin = @file.tell
-      @regions_begin = @db_begin + @db_items * @block_len
-      @cities_begin = @regions_begin + @region_size
+      @block_idx = @file.read(@block_idx_size * 4).unpack('N*')
+      @main_idx = @file.read(@main_idx_size * 4).scan(/.{1,4}/m)
+
+      @db_record_size = 3 + @id_size
+      @db = @file.read(@db_records_count * @db_record_size)
+      @regions_db = @file.read(@regions_db_size) if @regions_db_size > 0
+      @cities_db = @file.read(@cities_db_size) if @cities_db_size > 0
     end
 
     def search(ip)
-      ip1n = ip.to_i
+      octet = ip.to_i
+      return if octet == 0 or octet == 127 or octet >= @block_idx_size
 
-      return if ip1n == 0 or ip1n == 127 or ip1n >= 224
-
+      min, max = @block_idx[octet - 1], @block_idx[octet]
+      range = @range
       ipn = IPAddr.new(ip).hton
-      blocks_min, blocks_max = @b_idx_arr[ip1n - 1], @b_idx_arr[ip1n]
 
-      if blocks_max - blocks_min > @range
-        part = search_idx(ipn, blocks_min / @range, (blocks_max / @range) - 1)
-        min = part > 0 ? part * @range : 0
-        max = part > @m_idx_len ? @db_items : (part + 1) * @range
-        min = blocks_min if min < blocks_min
-        max = blocks_max if max > blocks_max
-      else
-        min = blocks_min
-        max = blocks_max
+      if max - min > range
+        min = range * main_idx_search(ipn, min / range, (max / range) - 1)
+        max = min + range
       end
 
-      search_db(ipn, min, max)
+      db_search(ipn, min, max)
     end
 
-    def search_idx(ipn, min, max)
-      idx = @m_idx_arr
+    def main_idx_search(ipn, min, max)
+      idx = @main_idx
 
-      while max - min > 8
-        offset = (min + max) >> 1
+      while min < max
+        mid = (min + max) / 2
 
-        if ipn > idx[offset]
-          min = offset
+        if ipn > idx[mid]
+          min = mid + 1
         else
-          max = offset
+          max = mid
         end
       end
 
-      while ipn > idx[min]
-        break if min >= max
-        min += 1
-      end
-
       min
     end
 
-    def search_db(ipn, min, max)
-      len = max - min
-      @file.pos = @db_begin + min * @block_len
-      search_db_chunk(@file.read(len * @block_len), ipn, 0, len - 1)
-    end
+    def db_search(ipn, min, max)
+      db = @db
+      db_record_size = @db_record_size
+      octets = ipn[1, 3]
 
-    def search_db_chunk(data, ipn, min, max)
-      block_len = @block_len
+      while min < max
+        mid = (min + max) / 2
 
-      if max - min > 1
-        ipn = ipn[1, 3]
-
-        while max - min > 8
-          offset = (min + max) >> 1
-
-          if ipn > data[offset * block_len, 3]
-            min = offset
-          else
-            max = offset
-          end
+        if octets > db[mid * db_record_size, 3]
+          min = mid + 1
+        else
+          max = mid
         end
-
-        while ipn >= data[min * block_len, 3]
-          min += 1
-          break if min >= max
-        end
-      else
-        min += 1
       end
 
-      data[min * block_len - @id_len, @id_len].unpack('H*').first.hex
-    end
-
-    def read_data(seek, limit, type)
-      @file.pos = (type == TYPE_REGION ? @regions_begin : @cities_begin) + seek
-      Pack.parse(@pack[type], @file.read(limit))
-    end
-
-    def read_country(seek)
-      read_data(seek, @max_country, TYPE_COUNTRY)
-    end
-
-    def read_region(seek)
-      read_data(seek, @max_region, TYPE_REGION)
-    end
-
-    def read_city(seek)
-      read_data(seek, @max_city, TYPE_CITY)
-    end
-
-    def read_location(seek, full = false)
-      region = nil
-      city = nil
-      country = nil
-
-      if seek < @country_size
-        country = read_country(seek)
-      elsif city = read_city(seek)
-        region_seek = city.delete(:region_seek)
-        country_id = city.delete(:country_id)
-        country = { id: country_id, iso: COUNTRY_CODES[country_id - 1] }
-      end
-
-      if full and region_seek
-        region = read_region(region_seek)
-        country = read_country(region.delete(:country_seek))
-      end
-
-      { city: city, region: region, country: country }
+      db[min * db_record_size - @id_size, @id_size].unpack('H*')[0].hex
     end
   end
 end