vm: compile short-circuit and/or for the dispatcher

lib/lua/compiler/bytecode.ex

@@ -145,6 +145,12 @@ defmodule Lua.Compiler.Bytecode do
   # `Lua.Compiler.Peephole` emits the instruction; codegen never does.
   @op_call_self 76
 
+  # Short-circuit `and` / `or`. Each carries the right operand's code as a
+  # nested body: the dispatcher runs it only when the left operand does not
+  # decide the result, and otherwise copies the left operand into `dest`.
+  @op_test_and 77
+  @op_test_or 78
+
   # The call opcodes whose tuple is `{tag, base, name_hint}` before
   # `annotate_line/2` bakes the source line in.
   @static_arity_calls [
@@ -188,8 +194,7 @@ defmodule Lua.Compiler.Bytecode do
   encoding — i.e. nothing in the tree fell back to the interpreter.
 
   Use this after `compile/1` to assert dispatcher coverage. A `false` result
-  means at least one prototype contains an opcode the encoder still rejects
-  (today: `:goto` / `:label`).
+  means at least one prototype contains an opcode the encoder rejects.
   """
   @spec fully_compiled?(Prototype.t()) :: boolean()
   def fully_compiled?(%Prototype{bytecode: nil}), do: false
@@ -293,6 +298,20 @@ defmodule Lua.Compiler.Bytecode do
     end
   end
 
+  defp encode({:test_and, dest, source, body}, current_line) do
+    case encode_list(body, [], current_line) do
+      {:ok, body_enc} -> {:ok, {@op_test_and, dest, source, List.to_tuple(body_enc)}}
+      :fallback -> :fallback
+    end
+  end
+
+  defp encode({:test_or, dest, source, body}, current_line) do
+    case encode_list(body, [], current_line) do
+      {:ok, body_enc} -> {:ok, {@op_test_or, dest, source, List.to_tuple(body_enc)}}
+      :fallback -> :fallback
+    end
+  end
+
   defp encode({:generic_for, base, var_regs, body}, current_line) when is_list(var_regs) do
     case encode_list(body, [], current_line) do
       {:ok, body_enc} ->
@@ -745,4 +764,6 @@ defmodule Lua.Compiler.Bytecode do
   def op_call_zero_1, do: @op_call_zero_1
   def op_call_zero_2, do: @op_call_zero_2
   def op_call_self, do: @op_call_self
+  def op_test_and, do: @op_test_and
+  def op_test_or, do: @op_test_or
 end

lib/lua/vm/dispatcher.ex

@@ -153,6 +153,12 @@ defmodule Lua.VM.Dispatcher do
   # the loop already holds everything the call needs.
   @op_call_self 76
 
+  # Short-circuit `and` / `or`: `{tag, dest, source, body}`. When `source`
+  # decides the result it is copied into `dest`; otherwise `body` (the right
+  # operand, ending in a write to `dest`) runs as a nested block.
+  @op_test_and 77
+  @op_test_or 78
+
   @doc """
   Execute a compiled prototype against `args` and `state`.
   """
@@ -893,6 +899,56 @@ defmodule Lua.VM.Dispatcher do
 
         dispatch(branch, 1, regs, upvalues, proto, state, [{code, pc + 1} | cont], frames, instruction_count, cs, cd, ou)
 
+      # `and` / `or` resume exactly like `:test`: the nested body ends and
+      # `finish_body` pops the `{code, pc + 1}` marker. When the left operand
+      # decides the result, no body runs and nothing is pushed.
+
+      {@op_test_and, dest, source, body} ->
+        case :erlang.element(source + 1, regs) do
+          v when v === nil or v === false ->
+            regs = :erlang.setelement(dest + 1, regs, v)
+            dispatch(code, pc + 1, regs, upvalues, proto, state, cont, frames, instruction_count, cs, cd, ou)
+
+          _ ->
+            dispatch(
+              body,
+              1,
+              regs,
+              upvalues,
+              proto,
+              state,
+              [{code, pc + 1} | cont],
+              frames,
+              instruction_count,
+              cs,
+              cd,
+              ou
+            )
+        end
+
+      {@op_test_or, dest, source, body} ->
+        case :erlang.element(source + 1, regs) do
+          v when v === nil or v === false ->
+            dispatch(
+              body,
+              1,
+              regs,
+              upvalues,
+              proto,
+              state,
+              [{code, pc + 1} | cont],
+              frames,
+              instruction_count,
+              cs,
+              cd,
+              ou
+            )
+
+          v ->
+            regs = :erlang.setelement(dest + 1, regs, v)
+            dispatch(code, pc + 1, regs, upvalues, proto, state, cont, frames, instruction_count, cs, cd, ou)
+        end
+
       # ── Calls ───────────────────────────────────────────────────────
       #
       # `:call_one` always asks for exactly one result placed at `base`.

test/lua/compiler/bytecode_test.exs

@@ -26,6 +26,18 @@ defmodule Lua.Compiler.BytecodeTest do
     proto
   end
 
+  # Current codegen emits no opcode the encoder rejects, so the fallback
+  # path is exercised with a synthetic prototype: `:set_list` with
+  # `count == 0` is the one shape the encoder refuses (see below).
+  defp uncovered(children \\ []) do
+    %Prototype{
+      instructions: [{:set_list, 0, 1, 0, 0}, {:return, 0, 0}],
+      prototypes: children,
+      max_registers: 2,
+      source: "test-synthetic"
+    }
+  end
+
   describe "supported-opcode coverage" do
     test "a pure-arithmetic function compiles to bytecode" do
       proto = compile!("function f(a, b) return a + b - 1 end")
@@ -176,13 +188,10 @@ defmodule Lua.Compiler.BytecodeTest do
       assert result.bytecode == nil
     end
 
-    test "short-circuit and/or falls back (test_and/test_or not covered)" do
-      # `:test_and` / `:test_or` carry a nested continuation body the encoder
-      # does not lower yet, so a function using short-circuit `and`/`or` keeps
-      # its prototype on the interpreter.
+    test "short-circuit and/or compiles (test_and/test_or)" do
       proto = compile!("function f(a, b) return a and b or 0 end")
       [fn_proto] = proto.prototypes
-      assert fn_proto.bytecode == nil
+      assert is_tuple(fn_proto.bytecode)
     end
   end
 
@@ -205,33 +214,21 @@ defmodule Lua.Compiler.BytecodeTest do
 
   describe "cascade independence" do
     test "child prototype compiles even when sibling falls back" do
-      # `pure` is pure arithmetic (covered). `impure` uses short-circuit
-      # `and`/`or` (`:test_and` / `:test_or`), which the encoder does not yet
-      # cover, so it stays on the interpreter.
-      proto =
-        compile!("""
-        function pure(a, b) return a + b end
-        function impure(a, b) return a and b or 0 end
-        """)
+      # `pure` is pure arithmetic (covered). Its sibling holds an opcode the
+      # encoder rejects, so it stays on the interpreter.
+      chunk = compile!("function pure(a, b) return a + b end")
+      proto = Bytecode.compile(%{chunk | prototypes: chunk.prototypes ++ [uncovered()]})
 
       [pure_proto, impure_proto] = proto.prototypes
       assert is_tuple(pure_proto.bytecode)
       assert impure_proto.bytecode == nil
     end
 
     test "deeply-nested function compiles even when its parent falls back" do
-      # The outer `make` uses short-circuit `and`/`or` (fallback), but the
-      # inner adder is a pure-arithmetic single-result function (compiles).
-      proto =
-        compile!("""
-        function make(a, b)
-          local guard = a and b or 0
-          local function add(x, y) return x + y + guard end
-          return add
-        end
-        """)
-
-      [make_proto] = proto.prototypes
+      # The outer `make` holds an opcode the encoder rejects (fallback), but
+      # the inner adder is a pure-arithmetic single-result function (compiles).
+      %Prototype{prototypes: [add_source]} = compile!("function add(x, y) return x + y end")
+      make_proto = Bytecode.compile(uncovered([%{add_source | bytecode: nil}]))
       [add_proto] = make_proto.prototypes
 
       assert make_proto.bytecode == nil
@@ -242,9 +239,7 @@ defmodule Lua.Compiler.BytecodeTest do
   describe "fully_compiled?/1 coverage guard" do
     # A representative corpus exercising every covered opcode family. Each
     # program must compile end-to-end — root chunk and every nested function —
-    # so the dispatcher never silently falls back to the interpreter. The only
-    # documented exception is short-circuit `and`/`or` (`:test_and` /
-    # `:test_or`), asserted separately below.
+    # so the dispatcher never silently falls back to the interpreter.
     @corpus [
       {"goto forward", "local x = 1 goto skip x = 99 ::skip:: return x"},
       {"goto backward loop", "local i = 0 ::top:: i = i + 1 if i < 5 then goto top end return i"},
@@ -264,7 +259,8 @@ defmodule Lua.Compiler.BytecodeTest do
        "local function counter() local n = 0 return function() n = n + 1 return n end end local c = counter() c() return c()"},
       {"varargs + multi-return", "local function f(...) return ... end return f(1, 2, 3)"},
       {"method call (self)", "local t = {v = 10} function t:get() return self.v end return t:get()"},
-      {"string concat", ~s{local function f(a, b) return a .. b end return f("x", "y")}}
+      {"string concat", ~s{local function f(a, b) return a .. b end return f("x", "y")}},
+      {"short-circuit and/or", "local function f(a, b) return a and b or 0 end return f(1, 2)"}
     ]
 
     for {label, src} <- @corpus do
@@ -273,12 +269,9 @@ defmodule Lua.Compiler.BytecodeTest do
       end
     end
 
-    test "short-circuit and/or is the remaining documented exception" do
-      # `:test_and` / `:test_or` (short-circuit `and`/`or`) are the only
-      # opcodes current codegen emits that the dispatcher does not yet cover,
-      # so a function using them still falls back to the interpreter. (goto /
-      # label are now covered — see the corpus above.)
-      refute Bytecode.fully_compiled?(compile!("local function f(a, b) return a and b or 0 end return f(1, 2)"))
+    test "a tree with one fallen-back prototype is not fully compiled" do
+      chunk = compile!("local function f(a, b) return a + b end return f(1, 2)")
+      refute Bytecode.fully_compiled?(Bytecode.compile(%{chunk | prototypes: chunk.prototypes ++ [uncovered()]}))
     end
   end
 
@@ -335,17 +328,8 @@ defmodule Lua.Compiler.BytecodeTest do
     end
 
     test "fallback returns a Prototype with bytecode: nil, never an error" do
-      # The encoder must not crash on any well-formed prototype. Short-circuit
-      # `and`/`or` (`:test_and` / `:test_or`) stays on the interpreter, so use
-      # it to exercise the fallback path.
-      proto =
-        compile!("""
-        function f(a, b)
-          return a and b or 0
-        end
-        """)
-
-      [fn_proto] = proto.prototypes
+      # The encoder must not crash on any well-formed prototype.
+      [fn_proto] = Bytecode.compile(uncovered([uncovered()])).prototypes
       assert %Prototype{} = fn_proto
       assert fn_proto.bytecode == nil
     end

test/lua/compiler/max_registers_invariant_test.exs

@@ -129,6 +129,9 @@ defmodule Lua.Compiler.MaxRegistersInvariantTest do
       op == Bytecode.op_equal_k() -> [1, 2]
       op == Bytecode.op_get_field_upvalue() -> [1]
       op == Bytecode.op_set_field_upvalue() -> [3]
+      # Short-circuit `and` / `or`: {tag, dest, source, body_bc}.
+      op == Bytecode.op_test_and() -> :short_circuit
+      op == Bytecode.op_test_or() -> :short_circuit
       true -> raise "register_positions/1 is missing a case for opcode #{inspect(op)}"
     end
   end
@@ -194,6 +197,12 @@ defmodule Lua.Compiler.MaxRegistersInvariantTest do
         var_max = Enum.reduce(Tuple.to_list(var_regs_tuple), -1, &max/2)
         Enum.max([base + 2, var_max, max_register_used(body_bc)])
 
+      :short_circuit ->
+        # Writes dest, reads source, recurses into the right operand's body.
+        dest = :erlang.element(2, instr)
+        source = :erlang.element(3, instr)
+        Enum.max([dest, source, max_register_used(:erlang.element(4, instr))])
+
       :call_arity_1 ->
         # {tag, base, hint, line}: reads base (the callee) and base + 1.
         :erlang.element(2, instr) + 1

test/lua/compiler/peephole_test.exs

@@ -503,8 +503,8 @@ defmodule Lua.Compiler.PeepholeTest do
     end
 
     test "the interpreter runs the fused opcode too" do
-      # `and` / `or` still fall back to the interpreter, so this child
-      # prototype carries `:call_self` with no bytecode behind it.
+      # Strip `f`'s bytecode so the interpreter, not the dispatcher, runs the
+      # `:call_self` the peephole left in its instructions.
       source = """
       local function f(n, flag)
         if n == 0 then return 0 end
@@ -516,11 +516,12 @@ defmodule Lua.Compiler.PeepholeTest do
 
       proto = compile!(source)
       [f] = proto.prototypes
-
-      assert f.bytecode == nil
       assert self_calls(proto) == 1
+
+      interpreted = %{proto | prototypes: [%{f | bytecode: nil}]}
+      state = Lua.VM.Stdlib.install(Lua.VM.State.new())
+      assert {:ok, [3], _state} = Lua.VM.execute(interpreted, state)
       assert run(source, peephole: false) == run(source, peephole: true)
-      assert {[3], _} = Lua.eval!(source)
     end
   end