run

#!/usr/bin/env ruby
#/ Usage: ./run [--local] [-p N] [[YYYY] DD | INPUT... | SCRIPT...]
#/ Run all scripts for the specified day of Advent of Code.
#/ Default is today.
#/ --local skips downloading the input.
#/ -p sets the max number of scripts that run in parallel. By default there is
#/    no limit.
#/
#/ For example, for day 1 of 2020, you might name inputs like:
#/  2020/examples/01.simple.txt
#/  2020/examples/01.complex.txt
#/  2020/input/01.txt
#/ And scripts like:
#/  2020/01.rb
#/  2020/py/01.py

require "fileutils"
require "net/http"
require "thread"
require "uri"

def main(year: nil, day: nil, inputs: nil, scripts: nil, local: false, p: nil)
  if year.nil?
    year, day = get_date(day: day, inputs: inputs, scripts: scripts)
  elsif day.nil?
    raise ArgumentError, "Year specified without day."
  end

  inputs, scripts = scan_files(inputs: inputs, scripts: scripts, local: local, year: year, day: day)

  if inputs.empty? || scripts.empty?
    puts "!!! Nothing to do!"
    exit 0
  end

  c = {}
  t = []
  q = []
  t << str_bg("======= Line counts =======")
  t << Thread.new { `wc -l #{scripts.join(" ")} 2>&1` }
  inputs.each do |input|
    t << str_bg("======= #{input} =======")
    scripts.each do |script|
      w = p.nil? ? nil : Queue.new
      q << w if w
      t << str_bg(">>> #{script} < #{input}")
      t << run_bg(script: script, input: input, c: c, w: w)
      t << w if w
      t << str_bg("")
    end
  end
  p.times { q.shift&.enq :go } unless p.nil?
  t.each do |x|
    case x
    when Queue
      q.shift&.enq :go
    else
      puts x.value
    end
  end
end

def get_date(day:, inputs:, scripts:)
  if day
    return Time.now.year, day
  end
  ((inputs || []) + (scripts || [])).each do |filename|
    if info = derive_date(filename)
      return info
    end
  end
  now = Time.now
  [now.year, now.day]
end

def derive_date(filename)
  year = nil
  day = nil
  filename.split("/").each do |part|
    if year.nil? && part =~ /\A\d{4}\z/
      year = part.to_i
    end
    if day.nil? && part =~ /\A\d{1,2}(\D|\z)/
      day = part.to_i
    end
  end
  if year && day
    [year.to_i, day.to_i]
  end
end

def scan_files(inputs:, scripts:, local:, year:, day:)
  return [inputs, scripts] if inputs && scripts
  found_inputs = []
  found_scripts = []
  Dir["**/*"].each do |file|
    case typeof(file, year: year, day: day)
    when :input
      found_inputs << file
    when :script
      found_scripts << file
    end
  end

  download_input = inputs.nil? && !local && !found_inputs.any? { |input| input =~ /input/ }
  if download_input
    begin
      token = File.read(".token").strip
      puts "Downloading input..."
      uri = URI("https://adventofcode.com/#{year}/day/#{day}/input")
      Net::HTTP.start(uri.host, uri.port, use_ssl: true) do |http|
        req = Net::HTTP::Get.new(uri)
        req["Cookie"] = "session=#{token}"
        resp = http.request(req)
        if resp.code != "200"
          puts "Error downloading input (HTTP #{resp.code}):"
          puts resp.body
        else
          input_file = sprintf("%04d/input/%02d.txt", year, day)
          FileUtils.mkdir_p(File.dirname(input_file))
          File.write(input_file, resp.body)
          found_inputs << input_file
        end
      end
    rescue Errno::ENOENT
      puts "!!!HINT!!! Create a .token file with your session token and I will"
      puts "!!!HINT!!! download the input file for you!"
    end
  end
  [inputs || found_inputs, scripts || found_scripts]
end

def str_bg(str)
  Thread.new { str }
end

def run_bg(script:, input:, c:, w:)
  case File.extname(script)
  when ".rb"
    Thread.new { w&.deq; capture_output "ruby", script, in: input }
  when ".py"
    Thread.new { w&.deq; capture_output "python3", script, in: input }
  when ".rs"
    c_th = c[script]
    if c_th.nil?
      c_th = c[script] = compile_bg(script) { |out_file| ["rustc", "-O", "-o", out_file, script] }
    end
    Thread.new do
      w&.deq
      res = c_th.value
      if res.ok?
        # warm up
        system res.bin, in: "/dev/null", out: "/dev/null", err: "/dev/null"
        capture_output res.bin, in: input
      else
        res.output
      end
    end
  when ".go"
    c_th = c[script]
    if c_th.nil?
      c_th = c[script] = compile_bg(script) { |out_file| ["go", "build", "-o", out_file, script] }
    end
    Thread.new do
      w&.deq
      res = c_th.value
      if res.ok?
        # warm up
        system res.bin, in: "/dev/null", out: "/dev/null", err: "/dev/null"
        capture_output res.bin, in: input
      else
        res.output
      end
    end
  else
    puts "Unrecognized script extension."
  end
end

def capture_output(*cmd, **opts)
  cmd = [{"IS_REAL_INPUT" => "Y"}] + cmd if opts[:in] =~ /input/
  capture_output2(*cmd, **opts).first
end

def capture_output2(*cmd, **opts)
  r, w = IO.pipe
  start = Time.now
  pid = spawn(*cmd, **opts, out: w, err: w)
  w.close
  res = r.read
  elapsed = Time.now - start

  r.close
  _, status = Process.wait2(pid)

  [
    sprintf("%s> elapsed: %d ms\n", res, 1000 * elapsed),
    status,
  ]
end

def compile_bg(script)
  bin = ".output/#{script}"
  cmd = yield bin

  Thread.new do
    system "mkdir", "-p", File.dirname(bin)

    output, status = capture_output2(*cmd)

    Result.new \
      status: status,
      output: output,
      bin: bin
  end
end

class Result
  def initialize(status:, output:, bin:)
    @status = status
    @output = output
    @bin = bin
  end

  attr_reader :output, :bin

  def ok?
    @status.success?
  end
end

def typeof(file, year:, day:)
  parts = file.split("/")
  return :wrong_year unless parts.any? { |part| part.to_i == year }
  return :wrong_day  unless parts.any? { |part| part.to_i == day  }
  return :input if parts.include?("input") || parts.include?("examples")
  return :script
end

options = {}
if ARGV.delete("--local")
  options[:local] = true
end
if i = ARGV.index("-p")
  ARGV.delete_at(i)
  options[:p] = ARGV.delete_at(i).to_i
  if options[:p] < 1
    raise "illegal value for -p"
  end
end

def usage
  File.readlines(__FILE__).grep(/^#\//).each { |line| puts line[3..-1] }
end

options = {}
until ARGV.empty?
  case arg = ARGV.shift
  when "--local"
    options[:local] = true
  when "-p"
    options[:p] = ARGV.shift.to_i
  when /\A-/
    usage
    exit 1
  when /\A\d{4}\z/
    options[:year] = arg.to_i
  when /\A\d{1,2}\z/
    options[:day] = arg.to_i
  when /example/, /input/
    (options[:inputs] ||= []) << arg
  else
    (options[:scripts] ||= []) << arg
  end
end

if options[:year] && options[:day].nil?
  usage
  exit 1
end

main(**options)