#!/usr/bin/ruby

require "fileutils"
require "shellwords"

GCC = "riscv64-elf-gcc -fPIE -fvisibility=hidden -I. -Itests/macros/scalar -ffreestanding -nostdlib -nostartfiles -Wl,--emit-relocs -frandom-seed=banana"
REGS = /\b(ra|sp|gp|tp|t[0-6]|s[0-9]|s10|s11|a[0-7])\b/
POLKAVM_REGS = ["ra", "sp", "t0", "t1", "t2", "s0", "s1", "a0", "a1", "a2", "a3", "a4", "a5"]
ALL_REGS = ["zero", "ra", "sp", "gp", "tp", "t0", "t1", "t2", "s0", "s1", "a0", "a1", "a2", "a3", "a4", "a5", "a6", "a7", "s2", "s3", "s4", "s5", "s6", "s7", "s8", "s9", "s10", "s11", "t3", "t4", "t5", "t6"]

rust = ["// This file was AUTOGENERATED. DO NOT EDIT."]

Dir["tests/*/*.S"].sort.each do |path|
    _, category, filename = path.split("/")
    name = filename.sub(".S", "")

    # This uses self modifying code; skip it.
    next if name == "fence_i"

    # This uses relocations with an addend; skip it.
    next if name == "auipc"

    # This assumes we have multiple cores; skip it.
    next if name == "lrsc"

    raw_output_path = "output-raw/#{category}/#{name}.elf"
    output_path = "output/#{category}/#{name}.elf"
    obj_path = "output/#{category}/#{name}.o"

    if path =~ /rv64/
        march = "rv64"
        mabi = "lp64"
    else
        march = "rv32"
        mabi = "ilp32"
    end

    march += "ima_zifencei_zbb"

    FileUtils.mkdir_p File.dirname(output_path)
    FileUtils.mkdir_p File.dirname(raw_output_path)

    # First compile the programs as is so that we can see which registers they use.
    if !File.exist?(raw_output_path) || ARGV.include?("--force")
        puts "Building test (raw): #{category}/#{name}"
        system "#{GCC} -march=#{march} -mabi=#{mabi} #{path.shellescape} -DTESTNUM=gp -o #{raw_output_path.shellescape}"
        if $?.exitstatus != 0
            puts "ERROR: building failed!"
            exit 1
        end
    else
        puts "Skipping building (already built): #{category}/#{name}"
    end

    # Check which existing registers the program uses.
    assembly = `riscv32-elf-objdump --no-show-raw-insn --visualize-jumps=off --no-addresses -M no-aliases -d -j .text #{raw_output_path.shellescape}`
    raise if $?.exitstatus != 0
    used_regs = assembly.split("\n").filter { |line| line.start_with?("\t") }.join("\n").scan(REGS).flatten.sort.uniq
    free_regs = POLKAVM_REGS - used_regs

    # Pick a register to hold the test case number.
    testnum = nil
    if free_regs.delete("s0") != nil
        testnum = "s0"
    else
        testnum = free_regs.pop
    end
    raise if testnum == nil

    # Replace any registers that PolkaVM doesn't support.
    regmap = {}

    missing_regs = used_regs - POLKAVM_REGS
    missing_regs.each do |reg|
        replacement_reg = free_regs.pop
        raise if replacement_reg == nil
        regmap[reg] = replacement_reg
        regmap["x#{ALL_REGS.index(reg)}"] = replacement_reg
    end

    # Recompile with the registers replaced.
    if !File.exist?(output_path) || ARGV.include?("--force")
        puts "Building test (final): #{category}/#{name} (testnum = #{testnum}, regmap: #{regmap.inspect})"
        extra_args = regmap.map { |src, dst| "-D#{src}=#{dst}" }.join(" ")

        # We need a two-step build to make the output deterministic; otherwise GCC will generate
        # a random filename for the '.o' file and embed that in the resulting ELF file.
        system "#{GCC} -march=#{march} -mabi=#{mabi} #{extra_args} #{path.shellescape} -DTESTNUM=#{testnum} -c -o #{obj_path.shellescape}"
        if $?.exitstatus != 0
            puts "ERROR: building failed!"
            exit 1
        end

        system "#{GCC} -march=#{march} -mabi=#{mabi} #{obj_path.shellescape} -o #{output_path.shellescape}"
        if $?.exitstatus != 0
            puts "ERROR: building failed!"
            exit 1
        end
    end

    rust << "riscv_test!(riscv_unoptimized_#{category}_#{name}, \"../../../guest-programs/riscv-tests/output/#{category}/#{name}.elf\", #{testnum.upcase}, false);"
    rust << "riscv_test!(riscv_optimized_#{category}_#{name}, \"../../../guest-programs/riscv-tests/output/#{category}/#{name}.elf\", #{testnum.upcase}, true);"
end

File.write "../../crates/polkavm/src/tests_riscv.rs", rust.join("\n")
