require 'voodoo/generators/common_code_generator'

module Voodoo
  # = NASM Code Generator
  #
  # The NASM code generator is a common base class for generators that output
  # assembly code for use with the {Netwide Assembler}[http://www.nasm.us/].
  #
  # This class is used by both the I386NasmGenerator and the
  # AMD64NasmGenerator, and contains the functionality that is common to
  # both.
  #
  # To use the functionality from this class, a subclass must define the
  # following methods and constants:
  # - #emit_function_prologue
  # - #load_arg
  # - #load_local
  # - @CODE_ALIGNMENT
  # - @DATA_ALIGNMENT
  # - @FUNCTION_ALIGNMENT
  # - @SCRATCH_REG
  # - @WORD_NAME
  # - @WORDSIZE
  # - @AX, @BX, @CX, and @DX
  class NasmGenerator < CommonCodeGenerator
    def initialize params = {}
      super params
      @if_labels = []
      @output_file_suffix = '.asm'
    end

    #
    # == Information About the Generator
    #

    # Returns the number of bits per word for this code generator.
    def wordsize
      @WORDSIZE * 8
    end

    #
    # == Labels
    #

    # Export symbols from the current section
    def export *symbols
      emit "global #{symbols.join ', '}\n"
    end

    # Continue execution at the given address
    def goto value
      emit "; goto #{value}\n"
      value_ref = load_value value, @SCRATCH_REG
      emit "jmp #{value_ref}\n"
    end

    # Import labels into the current section
    def import *symbols
      emit "extern #{symbols.join ', '}\n"
    end

    # Define a label in the current section
    def label name
      emit "#{name}:\n"
    end

    #
    # == Data definition
    #

    # Define a byte with the given value
    def byte value
      emit "db #{value}\n"
    end

    # Define a dword with the given value
    def dword value
      emit "dd #{value}\n"
    end

    # Define a qword with the given value
    def qword value
      emit "dq #{value}\n"
    end

    # Define a string with the given value
    def string value
      code = ''
      in_quote = false
      value.each_byte do |b|
        if b >= 32 && b < 128
          if in_quote
            code << b.chr
          else
            code << ',' unless code.empty?
            code << "'" + b.chr
            in_quote = true
          end
        else
          if in_quote
            code << "',#{b}"
            in_quote = false
          else
            code << ',' unless code.empty?
            code << "#{b}"
          end
        end
      end
      code << "'" if in_quote
      emit "db #{code}\n"
    end

    #
    # == Alignment
    #

    def align alignment = nil
      unless alignment
        # Get default alignment
        case @section
        when :code
          alignment = @CODE_ALIGNMENT
        when :data
          alignment = @DATA_ALIGNMENT
        when :function
          alignment = @FUNCTION_ALIGNMENT
        else
          # Use data alignment as default
          alignment = @DATA_ALIGNMENT
        end
      end
      emit "align #{alignment}\n" unless alignment == 0
    end

    #
    # == Functions
    #

    # Emit function preamble and declare _formals_ as function arguments
    def begin_function *formals
      emit "; function #{formals.join ' '}\n"
      environment = Environment.new @environment
      environment.add_args formals
      @environment = environment
      emit_function_prologue formals
    end

    # Emit function epilogue.
    def emit_function_epilogue formals = []
      emit "leave\n"
    end

    # End a function body
    def end_function
      emit "; end function\n\n"
      if @environment == @top_level
        raise "Cannot end function when not in a function"
      else
        @environment = @top_level
      end
    end

    # Return a from a function.
    # 
    # _words_ may contain an expression to be evaluated. The result
    # of the evaluation is returned from the function.
    def ret *words
      emit "; return #{words.join ' '}\n"
      eval_expr words
      emit_function_epilogue
      emit "ret\n"
    end

    #
    # == Conditionals
    #

    # End a conditional body
    def end_if
      label = @if_labels.pop
      emit "#{label}:\n"
    end

    #
    # == Value Classification
    #

    # Test if op is a binary operation
    def binop? op
      [:div, :mod, :sub].member?(op) || symmetric_operation?(op)
    end

    # Test if a value is an integer
    def integer? value
      value.kind_of? Integer
    end

    # Test if a symbol refers to a global
    def global? symbol
      symbol?(symbol) && @environment[symbol] == nil
    end

    # Tests if an operand is an immediate operand
    def immediate_operand? operand
      integer?(operand) || global?(operand)
    end

    # Tests if an operand is a memory operand
    def memory_operand? operand
      operand[0] == ?[
    end

    # Test if a value is a symbol
    def symbol? value
      value.kind_of? Symbol
    end

    # Test if op is a symmetric operation (i.e. it will yield the
    # same result if the order of its source operands is changed).
    def symmetric_operation? op
      [:add, :and, :mul, :or, :xor].member? op
    end

    #
    # == Loading Values
    #

    # Create a value reference to an address.
    # Invoking this code may clobber @BX and/or @CX
    def load_address base, offset, scale
      base_ref = load_value base, @BX
      offset_ref = load_value offset, @CX

      if offset_ref == 0
        if integer? base_ref
          # Only an integer base
          "[#{base_ref}]"
        else
          # Some complex base; load in @BX
          emit "mov #{@BX}, #{base_ref}\n"
          "[#{@BX}]"
        end
      elsif base_ref == 0
        if integer? offset_ref
          # Only a scaled offset
          "[#{offset_ref.to_i * scale}]"
        else
          # Some complex offset; load in @CX
          emit "mov #{@CX}, #{offset_ref}\n"
          "[#{@CX} * #{scale}]"
        end
      elsif integer? base_ref
        if integer? offset_ref
          # All integers, combine them together
          "[#{base_ref.to_i + (offset_ref.to_i * scale)}]"
        else
          # Complex offset; use @CX
          emit "mov #{@CX}, #{offset_ref}\n"
          "[#{base_ref} + #{@CX} * #{scale}]"
        end
      elsif integer? offset_ref
        # Complex base, integer offset; use @BX
        emit "mov #{@BX}, #{base_ref}\n"
        "[#{@BX} + #{offset_ref.to_i * scale}]"
      else
        # Both base and offset are complex
        # Use both @BX and @CX
        emit "mov #{@BX}, #{base_ref}\n"
        emit "mov #{@CX}, #{offset_ref}\n"
        "[#{@BX} + #{@CX} * #{scale}]"
      end
    end

    # Load the value associated with the given symbol.
    # Returns a string that can be used to refer to the loaded value.
    def load_symbol symbol, reg = @SCRATCH_REG
      x = @environment[symbol]
      if x
        case x[0]
        when :arg
          load_arg x[1], reg
        when :local
          load_local x[1], reg
        else
          raise "Invalid variable type: #{x[0]}"
        end
      else
        # Assume global
        symbol.to_s
      end
    end
    
    # Load a value.
    # Returns a string that can be used to refer to the loaded value.
    def load_value value, reg = @SCRATCH_REG
      if integer? value
        # Integers can be used as is
        value
      elsif symbol? value
        load_symbol value, reg
      end
    end

    #
    # == Storing Values
    #

    # Evaluate the expr in words and store the result in target
    def set target, *words
      if integer? target
        raise "Cannot change value of integer #{target}"
      elsif global?(target)
        raise "Cannot change value of global #{target}"
      end

      emit "; set #{target} #{words.join ' '}\n"
      if words.length == 1
        if words[0] == target
          emit "; nothing to do; destination equals source\n"
        else
          target_ref = load_value target, @BX
          if integer?(words[0])
            if words[0].to_i == 0
              # Set destination to 0
              emit "xor #{@AX}, #{@AX}\n"
              emit "mov #{target_ref}, #{@AX}\n"
            else
              # Load immediate
              emit "mov #{@WORD_NAME} #{target_ref}, #{words[0]}\n"
            end
          else
            # Copy source to destination
            eval_expr words, @RETURN_REG
            emit "mov #{target_ref}, #{@RETURN_REG}\n"
          end
        end
      else
        op = words[0]

        if words.length == 3 && binop?(op)
          # Binary operation
          binop op, target, words[1], words[2]
        else
          # Not a binary operation
          eval_expr words, @RETURN_REG
          target_ref = load_value target, @BX
          emit "mov #{target_ref}, #{@RETURN_REG}\n"
        end
      end
    end

    # Set the byte at _base_ + _offset_ to _value_
    def set_byte base, offset, value
      emit "; set-byte #{base} #{offset} #{value}\n"
      value_ref = load_value value, @RETURN_REG
      addr_ref = load_address base, offset, 1
      emit "mov byte #{addr_ref}, #{value_ref}\n"
    end

    # Set the word at _base_ + _offset_ * +@WORDSIZE+ to _value_
    def set_word base, offset, value
      emit "; set-word #{base} #{offset} #{value}\n"
      if immediate_operand?(value)
        value_ref = value
      else
        load_value_into_register value, @RETURN_REG
        value_ref = @RETURN_REG
      end
      addr_ref = load_address base, offset, @WORDSIZE
      emit "mov #{@WORD_NAME} #{addr_ref}, #{value_ref}\n"
    end

    #
    # == Binary Operations
    #

    # Emit code for a binary operation
    def binop op, target, x, y
      if target == x
        binop2 op, target, y
      elsif symmetric_operation?(op) && y == target
        binop2 op, target, x
      else
        # Cases that are handled specially
        return div(target, x, y) if op == :div
        return mod(target, x, y) if op == :mod
        return mul(target, x, y) if op == :mul

        target_ref = load_value target, @AX
        x_ref = load_value x, @DX
        y_ref = load_value y, @BX

        if memory_operand?(target_ref)
          if memory_operand?(x_ref) || memory_operand?(y_ref)
            emit "mov #{@CX}, #{x_ref}\n"
            emit "#{op} #{@CX}, #{y_ref}\n"
            emit "mov #{target_ref}, #{@CX}\n"
          else
            emit "mov #{@WORD_NAME} #{target_ref}, #{x_ref}\n"
            emit "#{op} #{@WORD_NAME} #{target_ref}, #{y_ref}\n"
          end
        else
          raise "Can't happen: target_ref is #{target_ref.inspect}"
        end
      end
    end

    # Emit code for a binary operation where the first operand
    # is also the target
    def binop2 op, target, y
      # Cases that are handled specially
      return div2(target, target, y) if op == :div
      return mod2(target, y) if op == :mod
      return mul2(target, y) if op == :mul

      target_ref = load_value target, @BX
      y_ref = load_value y, @DX
      if memory_operand?(target_ref) && memory_operand?(y_ref)
        emit "mov #{@AX}, #{y_ref}\n"
        emit "#{op} #{target_ref}, #{@AX}\n"
      else
        emit "#{op} #{@WORD_NAME} #{target_ref}, #{y_ref}\n"
      end
    end

    # Divide x by y and store the quotient in target
    def div target, x, y
      eval_div x, y
      target_ref = load_value target, @BX
      emit "mov #{target_ref}, #{@AX}\n"
    end

    # Divide target by x and store the quotient in target
    def div2 target, x
      div target, target, x
    end

    # Divide x by y and store the remainder in target
    def mod target, x, y
      eval_div x, y
      target_ref = load_value target, @BX
      emit "mov #{target_ref}, #{@DX}\n"
    end

    # Divide target by x and store the remainder in target
    def mod2 target, x
      mod target, target, x
    end

    # Multiply x by y and store the result in target
    def mul target, x, y
      eval_mul x, y
      target_ref = load_value target, @BX
      emit "mov #{target_ref}, #{@RETURN_REG}\n"      
    end

    # Multiply target by x and store the result in target
    def mul2 target, x
      mul target, target, x
    end

    #
    # == Expressions
    #

    # Perform division.
    # The quotient is stored in @AX, the remainder in @DX.
    def eval_div x, y
      x_ref = load_value_into_register x, @AX
      y_ref = load_value y, @SCRATCH_REG
      emit "mov #{@DX}, #{@AX}\n"
      emit "sar #{@DX}, #{@WORDSIZE * 8 - 1}\n"
      if immediate_operand?(y_ref)
        set_register @BX, y_ref
        emit "idiv #{@BX}\n"
      else
        emit "idiv #{@WORD_NAME} #{y_ref}\n"
      end
    end

    # Evaluate an expression.
    # The result is stored in _register_ (@RETURN_REG by default).
    def eval_expr words, register = @RETURN_REG
      if words.length == 1
        if words[0] == 0
          emit "xor #{register}, #{register}\n"
        else
          load_value_into_register words[0], register
        end
      else
        op = words[0]
        case op
        when :call
          call *words[1..-1]
        when :div
          eval_div words[1], words[2]
          set_register register, @AX
        when :'get-byte'
          # Clear register
          set_register register, 0
          # Get address reference
          address_ref = load_address words[1], words[2], 1
          # Load byte from address
          case register
          when 'eax', 'rax'
            set_register 'al', address_ref
          when 'ebx', 'rbx'
            set_register 'bl', address_ref
          when 'ecx', 'rcx'
            set_register 'cl', address_ref
          when 'edx', 'rdx'
            set_register 'dl', address_ref
          else
            set_register @BX, 0
            set_register 'bl', address_ref
            set_register register, @BX
          end
        when :'get-word'
          address_ref = load_address words[1], words[2], @WORDSIZE
          set_register register, address_ref
        when :mod
          eval_div words[1], words[2]
          set_register register, @DX
        when :mul
          eval_mul words[1], words[2], register
        when :not
          load_value_into_register words[1], register
          emit "not #{register}\n"
        else
          if binop?(op)
            x_ref = load_value words[1], @DX
            y_ref = load_value words[2], @BX
            emit "mov #{register}, #{x_ref}\n"
            emit "#{op} #{register}, #{y_ref}\n"
          else
            raise "Not a magic word: #{words[0]}"
          end
        end
      end
    end

    # Multiply x by y.
    # The result is stored in @AX by default, but
    # a different register can be specified by passing
    # a third argument.
    def eval_mul x, y, register = @AX
      x_ref = load_value x, @DX
      y_ref = load_value y, @BX

      if immediate_operand? x_ref
        if immediate_operand? y_ref
          set_register register, x_ref * y_ref
        else
          emit "imul #{register}, #{y_ref}, #{x_ref}\n"
        end
      elsif immediate_operand? y_ref
        emit "imul #{register}, #{x_ref}, #{y_ref}\n"
      else
        emit "mov #{register}, #{x_ref}\n"
        emit "imul #{register}, #{y_ref}\n"
      end
    end

    #
    # == Conditionals
    #

    # Start a conditional using the specified branch instruction
    # after the comparison.
    def common_if branch, x, y = nil
      load_value_into_register y, @DX if y
      load_value_into_register x, @AX
      truelabel = @environment.gensym
      falselabel = @environment.gensym
      @if_labels.push falselabel

      emit "cmp #{@AX}, #{@DX}\n"
      emit "#{branch} #{truelabel}\n"
      emit "jmp #{falselabel}\n"
      emit "#{truelabel}:\n"
    end

    # End a conditional.
    def end_if
      label = @if_labels.pop
      emit "#{label}:\n"
    end

    # Start the false path of a conditional.
    def ifelse
      emit "; else\n"
      newlabel = @environment.gensym
      emit "jmp #{newlabel}\n"
      label = @if_labels.pop
      emit "#{label}:\n"
      @if_labels.push newlabel
    end

    # Test if x is equal to y
    def ifeq x, y
      emit "; ifeq #{x} #{y}\n"
      common_if 'je', x, y
    end

    # Test if x is greater than or equal to y
    def ifge x, y
      emit "; ifge #{x} #{y}\n"
      common_if 'jge', x, y
    end

    # Test if x is strictly greater than y
    def ifgt x, y
      emit "; ifgt #{x} #{y}\n"
      common_if 'jg', x, y
    end

    # Test if x is less than or equal to y
    def ifle x, y
      emit "; ifle #{x} #{y}\n"
      common_if 'jle', x, y
    end

    # Test if x is strictly less than y
    def iflt x, y
      emit "; iflt #{x} #{y}\n"
      common_if 'jl', x, y
    end

    # Test if x different from y
    def ifne x, y
      emit "; ifne #{x} #{y}\n"
      common_if 'jne', x, y
    end

    #
    # == Miscellaneous
    #

    # Emit a comment
    def comment text
      emit ";#{text}\n"
    end

    # Load a value into a register
    def load_value_into_register value, register
      value_ref = load_value value, register
      set_register register, value_ref
    end

    # Set a register to a value.
    # The value must be a valid operand to the mov instruction.
    def set_register register, value_ref
      case value_ref
      when register
        # Nothing to do
      when 0
        emit "xor #{register}, #{register}\n"
      else
        emit "mov #{register}, #{value_ref}\n"
      end
    end

    #
    # == Output
    #

    # Write generated code to the given IO object.
    def write io
      io.puts "bits #{@WORDSIZE * 8}\n\n"
      @sections.each do |section,code|
        unless code.empty?
          io.puts "section #{section.to_s}"
          io.puts code
          io.puts
        end
      end
    end

  end
end
