rdb/lib/rdb/fast.rb

# frozen_string_literal: true

module Rdb
  # Fast-access methods for tables. Rows are found by byte offset, so a record
  # can be blanked or replaced in place without rewriting the file.
  #
  #   -ms  sequential (no index file)         -mb  binary search on a sorted table
  #   -mr  record numbers   table.r           -mi  inverted index   table.i
  #   -mh  hash             table.h
  #
  # table.r and table.h hold fixed-width entries (a head line, a dash line,
  # then "%8d\n" per entry), so one entry can be rewritten in place.
  module Fast
    ENTRY_WIDTH = 9
    SUFFIX = { 'h' => '.h', 'i' => '.i', 'r' => '.r' }.freeze
    FIXED_HEAD = "Offset\n------\n"

    Row = Struct.new(:from, :to, :line) do
      def cells
        @cells ||= line.split("\t", -1)
      end

      def blank?
        line.delete("\t ").empty?
      end
    end

    Hit = Struct.new(:row, :xfrom, :xto)

    # A table file. The head is read at once; the rows only when a method
    # needs all of them (sequential search, -mb, update.rdb). A lookup through
    # an index reads just its row, by byte offset.
    class Data
      attr_reader :path, :head

      def initialize(path)
        @path = path
        File.exist?(path) or Rdb.die("#{path}: no such table")
        File.open(path, 'rb') do |f|
          first = f.gets.to_s
          Rdb.die("#{path}: is a list; fast access needs a table") if first == "\n"
          @head = first.chomp("\n").split("\t", -1).map { |h| h.dup.force_encoding(Encoding::UTF_8).strip }
          @body = f.pos
          second = f.gets
          @body = f.pos if second && Table.dash_line?(second.chomp("\n").split("\t", -1))
        end
      end

      def rows
        @rows ||= begin
          bytes = File.binread(path).byteslice(@body..).to_s
          off = @body
          bytes.each_line("\n").map do |l|
            row = Row.new(off, off + l.bytesize, l.chomp("\n"))
            off += l.bytesize
            row
          end
        end
      end

      def live
        @live ||= rows.reject(&:blank?)
      end

      def column(name)
        Rdb.column_index(@head, name) || Rdb.die("#{path}: no column named '#{name}'")
      end

      def row_at(offset)
        return @by_from[offset] if @rows && (@by_from ||= @rows.to_h { |r| [r.from, r] })
        return nil if offset < @body

        File.open(path, 'rb') do |f|
          f.seek(offset)
          line = f.gets or return nil
          Row.new(offset, offset + line.bytesize, line.chomp("\n"))
        end
      end
    end

    module_function

    def index_path(path, method)
      SUFFIX[method] ? path + SUFFIX[method] : nil
    end

    def key_of(cells, idx)
      cells[idx].to_s
    end

    def norm(str, fold)
      s = str.to_s.strip
      fold ? s.downcase : s
    end

    def fmt_entry(num)
      format("%8d\n", num)
    end

    # FNV-1a over the key without blanks. (A plain sum of the bytes put
    # 3000 five-digit numbers into 30 slots.)
    def hash_slot(slots, key)
      key.b.delete(" \t").each_byte.reduce(0xcbf29ce484222325) { |h, b| ((h ^ b) * 0x100000001b3) & 0xffffffffffffffff } % slots
    end

    def live_rows(data)
      data.live
    end

    # ---- building indexes -------------------------------------------------

    def build(path, method, keycol, fold: false)
      data = Data.new(path)
      rows = live_rows(data)
      case method
      when 'r'
        Rdb.write_file(path + '.r', FIXED_HEAD + rows.map { |r| fmt_entry(r.from) }.join)
      when 'i'
        k = data.column(keycol)
        ents = rows.map { |r| [r.cells[k].to_s, r.from] }
        ents = ents.sort_by { |key, off| [fold ? key.strip.downcase : key.strip, off] }
        Rdb.write_file(path + '.i', "#{keycol}\tOffset\n#{'-' * keycol.size}\t------\n" + ents.map { |key, off| "#{key}\t#{off}\n" }.join)
      when 'h'
        k = data.column(keycol)
        slots = [rows.size * 2, 2].max
        table = Array.new(slots, 0)
        rows.each { |r| table[free_slot(table, hash_slot(slots, r.cells[k].to_s))] = r.from }
        Rdb.write_file(path + '.h', FIXED_HEAD + table.map { |o| fmt_entry(o) }.join)
      when 'b'
        k = data.column(keycol)
        body = live_rows(data).sort_by { |r| r.cells[k].to_s.strip }.map { |r| r.line + "\n" }
        Rdb.write_file(path, Table.head_text(data.head) + body.join)
      end
    end

    def free_slot(table, start)
      i = start
      i = (i + 1) % table.size while table[i] != 0 && table[i] != -1
      i
    end

    # ---- searching ----------------------------------------------------------

    # One entry of a fixed-width index (.r, .h) without reading the rest.
    def fixed_count(path)
      [(File.size(path) - FIXED_HEAD.bytesize) / ENTRY_WIDTH, 0].max
    end

    def fixed_entry(io, i)
      io.seek(FIXED_HEAD.bytesize + ENTRY_WIDTH * i)
      io.read(ENTRY_WIDTH).to_i
    end

    # Entries of an inverted index (.i) from the first key that is not below
    # target, found by bisecting on byte positions. Yields hashes like
    # read_inverted does.
    def each_inverted_from(path, below)
      File.open(path, 'rb') do |f|
        f.gets
        f.gets
        lo = f.pos
        hi = f.size
        while lo < hi
          mid = (lo + hi) / 2
          if mid == lo
            at = lo
          else
            f.seek(mid - 1)
            f.gets
            at = f.pos
          end
          if at >= hi
            hi = mid
            next
          end
          f.seek(at)
          line = f.gets
          if below.call(line.split("\t", 2)[0])
            lo = at + line.bytesize
          else
            hi = at
          end
        end
        f.seek(lo)
        loop do
          from = f.pos
          line = f.gets or break
          key, off = line.chomp("\n").split("\t", 2)
          yield({ from: from, to: from + line.bytesize, key: key, offset: off.to_i })
        end
      end
    end

    def read_fixed(path)
      File.binread(path).byteslice(FIXED_HEAD.bytesize..).to_s.lines.map { |l| l.to_i }
    end

    def read_inverted(path)
      off = 0
      ents = []
      File.binread(path).each_line("\n").with_index do |l, n|
        ents << { from: off, to: off + l.bytesize, key: l.chomp("\n").split("\t", 2)[0], offset: l.chomp("\n").split("\t", 2)[1].to_i } if n >= 2
        off += l.bytesize
      end
      ents
    end

    def matches?(cell, key, partial, exact)
      if partial
        norm(cell, true).start_with?(norm(key, true))
      elsif exact
        cell == key
      else
        cell.strip == key.strip
      end
    end

    # Find rows for one key. method is one of s b h i r; keycols are column names.
    def search(data, method, keycols, key, first: false, partial: false, exact: false)
      path = data.path
      idxs = keycols.map { |c| data.column(c) }
      keyvals = Array(key)
      hits = []
      case method
      when 'b'
        rows = live_rows(data)
        k = idxs[0]
        lo = rows.bsearch_index { |r| (partial ? norm(r.cells[k], true)[0, keyvals[0].strip.size] : r.cells[k].strip) >= (partial ? norm(keyvals[0], true) : keyvals[0].strip) }
        if lo
          rows[lo..].each do |r|
            break unless matches?(r.cells[k], keyvals[0], partial, exact)

            hits << Hit.new(r, 0, 8)
            break if first
          end
        end
      when 'r'
        ip = path + '.r'
        n = keyvals[0].to_i
        if n.positive? && n <= fixed_count(ip)
          off = File.open(ip, 'rb') { |f| fixed_entry(f, n - 1) }
          r = data.row_at(off) and hits << Hit.new(r, FIXED_HEAD.bytesize + ENTRY_WIDTH * (n - 1), FIXED_HEAD.bytesize + ENTRY_WIDTH * n)
        end
      when 'i'
        target = partial ? norm(keyvals[0], true) : keyvals[0].strip
        below = ->(key) { (partial ? norm(key, true)[0, keyvals[0].strip.size] : key.strip) < target }
        each_inverted_from(path + '.i', below) do |e|
          break unless matches?(e[:key], keyvals[0], partial, exact)

          r = data.row_at(e[:offset])
          hits << Hit.new(r, e[:from], e[:to]) if r
          break if first
        end
      when 'h'
        ip = path + '.h'
        slots = fixed_count(ip)
        File.open(ip, 'rb') do |f|
          i = hash_slot(slots, keyvals[0])
          slots.times do
            o = fixed_entry(f, i)
            break if o.zero?

            if o.positive? && (r = data.row_at(o)) && matches?(r.cells[idxs[0]], keyvals[0], false, exact)
              hits << Hit.new(r, FIXED_HEAD.bytesize + ENTRY_WIDTH * i, FIXED_HEAD.bytesize + ENTRY_WIDTH * (i + 1))
              break if first
            end
            i = (i + 1) % slots
          end
        end
      else
        live_rows(data).each do |r|
          c = r.cells
          next unless idxs.each_with_index.all? { |ci, n| matches?(c[ci].to_s, keyvals[n].to_s, partial, exact) }

          hits << Hit.new(r, 0, 8)
          break if first
        end
      end
      hits
    end

    # ---- changing rows --------------------------------------------------------

    def blanked(line)
      line.gsub(/[^\t]/, ' ')
    end

    def splice(path, from, to, text)
      b = File.binread(path)
      Rdb.write_file(path, b.byteslice(0, from) + text + b.byteslice(to..).to_s)
    end

    # Remove a record's entry from the secondary index for method.
    def index_remove(path, method, row_from)
      case method
      when 'r'
        ip = path + '.r'
        ents = read_fixed(ip)
        ents.delete(row_from)
        Rdb.write_file(ip, FIXED_HEAD + ents.map { |o| fmt_entry(o) }.join)
      when 'h'
        ip = path + '.h'
        ents = read_fixed(ip).map { |o| o == row_from ? -1 : o }
        Rdb.write_file(ip, FIXED_HEAD + ents.map { |o| fmt_entry(o) }.join)
      when 'i'
        ip = path + '.i'
        e = read_inverted(ip).find { |x| x[:offset] == row_from }
        splice(ip, e[:from], e[:to], '') if e
      end
    end

    def index_add(path, method, row_from, cells, data_head, keycol)
      case method
      when 'r'
        Rdb.write_file(path + '.r', File.binread(path + '.r') + fmt_entry(row_from))
      when 'i'
        ip = path + '.i'
        name = File.binread(ip).lines.first.split("\t")[0]
        key = cells[data_head.index(name)].to_s
        ents = read_inverted(ip)
        pos = ents.find { |e| e[:key].strip > key.strip }
        line = "#{key}\t#{row_from}\n"
        if pos then splice(ip, pos[:from], pos[:from], line)
        else Rdb.write_file(ip, File.binread(ip) + line)
        end
      when 'h'
        ip = path + '.h'
        ents = read_fixed(ip)
        if ents.any? { |o| o <= 0 }
          key = cells[data_head.index(keycol)].to_s
          ents[free_slot(ents, hash_slot(ents.size, key))] = row_from
          Rdb.write_file(ip, FIXED_HEAD + ents.map { |o| fmt_entry(o) }.join)
        else
          build(path, 'h', keycol)
        end
      when 'b'
        build(path, 'b', keycol)
      end
    end

    def append_row(path, cells, method, keycol)
      Rdb.check_cells!(cells)
      line = cells.join("\t")
      b = File.binread(path)
      b += "\n" unless b.empty? || b.end_with?("\n")
      from = b.bytesize
      Rdb.write_file(path, b + line + "\n")
      index_add(path, method, from, cells, Data.new(path).head, keycol) if SUFFIX[method] || method == 'b'
      from
    end

    def delete_hit(path, hit, method)
      r = hit.row
      splice(path, r.from, r.to, blanked(r.line) + "\n")
      index_remove(path, method, r.from)
    end

    # Put cells where the record at hit was. If the new record is too long
    # even after trimming blanks, append it (unless append is false).
    def replace_row(path, from, to, cells, method, xfrom, xto, keycol, allow_append: true)
      Rdb.check_cells!(cells)
      room = to - from - 1
      line = cells.join("\t")
      line = cells.map(&:strip).join("\t") if line.bytesize > room
      if line.bytesize <= room
        splice(path, from, to, line + ' ' * (room - line.bytesize) + "\n")
        return from
      end
      return nil unless allow_append

      old = File.binread(path).byteslice(from, room)
      splice(path, from, to, blanked(old) + "\n")
      newfrom = append_row(path, cells.map(&:strip), 'x', keycol)
      reindex_entry(path, method, xfrom, xto, newfrom, cells, keycol)
      newfrom
    end

    def reindex_entry(path, method, xfrom, xto, newfrom, cells, keycol)
      case method
      when 'r', 'h'
        splice(path + SUFFIX[method], xfrom, xto, fmt_entry(newfrom))
      when 'i'
        ip = path + '.i'
        name = File.binread(ip).lines.first.split("\t")[0]
        head = Data.new(path).head
        splice(ip, xfrom, xto, "#{cells[head.index(name)].to_s.strip}\t#{newfrom}\n")
      end
    end

    # Read keys: a table of keys, or one key per line.
    def read_keys(io, ncols)
      text = io.read.to_s
      return [] if text.empty?

      lines = text.lines.map { |l| l.chomp("\n") }
      if Table.headed?(lines)
        lines.drop(2).reject(&:empty?).map { |l| l.split("\t", -1).first(ncols) }
      else
        lines.reject(&:empty?).map { |l| l.split("\t", -1).first(ncols) }
      end
    end
  end
end