compiler: peephole pass with fused and constant opcodes

lib/lua/compiler.ex

@@ -11,22 +11,31 @@ defmodule Lua.Compiler do
   alias Lua.Compiler.Codegen
   alias Lua.Compiler.GotoResolution
   alias Lua.Compiler.GotoValidation
+  alias Lua.Compiler.Peephole
   alias Lua.Compiler.Prototype
   alias Lua.Compiler.Scope
 
   @type compile_opts :: [
-          source: binary()
+          source: binary(),
+          peephole: boolean()
         ]
 
   @doc """
   Compiles a Lua AST chunk into a prototype.
 
-  After codegen, the prototype is offered to `Lua.Compiler.Bytecode` for
-  dense encoding. Sub-prototypes are encoded independently — the dispatcher
-  takes over per-prototype wherever every opcode in that prototype falls
-  within its coverage; anything else stays on the interpreter. The
-  original instruction stream is preserved either way, so error reporting
-  and tooling continue to work unchanged.
+  Codegen's output first goes through `Lua.Compiler.Peephole`, which elides
+  redundant moves, folds literals into `_k` opcode variants, fuses upvalue
+  field access, and re-derives `max_registers` from the result. Both engines
+  run the rewritten stream. Pass `peephole: false` to skip it — the
+  unoptimised stream is semantically identical and the differential tests
+  compare the two.
+
+  The prototype is then offered to `Lua.Compiler.Bytecode` for dense
+  encoding. Sub-prototypes are encoded independently — the dispatcher takes
+  over per-prototype wherever every opcode in that prototype falls within
+  its coverage; anything else stays on the interpreter. The instruction
+  stream is preserved either way, so error reporting and tooling continue to
+  work unchanged.
   """
   @spec compile(Chunk.t(), compile_opts()) :: {:ok, Prototype.t()} | {:error, term()}
   def compile(%Chunk{} = chunk, opts \\ []) do
@@ -41,19 +50,30 @@ defmodule Lua.Compiler do
     with :ok <- GotoValidation.validate(chunk),
          {:ok, scope_state} <- Scope.resolve(chunk, opts),
          {:ok, prototype} <- Codegen.generate(chunk, scope_state, opts) do
-      # Encode bytecode first (it reads the raw `:goto` / `:label` stream),
-      # then resolve gotos for the list interpreter. The two passes are
-      # independent: the dispatcher runs `bytecode`, the interpreter runs the
-      # resolved `instructions` plus `goto_targets`.
+      # Peephole first — it rewrites the raw instruction stream, so both the
+      # bytecode encoding and the interpreter's list see the same code. Then
+      # encode bytecode (it reads the raw `:goto` / `:label` stream) and
+      # finally resolve gotos for the list interpreter. The last two passes
+      # are independent: the dispatcher runs `bytecode`, the interpreter runs
+      # the resolved `instructions` plus `goto_targets`.
       prototype =
         prototype
+        |> maybe_peephole(opts)
         |> Bytecode.compile()
         |> GotoResolution.resolve()
 
       {:ok, prototype}
     end
   end
 
+  defp maybe_peephole(prototype, opts) do
+    if Keyword.get(opts, :peephole, true) do
+      Peephole.optimize(prototype)
+    else
+      prototype
+    end
+  end
+
   @doc """
   Compiles a Lua AST chunk, raising on error.
   """

lib/lua/compiler/bytecode.ex

@@ -112,6 +112,19 @@ defmodule Lua.Compiler.Bytecode do
   @op_label 60
   @op_goto 61
 
