require 'voodoo/generators/nasm_generator'

module Voodoo
  # = AMD64 NASM Code Generator
  #
  # Code generator that emits NASM assembly code for AMD64 processors.
  # 
  # == Calling Convention
  #
  # The calling convention implemented by this code generator is
  # compatible with the System V ABI for AMD64, provided that all
  # arguments are integers or pointers.
  #
  # Arguments are passed in registers. The registers are used in the
  # following order:
  #
  # 1. +rdi+
  # 2. +rsi+
  # 3. +rdx+
  # 4. +rcx+
  # 5. +r8+
  # 6. +r9+
  #
  # Additional arguments are pushed on the stack, starting with the last
  # argument and working backwards. These arguments are removed from the
  # stack by the caller, after the called function returns.
  #
  # The return value is passed in +rax+.
  #
  # For varargs functions, +rax+ must be set to an upper bound on the
  # number of vector arguments. Since the code generator does not know
  # whether the called function is a varargs function, this is always
  # done. Since the code generator never passes any vector arguments,
  # this means +rax+ is set to +0+ before each call. 
  #
  # == Call Frames
  #
  #   arg_n
  #   :
  #   arg_7
  #   arg_6
  #   saved_rip
  #   saved_rbp <-- rbp
  #   arg_0
  #   arg_1
  #   :
  #   arg_5
  #   local_0
  #   local_1
  #   :
  #   local_n   <-- rsp
  #
  class AMD64NasmGenerator < NasmGenerator
    def initialize params = {}
      # Number of bytes in a word
      @WORDSIZE = 8
      # Word name in NASM lingo
      @WORD_NAME = 'qword'
      # Default alignment for code
      @CODE_ALIGNMENT = 0
      # Default alignment for data
      @DATA_ALIGNMENT = @WORDSIZE
      # Default alignment for functions
      @FUNCTION_ALIGNMENT = 16
      # Register used for return values
      @RETURN_REG = 'rax'
      # Register used as scratch register
      @SCRATCH_REG = 'r11'
      # Registers used for argument passing
      @ARG_REGS = ['rdi', 'rsi', 'rdx', 'rcx', 'r8', 'r9']
      # Accumulator index
      @AX = 'rax'
      # Base index
      @BX = 'rbx'
      # Count index
      @CX = 'rcx'
      # Data index
      @DX = 'rdx'
      super params
    end

    #
    # == Data Definition
    #

    # Define a machine word with the given value.
    def word value
      qword value
    end

    #
    # == Functions
    #

    # Call a function.
    def call func, *args
      emit "; call #{func} #{args.join ' '}\n"
      # First couple of arguments go in registers
      register_args = args[0..(number_of_register_arguments - 1)] || []
      # Rest of arguments go on the stack
      stack_args = args[number_of_register_arguments..-1] || []
      emit "; register_args: #{register_args.inspect}\n"
      emit "; stack_args: #{stack_args.inspect}\n"
      # Push stack arguments
      stack_args.reverse.each { |arg| push_qword arg }
      # Load register arguments
      register_args.each_with_index do |arg,i|
        register = @ARG_REGS[i]
        value_ref = load_value arg, register
        if value_ref != register
          emit "mov #{register}, #{value_ref}\n"
        end
      end
      # Call function
      value_ref = load_value func, @SCRATCH_REG
      emit "xor rax, rax\n"
      # If value_ref is a symbol, use PLT-relative addressing
      if global?(func)
        emit "call #{value_ref} wrt ..plt\n"
      else
        emit "call #{value_ref}\n"
      end
      # Clean up stack
      unless stack_args.empty?
        emit "add rsp, #{stack_args.length * @WORDSIZE}\n"
      end
    end

    # Emit function prologue.
    def emit_function_prologue formals = []
      emit "push rbp\nmov rbp, rsp\n"
      unless formals.empty?
        register_args = formals[0...number_of_register_arguments]
        register_args.each_with_index do |arg,i|
          emit "push #{@ARG_REGS[i]}\n"
        end
      end
    end

    # Call a function, re-using the current call frame if possible.
    def tail_call func, *args
      emit "; tail-call #{func} #{args.join ' '}\n"
      # Compute required number of stack words
      nstackargs = number_of_stack_arguments args.length
      # If we need more stack arguments than we have now,
      # perform a normal call and return
      if nstackargs > number_of_stack_arguments(@environment.args)
        emit "; Not enough space for proper tail call; using regular call\n"
        ret :call, func, *args
      end

      # If any arguments are going to be overwritten before they are
      # used, save them to new local variables and use those instead.
      i = args.length - 1
      while i >= -1
        arg = (i >= 0) ? args[i] : func

        if symbol?(arg)
          x = @environment[arg]
          if x && x[0] == :arg && x[1] < args.length && x[1] > i &&
              (i >= 0 || func != args[x[1]])
            # Save value
            newsym = @environment.gensym
            let newsym, arg
            # Change reference
            if i >= 0
              args[i] = newsym
            else
              func = newsym
            end
          end
        end
        i = i - 1
      end

      # Set stack arguments
      if args.length > number_of_register_arguments
        (args.length - 1).downto(number_of_register_arguments).each do |i|
          arg = args[i]
          
          value_ref = load_value arg, @SCRATCH_REG
          newarg_ref = load_arg i
          # Elide code if source is same as destination
          unless value_ref == newarg_ref
            emit "mov #{@SCRATCH_REG}, #{value_ref}\n"
            emit "mov #{newarg_ref}, #{@SCRATCH_REG}\n"
          end
        end
      end

      # Set register arguments
      number_of_register_arguments(args.length).times do |i|
        register = @ARG_REGS[i]
        load_value_into_register args[i], register
      end

      # Tail call
      func_ref = load_value func, @BX
      emit "leave\n"
      set_register @AX, 0
      # If func_ref is a symbol, use PLT-relative addressing
      if global?(func)
        emit "jmp #{func_ref} wrt ..plt\n"
      else
        emit "jmp #{func_ref}\n"
      end
    end

    #
    # == Loading Values
    #

    # Load the value of the nth argument
    def load_arg n, reg = @SCRATCH_REG
      if register_argument?(n)
        # Arguments that were originally passed in a register
        # are now below rbp
        "[rbp - #{(n + 1) * @WORDSIZE}]"
      else
        # Arguments that were originally passed on the stack
        # are now above rbp, starting 2 words above it
        "[rbp + #{(n + 2 - number_of_register_arguments) * @WORDSIZE}]"
      end
    end

    # Load the value of the nth local variable
    def load_local n, reg = @SCRATCH_REG
      # If there current function has any arguments,
      # local variables are offset by
      # number_of_register_arguments(number_of_arguments)
      # words.
      offset = number_of_register_arguments(@environment.args) * @WORDSIZE
      "[rbp - #{offset + (n + 1) * @WORDSIZE}]"
    end

    #
    # == Variables
    #

    # Introduce a new local variable
    def let symbol, *words
      emit "; let #{symbol} #{words.join ' '}\n"
      @environment.add_local symbol
      eval_expr words, @RETURN_REG
      emit "push #{@RETURN_REG}\n"
    end

    #
    # == Miscellaneous
    #

    # Load a value and push it on the stack.
    def push_qword value
      value_ref = load_value value, @SCRATCH_REG
      emit "push qword #{value_ref}\n"
    end

    # Calculate the number of register arguments,
    # given the total number of arguments.
    # If _n_ is +nil+, returns the maximum number of
    # register arguments.
    def number_of_register_arguments n = nil
      if n.nil?
        @ARG_REGS.length
      else
        [@ARG_REGS.length, n].min
      end
    end

    # Calculate the number of stack arguments,
    # given the total number of arguments.
    def number_of_stack_arguments n
      [0, n - number_of_register_arguments].max
    end

    # Tests if the nth argument is a register argument.
    def register_argument? n
      n < number_of_register_arguments
    end

  end

  # Register class
  Voodoo::CodeGenerator.register_generator AMD64NasmGenerator,
                                           :architecture => :amd64,
                                           :format => :nasm
end
