#!/usr/bin/env ruby

# Low-effort benchmark comparing Lite3::DBM in various modes against
# other equivalent Hash-like Ruby storage mechanisms.

require 'dbmlite3'
require 'yaml/dbm' unless RUBY_PLATFORM == "java"
require 'fileutils'
require 'optparse'

COUNT = 10000

Opts = proc do
  opt_defaults = {
    count: COUNT,
    transaction_only: true,
  }

  opts = Struct.new(*opt_defaults.keys).new(*opt_defaults.values)

  OptionParser.new do |opo|
    opo.banner = "Usage: bench_1.rb [options]"

    opo.on("-t", "--multi-transaction",
           "Don't batch Lite3 accesses in one big transaction.") {
      opts.transaction_only = false
    }

    opo.on("-c", "--count N", Integer, "Set test count.") { |c|
      opts.count = c
    }
  end.parse!

  next opts
end.call


module Tmp
  @root = File.join( File.dirname(__FILE__), "tmpdata")
  @count = 0

  def self.file
    FileUtils.mkdir(@root) unless File.directory?(@root)

    file = "testfile_#{@count}_#{$$}.sqlite3"
    @count += 1

    return File.join(@root, file)
  end

  def self.cleanup
    return unless File.directory?(@root)
    FileUtils.rm_rf(@root)
  end
end



def insert(count, db, offset)
  for n in 0 .. count
    db["k_#{n}"] = "#{n + offset}"
  end
end


def lookup(count, db)
  rnd = Random.new(69_420)
  sz = db.size

  for _ in 0 .. count * 3
    idx = rnd.rand(sz)
    v = db["k_#{idx}"]
    raise "Invalid value: #{v}" unless v.to_s == "#{idx + 1}"
  end
end


Times = {}

def time_it(type, task, db, use_tr, &block)
  print "#{type} #{task} - "
  STDOUT.flush

  start = Time.now

  if use_tr && db.respond_to?(:transaction)
    db.transaction { block.call }
  else
    block.call
  end

  finish = Time.now

  elapsed = (finish - start).to_f
  Times[type] = Times.fetch(type, 0) + elapsed

  puts "#{elapsed.round(4)}"
end

def bench(count, desc, db, use_tr)
  time_it(desc, "insert", db, use_tr) {
    insert(count, db, 0)
  }
  time_it(desc, "upsert", db, use_tr) {
    insert(count, db, 1)
  }
  time_it(desc, "lookup", db, use_tr) {
    lookup(count, db)
  }
  time_it(desc, "delete_if", db, use_tr) {
    rnd = Random.new(69_420)
    db.delete_if{|k,v| rnd.rand(2) == 0}
  }
  puts
end



def main
  puts "Count = #{Opts.count}\n"

  Tmp.cleanup

  bench(Opts.count, "hash", {}, false)

  if RUBY_PLATFORM != "java"
    DBM.open(Tmp.file) {|dbm|
      bench(Opts.count, "DBM", dbm, false)
    }
  
    YAML::DBM.open(Tmp.file) {|dbm|
      bench(Opts.count, "YAML::DBM", dbm, false)
    }
  end

  Lite3::DBM.open(Tmp.file, "benchmark", :yaml) { |dbm|
    bench(Opts.count, "Lite3::DBM(yaml)", dbm, true)
  }

  Lite3::DBM.open(Tmp.file, "benchmark", :marshal) { |dbm|
    bench(Opts.count, "Lite3::DBM(marshal)", dbm, true)
  }

  Lite3::DBM.open(Tmp.file, "benchmark", :string) { |dbm|
    bench(Opts.count, "Lite3::DBM(string)", dbm, true)
  }

  if !Opts.transaction_only
    Lite3::DBM.open(Tmp.file, "benchmark", :yaml) { |dbm|
      bench(Opts.count, "Lite3::DBM(yaml, single-trans)", dbm, false)
    }

    Lite3::DBM.open(Tmp.file, "benchmark", :marshal) { |dbm|
      bench(Opts.count, "Lite3::DBM(marshal, single-trans)", dbm, false)
    }

    Lite3::DBM.open(Tmp.file, "benchmark", :string) { |dbm|
      bench(Opts.count, "Lite3::DBM(string, single-trans)", dbm, false)
    }
  end


  puts
  puts "Totals:"
  Times.each{|k,v|
    puts "    #{k} - #{v.round(4)}"
  }

  Tmp.cleanup
end

main