#!/usr/bin/env ruby
# vim: set sw=3 tw=80:
require 'base64'

POS_INT_REGEX = /(?:0[0-7]+|0b[01]+|0x[0-9a-fA-F]+|[0-9]+)/

# Make sure input is treated as a signed 32-bit number
def signed32(n)
   if n and (n < 0 or n >= 0x80000000)
      -(-n & 0xffffffff)
   else
      n
   end
end

def split_immed(n)
   result = []

   while (n < -16) || (n > 15)
      result = [ 0b10000000 | (n & 0x7f) ] + result
      n >>= 7
   end

   [ 0b00100000 | (n & 0x1f) ] + result
end

$labels = {}
$forward_labels = {}

class Label
   attr_accessor :name
   def initialize(name, pc=nil)
      unless name =~ /^\.*[a-zA-Z_]\w*$/
         raise StandardError.new("Bad label name '#{name}' at #{$file}:#{$line}")
      end
      @name, @pc = name, pc
      @export, @global = false, false
      @alias = nil
   end

   def export=(x) @export = x ? true : false; end
   def export?() @export; end
   
   def global=(x) @global = x ? true : false; end
   def global?() @global; end

   def local?() (@name =~ /^(\.+)/) ? $1.length : false; end
   def level() self.local? || 0; end

   def pc=(x) @pc = x; end
   def pc() @alias ? @alias.pc : @pc; end
   def alias() @alias; end
   def alias?() @alias ? true : false; end

   def Label.get(name)
      lbl = $labels[name] || $forward_labels[name]
      if not lbl
         lbl = Label.new(name)
         $forward_labels[name] = lbl
      end
      return lbl
   end

   def Label.define(name)
      if $labels.has_key? name
         $stderr.puts "WARNING: Redefining label '#{name}'."
      end

      lbl = $forward_labels[name] || Label.new(name)
      $forward_labels.delete name
      $labels[name] = lbl
      return lbl
   end

   def Label.alias(label, name)
      if $labels.has_key? name
         $stderr.puts "WARNING: Redefining label '#{name}'."
      end

      lbl = $forward_labels[name] || Label.new(name)
      lbl.instance_variable_set(:@alias, label)
      $forward_labels.delete name
      $labels[name] = lbl
      return lbl
   end

   def Label.in_file_scope()
      old_labels, old_forwards = $labels.dup, $forward_labels.dup
      
      # Local references end at file boundaries
      $labels.delete_if { |k,v| not v.global? }
      $forward_labels.delete_if { |k,v| not v.global? }

      yield
      
      # Local references end at file boundaries
      $labels.delete_if { |k,v| not v.global? }
      $forward_labels.delete_if { |k,v| not v.global? }

      $labels = old_labels.merge($labels)
      $forward_labels = old_forwards.merge($forward_labels)
   end
end

