module Rufus::Lua
  class Table
    # Method to behave like Hash#fetch.
    def fetch(key, default=nil)
      value = self[key]
      if value.nil?
        default
      else
        value
      end
    end

    def has_key?(key)
      self.fetch(key) != nil
    end

    # Fix memory leak
    undef :"[]"
    def [](k)
      # Push table onto stack
      load_onto_stack

      # Push key onto stack
      stack_push(k)

      # Fetches the value from the table, replaces the key at the top of the
      # stack with the value. At the end, the stack still has the table at -2.
      Lib.lua_gettable(@pointer, -2)

      # Pop value off stack
      v = stack_pop

      # Pop table off stack
      stack_pop # <= Memory leak fixed here

      v
    end

    # Fix memory leak
    undef :"[]="
    def []=(k, v)
      # Push table onto stack
      load_onto_stack

      stack_push(k)
      stack_push(v)

      # For the VM heavy thread test, we need to cause out-of-order execution in
      # methods modifying the Lua stack. Here we tell Ruby to swap out threads
      # if we are in the thread test.
      Thread.pass if $IMMUNIO_IN_THREAD_TEST

      # Sets the key-value pair onto the table, then pops the key and value but
      # leaves the table on the stack
      Lib.lua_settable(@pointer, -3)

      # Pop table off stack
      stack_pop # <= Memory leak fixed here

      v
    end

    # Fix memory leak
    undef :objlen
    def objlen
      # Push table onto stack
      load_onto_stack

      # Gets the length of the table, but leaves the table on the stack
      len = Lib.lua_objlen(@pointer, -1)

      # Pop table off stack
      stack_pop # <= Memory leak fixed here

      len
    end

    # Override original to_h to fix bad implementation and convert nested
    # collections.
    # See http://www.lua.org/manual/5.1/manual.html#lua%5Fnext for details on
    # this crazy complex iteration system in lua...
    undef :to_h
    def to_h
      # Load table onto stack
      load_onto_stack

      # Get position of table on stack
      table_pos = stack_top

      # This is the starter for the lua_next iterator. Push a nil value as a
      # sentinel onto the stack. This tells lua_next to fetch the first key-
      # value pair from the table.
      Lib.lua_pushnil(@pointer)

      h = {}

      # Fetch the next key-value pair. The last key must be on the top of the
      # stack. The key on the stack will be overwritten with the next key
      # fetched from the table. The value will be pushed onto the stack after
      # the key.
      #
      # In order to iterate, the caller must pop the value off the stack.
      # lua_next then uses the current key, which is once again at the top of
      # the stack, to locate the next key.
      #
      # But! When there are no more key-value pairs, this call will remove the
      # key from the stack, leaving the table at the top.
      #
      # WARNING: Lua keys can be anything (yes, that includes a table). We do
      # not support tables as keys. Keys must be nice numbers or strings.
      while Lib.lua_next(@pointer, table_pos) != 0 do
        # Grab key and value from top of stack
        value = stack_fetch(-1)
        key = stack_fetch(-2)

        # Pop the value off the stack leaving the key on top
        stack_unstack

        # Recurse into table if needed
        if Table === value
          if value[1]
            h[key] = value.to_a
          else
            h[key] = value.to_h
          end
        else
          h[key] = value
        end
      end

      # Remove the table pushed onto the stack at the very beginning
      Lib.lua_remove(@pointer, table_pos)

      h
    end
  end
end