+  # Fused opcodes produced by `Lua.Compiler.Peephole`. The `_k` family
+  # carries its right operand as an inline literal instead of a register;
+  # the upvalue-field pair folds a `get_upvalue` into the field access that
+  # consumes it. Codegen never emits any of them directly.
+  @op_add_k 62
+  @op_subtract_k 63
+  @op_multiply_k 64
+  @op_less_than_k 65
+  @op_less_equal_k 66
+  @op_equal_k 67
+  @op_get_field_upvalue 68
+  @op_set_field_upvalue 69
+
   @doc """
   Compile a prototype, populating its `bytecode` field on success.
 
@@ -306,6 +319,26 @@ defmodule Lua.Compiler.Bytecode do
   defp encode({:shift_right, dest, a, b, hint_a, hint_b}), do: {:ok, {@op_shift_right, dest, a, b, hint_a, hint_b}}
   defp encode({:bitwise_not, dest, src, hint}), do: {:ok, {@op_bitwise_not, dest, src, hint}}
 
+  # Constant-folded arithmetic and comparison. Slot 4 is a literal Lua
+  # value, not a register index — the dispatcher and the interpreter both
+  # use it directly as the right operand. `hint_a` still rides along so
+  # `attempt to perform arithmetic` errors keep their `(local 'n')` suffix;
+  # the constant side never had a hint.
+  defp encode({:add_k, dest, a, constant, hint_a}), do: {:ok, {@op_add_k, dest, a, constant, hint_a}}
+  defp encode({:subtract_k, dest, a, constant, hint_a}), do: {:ok, {@op_subtract_k, dest, a, constant, hint_a}}
+  defp encode({:multiply_k, dest, a, constant, hint_a}), do: {:ok, {@op_multiply_k, dest, a, constant, hint_a}}
+  defp encode({:less_than_k, dest, a, constant}), do: {:ok, {@op_less_than_k, dest, a, constant}}
+  defp encode({:less_equal_k, dest, a, constant}), do: {:ok, {@op_less_equal_k, dest, a, constant}}
+  defp encode({:equal_k, dest, a, constant}), do: {:ok, {@op_equal_k, dest, a, constant}}
+
+  # Field access through an upvalue-held table — the shape of every global
+  # read and write outside the chunk itself.
+  defp encode({:get_field_upvalue, dest, index, name, name_hint}),
+    do: {:ok, {@op_get_field_upvalue, dest, index, name, name_hint}}
+
+  defp encode({:set_field_upvalue, index, name, value_reg, name_hint}),
+    do: {:ok, {@op_set_field_upvalue, index, name, value_reg, name_hint}}
+
   defp encode({:less_than, dest, a, b}), do: {:ok, {@op_less_than, dest, a, b}}
   defp encode({:less_equal, dest, a, b}), do: {:ok, {@op_less_equal, dest, a, b}}
   defp encode({:greater_than, dest, a, b}), do: {:ok, {@op_greater_than, dest, a, b}}
@@ -642,4 +675,12 @@ defmodule Lua.Compiler.Bytecode do
   def op_set_list_multi, do: @op_set_list_multi
   def op_label, do: @op_label
   def op_goto, do: @op_goto
+  def op_add_k, do: @op_add_k
+  def op_subtract_k, do: @op_subtract_k
+  def op_multiply_k, do: @op_multiply_k
+  def op_less_than_k, do: @op_less_than_k
+  def op_less_equal_k, do: @op_less_equal_k
+  def op_equal_k, do: @op_equal_k
+  def op_get_field_upvalue, do: @op_get_field_upvalue
+  def op_set_field_upvalue, do: @op_set_field_upvalue
 end

lib/lua/compiler/codegen.ex

@@ -120,7 +120,8 @@ defmodule Lua.Compiler.Codegen do
 
   # Returns the register-slot count an instruction proves is needed (its
   # highest written register index + 1), recursing into nested bodies.
-  defp instruction_size({:load_nil, dest, count}), do: dest + count
+  # `load_nil` clears `count + 1` registers, `dest..dest + count`.
+  defp instruction_size({:load_nil, dest, count}), do: dest + count + 1
   defp instruction_size({:vararg, base, count}) when is_integer(count) and count > 0, do: base + count
   defp instruction_size({:vararg, base, _}), do: base + 1
   defp instruction_size({:self, base, _obj, _name, _hint}), do: base + 2
@@ -160,6 +161,10 @@ defmodule Lua.Compiler.Codegen do
   defp instruction_size({:set_upvalue, _index, _source}), do: 0
   defp instruction_size({:set_open_upvalue, _reg, _source}), do: 0
 
+  # `Lua.Compiler.Peephole` emits this: operand 1 is an *upvalue* index, not
+  # a register, so it must not reach the default clause below.
+  defp instruction_size({:set_field_upvalue, _index, _name, _value, _hint}), do: 0
+
   # Everything that reaches here is an ordinary value-producing opcode —
   # `{tag, dest, ...}` whose destination is operand 1. That is the rule for
   # every load / move / arithmetic / comparison / bitwise / table-read /

lib/lua/compiler/instruction.ex

@@ -39,6 +39,16 @@ defmodule Lua.Compiler.Instruction do
   def set_table(table, key, value, name_hint \\ nil), do: {:set_table, table, key, value, name_hint}
   def get_field(dest, table, name, name_hint \\ nil), do: {:get_field, dest, table, name, name_hint}
   def set_field(table, name, value, name_hint \\ nil), do: {:set_field, table, name, value, name_hint}
+
+  # Field access through an upvalue-held table, fusing a `get_upvalue` with
+  # the `get_field` / `set_field` that consumes it. Every read or write of a
+  # global is exactly that pair (`_ENV` is an upvalue in every function but
+  # the chunk), so the fused form halves their instruction count.
+  # `Lua.Compiler.Peephole` emits these; codegen never does.
+  def get_field_upvalue(dest, index, name, name_hint \\ nil), do: {:get_field_upvalue, dest, index, name, name_hint}
+
+  def set_field_upvalue(index, name, value, name_hint \\ nil), do: {:set_field_upvalue, index, name, value, name_hint}
+
   def set_list(table, start, count, offset), do: {:set_list, table, start, count, offset}
 
   # Arithmetic.
@@ -78,6 +88,19 @@ defmodule Lua.Compiler.Instruction do
   def less_than(dest, a, b), do: {:less_than, dest, a, b}
   def less_equal(dest, a, b), do: {:less_equal, dest, a, b}
 
+  # Constant-folded variants. The right operand is an inline literal rather
+  # than a register, so the `load_constant` that materialised it disappears
+  # along with the register it occupied. Only `hint_a` survives: the
+  # constant side never carried a name hint to begin with, so error
+  # rendering is unchanged. `Lua.Compiler.Peephole` emits these; codegen
+  # never does.
+  def add_k(dest, a, constant, hint_a \\ nil), do: {:add_k, dest, a, constant, hint_a}
+  def subtract_k(dest, a, constant, hint_a \\ nil), do: {:subtract_k, dest, a, constant, hint_a}
+  def multiply_k(dest, a, constant, hint_a \\ nil), do: {:multiply_k, dest, a, constant, hint_a}
+  def equal_k(dest, a, constant), do: {:equal_k, dest, a, constant}
+  def less_than_k(dest, a, constant), do: {:less_than_k, dest, a, constant}
+  def less_equal_k(dest, a, constant), do: {:less_equal_k, dest, a, constant}
+
   # Unary / logical
   def logical_not(dest, source), do: {:not, dest, source}
   def length(dest, source), do: {:length, dest, source}