class Instruction
   def initialize(labels, pc, opcode, args)
      @labels = labels.map do |name|
         lbl = Label.define(name)
         $labels.delete_if { |k,v| v.level > lbl.level }
         lbl
      end

      @type = opcode.downcase.to_sym
      @args = args
      @align = 1

      valid = true

      case @type
      when :nop, :rshift, :nip, :and, :or,
           :xor, :swap, :dup, :add, :sub
         valid = args.empty? || ((args.length == 1) && (args[0] =~ /return/i))
         @return = !args.empty?
      when :load, :store
         valid = args.empty? || ((args.length == 1) && (args[0] =~ /byte/i))
         @byte_op = !args.empty?
      when :return
         valid = args.empty?
         @type = :nop
         @return = true
      when :rdrop, :pop, :pushpc, :push, :rdup, :mult, :over, :nlz
         valid = args.empty?
      when :immed
         if (args.length == 1) and (args[0] =~ /^-?#{POS_INT_REGEX}$/)
            # Integer constant
            valid, @value = true, signed32(args[0].to_i(0))
         elsif (args.length == 1) and (args[0] =~ /^'.'$/)
            # Character constant
            valid, @value = true, args[0][1].ord
         else
            valid, @target = true, Label.get(args[0])
         end
      when :jump, :call
         @rel, @target, @cc, @drop = false, nil, :always, false

         args.each do |arg|
            case arg
            when /^rel$/i
               @rel = true
            when /^(never|eq|lt|gt|always|ne|ge|le)$/i
               @cc = $1.downcase.to_sym
            when /^drop$/i
               @drop = true;
            else
               if @target
                  raise StandardError.new("Syntax error: #{opcode} #{args.join(' ')}")
               else
                  @target = Label.get(arg)
               end
            end
         end
      when :drop
         @type = :jump
         @rel = false
         @target = nil
         @cc = :never
         @drop = false
      when :drop2
         @type = :jump
         @target = nil
         @rel = false
         @cc = :never
         @drop = true
      when :merge
         valid = (args.length == 1) && (args[0] =~ /^#{POS_INT_REGEX}$/)
         @value = args[0].to_i(0) if valid
         valid &&= (@value >= 0 and @value <= 127)
      when :resb
         valid = (args.length == 1) && (args[0] =~ /^#{POS_INT_REGEX}$/)
         @immed = [ 0 ] * args[0].to_i(0) if valid
      when :align
         valid = (args.length == 1) && (args[0] =~ /^#{POS_INT_REGEX}$/)
         @align = args[0].to_i(0) if valid
         valid &&= (@align >= 1)
      when :byte
         valid = true
         @immed = []
         args.each do |arg|
            case arg
            when /^-?#{POS_INT_REGEX}$/
               v = arg.to_i(0)
               valid &&= (-128..255).member? v
               @immed << (v & 0xff) if valid
            when /^'(.)'$/
               @immed << (arg[1].ord & 0xff) if valid
            when /^".*"$/
               arg[1..-2].scan(/([^\\"])|(\\")|(\\\\)/) do |a,b,c|
                  a = '"' if b
                  a = '\\' if c
                  @immed << (a[0].ord & 0xff)
               end
            else
               raise StandardError.new("Unsupported BYTE operand.")
            end
         end
      when :word
         valid = true
         @immed = []
         @align = 4
         args.each do |arg|
            case arg
            when /^-?#{POS_INT_REGEX}$/
               v = arg.to_i(0)
               # Convert long to unsigned chars w/ big-endian byte order
               @immed += [v].pack("N").unpack("C*")
            else
               raise StandardError.new("WORD only supports integer operands.")
            end
         end
      else
         raise StandardError.new("Syntax error: #{opcode} #{args.join(' ')}")
      end

      if not valid
         raise StandardError.new("Invalid instruction: #{opcode} #{args.join(' ')}")
      end

      reflow(pc)
   end

   def reflow(pc)
      if pc >= $options[:ram_size]
         raise StandardError.new("No more room in Block RAM.")
      end

      @align_bytes = (-pc) % @align
      @pc = pc + @align_bytes
      old_bytes = @bytes
      @labels.each { |lbl| lbl.pc = @pc }

      if @target
         @value = signed32(@target.pc)

         if @value
            if @rel
               # Known relative; PC of jump/call depends on immed. length
               diff = @value - (@pc + 6)
               im_size = split_immed(diff).length
               diff = @value - (@pc + im_size + 1)
               @immed = split_immed(diff)
               
               if @immed.length > im_size
                  im_size = @immed.length
                  diff = @value - (@pc + im_size + 1)
                  @immed = split_immed(diff)
               end
               
               if @immed.length > im_size
                  raise StandardError.new("Relative target offset won't converge!")
               end
            else
               # Absolute known address
               @immed = split_immed(@value)
            end
         else
            # Unknown address; reserve maximum space that may be required
            @immed = split_immed(-($max_pc+1))
         end
      elsif @value
         @immed = split_immed(@value)
      end

      if @type == :align
         @bytes = 0
      elsif not @immed
         @bytes = 1
      elsif @type == :jump or @type == :call
         @bytes = @immed.length + 1
      else
         @bytes = @immed.length
      end

      @bytes += @align_bytes

      if old_bytes != @bytes
         $layout_changed = true
      end
   end

   def bytes() @bytes; end

   def finalize_target
      if @target
         @value = signed32(@target.pc)

         if not @value
            raise StandardError.new("Unknown label: '#{@target.name}'")
         else
            if @rel
               @immed = split_immed(@value - (@pc + @bytes))
            else
               @immed = split_immed(@value)
            end
         end

         if @type == :immed
            bytes = @immed.length
         else
            bytes = @immed.length + 1
         end

         if @bytes < bytes
            raise StandardError.new("Insufficient space for target address.")
         end
      end
   end

   OPCODES = {
      :nop    => 0b00000000,
      :rshift => 0b00000010,
      :add    => 0b00000100,
      :sub    => 0b00000110,
      :nip    => 0b00001000,
      :and    => 0b00001010,
      :or     => 0b00001100,
      :xor    => 0b00001110,
      :swap   => 0b00010000,
      :dup    => 0b00010010,
      :over   => 0b00010100,
      :rdup   => 0b00010101,
      :nlz    => 0b00010110,
      :mult   => 0b00010111,
      :load   => 0b00011000,
      :store  => 0b00011010,
      :rdrop  => 0b00011100,
      :pop    => 0b00011101,
      :pushpc => 0b00011110,
      :push   => 0b00011111,
      :jump   => 0b01000000,
      :call   => 0b01100000
   }.freeze

   CC_MAP = {
      :never  => 0b000,
      :eq     => 0b001,
      :lt     => 0b010,
      :gt     => 0b011,
      :always => 0b100,
      :ne     => 0b101,
      :ge     => 0b110,
      :le     => 0b111
   }.freeze

   def to_bytes
      finalize_target
      bytes = [0] * @align_bytes

      case @type
      when :align
         # no opcode
         nil
      when :nop, :rshift, :add, :sub, :nip,
           :and, :or, :xor, :swap, :dup
         bytes << (OPCODES[@type] | (@return ? 1 : 0))
      when :load, :store
         bytes << (OPCODES[@type] | (@byte_op ? 1 : 0))
      when :rdrop, :pop, :pushpc, :push, :rdup, :mult, :over, :nlz
         bytes << OPCODES[@type]
      when :immed, :resb, :byte, :word
         bytes += @immed
      when :jump, :call
         result = OPCODES[@type]
         result |= 0b10000 if @rel
         result |= (CC_MAP[@cc]<<1) if @cc
         result |= 0b00001 if @drop
         # Insert NOPs if immed is shorter than allocated
         pad = @immed && ([0] * (@bytes - (@immed.length + 1)))
         bytes += pad if pad
         bytes += @immed if @immed
         bytes << result
      when :merge
         bytes << (0b10000000 | (@value & 0x7f))
      else
         raise StandardError.new("Don't know how to output a #{@type} opcode.")
      end
      bytes
   end

   def to_s
      finalize_target

      s = @labels.map {|lbl| lbl.name + ":\n" }.join('')
      s << "  #{ADDR_FMT % @pc}  "

      case @type
      when :align
         s << ('%s %d' % [ @type, @align ])
      when :nop, :rshift, :add, :sub, :nip,
           :and, :or, :xor, :swap, :dup
         s << @type.to_s
         s << ' return' if @return
      when :load, :store
         s << @type.to_s
         s << ' byte' if @byte_op
      when :rdrop, :pop, :pushpc, :push, :rdup, :mult, :over, :nlz
         s << @type.to_s
      when :immed
         s << @type.to_s
         s << (' 0x%-8X' % (@value & 0xffffffff))
         s << '  # ' << @value.to_s << ' decimal'
      when :jump, :call
         s << @type.to_s
         s << ' ' << @target.name if @target
         s << ' (' << (ADDR_FMT % @value) << ')' if @value
         s << ' rel' if @rel
         s << ' ' << @cc.to_s if @cc
         s << ' drop' if @drop
      when :merge
         s << @type.to_s << (' 0b%07b' % @value)
      when :resb
         s << @type.to_s << ' ' << @immed.length.to_s
      when :byte
         s << @type.to_s << ' ' << @immed.join(' ')
      when :word
         words = @immed.pack('C*').unpack('N*').map { |w| '0x%08x' % w }
         s << @type.to_s << ' ' << words.join(' ')
      else
         raise StandardError.new("Don't know how to output a #{@type} opcode.")
      end

      s
   end
end

def parse_options(argv)
   options = {
      :include_path => [],
      :files => []
   }

   ARGV.each_with_index do |option,idx|
      case option
      when /^--origin=(#{POS_INT_REGEX})$/
         options[:origin] = $1.to_i(0)
      when /^--format=(binary|hex)$/
         options[:format] = $1.to_sym
         options[:pad] = true unless options.has_key? :pad
      when /^--format=(raw|write|base64)$/
         options[:format] = $1.to_sym
      when /^-p|--pad$/
         options[:pad] = true
      when /^-r|--no-pad|--raw$/
         options[:pad] = false
      when /^--(?:code|disassembly)=(.*)$/
         options[:code_file] = ($1=='-') ? $stdout : File.open($1, "w")
      when /^--sym(?:bol)?s=(.*)$/
         options[:symbol_file] = ($1=='-') ? $stdout : File.open($1, "w")
      when /^--ram(?:-size)?=(#{POS_INT_REGEX})$/
         options[:ram_size] = $1.to_i(0)
      when /^--partial$/
         options[:code_file] ||= $stdout
         options[:partial] = true
      when /^(?:--includes=|-I)(.*)$/
         options[:include_path] << File.expand_path($1)
      when /^--depends$/
         options[:depends_file] = $stdout
      when /^--$/
         options[:files] += ARGV[(idx+1) .. -1]
         break
      when /^-/
         raise StandardError.new("Unknown option '#{option}'.")
      else
         options[:files] << option
      end
   end

   return options
end

$options = {
   # Defaults
   :origin       => 0x400,
   :format       => :base64,
   :pad          => false,
   :code_file    => nil,
   :symbol_file  => nil,
   :ram_size     => (4 * 2048),
   :partial      => false,
   :depends_file => false
}.merge(parse_options(ARGV)).freeze

$pc = $options[:origin]
$max_pc = $options[:ram_size] - 1

$instructions = []
$accum_labels = []

OPERAND = /(?:'.'|"(?:[^\\"]|\\"|\\\\)*"|\S+)/
ADDR_FMT = "0x%0#{$max_pc.to_s(16).length}x"

$included_files = []

def read_asm_file(fname)
   old_file, old_line, $file, $line = $file, $line, fname, 1

   Label.in_file_scope do
      File.open(fname).each_line do |ln|
         ln.sub!(/#.*$/,'')

         ln.split(';').each do |part|
            while part.sub!(/^\s*(\S+)\s*:\s*/, '')
               $accum_labels << $1
            end

            case part
            when /^\s*$/
               nil # Empty line or comment

            when /^\s*include\s+"([^"]+)"\s*$/
               inc_name = $1
               search_path = [File.dirname(fname)] + $options[:include_path]
               found = nil

               search_path.each do |dir|
                  try_name = File.expand_path(inc_name, dir)
                  if File.exist? try_name
                     found = try_name
                     break
                  end
               end

               if found
                  $included_files << found
                  read_asm_file(found)
               else
                  raise StandardError.new("Unable to locate include file '#{inc_name}'.")
               end

            when /^\s*export\s+(#{OPERAND}(?:\s*#{OPERAND})*)\s*$/
               $1.scan(OPERAND).each do |sym|
                  if sym =~ /^\./
                     raise StandardError.new("Can't export local label '#{sym}'.")
                  else
                     lbl = Label.get(sym)
                     lbl.export = true
                     lbl.global = true
                  end
               end

            when /^\s*import\s+(#{OPERAND}(?:\s*#{OPERAND})*)\s*$/
               $1.scan(OPERAND).each do |sym|
                  if sym =~ /^\./
                     raise StandardError.new("Can't import local label '#{sym}'.")
                  else
                     Label.get(sym).global = true
                  end
               end

            when /^\s*(#{OPERAND})\s+eql\s+(#{OPERAND})\s*$/i
               name, value = $1, $2

               if value =~ /^-?#{POS_INT_REGEX}$/
                  value = value.to_i(0)

                  unless $labels[name] and $labels[name].pc == value
                     lbl = Label.define(name)
                     lbl.pc = value
                     lbl.global = !lbl.local?
                  end
               else
                  lbl = Label.alias(Label.get(value), name)
                  lbl.global = !lbl.local?
               end

            when /^\s*(\S+)((?:\s+#{OPERAND})*)\s*$/
               ins = Instruction.new($accum_labels,$pc,$1,$2.scan(/#{OPERAND}/))
               $accum_labels = []
               $instructions << ins
               $pc += ins.bytes
            end
         end

         $line += 1
      end
   end

   $file, $line = old_file, old_line
end

begin
   $options[:files].each { |fname| read_asm_file(fname) }
rescue => e
   $stderr.puts "Error at #{$file} line #{$line}:"
   $stderr.puts e
   $stderr.puts e.backtrace
end

if $options[:depends_file]
   $options[:depends_file].puts $included_files
   exit
end

if $instructions.empty? or not $accum_labels.empty?
   ins = Instruction.new($accum_labels, $pc, "align", ["1"])
   $instructions << ins
   $pc += ins.bytes
end

5.times do
   $max_pc = $pc
   $pc = $options[:origin]
   $layout_changed = false
   $instructions.each { |ins| ins.reflow($pc); $pc += ins.bytes }
   break unless $layout_changed
end

if $layout_changed
   $stderr.puts "WARNING: Opcode layout may be suboptimal."
end

if $options[:code_file]
   $instructions.each do |ins|
      $options[:code_file].puts ins
   end
end

$max_length = $labels.values.inject(0) do |max, lbl|
   len = lbl.name.length
   if lbl.export? and (len > max)
      len
   else
      max
   end
end

if $options[:symbol_file]
   file = $options[:symbol_file]
   ($labels.values.sort {|a,b| a.pc <=> b.pc }).each do |lbl|
      if lbl.pc and lbl.export?
         file.printf "%-*s EQL 0x%08x\n", $max_length, lbl.name, lbl.pc
      end
   end
end

unless $options[:partial]
   bytes = ($instructions.map { |ins| ins.to_bytes }).flatten

   if $options[:pad] and [ :raw, :binary, :hex ].member? $options[:format]
      bytes = ([0] * $options[:origin]) + bytes
      bytes += ([0] * ($options[:ram_size] - bytes.length))
   end

   file = $stdout

   case $options[:format]
   when :raw
      file.write bytes.pack('C*')
   when :base64
      file.puts('STORE %X' % $options[:origin])
      file.print Base64.encode64(bytes.pack('C*'))
      file.puts '.'
      ($labels.values.sort {|a,b| a.name <=> b.name }).each do |lbl|
         if lbl.pc and lbl.export?
            file.printf "# %-*s EQL 0x%08x\n", $max_length, lbl.name, lbl.pc
         end
      end
   when :binary, :hex, :write
      pc = $options[:origin]
      bytes = ([0] * (pc & 0x3)) + bytes
      pc &= ~0x3
      words = (bytes+[0,0,0]).pack('C*').unpack('N*')
      addr_len = $max_pc.to_s(16).length
      words.each do |w|
         if $options[:format] == :write
            file.printf "WRITE %08X %0*X\n", w, addr_len, pc
         elsif $options[:format] == :binary
            file.printf "%032b\n", w
         else # hex
            file.printf "%08X\n", w
         end

         pc += 4
      end
   else
      raise StandardError.new("Unknown output format")
   end
end