require 'voodoo/generators/common_code_generator'

module Voodoo
  # = MIPS GNU Assembler Code Generator
  #
  # The MIPS code generator generates i386 assembly code for use with
  # the GNU assembler.
  #
  # == Calling Convention
  #
  # The first four arguments are passed in the registers $4 through $7.
  # Any additional arguments are passed on the stack, starting at
  # $sp + 16. Words $sp through $sp + 12 will be available for the called
  # function to use. $sp will always be a multiple of 8.
  #
  # The return address for the called function is passed in $31.
  #
  # When performing a position-independent call, the address of the called
  # function is passed in $25.
  #
  # The called function will store its return value in $2.
  #
  # The called function is required to preserve the values of registers
  # $16 through $23 and register $30.
  #
  # This calling convention is compatible with the System V ELF ABI.
  #
  # == Call Frames
  #
  # Call frames have the following layout:
  #
  # When a function is called, it receives a stack frame that looks like
  # the following:
  #
  #   :
  #   old frame
  #   argn
  #   :
  #   arg4
  #   empty3
  #   empty2
  #   empty1
  #   empty0    <-- $sp points here
  #
  # The function prologue of functions generated by this code generator
  # creates activiation frames that look as follows:
  #
  # :
  # old frame
  # argn
  # :
  # arg4
  # arg3
  # arg2
  # arg1
  # arg0
  # return address
  # pointer to Global Offset Table
  # local0
  # :
  # local7
  # local8
  # :
  # localn
  # padding    <-- $sp points here
  #
  # In words:
  #
  # The four empty slots provided by the caller are used to store
  # arguments 0 through 3 (originally passed in $4 through $7),
  # if necessary.
  #
  # The stack frame created below the four slots provided by the caller
  # contains the following data, in order:
  #
  # - Saved return address (originally passed in $31), if necessary.
  #
  # - Saved pointer to global offset table (computed from $25), if necessary.
  #
  # - Saved values of caller's locals 0 through 7
  #   (originally in $16 through $23), if necessary.
  #
  # - Values of our locals > 8, if necessary.
  #
  # - Padding to align $sp on a multiple of 8, if necessary.
  #
  #
  # In accordance with the System V ELF ABI for MIPS, the pointer to
  # the global offset table is calculated as follows:
  #
  # $gp = _gp_disp + $25
  #
  # where $gp is the register to store the pointer, $25 is the register
  # holding the address of the called function, and _gp_disp is the
  # offset between the beginning of the function and the global offset table.
  #
  class MIPSGasGenerator < CommonCodeGenerator
    def initialize params
      @WORDSIZE = 4
      @CODE_ALIGNMENT = 0
      @DATA_ALIGNMENT = @WORDSIZE
      @FUNCTION_ALIGNMENT = @WORDSIZE

      @INITIAL_FRAME_SIZE = 2 * @WORDSIZE
      @REGISTER_ARG_BASE = 4
      @NREGISTER_ARGS = 4
      @REGISTER_LOCAL_BASE = 16
      @NREGISTER_LOCALS = 8
      @RETURN = :'$2'
      @TEMPORARY = :'$1'
      @FUNCTION = :'$25'
      @GOT = :'$28'
      @TEMPORARIES = [:'$8', :'$9', :'$10', :'$11',
                      :'$12', :'$13', :'$14', :'$15']
      @function_end_label = nil
      @imports = {}
      @if_labels = []
      super params
      @output_file_suffix = '.s'
    end

    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

    # Return the offset from $sp at which the nth (0-based) argument is
    # stored.
    def arg_offset n
      @frame_size + n * @WORDSIZE
    end

    # Return an $sp-relative reference for the nth (0-based) argument
    def arg_reference n
      "#{arg_offset n}($sp)"
    end

    # Return the register in which the nth (0-based) argument is stored, or
    # nil if not stored in a register
    def arg_register n
      if register_arg? n
        "$#{@REGISTER_ARG_BASE + n}"
      else
        nil
      end
    end

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

    # Emit function prologue and declare _formals_ as function arguments
    def begin_function formals, nlocals
      if @environment != @top_level
        raise "Cannot begin function when already in a function"
      end

      emit "# function #{formals.join ' '}\n"
      environment = Environment.new @environment
      environment.add_args formals
      @environment = environment
      @function_end_label = gensym
      emit_function_prologue formals, nlocals
    end

    # Test if op is a binary operation
    def binop? op
      assymetric_binop?(op) || symmetric_binop?(op)
    end

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

    # Call a function.
    def call func, *args
      emit "# call #{func} #{args.join ' '}\n"
      # Grow the stack frame, subject to 3 conditions:
      # 1. There must be enough space to hold all arguments
      # 2. $sp must remain a multiple of 8
      # 3. We must provide space for at least 4 words
      increment = ([args.length, 4].max * @WORDSIZE + 7) / 8 * 8
      emit "addiu $sp, $sp, -#{increment}\n"
      @frame_size = @frame_size + increment

      # Put arguments in the right places
      n = 0
      while n < args.length
        if n < @NREGISTER_ARGS
          # Put arguments up to @NREGISTER_ARGS in registers
          load_value_into_register args[n], arg_register(n)
        else
          # Put other arguments on the stack
          load_value_into_register args[n], @TEMPORARY
          emit "sw #{@TEMPORARY}, #{n * @WORDSIZE}($sp)\n"
        end
        n = n + 1
      end

      # Load function address
      if global? func
        if @imports.has_key? func
          # Load address from global offset table
          emit "lw #{@FUNCTION}, %call16(#{func})(#{@GOT})\n"
        else
          # Assume label defined in this module
          emit "lui #{@FUNCTION}, %hi(#{func})\n"
          emit "addiu #{@FUNCTION}, %lo(#{func})\n"
        end
      else
        # Load address
        load_value_into_register func, "#{@FUNCTION}"
      end

      # Perform call
      emit "jalr #{@FUNCTION}\n"
      # brach delay slot
      emit "nop\n"

      # Restore original stack frame
      emit "addiu $sp, $sp, #{increment}\n"
      @frame_size = @frame_size - increment

      # restore gp
      emit "lw #{@GOT}, #{@frame_size - 2 * @WORDSIZE}($sp)\n"
    end

    # Start a conditional using the specified branch instruction
    # after the comparison.
    def common_if comp, x, y = nil
      emit "# #{comp} #{x} #{y}\n"

      temporaries = @TEMPORARIES.dup
      xreg = load_value x, temporaries[0]
      temporaries.delete xreg
      yreg = load_value y, temporaries[0]
      temporaries.delete yreg

      falselabel = @environment.gensym
      @if_labels.push falselabel

      case comp
      when :ifeq
        emit "bne #{xreg}, #{yreg}, #{falselabel}\n"
      when :ifge
        emit "slt #{@TEMPORARY}, #{xreg}, #{yreg}\n"
        emit "bne #{@TEMPORARY}, $0, #{falselabel}\n"
      when :ifgt
        emit "slt #{@TEMPORARY}, #{yreg}, #{xreg}\n"
        emit "beq #{@TEMPORARY}, $0, #{falselabel}\n"
      when :ifle
        emit "slt #{@TEMPORARY}, #{yreg}, #{xreg}\n"
        emit "bne #{@TEMPORARY}, $0, #{falselabel}\n"
      when :iflt
        emit "slt #{@TEMPORARY}, #{xreg}, #{yreg}\n"
        emit "beq #{@TEMPORARY}, $0, #{falselabel}\n"
      when :ifne
        emit "beq #{xreg}, #{yreg}, #{falselabel}\n"
      else
        raise "Unknown conditional: #{comp}"
      end
      emit "nop\n"
    end

    # Emit function prologue.
    def emit_function_prologue formals = [], nlocals = 0
      # Calculate new value for $gp
      emit "lui #{@GOT}, %hi(_gp_disp)\n"
      emit "addiu #{@GOT}, %lo(_gp_disp)\n"
      emit "addu #{@GOT}, #{@GOT}, #{@FUNCTION}\n"

      # Save parameters 0 .. 3 in the stack space that the caller
      # should have allocated for them.
      [@NREGISTER_ARGS, formals.length].min.times do |n|
        emit "sw $#{n + @REGISTER_ARG_BASE}, #{n * @WORDSIZE}($29)\n"
      end

      # Create initial frame
      @frame_size = (@INITIAL_FRAME_SIZE + nlocals * @WORDSIZE + 7) / 8 * 8
      emit "addiu $sp, $sp, -#{@frame_size}\n"

      # Save return address
      emit "sw $31, #{@frame_size - @WORDSIZE}($sp)\n"

      # Save gp
      emit "sw #{@GOT}, #{@frame_size - 2 * @WORDSIZE}($sp)\n"
    end

    # End a function body
    def end_function
      if @environment == @top_level
        raise "Cannot end function when not in a function"
      end

      emit "# function epilogue\n"

      label @function_end_label
      # Restore saved locals
      [@NREGISTER_LOCALS, @environment.locals].min.times do |n|
        emit "lw #{local_register n}, #{local_reference n}\n"
      end
      # Load return address
      emit "lw $31, #{@frame_size - @WORDSIZE}($sp)\n"
      # Restore stack pointer
      emit "addiu $sp, $sp, #{@frame_size}\n"
      # Return
      emit "jr $31\n"
      emit "nop\n"

      emit "# end function\n\n"

      @function_end_label = nil
      @environment = @top_level
    end

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

    # Evaluate the binary operation expr and store the result in register
    def eval_binop expr, register
      temporaries = @TEMPORARIES.dup
      x = load_value expr[1], temporaries[0]
      temporaries.delete x
      y = load_value expr[2], temporaries[0]
      temporaries.delete y

      case expr[0]
      when :div
        emit "div $0, #{x}, #{y}\n"
        emit "mflo #{register}\n"
      when :mod
        emit "div $0, #{x}, #{y}\n"
        emit "mfhi #{register}\n"
      when :mul
        emit "mult #{x}, #{y}\n"
        emit "mflo #{register}\n"
      else
        emit "#{expr[0]} #{register}, #{x}, #{y}\n"
      end
    end

    # Evaluate the expression expr and store the result in register
    def eval_expr expr, register
      if expr.length == 1
        # Load value
        load_value_into_register expr[0], register
      else
        # Evaluate expression
        op = expr[0]
        case op
        when :call
          call *expr[1..-1]
          emit "move #{register}, #{@RETURN}\n" if register != @RETURN
        when :'get-byte'
          get_byte expr[1], expr[2], register
        when :'get-word'
          get_word expr[1], expr[2], register
        when :not
            eval_binop [:nor, 0, expr[1]], register
        else
          if binop? op
            eval_binop expr, register
          else
            raise "Not a magic word: #{op}"
          end
        end
      end
    end

    # Export symbols from the current section
    def export *symbols
      symbols.each { |sym| emit ".globl #{sym}\n" }
    end

    # Add a function to the current section
    def function formals, *code
      # Set nlocals to the number of local variables in the function
      nlocals = 0
      code.each { |action| nlocals = nlocals + 1 if action[0] == :let }
      begin_function formals, nlocals
      code.each { |action| add section, action }
      end_function
    end

    # Load byte from _base_ + _offset_ into _register_
    def get_byte base, offset, register
      # If base is an integer, but offset isn't, swap them
      if !integer?(offset) && integer?(base)
        base, offset = [offset, base]
      end

      if integer? offset
        base_reg = load_value base
        emit "lb #{register}, #{offset}(#{base_reg})\n"
      else
        eval_binop [:add, base, offset], @TEMPORARY
        emit "lb #{register}, 0(#{@TEMPORARY})\n"
      end
    end

    # Load word from _base_ + _offset_ * _@WORDSIZE_ into _register_
    def get_word base, offset, register
      if integer? offset
        base_reg = load_value base
        emit "lw #{register}, #{offset * @WORDSIZE}(#{base_reg})\n"
      else
        offset_reg = @TEMPORARIES.pop
        begin
          load_value_into_register offset, offset_reg
          base_reg = load_value base
          emit "sll #{offset_reg}, #{offset_reg}, 2\n" # multiply by @WORDSIZE
          emit "add #{@TEMPORARY}, #{base_reg}, #{offset_reg}\n"
          emit "lw #{register}, 0(#{@TEMPORARY})\n"
        ensure
          @TEMPORARIES.push offset_reg
        end
      end
    end

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

    # Jump to a label.
    def goto label
      emit "j #{label}\n"
      emit "nop\n"
    end

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

    # Test if x is equal to y
    def ifeq x, y
      common_if :ifeq, x, y
    end

    # Test if x is greater than or equal to y
    def ifge x, y
      common_if :ifge, x, y
    end

    # Test if x is strictly greater than y
    def ifgt x, y
      common_if :ifgt, x, y
    end

    # Test if x is less than or equal to y
    def ifle x, y
      common_if :ifle, x, y
    end

    # Test if x is strictly less than y
    def iflt x, y
      common_if :iflt, x, y
    end

    # Test if x different from y
    def ifne x, y
      common_if :ifne, x, y
    end

    # Import labels into the current section
    def import *symbols
      # Record imported labels in @imports
      symbols.each { |sym| @imports[sym] = sym }
    end

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

    # Emit a label
    def label name
      emit "#{name}:\n"
    end

    # Introduce a new local variable
    def let symbol, *expr
      emit "# let #{symbol} #{expr.join ' '}\n"
      n = @environment.locals
      @environment.add_local symbol

      # Extend stack frame, if necessary
      if @environment.locals > max_locals
        $stderr.puts "WARNING: Stack frame too small. Possible BUG"
        # Extend frame by 8 bytes, to keep $sp a multiple of 8
        emit "addiu $sp, $sp, -8\n"
        @frame_size = @frame_size + 8
      end

      register = local_register n
      ref = local_reference n
      if register
        # Save current value of register
        emit "sw #{register}, #{ref}\n"
        # Set new value
        eval_expr expr, register
      else
        eval_expr expr, @TEMPORARY
        emit "sw #{@TEMPORARY}, #{ref}\n"
      end
    end

    # Load a value into a register.
    # Returns the name of the register.
    # If the value was already in a register, the name of that
    # register is returned.
    # Else, the value is loaded into a register and the name of
    # that register is returned. The register to use in that case
    # may be specified using the optional second argument.
    def load_value x, register = @TEMPORARY
      if x == 0
        return "$0"
      elsif integer? x
        if x >= -32768 && x <= 32767
          emit "addiu #{register}, $0, #{x}\n"
          return register
        else
          emit "lui #{register}, #{x >> 16}\n"
          emit "ori #{register}, #{register}, #{x & 0xffff}\n"
          return register
        end
      elsif symbol? x
        binding = @environment[x]
        if binding
          case binding[0]
          when :arg
            emit "lw #{register}, #{arg_reference binding[1]}\n"
            return register
          when :local
            n = binding[1]
            if register_local? n
              return local_register(n)
            else
              emit "lw #{register}, #{local_reference n}\n"
              return register
            end
          else
            raise "Don't know how to load #{x.inspect}"
          end
        else
          # Assume global
          emit "lui #{register}, %hi(#{x})\n"
          emit "addiu #{register}, %lo(#{x})\n"
          return register
        end
      else
        raise "Don't know how to load #{x.inspect}"
      end
    end

    # Load a value into a specific register
    def load_value_into_register x, register
      reg = load_value x, register
      if reg != register
        emit "move #{register}, #{reg}\n"
      end
    end

    # Return the offset from $sp at which the nth (0-based) local variable is
    # stored. For register locals this is the offset at which the saved
    # value (from the calling function) is stored.
    def local_offset n
      @frame_size - @INITIAL_FRAME_SIZE - (n + 1) * @WORDSIZE
    end

    # Return an $sp-relative reference for the nth (0-based) local
    def local_reference n
      "#{local_offset n}($sp)"
    end

    # Return the register in which the nth local (0-based) is stored, or
    # nil if not stored in a register
    def local_register n
      if register_local? n
        "$#{@REGISTER_LOCAL_BASE + n}"
      else
        nil
      end
    end

    # Compute the maximum number of locals that would fit in the current
    # frame size.
    def max_locals
      (@frame_size - @INITIAL_FRAME_SIZE) / @WORDSIZE
    end

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

    # Returns true if the nth (0-based) argument is stored in a register
    def register_arg? n
      n < @NREGISTER_ARGS
    end

    # Returns true if the nth (0-based) local is stored in a register
    def register_local? n
      n < @NREGISTER_LOCALS
    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"
      # Compute return value and store it in @RETURN
      eval_expr words, @RETURN
      # Go to epilogue
      goto @function_end_label
    end
    
    # Set a variable to the result of evaluating an expression
    def set symbol, *expr
      emit "# set #{symbol} #{expr.join ' '}\n"

      x = @environment[symbol]
      if x == nil
        raise "Cannot change value of constant #{symbol}"
      end

      if x[0] == :local
        register = local_register x[1]
      else
        register = nil
      end

      if register
        # Set new value
        eval_expr expr, register
      else
        case x[0]
        when :local
          ref = local_reference x[1]
        when :arg
          ref = arg_reference x[1]
        else
          raise "??? #{sym} is neither a local nor an argument"
        end
        eval_expr expr, @TEMPORARY
        emit "sw #{@TEMPORARY}, #{ref}\n"
      end
    end

    # Set the byte at _base_ + _offset_ to _value_
    def set_byte base, offset, value
      emit "# set-byte #{base} #{offset} #{value}\n"
      # If base is an integer, but offset isn't, swap them
      if !integer?(offset) && integer?(base)
        base, offset = [offset, base]
      end

      value_reg = @TEMPORARIES.pop
      load_value_into_register value, value_reg
      begin
        if integer? offset
          base_reg = load_value base
          emit "sb #{value_reg}, #{offset}(#{base_reg})\n"
        else
          eval_binop [:add, base, offset], @TEMPORARY
          emit "sb #{value_reg}, 0(#{@TEMPORARY})\n"
        end
      ensure
        @TEMPORARIES.push value_reg
      end
    end

    # Set the word at _base_ + _offset_ * +@WORDSIZE+ to _value_
    def set_word base, offset, value
      emit "# set-word #{base} #{offset} #{value}\n"

      value_reg = @TEMPORARIES.pop
      load_value_into_register value, value_reg
      begin
        if integer? offset
          base_reg = load_value base
          emit "sw #{value_reg}, #{offset * @WORDSIZE}(#{base_reg})\n"
        else
          offset_reg = @TEMPORARIES.pop
          begin
            load_value_into_register offset, offset_reg
            base_reg = load_value base
            emit "sll #{offset_reg}, #{offset_reg}, 2\n"
            emit "add #{@TEMPORARY}, #{base_reg}, #{offset_reg}\n"
            emit "sw #{value_reg}, 0(#{@TEMPORARY})\n"
          ensure
            @TEMPORARIES.push offset_reg
          end
        end
      ensure
        @TEMPORARIES.push value_reg
      end
    end

    # Define a string with the given value
    def string value
      code = ''
      value.each_byte do |b|
        if b >= 32 && b < 128
          code << b.chr
        else
          code << sprintf("\\%03o", b)
        end
      end
      emit ".ascii \"#{code}\"\n"
    end

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

    # Test if op is a symmetric binary operation (i.e. it will yield the
    # same result if the order of its source operands is changed).
    def symmetric_binop? op
      [:add, :and, :mul, :or, :xor].member? op
    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 number of stack arguments
      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

      # TODO

      # Back up any stack arguments we will need after we overwrite them
      temporaries = @TEMPORARIES.dup
      nlocals = @environment.locals
      (@NREGISTER_ARGS...args.length).each do |i|
        x = @environment[args[i]]
        if x && x[0] == :arg && x[1] >= @NREGISTER_ARGS && x[1] < i
          # args[i] is a stack arg, but has an index < i
          # That means that, by the time we get to pass arg[i] to the called
          # function, it will have been overwritten. So we make a backup.
          if temporaries.empty?
            # Oh dear, we're out of temporaries.
            # Store the value in the current stack frame, extending it
            # as necessary.
            if nlocals >= max_locals
              emit "addiu $sp, -8\n"
              @frame_size = @frame_size + 8
            end
            reg = load_value args[i]
            emit "sw #{reg}, #{local_reference nlocals}\n"
            args[i] = [:local, nlocals]
            nlocals = nlocals + 1
          else
            # Load the argument into a temporary
            reg = temporaries.shift
            load_value_into_register args[i], reg
            # Update args[i] so we know how to get to it later
            args[i] = [:reg, reg]
          end
        end
      end

      # Set stack arguments
      (@NREGISTER_ARGS...args.length).each do |i|
        arg = args[i]
        reg = @TEMPORARY

        # Load value
        if arg.kind_of? Array
          # Special cases, created above
          case arg[0]
          when :reg
            reg = arg[1]
          when :local
            emit "lw #{reg}, #{local_reference arg[1]}\n"
          end
        else
          load_value_into_register arg, reg
        end

        # Store value in designated place
        emit "sw #{reg}, #{arg_reference i}\n"
      end

      # Set register arguments
      [@NREGISTER_ARGS,args.length].min.times do |i|
        reg = arg_register i
        load_value_into_register args[i], reg
      end

      # Restore saved registers
      [@NREGISTER_LOCALS, @environment.locals].min.times do |n|
        emit "lw #{local_register n}, #{local_reference n}\n"
      end
      # Restore return address
      emit "lw $31, #{@frame_size - @WORDSIZE}($sp)\n"

      # Adjust stack pointer
      emit "addiu $sp, #{@frame_size}\n"

      # Load function address
      if global? func
        if @imports.has_key? func
          # Load address from global offset table
          emit "lw #{@FUNCTION}, %call16(#{func})(#{@GOT})\n"
        else
          # Assume label defined in this module
          emit "lui #{@FUNCTION}, %hi(#{func})\n"
          emit "addiu #{@FUNCTION}, %lo(#{func})\n"
        end
      else
        # Load address
        load_value_into_register func, "#{@FUNCTION}"
      end

      # Perform call
      emit "jr #{@FUNCTION}\n"
      # brach delay slot
      emit "nop\n"
    end

    # Define a word with the given value
    def word value
      emit ".int #{value}\n"
    end

    # Write generated code to the given IO object.
    def write io
      io.puts ".set noat"
      @sections.each do |section,code|
        unless code.empty?
          io.puts ".section #{section.to_s}"
          io.puts code
          io.puts
        end
      end
    end

  end

  # Register class for big endian MIPS
  Voodoo::CodeGenerator.register_generator MIPSGasGenerator,
                                           :architecture => :mips,
                                           :format => :gas


  # Register class for little endian MIPS
  Voodoo::CodeGenerator.register_generator MIPSGasGenerator,
                                           :architecture => :mipsel,
                                           :format => :gas

end
